[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

@@ -179,7 +179,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 +245,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 +279,6 @@ def close_session(self, session):
 
     return response
 
-
   def close_session_idle(self, notebook, session):
     idle = True
     response = {'result': []}
@@ -317,16 +314,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 +351,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 +387,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 +413,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 +441,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 +449,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 +470,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 +481,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 +495,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 +547,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 +571,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 +579,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 +605,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 +617,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 +644,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 +673,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 +696,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 +710,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 +732,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 +760,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 +773,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 +788,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 +810,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 +845,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 +859,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 +871,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 +883,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 +896,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 +907,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 +921,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)