def test_call_session_single(self): finish = ( MAX_NUMBER_OF_SESSIONS.set_for_testing(1), CLOSE_SESSIONS.set_for_testing(False) ) try: with patch('beeswax.server.hive_server2_lib.thrift_util.get_client') as get_client: with patch('beeswax.server.hive_server2_lib.HiveServerClient.open_session') as open_session: with patch('beeswax.server.hive_server2_lib.Session.objects.get_session') as get_session: open_session.return_value = MagicMock(status_code=0) get_session.return_value = None fn = MagicMock(attr='test') req = MagicMock() server_config = get_query_server_config(name='beeswax') client = HiveServerClient(server_config, self.user) (res, session1) = client.call(fn, req, status=None) open_session.assert_called_once() # Reuse session from argument (res, session2) = client.call(fn, req, status=None, session=session1) open_session.assert_called_once() # open_session should not be called again, because we're reusing session assert_equal(session1, session2) # Reuse session from get_session get_session.return_value = session1 (res, session3) = client.call(fn, req, status=None) open_session.assert_called_once() # open_session should not be called again, because we're reusing session assert_equal(session1, session3) finally: for f in finish: f()
def test_call_session_close_idle(self): finish = ( MAX_NUMBER_OF_SESSIONS.set_for_testing(-1), CLOSE_SESSIONS.set_for_testing(True) ) try: with patch('beeswax.server.hive_server2_lib.thrift_util.get_client') as get_client: with patch('beeswax.server.hive_server2_lib.HiveServerClient.open_session') as open_session: open_session.return_value = MagicMock(status_code=0) fn = MagicMock(return_value=MagicMock(status=MagicMock(statusCode=0))) req = MagicMock() server_config = get_query_server_config(name='beeswax') client = HiveServerClient(server_config, self.user) (res, session1) = client.call(fn, req, status=None) open_session.assert_called_once() # Reuse session from argument (res, session2) = client.call(fn, req, status=None, session=session1) open_session.assert_called_once() # open_session should not be called again, because we're reusing session assert_equal(session1, session2) # Create new session open_session.return_value = MagicMock(status_code=0) (res, session3) = client.call(fn, req, status=None) assert_equal(open_session.call_count, 2) assert_not_equal(session1, session3) finally: for f in finish: f()
def test_call_session_close_idle_limit(self): finish = ( MAX_NUMBER_OF_SESSIONS.set_for_testing(2), CLOSE_SESSIONS.set_for_testing(True) ) try: with patch('beeswax.server.hive_server2_lib.thrift_util.get_client') as get_client: with patch('beeswax.server.hive_server2_lib.HiveServerClient.open_session') as open_session: with patch('beeswax.server.hive_server2_lib.Session.objects.get_n_sessions') as get_n_sessions: get_n_sessions.return_value = MagicMock(count=MagicMock(return_value=2)) open_session.return_value = MagicMock(status_code=0) fn = MagicMock(return_value=MagicMock(status=MagicMock(statusCode=0))) req = MagicMock() server_config = get_query_server_config(name='beeswax') client = HiveServerClient(server_config, self.user) assert_raises(Exception, client.call, fn, req, status=None) get_n_sessions.return_value = MagicMock(count=MagicMock(return_value=1)) (res, session1) = client.call(fn, req, status=None) open_session.assert_called_once() finally: for f in finish: f()
def test_get_databases_impala_specific(self): query = Mock( get_query_statement=Mock(return_value=['SELECT 1']), settings=[] ) with patch('beeswax.server.hive_server2_lib.HiveServerTRowSet') as HiveServerTRowSet: client = HiveServerClient(self.query_server, self.user) client.call = Mock(return_value=(Mock(), Mock())) client.fetch_result = Mock(return_value=(Mock(), Mock())) client._close = Mock() client.get_databases(query) assert_not_equal( None, client.call.call_args[0][1].schemaName, client.call.call_args.args ) with patch.dict(self.query_server, {'dialect': 'impala'}, clear=True): client.get_databases(query) assert_equal( None, # Should be empty and not '*' with Impala client.call.call_args[0][1].schemaName, client.call.call_args.args )