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()