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_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_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_open_session(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:
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'
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
)
),
get_coordinator_host=Mock(return_value='hive-host')
)
session_count = Session.objects.filter(owner=self.user, application=self.query_server['server_name']).count()
# Send open session
session = HiveServerClient(self.query_server, self.user).open_session(self.user)
assert_equal(
session_count + 1, # +1 as setUp resets the user which deletes cascade the sessions
Session.objects.filter(owner=self.user, application=self.query_server['server_name']).count()
)
session = Session.objects.get_session(self.user, self.query_server['server_name'])
secret, guid = session.get_adjusted_guid_secret()
secret, guid = HiveServerQueryHandle.get_decoded(secret, guid)
assert_equal(
original_secret,
secret
)
assert_equal(
original_guid,
guid
)
handle = session.get_handle()
assert_equal(
original_secret,
handle.sessionId.secret
)
assert_equal(
original_guid,
handle.sessionId.guid
)
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_explain(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), )),
FetchResults=Mock(return_value=Mock(
status=Mock(statusCode=TStatusCode.SUCCESS_STATUS),
results=Mock(columns=[
# Dump of `EXPLAIN SELECT 1`
Mock(stringVal=Mock(values=[
'Plan optimized by CBO.', '', 'Stage-0',
' Fetch Operator', '5 limit:-1'
],
nulls='')),
]),
schema=Mock(columns=[
Mock(columnName='Explain'),
]))),
GetResultSetMetadata=Mock(return_value=Mock(
status=Mock(statusCode=TStatusCode.SUCCESS_STATUS),
results=Mock(columns=[
Mock(stringVal=Mock(values=[
'Explain',
], nulls='')), # Fake but ok
]),
schema=Mock(columns=[
Mock(columnName='primitiveEntry 7'),
]))),
get_coordinator_host=Mock(return_value='hive-host'))
session_count = Session.objects.filter(
owner=self.user,
application=self.query_server['server_name']).count()
# Send explain
explain = HiveServerClient(self.query_server,
self.user).explain(query)
assert_equal([['Plan optimized by CBO.'], [''], ['Stage-0'],
[' Fetch Operator'], ['5 limit:-1']],
list(explain.rows()))
assert_equal(
session_count + 1,
Session.objects.filter(
owner=self.user,
application=self.query_server['server_name']).count())
def test_call_session_pool_limit(self):
finish = (MAX_NUMBER_OF_SESSIONS.set_for_testing(2),
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_tez_session'
) as get_tez_session:
get_tez_session.side_effect = Exception('')
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)
finally:
for f in finish:
f()
def test_open_session(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)),
get_coordinator_host=Mock(
return_value='hive-host'))
session_count = Session.objects.filter(
owner=self.user,
application=self.query_server['server_name']).count()
# Send open session
session = HiveServerClient(self.query_server,
self.user).open_session(self.user)
assert_equal(
session_count +
1, # +1 as setUp resets the user which deletes cascade the sessions
Session.objects.filter(
owner=self.user,
application=self.query_server['server_name']).count())
assert_equal(
session.guid,
Session.objects.get_session(
self.user, self.query_server['server_name']).guid.encode())
def get(user, query_server=None):
global DBMS_CACHE
global DBMS_CACHE_LOCK
if query_server is None:
query_server = get_query_server_config()
DBMS_CACHE_LOCK.acquire()
try:
DBMS_CACHE.setdefault(user.username, {})
if query_server['server_name'] not in DBMS_CACHE[user.username]:
# Avoid circular dependency
from beeswax.server.hive_server2_lib import HiveServerClientCompatible
if query_server['server_name'] == 'impala':
from impala.dbms import ImpalaDbms
from impala.server import ImpalaServerClient
DBMS_CACHE[user.username][
query_server['server_name']] = ImpalaDbms(
HiveServerClientCompatible(
ImpalaServerClient(query_server, user)),
QueryHistory.SERVER_TYPE[1][0])
else:
from beeswax.server.hive_server2_lib import HiveServerClient
DBMS_CACHE[user.username][
query_server['server_name']] = HiveServer2Dbms(
HiveServerClientCompatible(
HiveServerClient(query_server, user)),
QueryHistory.SERVER_TYPE[1][0])
return DBMS_CACHE[user.username][query_server['server_name']]
finally:
DBMS_CACHE_LOCK.release()
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 get(user, query_server=None, cluster=None):
global DBMS_CACHE
global DBMS_CACHE_LOCK
global RESET_HS2_QUERY_SERVER
if query_server is None:
query_server = get_query_server_config(connector=cluster)
DBMS_CACHE_LOCK.acquire()
try:
DBMS_CACHE.setdefault(user.id, {})
if query_server['server_name'] not in DBMS_CACHE[user.id]:
# Avoid circular dependency
from beeswax.server.hive_server2_lib import HiveServerClientCompatible
if query_server.get('dialect') == 'impala':
from impala.dbms import ImpalaDbms
from impala.server import ImpalaServerClient
DBMS_CACHE[user.id][query_server['server_name']] = ImpalaDbms(
HiveServerClientCompatible(ImpalaServerClient(query_server, user)),
QueryHistory.SERVER_TYPE[1][0]
)
elif query_server['server_name'] == 'hms':
from beeswax.server.hive_metastore_server import HiveMetastoreClient
DBMS_CACHE[user.id][query_server['server_name']] = HiveServer2Dbms(
HiveMetastoreClient(query_server, user),
QueryHistory.SERVER_TYPE[1][0]
)
else:
from beeswax.server.hive_server2_lib import HiveServerClient
DBMS_CACHE[user.id][query_server['server_name']] = HiveServer2Dbms(
HiveServerClientCompatible(HiveServerClient(query_server, user)),
QueryHistory.SERVER_TYPE[1][0]
)
elif RESET_HS2_QUERY_SERVER:
from beeswax.server.hive_server2_lib import HiveServerClient, HiveServerClientCompatible
RESET_HS2_QUERY_SERVER = False
LOG.debug('Setting DBMS cache for the new hs2')
DBMS_CACHE[user.id].clear()
DBMS_CACHE[user.id][query_server['server_name']] = HiveServer2Dbms(
HiveServerClientCompatible(HiveServerClient(query_server, user)),
QueryHistory.SERVER_TYPE[1][0]
)
return DBMS_CACHE[user.id][query_server['server_name']]
finally:
DBMS_CACHE_LOCK.release()
def create_from_request(request, id):
query_history = HiveServerQueryHistory.objects.get(id=id)
handle, state = _get_query_handle_and_state(query_history)
return QueryResult(
query_history,
HiveServerClientCompatible(
HiveServerClient(query_history.get_query_server_config(),
request.user)), request.fs, handle, state)
def get(user, query_server=None):
# Avoid circular dependency
from beeswax.server.hive_server2_lib import HiveServerClientCompatible, HiveServerClient
if query_server is None:
query_server = get_query_server_config()
return Dbms(HiveServerClientCompatible(HiveServerClient(query_server, user)), QueryHistory.SERVER_TYPE[1][0])
def get(user, query_server=None):
# Avoid circular dependency
from beeswax.server.hive_server2_lib import HiveServerClientCompatible, HiveServerClient
from beeswax.server.beeswax_lib import BeeswaxClient
if query_server is None:
query_server = get_query_server(support_ddl=True)
if SERVER_INTERFACE.get() == HIVE_SERVER2:
return Dbms(HiveServerClientCompatible(HiveServerClient(query_server, user)), QueryHistory.SERVER_TYPE[1][0])
else:
return Dbms(BeeswaxClient(query_server, user), QueryHistory.SERVER_TYPE[0][0])
def get(user, query_server=None):
# Avoid circular dependency
from beeswax.server.hive_server2_lib import HiveServerClientCompatible, HiveServerClient
from beeswax.server.beeswax_lib import BeeswaxClient
if query_server is None:
query_server = get_query_server_config()
if query_server['server_interface'] == HIVE_SERVER2:
return Dbms(HiveServerClientCompatible(HiveServerClient(query_server, user)), QueryHistory.SERVER_TYPE[1][0])
else:
return Dbms(BeeswaxClient(query_server, user), QueryHistory.SERVER_TYPE[0][0])
def test_get_configuration(self):
with patch('beeswax.server.hive_server2_lib.HiveServerClient.execute_query_statement') as execute_query_statement:
with patch('beeswax.server.hive_server2_lib.CONFIG_WHITELIST.get') as CONFIG_WHITELIST:
execute_query_statement.return_value = Mock(
rows=Mock(
return_value=[
['hive.server2.tez.default.queues=gethue'],
['hive.server2.tez.initialize.default.sessions=true']
]
)
)
CONFIG_WHITELIST.return_value = ['hive.server2.tez.default.queues']
configuration = HiveServerClient(self.query_server, self.user).get_configuration()
assert_equal(
configuration,
{'hive.server2.tez.default.queues': 'gethue'}
)
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)
)
def create(query_history, fs):
client = HiveServerClientCompatible(
HiveServerClient(query_history.get_query_server_config(),
query_history.owner))
handle, state = _get_query_handle_and_state(query_history)
return QueryResult(query_history, client, fs, handle, state)
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)
)
