Browse Source

PR436 [hive] Multiple Hive sessions for Tez

https://github.com/cloudera/hue/pull/436

* Add session_guid information to operations

* Mark and save hiveserver2 session as invalid on error

* Better error message on missing handle fields for snippet

* Allow multiple concurrent Hive operations by using multiple sessions

* Fixed help string for max_number_of_sessions configuration parameter

* Use 40 user documents instead of 100 when checking busy HS2 sessions
antbell 9 years ago
parent
commit
874d728

+ 7 - 0
apps/beeswax/src/beeswax/conf.py

@@ -119,6 +119,13 @@ CLOSE_QUERIES = Config(
   default=False
 )
 
+MAX_NUMBER_OF_SESSIONS = Config(
+  key="max_number_of_sessions",
+  help=_t("Hue will use at most this many HiveServer2 sessions per user at a time"),
+  type=int,
+  default=1
+)
+
 THRIFT_VERSION = Config(
   key="thrift_version",
   help=_t("Thrift version to use when communicating with HiveServer2."),

+ 10 - 1
apps/beeswax/src/beeswax/models.py

@@ -395,6 +395,12 @@ class SessionManager(models.Manager):
     except Session.DoesNotExist, e:
       return None
 
+  def get_n_sessions(self, user, n, application='beeswax', filter_open=True):
+    q = self.filter(owner=user, application=application).exclude(guid='').exclude(secret='')
+    if filter_open:
+      q = q.filter(status_code=0)
+    return q.order_by("-last_used")[0:n]
+
 
 class Session(models.Model):
   """
@@ -428,7 +434,7 @@ class Session(models.Model):
 
 
 class QueryHandle(object):
-  def __init__(self, secret=None, guid=None, operation_type=None, has_result_set=None, modified_row_count=None, log_context=None):
+  def __init__(self, secret=None, guid=None, operation_type=None, has_result_set=None, modified_row_count=None, log_context=None, session_guid=None):
     self.secret = secret
     self.guid = guid
     self.operation_type = operation_type
@@ -449,10 +455,13 @@ class HiveServerQueryHandle(QueryHandle):
   QueryHandle for Hive Server 2.
 
   Store THandleIdentifier base64 encoded in order to be unicode compatible with Django.
+
+  Also store session handle if provided.
   """
   def __init__(self, **kwargs):
     super(HiveServerQueryHandle, self).__init__(**kwargs)
     self.secret, self.guid = self.get_encoded()
+    self.session_guid = kwargs.get('session_guid')
 
   def get(self):
     return self.secret, self.guid

+ 74 - 11
apps/beeswax/src/beeswax/server/hive_server2_lib.py

@@ -27,6 +27,8 @@ from django.utils.translation import ugettext as _
 
 from desktop.lib import thrift_util
 from desktop.conf import DEFAULT_USER
+from desktop.models import Document2
+from beeswax import conf
 from hadoop import cluster
 
 from TCLIService import TCLIService
@@ -40,7 +42,7 @@ from beeswax import conf as beeswax_conf
 from beeswax import hive_site
 from beeswax.hive_site import hiveserver2_use_ssl
 from beeswax.conf import CONFIG_WHITELIST, LIST_PARTITIONS_LIMIT
-from beeswax.models import Session, HiveServerQueryHandle, HiveServerQueryHistory
+from beeswax.models import Session, HiveServerQueryHandle, HiveServerQueryHistory, QueryHistory
 from beeswax.server.dbms import Table, NoSuchObjectException, DataTable,\
                                 QueryServerException
 
@@ -617,8 +619,64 @@ class HiveServerClient:
     return session
 
 
-  def call(self, fn, req, status=TStatusCode.SUCCESS_STATUS):
-    session = Session.objects.get_session(self.user, self.query_server['server_name'])
+  def call(self, fn, req, status=TStatusCode.SUCCESS_STATUS,
+           withMultipleSession=False):
+    (res, session) = self.call_return_result_and_session(fn, req, status, withMultipleSession)
+    return res
+
+
+  def call_return_result_and_session(self, fn, req, status=TStatusCode.SUCCESS_STATUS,
+                                     withMultipleSession=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
+
+    session = None
+
+    if not withMultipleSession:
+
+      # 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]:
+          snippet_data = json.loads(doc.data)['snippets'][0]
+          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
 
     if session is None:
       session = self.open_session(self.user)
@@ -632,8 +690,11 @@ class HiveServerClient:
     if res.status.statusCode == TStatusCode.ERROR_STATUS and \
         re.search('Invalid SessionHandle|Invalid session|Client session expired', res.status.errorMessage or '', re.I):
       LOG.info('Retrying with a new session because for %s of %s' % (self.user, res))
+      session.status_code = TStatusCode.INVALID_HANDLE_STATUS
+      session.save()
 
       session = self.open_session(self.user)
+
       req.sessionHandle = session.get_handle()
 
       # Get back the name of the function to call
@@ -647,7 +708,7 @@ class HiveServerClient:
         message = ''
       raise QueryServerException(Exception('Bad status for request %s:\n%s' % (req, res)), message=message)
     else:
-      return res
+      return (res, session)
 
 
   def close_session(self, sessionHandle):
@@ -762,7 +823,7 @@ class HiveServerClient:
     return HiveServerDataTable(results, schema, operation_handle, self.query_server)
 
 
-  def execute_async_query(self, query, statement=0):
+  def execute_async_query(self, query, statement=0, withMultipleSession=False):
     if statement == 0:
       # Impala just has settings currently
       if self.query_server['server_name'] == 'beeswax':
@@ -778,7 +839,8 @@ 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)
+    return self.execute_async_statement(statement=query_statement, confOverlay=configuration,
+                                        withMultipleSession=withMultipleSession)
 
 
   def execute_statement(self, statement, max_rows=1000, configuration={}, orientation=TFetchOrientation.FETCH_NEXT):
@@ -791,18 +853,19 @@ class HiveServerClient:
     return self.fetch_result(res.operationHandle, max_rows=max_rows, orientation=orientation), res.operationHandle
 
 
-  def execute_async_statement(self, statement, confOverlay):
+  def execute_async_statement(self, statement, confOverlay, withMultipleSession=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 = self.call(self._client.ExecuteStatement, req)
+    (res, session) = self.call_return_result_and_session(self._client.ExecuteStatement, req, withMultipleSession=withMultipleSession)
 
     return HiveServerQueryHandle(secret=res.operationHandle.operationId.secret,
                                  guid=res.operationHandle.operationId.guid,
                                  operation_type=res.operationHandle.operationType,
                                  has_result_set=res.operationHandle.hasResultSet,
-                                 modified_row_count=res.operationHandle.modifiedRowCount)
+                                 modified_row_count=res.operationHandle.modifiedRowCount,
+                                 session_guid=session.guid)
 
 
   def fetch_data(self, operation_handle, orientation=TFetchOrientation.FETCH_NEXT, max_rows=1000):
@@ -1046,8 +1109,8 @@ class HiveServerClientCompatible(object):
     self.query_server = client.query_server
 
 
-  def query(self, query, statement=0):
-    return self._client.execute_async_query(query, statement)
+  def query(self, query, statement=0, withMultipleSession=False):
+    return self._client.execute_async_query(query, statement, withMultipleSession=withMultipleSession)
 
 
   def get_state(self, handle):

+ 3 - 0
desktop/conf/pseudo-distributed.ini.tmpl

@@ -897,6 +897,9 @@
   # This will free all the query resources in HiveServer2, but also make its results inaccessible.
   ## close_queries=false
 
+  # Hue will use at most this many HiveServer2 sessions per user at a time
+  ## max_number_of_sessions=1
+
   # Thrift version to use when communicating with HiveServer2.
   # New column format is from version 7.
   ## thrift_version=7

+ 2 - 2
desktop/core/src/desktop/models.py

@@ -954,8 +954,8 @@ class Document2Manager(models.Manager, Document2QueryMixin):
 
     return latest_doc
 
-  def get_history(self, user, doc_type):
-    return self.documents(user, perms='owned', include_history=True).filter(type=doc_type, is_history=True)
+  def get_history(self, user, doc_type, include_trashed=False):
+    return self.documents(user, perms='owned', include_history=True, include_trashed=include_trashed).filter(type=doc_type, is_history=True)
 
   def get_home_directory(self, user):
     try:

+ 6 - 2
desktop/libs/notebook/src/notebook/connectors/hiveserver2.py

@@ -229,7 +229,7 @@ class HS2Api(Api):
     try:
       if statement.get('statement_id') == 0:
         db.use(query.database)
-      handle = db.client.query(query)
+      handle = db.client.query(query, withMultipleSession=True)
     except QueryServerException, ex:
       raise QueryError(ex.message, handle=statement)
 
@@ -242,6 +242,7 @@ class HS2Api(Api):
       'has_result_set': handle.has_result_set,
       'modified_row_count': handle.modified_row_count,
       'log_context': handle.log_context,
+      'session_guid': handle.session_guid
     }
     response.update(statement)
 
@@ -618,7 +619,10 @@ class HS2Api(Api):
 
 
   def _get_handle(self, snippet):
-    snippet['result']['handle']['secret'], snippet['result']['handle']['guid'] = HiveServerQueryHandle.get_decoded(snippet['result']['handle']['secret'], snippet['result']['handle']['guid'])
+    try:
+      snippet['result']['handle']['secret'], snippet['result']['handle']['guid'] = HiveServerQueryHandle.get_decoded(snippet['result']['handle']['secret'], snippet['result']['handle']['guid'])
+    except KeyError, ex:
+      raise Exception('Operation has no valid handle attached')
 
     for key in snippet['result']['handle'].keys():
       if key not in ('log_context', 'secret', 'has_result_set', 'operation_type', 'modified_row_count', 'guid'):