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