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()
def test_get_table_with_error(self):
query = Mock(
get_query_statement=Mock(return_value=['SELECT 1']),
settings=[]
)
with patch('beeswax.server.hive_server2_lib.thrift_util.get_client') as get_client:
get_client.return_value = Mock(
OpenSession=Mock(
return_value=Mock(
status=Mock(
statusCode=TStatusCode.SUCCESS_STATUS
),
configuration={},
sessionHandle=Mock(
sessionId=Mock(
secret=b'1',
guid=b'1'
)
),
serverProtocolVersion=11
)
),
ExecuteStatement=Mock(
return_value=Mock(
status=Mock(
statusCode=TStatusCode.SUCCESS_STATUS
),
)
),
get_coordinator_host=Mock(return_value='hive-host')
)
client = HiveServerClient(self.query_server, self.user)
# Non empty error message from HS2
client._client.GetTables = Mock(
return_value=Mock(
status=Mock(
errorMessage='Error while compiling statement: FAILED: HiveAccessControlException Permission denied'
)
)
)
assert_raises(QueryServerException, client.get_table, database='database', table_name='table_name')
try:
client.get_table(database='database', table_name='table_name')
except QueryServerException as e:
assert_equal(
'Error while compiling statement: FAILED: HiveAccessControlException Permission denied',
str(e)
)
# Empty error message from HS2
get_tables_res = Mock(
status=Mock(
errorMessage=None
)
)
client._client.GetTables = Mock(
return_value=get_tables_res
)
try:
client.get_table(database='database', table_name='table_name')
except QueryServerException as e:
if sys.version_info[0] > 2:
req_string = ("TGetTablesReq(sessionHandle=TSessionHandle(sessionId=THandleIdentifier(guid=b'l\\xc4', secret=b'l\\xc4')), "
"catalogName=None, schemaName='database', tableName='table_name', tableTypes=None)")
else:
req_string = ("TGetTablesReq(schemaName='database', sessionHandle=TSessionHandle(sessionId=THandleIdentifier"
"(secret='1', guid='1')), tableName='table_name', tableTypes=None, catalogName=None)")
assert_equal(
"Bad status for request %s:\n%s" % (req_string, get_tables_res),
str(e)
)
def test_get_table_with_error(self):
query = Mock(
get_query_statement=Mock(return_value=['SELECT 1']),
settings=[]
)
original_secret = b's\xb6\x0ePP\xbdL\x17\xa3\x0f\\\xf7K\xe8Y\x1d'
original_guid = b'\xd9\xe0hT\xd6wO\xe1\xa3S\xfb\x04\xca\x93V\x01'
with patch('beeswax.server.hive_server2_lib.thrift_util.get_client') as get_client:
get_client.return_value = Mock(
OpenSession=Mock(
return_value=Mock(
status=Mock(
statusCode=TStatusCode.SUCCESS_STATUS
),
configuration={},
sessionHandle=Mock(
sessionId=Mock(
secret=original_secret,
guid=original_guid
)
),
serverProtocolVersion=11
)
),
ExecuteStatement=Mock(
return_value=Mock(
status=Mock(
statusCode=TStatusCode.SUCCESS_STATUS
),
)
),
get_coordinator_host=Mock(return_value='hive-host')
)
client = HiveServerClient(self.query_server, self.user)
# Non empty error message from HS2
client._client.GetTables = Mock(
return_value=Mock(
status=Mock(
errorMessage='Error while compiling statement: FAILED: HiveAccessControlException Permission denied'
)
)
)
assert_raises(QueryServerException, client.get_table, database='database', table_name='table_name')
try:
client.get_table(database='database', table_name='table_name')
except QueryServerException as e:
assert_equal(
'Error while compiling statement: FAILED: HiveAccessControlException Permission denied',
str(e)
)
# Empty error message from HS2
get_tables_res = Mock(
status=Mock(
errorMessage=None
)
)
client._client.GetTables = Mock(
return_value=get_tables_res
)
try:
client.get_table(database='database', table_name='table_name')
except QueryServerException as e:
if sys.version_info[0] > 2:
req_string = ("TGetTablesReq(sessionHandle=TSessionHandle(sessionId=THandleIdentifier(guid=%s, secret=%s)), "
"catalogName=None, schemaName='database', tableName='table_name', tableTypes=None)")\
% (str(original_guid), str(original_secret))
else:
req_string = ("TGetTablesReq(schemaName='database', sessionHandle=TSessionHandle(sessionId=THandleIdentifier"
"(secret='%s', guid='%s')), tableName='table_name', tableTypes=None, catalogName=None)")\
% ('s\\xb6\\x0ePP\\xbdL\\x17\\xa3\\x0f\\\\\\xf7K\\xe8Y\\x1d',
'\\xd9\\xe0hT\\xd6wO\\xe1\\xa3S\\xfb\\x04\\xca\\x93V\\x01') # manually adding '\'
assert_equal(
"Bad status for request %s:\n%s" % (req_string, get_tables_res),
str(e)
)