浏览代码

HUE-8735 [jb] Handle null bytes in query id

jdesjean 6 年之前
父节点
当前提交
cdab63f

+ 2 - 1
apps/impala/src/impala/api.py

@@ -32,6 +32,7 @@ from beeswax.server.dbms import get_cluster_config
 from beeswax.views import authorized_get_query_history
 
 from desktop.lib.django_util import JsonResponse
+from desktop.lib.thrift_util import unpack_guid
 from desktop.models import Document2
 
 from jobbrowser.apis.query_api import _get_api
@@ -137,7 +138,7 @@ def alanize(request):
     doc = Document2.objects.get(id=query_id)
     snippets = doc.data_dict.get('snippets', [])
     secret = snippets[0]['result']['handle']['secret']
-    impala_query_id = "%x:%x" % struct.unpack(b"QQ", base64.decodestring(secret))
+    impala_query_id = unpack_guid(base64.decodestring(secret))
     query_profile = api.get_query_profile_encoded(impala_query_id)
     profile = analyzer.analyze(analyzer.parse_data(query_profile))
     ANALYZER.pre_process(profile)

+ 6 - 3
desktop/core/src/desktop/lib/thrift_util.py

@@ -528,16 +528,19 @@ def _unpack_guid_secret_in_handle(str_args):
 
     if secret and guid:
       try:
-        encoded_secret = eval(secret.group(1))
+        encoded_secret = eval(secret.group(1)) # Does not take null bytes, but don't know how to fix.
         encoded_guid = eval(guid.group(1))
 
-        str_args = str_args.replace(secret.group(1), "%x:%x" % struct.unpack(b"QQ", encoded_secret))
-        str_args = str_args.replace(guid.group(1), "%x:%x" % struct.unpack(b"QQ", encoded_guid))
+        str_args = str_args.replace(secret.group(1), unpack_guid(encoded_secret))
+        str_args = str_args.replace(guid.group(1), unpack_guid(encoded_guid))
       except Exception, e:
         logging.warn("Unable to unpack the secret and guid in Thrift Handle: %s" % e)
 
   return str_args
 
+def unpack_guid(guid):
+  return "%016x:%016x" % struct.unpack(b"QQ", guid)
+
 def simpler_string(thrift_obj):
   """
   Strips out nulls and empty arrays from the string representation.

+ 4 - 0
desktop/core/src/desktop/lib/thrift_util_test.py

@@ -225,6 +225,10 @@ class ThriftUtilTest(unittest.TestCase):
     impala_handle = """(TGetTablesReq(schemaName=u\'default\', sessionHandle=TSessionHandle(sessionId=THandleIdentifier(secret=\'\x7f\x98\x97s\xe1\xa8G\xf4\x8a\x8a\\r\x0e6\xc2\xee\xf0\', guid=\'\xfa\xb0/\x04 \xfeDX\x99\xfcq\xff2\x07\x02\xfe\')), tableName=u\'customers\', tableTypes=None, catalogName=None),)"""
     self.assertEqual(_unpack_guid_secret_in_handle(impala_handle), """(TGetTablesReq(schemaName=u\'default\', sessionHandle=TSessionHandle(sessionId=THandleIdentifier(secret=f447a8e17397987f:f0eec2360e0d8a8a, guid=5844fe20042fb0fa:fe020732ff71fc99)), tableName=u\'customers\', tableTypes=None, catalogName=None),)""")
 
+    # Following should be added to test, but fails because eval doesn't handle null bytes
+    #impala_handle = """(TGetTablesReq(schemaName=u\'default\', sessionHandle=TSessionHandle(sessionId=THandleIdentifier(secret=\'\x7f\x98\x97s\xe1\xa8G\xf4\x8a\x8a\\r\x0e6\xc2\xee\xf0\', guid=\'\xd23\xfa\x150\xf5D\x91\x00\x00\x00\x00\xd7\xef\x91\x00\')), tableName=u\'customers\', tableTypes=None, catalogName=None),)"""
+    #self.assertEqual(_unpack_guid_secret_in_handle(impala_handle), """(TGetTablesReq(schemaName=u\'default\', sessionHandle=TSessionHandle(sessionId=THandleIdentifier(secret=f447a8e17397987f:f0eec2360e0d8a8a, guid=9144f53015fa33d2:0091efd700000000)), tableName=u\'customers\', tableTypes=None, catalogName=None),)""")
+
 class TestJsonable2Thrift(unittest.TestCase):
   """
   Tests a handful of permutations of jsonable2thrift.

+ 2 - 1
desktop/libs/metadata/src/metadata/optimizer_api.py

@@ -26,6 +26,7 @@ from django.views.decorators.http import require_POST
 
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.i18n import force_unicode
+from desktop.lib.thrift_util import unpack_guid
 from desktop.models import Document2
 from libsentry.privilege_checker import MissingSentryPrivilegeException
 from notebook.api import _get_statement
@@ -311,7 +312,7 @@ def _convert_queries(queries_data):
     try:
       snippet = query_data['snippets'][0]
       if 'guid' in snippet['result']['handle']: # Not failed query
-        original_query_id = '%s:%s' % struct.unpack(b"QQ", base64.decodestring(snippet['result']['handle']['guid']))
+        original_query_id = unpack_guid(base64.decodestring(snippet['result']['handle']['guid']))
         execution_time = snippet['result']['executionTime'] * 100 if snippet['status'] in ('available', 'expired') else -1
         statement = _clean_query(_get_statement(query_data))
         queries.append((original_query_id, execution_time, statement, snippet.get('database', 'default').strip()))

+ 5 - 5
desktop/libs/notebook/src/notebook/connectors/hiveserver2.py

@@ -35,6 +35,7 @@ from desktop.lib.exceptions import StructuredException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.i18n import force_unicode, smart_str
 from desktop.lib.rest.http_client import RestException
+from desktop.lib.thrift_util import unpack_guid
 from desktop.models import DefaultConfiguration, Document2
 from metadata.optimizer_client import OptimizerApi
 
@@ -194,7 +195,7 @@ class HS2Api(Api):
 
     try:
       decoded_guid = session.get_handle().sessionId.guid
-      response['session_id'] = "%x:%x" % struct.unpack(b"QQ", decoded_guid)
+      response['session_id'] = unpack_guid(decoded_guid)
     except Exception, e:
       LOG.warn('Failed to decode session handle: %s' % e)
 
@@ -214,7 +215,7 @@ class HS2Api(Api):
     if not session_id:
       session = Session.objects.get_session(self.user, application=app_name)
       decoded_guid = session.get_handle().sessionId.guid
-      session_decoded_id = "%x:%x" % struct.unpack(b"QQ", decoded_guid)
+      session_decoded_id = unpack_guid(decoded_guid)
       if source_method == "dt_logout":
         LOG.debug("Closing Impala session id %s on logout for user %s" % (session_decoded_id, self.user.username))
 
@@ -441,7 +442,7 @@ class HS2Api(Api):
         'finished': job.get('finished', False)
       } for job in jobs_with_state]
     elif snippet['type'] == 'impala' and ENABLE_QUERY_BROWSER.get():
-      query_id = "%x:%x" % struct.unpack(b"QQ", snippet['result']['handle']['guid'])
+      query_id = unpack_guid(snippet['result']['handle']['guid'])
       progress = min(self.progress(snippet, logs), 99) if snippet['status'] != 'available' and snippet['status'] != 'success' else 100
       jobs = [{
         'name': query_id,
@@ -876,8 +877,7 @@ DROP TABLE IF EXISTS `%(table)s`;
     guid = None
     if 'result' in snippet and 'handle' in snippet['result'] and 'guid' in snippet['result']['handle']:
       try:
-        decoded_guid = base64.decodestring(snippet['result']['handle']['guid'])
-        guid = "%x:%x" % struct.unpack(b"QQ", decoded_guid)
+        guid = unpack_guid(base64.decodestring(snippet['result']['handle']['guid']))
       except Exception, e:
         LOG.warn('Failed to decode operation handle guid: %s' % e)
     else: