Эх сурвалжийг харах

HUE-6264 [hive] Refactor and componentize multiple Tez session logic

To simplify the core logic.
https://github.com/cloudera/hue/pull/436
Romain Rigaux 8 жил өмнө
parent
commit
10d7597dbb

+ 56 - 56
apps/beeswax/src/beeswax/server/hive_server2_lib.py

@@ -631,67 +631,24 @@ class HiveServerClient:
     return session
 
 
-  def call(self, fn, req, status=TStatusCode.SUCCESS_STATUS,
-           withMultipleSession=False):
-    (res, session) = self.call_return_result_and_session(fn, req, status, withMultipleSession)
+  def call(self, fn, req, status=TStatusCode.SUCCESS_STATUS, with_multiple_session=False):
+    (res, session) = self.call_return_result_and_session(fn, req, status, with_multiple_session)
     return res
 
 
-  def call_return_result_and_session(self, fn, req, status=TStatusCode.SUCCESS_STATUS,
-                                     withMultipleSession=False):
-
+  def call_return_result_and_session(self, fn, req, status=TStatusCode.SUCCESS_STATUS, with_multiple_session=False):
     n_sessions = conf.MAX_NUMBER_OF_SESSIONS.get()
 
     # When a single session is allowed, avoid multiple session logic
-    if n_sessions == 1:
-      withMultipleSession = False
+    with_multiple_session = n_sessions > 1
 
     session = None
 
-    if not withMultipleSession:
+    if not with_multiple_session:
       # Default behaviour: get one session
       session = Session.objects.get_session(self.user, self.query_server['server_name'])
-
     else:
-      # Get 2 + n_sessions sessions and filter out the busy ones
-      sessions = Session.objects.get_n_sessions(self.user, n=2 + n_sessions, application=self.query_server['server_name'])
-      LOG.debug('%s sessions found' % len(sessions))
-      if sessions:
-        # Include trashed documents to keep the query lazy
-        # and avoid retrieving all documents
-        docs = Document2.objects.get_history(doc_type='query-hive', user=self.user, include_trashed=True)
-        busy_sessions = set()
-
-        # Only check last 40 documents for performance
-        for doc in docs[:40]:
-          try:
-            snippet_data = json.loads(doc.data)['snippets'][0]
-          except (KeyError, IndexError):
-            # data might not contain a 'snippets' field or it might be empty
-            LOG.warn('No snippets in Document2 object of type query-hive')
-            continue
-          session_guid = snippet_data.get('result', {}).get('handle', {}).get('session_guid')
-          status = snippet_data.get('status')
-
-          if status in [str(QueryHistory.STATE.submitted), str(QueryHistory.STATE.running)]:
-            if session_guid is not None and session_guid not in busy_sessions:
-              busy_sessions.add(session_guid)
-
-        n_busy_sessions = 0
-        available_sessions = []
-        for session in sessions:
-          if session.guid not in busy_sessions:
-            available_sessions.append(session)
-          else:
-            n_busy_sessions += 1
-
-        if n_busy_sessions == n_sessions:
-          raise Exception('Too many open sessions. Stop a running query before starting a new one')
-
-        if available_sessions:
-          session = available_sessions[0]
-        else:
-          session = None # No available session found
+      session = self._get_tez_session(n_sessions)
 
     if session is None:
       session = self.open_session(self.user)
@@ -726,6 +683,50 @@ class HiveServerClient:
       return (res, session)
 
 
+  def _get_tez_session(self, n_sessions):
+    # Get 2 + n_sessions sessions and filter out the busy ones
+    sessions = Session.objects.get_n_sessions(self.user, n=2 + n_sessions, application=self.query_server['server_name'])
+    LOG.debug('%s sessions found' % len(sessions))
+    if sessions:
+      # Include trashed documents to keep the query lazy
+      # and avoid retrieving all documents
+      docs = Document2.objects.get_history(doc_type='query-hive', user=self.user, include_trashed=True)
+      busy_sessions = set()
+
+      # Only check last 40 documents for performance
+      for doc in docs[:40]:
+        try:
+          snippet_data = json.loads(doc.data)['snippets'][0]
+        except (KeyError, IndexError):
+          # data might not contain a 'snippets' field or it might be empty
+          LOG.warn('No snippets in Document2 object of type query-hive')
+          continue
+        session_guid = snippet_data.get('result', {}).get('handle', {}).get('session_guid')
+        status = snippet_data.get('status')
+
+        if status in [str(QueryHistory.STATE.submitted), str(QueryHistory.STATE.running)]:
+          if session_guid is not None and session_guid not in busy_sessions:
+            busy_sessions.add(session_guid)
+
+      n_busy_sessions = 0
+      available_sessions = []
+      for session in sessions:
+        if session.guid not in busy_sessions:
+          available_sessions.append(session)
+        else:
+          n_busy_sessions += 1
+
+      if n_busy_sessions == n_sessions:
+        raise Exception('Too many open sessions. Stop a running query before starting a new one')
+
+      if available_sessions:
+        session = available_sessions[0]
+      else:
+        session = None # No available session found
+
+      return session
+
+
   def close_session(self, sessionHandle):
     req = TCloseSessionReq(sessionHandle=sessionHandle)
     return self._client.CloseSession(req)
@@ -864,7 +865,7 @@ class HiveServerClient:
     return HiveServerDataTable(results, schema, operation_handle, self.query_server)
 
 
-  def execute_async_query(self, query, statement=0, withMultipleSession=False):
+  def execute_async_query(self, query, statement=0, with_multiple_session=False):
     if statement == 0:
       # Impala just has settings currently
       if self.query_server['server_name'] == 'beeswax':
@@ -880,8 +881,7 @@ class HiveServerClient:
     configuration.update(self._get_query_configuration(query))
     query_statement = query.get_query_statement(statement)
 
-    return self.execute_async_statement(statement=query_statement, confOverlay=configuration,
-                                        withMultipleSession=withMultipleSession)
+    return self.execute_async_statement(statement=query_statement, confOverlay=configuration, with_multiple_session=with_multiple_session)
 
 
   def execute_statement(self, statement, max_rows=1000, configuration={}, orientation=TFetchOrientation.FETCH_NEXT):
@@ -894,12 +894,12 @@ class HiveServerClient:
     return self.fetch_result(res.operationHandle, max_rows=max_rows, orientation=orientation), res.operationHandle
 
 
-  def execute_async_statement(self, statement, confOverlay, withMultipleSession=False):
+  def execute_async_statement(self, statement, confOverlay, with_multiple_session=False):
     if self.query_server['server_name'] == 'impala' and self.query_server['QUERY_TIMEOUT_S'] > 0:
       confOverlay['QUERY_TIMEOUT_S'] = str(self.query_server['QUERY_TIMEOUT_S'])
 
     req = TExecuteStatementReq(statement=statement.encode('utf-8'), confOverlay=confOverlay, runAsync=True)
-    (res, session) = self.call_return_result_and_session(self._client.ExecuteStatement, req, withMultipleSession=withMultipleSession)
+    (res, session) = self.call_return_result_and_session(self._client.ExecuteStatement, req, with_multiple_session=with_multiple_session)
 
     return HiveServerQueryHandle(secret=res.operationHandle.operationId.secret,
                                  guid=res.operationHandle.operationId.guid,
@@ -1161,8 +1161,8 @@ class HiveServerClientCompatible(object):
     self.query_server = client.query_server
 
 
-  def query(self, query, statement=0, withMultipleSession=False):
-    return self._client.execute_async_query(query, statement, withMultipleSession=withMultipleSession)
+  def query(self, query, statement=0, with_multiple_session=False):
+    return self._client.execute_async_query(query, statement, with_multiple_session=with_multiple_session)
 
 
   def get_state(self, handle):

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

@@ -231,7 +231,7 @@ class HS2Api(Api):
       if statement.get('statement_id') == 0:
         if query.database:
           db.use(query.database)
-      handle = db.client.query(query, withMultipleSession=True)
+      handle = db.client.query(query, with_multiple_session=True)
     except QueryServerException, ex:
       raise QueryError(ex.message, handle=statement)