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 )
def test_call_session_close_idle_managed_queries(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: with patch('beeswax.server.hive_server2_lib.HiveServerClient.close_session') as close_session: with patch('beeswax.server.hive_server2_lib.HiveServerTRowSet') as HiveServerTRowSet: status = MagicMock(status=MagicMock(statusCode=0)) status_return = MagicMock(return_value=status) get_client.return_value = MagicMock( return_value=status, GetSchemas=status_return, FetchResults=status_return, GetResultSetMetadata=status_return, CloseOperation=status_return, ExecuteStatement=status_return, GetTables=status_return, GetColumns=status_return ) open_session.return_value = MagicMock(status_code=0) server_config = get_query_server_config(name='beeswax') client = HiveServerClient(server_config, self.user) res = client.get_databases() assert_equal(open_session.call_count, 1) assert_equal(close_session.call_count, 1) res = client.get_database(MagicMock()) assert_equal(open_session.call_count, 2) assert_equal(close_session.call_count, 2) res = client.get_tables_meta(MagicMock(), MagicMock()) assert_equal(open_session.call_count, 3) assert_equal(close_session.call_count, 3) res = client.get_tables(MagicMock(), MagicMock()) assert_equal(open_session.call_count, 4) assert_equal(close_session.call_count, 4) res = client.get_table(MagicMock(), MagicMock()) assert_equal(open_session.call_count, 5) assert_equal(close_session.call_count, 5) res = client.get_columns(MagicMock(), MagicMock()) assert_equal(open_session.call_count, 6) assert_equal(close_session.call_count, 6) res = client.get_partitions(MagicMock(), MagicMock()) # get_partitions does 2 requests with 1 session each assert_equal(open_session.call_count, 8) assert_equal(close_session.call_count, 8) finally: for f in finish: f()