[computes] fixed session-type extraction for connectors

The problem is that connector based query execution is not able to
reuse session to fetch results. The frontend is sending the correct
session_id but our session fetching logic got broken when the computes
was implemented. we are now looking for the session_type from
compute['name'] for computes, connector['name'] for connector and
then snippets['type'] for old config file based hive/impala sessions.

Change-Id: I386f3af1b0057e078e7b804f447f401e976b433a

Author: amitsrivastava

desktop/libs/notebook/src/notebook/connectors/hiveserver2.py

@@ -16,8 +16,6 @@
 # limitations under the License.
 
 from __future__ import division
-from future import standard_library
-standard_library.install_aliases()
 from builtins import next, object
 import binascii
 import copy
@@ -44,6 +42,9 @@
 from notebook.connectors.base import Api, QueryError, QueryExpired, OperationTimeout, OperationNotSupported, _get_snippet_name, Notebook, \
     get_interpreter, patch_snippet_for_connector
 
+from future import standard_library
+standard_library.install_aliases()
+
 if sys.version_info[0] > 2:
   from urllib.parse import quote as urllib_quote, unquote as urllib_unquote
   from django.utils.translation import gettext as _
@@ -179,7 +180,6 @@ class HS2Api(Api):
   def get_properties(lang='hive'):
     return ImpalaConfiguration.PROPERTIES if lang == 'impala' else HiveConfiguration.PROPERTIES
 
-
   @query_error_handler
   def create_session(self, lang='hive', properties=None):
     application = 'beeswax' if lang == 'hive' or lang == 'llap' else lang
@@ -246,7 +246,6 @@ def create_session(self, lang='hive', properties=None):
 
     return response
 
-
   @query_error_handler
   def close_session(self, session):
     app_name = session.get('type')
@@ -281,7 +280,6 @@ def close_session(self, session):
 
     return response
 
-
   def close_session_idle(self, notebook, session):
     idle = True
     response = {'result': []}
@@ -317,16 +315,14 @@ def execute(self, notebook, snippet):
     db = self._get_db(snippet, interpreter=self.interpreter)
 
     statement = self._get_current_statement(notebook, snippet)
-    compute = snippet.get('compute', {})
-    session_type = compute['name'] if is_compute(snippet) and compute.get('name') else snippet['type']
+    session_type = self._find_session_type(snippet)
     session = self._get_session(notebook, session_type)
 
     query = self._prepare_hql_query(snippet, statement['statement'], session)
     _session = self._get_session_by_id(notebook, session_type)
 
-
     try:
-      if statement.get('statement_id') == 0: # TODO: move this to client
+      if statement.get('statement_id') == 0:  # TODO: move this to client
         if query.database and not statement['statement'].lower().startswith('set'):
           result = db.use(query.database, session=_session)
           if result.session:
@@ -356,6 +352,14 @@ def execute(self, notebook, snippet):
 
     return response
 
+  def _find_session_type(self, snippet):
+    compute = snippet.get('compute', {})
+    if is_compute(snippet) and compute.get('name'):
+      return compute['name']
+    connector = snippet.get('connector', {})
+    if connector and connector.get('name'):
+      return connector['name']
+    return snippet.get('type')
 
   @query_error_handler
   def check_status(self, notebook, snippet):
@@ -384,7 +388,6 @@ def check_status(self, notebook, snippet):
 
     return response
 
-
   @query_error_handler
   def fetch_result(self, notebook, snippet, rows, start_over):
     db = self._get_db(snippet, interpreter=self.interpreter)
@@ -411,7 +414,6 @@ def fetch_result(self, notebook, snippet, rows, start_over):
         'type': 'table'
     }
 
-
   @query_error_handler
   def fetch_result_size(self, notebook, snippet):
     resp = {
@@ -440,7 +442,6 @@ def fetch_result_size(self, notebook, snippet):
 
     return resp
 
-
   @query_error_handler
   def cancel(self, notebook, snippet):
     db = self._get_db(snippet, interpreter=self.interpreter)
@@ -449,15 +450,13 @@ def cancel(self, notebook, snippet):
     db.cancel_operation(handle)
     return {'status': 0}
 
-
   @query_error_handler
   def get_log(self, notebook, snippet, startFrom=None, size=None):
     db = self._get_db(snippet, interpreter=self.interpreter)
 
     handle = self._get_handle(snippet)
     return db.get_log(handle, start_over=startFrom == 0)
 
-
   @query_error_handler
   def close_statement(self, notebook, snippet):
     db = self._get_db(snippet, interpreter=self.interpreter)
@@ -472,7 +471,6 @@ def close_statement(self, notebook, snippet):
         raise e
     return {'status': 0}
 
-
   def can_start_over(self, notebook, snippet):
     try:
       db = self._get_db(snippet, interpreter=self.interpreter)
@@ -484,13 +482,12 @@ def can_start_over(self, notebook, snippet):
       raise e
     return can_start_over
 
-
   @query_error_handler
   def progress(self, notebook, snippet, logs=''):
     patch_snippet_for_connector(snippet)
 
     if snippet['dialect'] == 'hive':
-      match = re.search('Total jobs = (\d+)', logs, re.MULTILINE)
+      match = re.search(r'Total jobs = (\d+)', logs, re.MULTILINE)
       total = int(match.group(1)) if match else 1
 
       started = logs.count('Starting Job')
@@ -499,13 +496,12 @@ def progress(self, notebook, snippet, logs=''):
       progress = int((started + ended) * 100 / (total * 2))
       return max(progress, 5)  # Return 5% progress as a minimum
     elif snippet['dialect'] == 'impala':
-      match = re.findall('(\d+)% Complete', logs, re.MULTILINE)
+      match = re.findall(r'(\d+)% Complete', logs, re.MULTILINE)
       # Retrieve the last reported progress percentage if it exists
       return int(match[-1]) if match and isinstance(match, list) else 0
     else:
       return 50
 
-
   @query_error_handler
   def get_jobs(self, notebook, snippet, logs):
     jobs = []
@@ -552,7 +548,6 @@ def get_jobs(self, notebook, snippet, logs):
 
     return jobs
 
-
   @query_error_handler
   def autocomplete(self, snippet, database=None, table=None, column=None, nested=None, operation=None):
     db = self._get_db(snippet, interpreter=self.interpreter)
@@ -577,7 +572,6 @@ def autocomplete(self, snippet, database=None, table=None, column=None, nested=N
 
     return resp
 
-
   @query_error_handler
   def get_sample_data(self, snippet, database=None, table=None, column=None, is_async=False, operation=None):
     try:
@@ -586,7 +580,6 @@ def get_sample_data(self, snippet, database=None, table=None, column=None, is_as
     except QueryServerException as ex:
       raise QueryError(ex.message)
 
-
   @query_error_handler
   def explain(self, notebook, snippet):
     db = self._get_db(snippet, interpreter=self.interpreter)
@@ -613,7 +606,6 @@ def explain(self, notebook, snippet):
       'statement': statement,
     }
 
-
   @query_error_handler
   def export_data_as_hdfs_file(self, snippet, target_file, overwrite):
     db = self._get_db(snippet, interpreter=self.interpreter)
@@ -626,8 +618,7 @@ def export_data_as_hdfs_file(self, snippet, target_file, overwrite):
 
     return '/filebrowser/view=%s' % urllib_quote(
         urllib_quote(target_file.encode('utf-8'), safe=SAFE_CHARACTERS_URI_COMPONENTS)
-    ) # Quote twice, because of issue in the routing on client
-
+    )  # Quote twice, because of issue in the routing on client
 
   def export_data_as_table(self, notebook, snippet, destination, is_temporary=False, location=None):
     db = self._get_db(snippet, interpreter=self.interpreter)
@@ -654,7 +645,6 @@ def export_data_as_table(self, notebook, snippet, destination, is_temporary=Fals
 
     return hql, success_url
 
-
   def export_large_data_to_hdfs(self, notebook, snippet, destination):
     response = self._get_current_statement(notebook, snippet)
     session = self._get_session(notebook, snippet['type'])
@@ -684,7 +674,6 @@ def export_large_data_to_hdfs(self, notebook, snippet, destination):
 
     return hql, success_url
 
-
   def upgrade_properties(self, lang='hive', properties=None):
     upgraded_properties = copy.deepcopy(self.get_properties(lang))
 
@@ -708,7 +697,6 @@ def upgrade_properties(self, lang='hive', properties=None):
 
     return upgraded_properties
 
-
   def _get_session(self, notebook, type='hive'):
     session = next((session for session in notebook['sessions'] if session['type'] == type), None)
     return session
@@ -723,7 +711,6 @@ def _get_session_by_id(self, notebook, type='hive'):
           filters['owner'] = self.user
         return Session.objects.get(**filters)
 
-
   def _get_hive_execution_engine(self, notebook, snippet):
     # Get hive.execution.engine from snippet properties, if none, then get from session
     properties = snippet['properties']
@@ -746,7 +733,6 @@ def _get_hive_execution_engine(self, notebook, snippet):
 
     return engine
 
-
   def _prepare_hql_query(self, snippet, statement, session):
     settings = snippet['properties'].get('settings', None)
     file_resources = snippet['properties'].get('files', None)
@@ -775,7 +761,6 @@ def _prepare_hql_query(self, snippet, statement, session):
       database=database
     )
 
-
   def get_browse_query(self, snippet, database, table, partition_spec=None):
     db = self._get_db(snippet, interpreter=self.interpreter)
     table = db.get_table(database, table)
@@ -789,7 +774,6 @@ def get_browse_query(self, snippet, database, table, partition_spec=None):
     else:
       return db.get_select_star_query(database, table, limit=100)
 
-
   def _get_handle(self, snippet):
     try:
       handle = snippet['result']['handle'].copy()
@@ -805,7 +789,6 @@ def _get_handle(self, snippet):
 
     return HiveServerQueryHandle(**handle)
 
-
   def _get_db(self, snippet, is_async=False, interpreter=None):
     if interpreter and interpreter.get('dialect'):
       dialect = interpreter['dialect']
@@ -828,7 +811,6 @@ def _get_db(self, snippet, is_async=False, interpreter=None):
     # Note: name is not used if interpreter is present
     return dbms.get(self.user, query_server=get_query_server_config(name=name, connector=interpreter))
 
-
   def _parse_job_counters(self, job_id):
     # Attempt to fetch total records from the job's Hive counter
     total_records, total_size = None, None
@@ -864,7 +846,6 @@ def _parse_job_counters(self, job_id):
 
     return total_records, total_size
 
-
   def _get_hive_result_size(self, notebook, snippet):
     total_records, total_size, msg = None, None, None
     engine = self._get_hive_execution_engine(notebook, snippet).lower()
@@ -879,8 +860,8 @@ def _get_hive_result_size(self, notebook, snippet):
       else:
         msg = _('Hive query did not execute any jobs.')
     elif engine == 'spark':
-      total_records_re = "RECORDS_OUT_0: (?P<total_records>\d+)"
-      total_size_re = "Spark Job\[[a-z0-9-]+\] Metrics[A-Za-z0-9:\s]+ResultSize: (?P<total_size>\d+)"
+      total_records_re = r"RECORDS_OUT_0: (?P<total_records>\d+)"
+      total_size_re = r"Spark Job\[[a-z0-9-]+\] Metrics[A-Za-z0-9:\s]+ResultSize: (?P<total_size>\d+)"
       total_records_match = re.search(total_records_re, logs, re.MULTILINE)
       total_size_match = re.search(total_size_re, logs, re.MULTILINE)
 
@@ -891,7 +872,6 @@ def _get_hive_result_size(self, notebook, snippet):
 
     return total_records, total_size, msg
 
-
   def _get_impala_result_size(self, notebook, snippet):
     total_records_match = None
     total_records, total_size, msg = None, None, None
@@ -904,7 +884,7 @@ def _get_impala_result_size(self, notebook, snippet):
 
       fragment = self._get_impala_query_profile(server_url, query_id=query_id)
       total_records_re = \
-          "Coordinator Fragment F\d\d.+?RowsReturned: \d+(?:.\d+[KMB])? \((?P<total_records>\d+)\).*?(Averaged Fragment F\d\d)"
+          r"Coordinator Fragment F\d\d.+?RowsReturned: \d+(?:.\d+[KMB])? \((?P<total_records>\d+)\).*?(Averaged Fragment F\d\d)"
       total_records_match = re.search(total_records_re, fragment, re.MULTILINE | re.DOTALL)
 
     if total_records_match:
@@ -917,7 +897,6 @@ def _get_impala_result_size(self, notebook, snippet):
 
     return total_records, total_size, msg
 
-
   def _get_impala_query_id(self, snippet):
     guid = None
     if 'result' in snippet and 'handle' in snippet['result'] and 'guid' in snippet['result']['handle']:
@@ -929,7 +908,6 @@ def _get_impala_query_id(self, snippet):
       LOG.warning('Snippet does not contain a valid result handle, cannot extract Impala query ID.')
     return guid
 
-
   def _get_impala_query_profile(self, server_url, query_id):
     api = get_impalad_api(user=self.user, url=server_url)
 
@@ -944,18 +922,15 @@ def _get_impala_query_profile(self, server_url, query_id):
 
     return profile
 
-
   def _get_impala_profile_plan(self, query_id, profile):
-    query_plan_re = "Query \(id=%(query_id)s\):.+?Execution Profile %(query_id)s" % {'query_id': query_id}
+    query_plan_re = r"Query \(id=%(query_id)s\):.+?Execution Profile %(query_id)s" % {'query_id': query_id}
     query_plan_match = re.search(query_plan_re, profile, re.MULTILINE | re.DOTALL)
     return query_plan_match.group() if query_plan_match else None
 
-
   def describe_column(self, notebook, snippet, database=None, table=None, column=None):
     db = self._get_db(snippet, interpreter=self.interpreter)
     return db.get_table_columns_stats(database, table, column)
 
-
   def describe_table(self, notebook, snippet, database=None, table=None):
     db = self._get_db(snippet, interpreter=self.interpreter)
     tb = db.get_table(database, table)