Browse Source

[sparksql] Add /describe APIs, fix column /autocomplete API and skip /sample_data for transactional table (#2913)

- Added the /describe API support table and database.
- Fixed /autocomplete API for table columns for partition tables. Previously, column details in table browser was buggy when table has partition.
- Skipped /sample_data for transactional tables to suppress unnecessary error popup in the UI for transactional tables.
- Added units tests.
Harsh Gupta 3 years ago
parent
commit
4a7731d4ff

+ 222 - 13
desktop/libs/notebook/src/notebook/connectors/spark_shell.py

@@ -23,6 +23,8 @@ import time
 import textwrap
 import textwrap
 import json
 import json
 
 
+from beeswax.server.dbms import Table
+
 from desktop.conf import USE_DEFAULT_CONFIGURATION
 from desktop.conf import USE_DEFAULT_CONFIGURATION
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.i18n import force_unicode
 from desktop.lib.i18n import force_unicode
@@ -183,7 +185,6 @@ class SparkApi(Api):
 
 
 
 
   def _execute(self, api, session, snippet_type, statement):
   def _execute(self, api, session, snippet_type, statement):
-
     if not session or not self._check_session(session):
     if not session or not self._check_session(session):
       stored_session_info = self._get_session_info_from_user()
       stored_session_info = self._get_session_info_from_user()
       if stored_session_info and self._check_session(stored_session_info):
       if stored_session_info and self._check_session(stored_session_info):
@@ -247,11 +248,11 @@ class SparkApi(Api):
     try:
     try:
       response = api.fetch_data(session['id'], cell)
       response = api.fetch_data(session['id'], cell)
     except Exception as e:
     except Exception as e:
-      message = force_unicode(str(e)).lower()
+      message = force_unicode(str(e))
       if re.search("session ('\d+' )?not found", message):
       if re.search("session ('\d+' )?not found", message):
         raise SessionExpired(e)
         raise SessionExpired(e)
       else:
       else:
-        raise e
+        raise PopupException(_(message))
 
 
     content = response['output']
     content = response['output']
 
 
@@ -337,7 +338,6 @@ class SparkApi(Api):
   
   
 
 
   def _handle_session_health_check(self, session):
   def _handle_session_health_check(self, session):
-
     if not session or not self._check_session(session):
     if not session or not self._check_session(session):
       stored_session_info = self._get_session_info_from_user()
       stored_session_info = self._get_session_info_from_user()
       if stored_session_info and self._check_session(stored_session_info):
       if stored_session_info and self._check_session(stored_session_info):
@@ -480,15 +480,30 @@ class SparkApi(Api):
     columns_list = self._check_status_and_fetch_result(api, session, describe_tables_execute)
     columns_list = self._check_status_and_fetch_result(api, session, describe_tables_execute)
 
 
     if columns_list:
     if columns_list:
-      return [{
-        'name': col[0],
-        'type': col[1],
-        'comment': '',
-      } for col in columns_list['data']]
+      cols = []
+      cols_flag = True
+      for col in columns_list['data']:
+        # Columns are at the start of the result data and then comes all other info like partitions, we want to skip other info.
+        if col[0] == '' or col[0].startswith('#'):
+          cols_flag = False
+
+        if cols_flag:
+          cols.append({
+            'name': col[0],
+            'type': col[1],
+            'comment': col[2],
+          })
+
+      return cols
 
 
 
 
   def get_sample_data(self, snippet, database=None, table=None, column=None, is_async=False, operation=None):
   def get_sample_data(self, snippet, database=None, table=None, column=None, is_async=False, operation=None):
     api = self.get_api()
     api = self.get_api()
+    response = {
+      'status': 0,
+      'rows': [],
+      'full_headers': []
+    }
 
 
     # Trying to close unused sessions if there are any.
     # Trying to close unused sessions if there are any.
     # Calling the method here since this /sample_data call can be frequent enough and we dont need dedicated one.
     # Calling the method here since this /sample_data call can be frequent enough and we dont need dedicated one.
@@ -501,15 +516,20 @@ class SparkApi(Api):
     else:
     else:
       session = self.create_session(snippet.get('type'))
       session = self.create_session(snippet.get('type'))
 
 
+    # Skip sample data for transactional tables
+    notebook = {}
+    tb_describe = self.describe_table(notebook, snippet, database, table)
+
+    for stat in tb_describe['stats']:
+      if stat.get('data_type') and stat['data_type'] == 'transactional' and stat.get('col_name'):
+        return response
+
+
     statement = self._get_select_query(database, table, column, operation)
     statement = self._get_select_query(database, table, column, operation)
 
 
     sample_execute = self._execute(api, session, snippet.get('type'), statement)
     sample_execute = self._execute(api, session, snippet.get('type'), statement)
     sample_result = self._check_status_and_fetch_result(api, session, sample_execute)
     sample_result = self._check_status_and_fetch_result(api, session, sample_execute)
 
 
-    response = {
-      'status': 0,
-      'result': {}
-    }
     response['rows'] = sample_result['data']
     response['rows'] = sample_result['data']
     response['full_headers'] = sample_result['meta']
     response['full_headers'] = sample_result['meta']
 
 
@@ -539,6 +559,69 @@ class SparkApi(Api):
     return statement
     return statement
 
 
 
 
+  def describe_table(self, notebook, snippet, database=None, table=None):
+    api = self.get_api()
+
+    stored_session_info = self._get_session_info_from_user()
+    if stored_session_info and self._check_session(stored_session_info):
+      session = stored_session_info
+    else:
+      session = self.create_session(snippet.get('type'))
+
+    describe_query = 'DESCRIBE FORMATTED %(db)s.%(tb)s' % {'db': database, 'tb': table}
+    table_execute = self._execute(api, session, snippet.get('type'), describe_query)
+    table_result = self._check_status_and_fetch_result(api, session, table_execute)
+
+    tb = SparkDescribeTable(table_result)
+    tb.handle_describe_format()
+
+    return {
+      'status': 0,
+      'name': tb.name,
+      'partition_keys': tb.partition_keys,
+      'primary_keys': tb.primary_keys,
+      'cols': tb.cols,
+      'path_location': tb.path_location,
+      'hdfs_link': tb.hdfs_link,
+      'comment': tb.comment,
+      'is_view': tb.is_view,
+      'properties': tb.properties,
+      'details': tb.details,
+      'stats': tb.stats
+    }
+
+
+  def describe_database(self, notebook, snippet, database=None):
+    response = {'status': 0}
+    api = self.get_api()
+
+    stored_session_info = self._get_session_info_from_user()
+    if stored_session_info and self._check_session(stored_session_info):
+      session = stored_session_info
+    else:
+      session = self.create_session(snippet.get('type'))
+
+    describe_query = 'DESCRIBE DATABASE EXTENDED %(db)s' % {'db': database}
+    db_execute = self._execute(api, session, snippet.get('type'), describe_query)
+    db_result = self._check_status_and_fetch_result(api, session, db_execute)
+
+    for d in db_result.get('data', []):
+      if d[0] in ('Database Name', 'Namespace Name'):
+        response['db_name'] = d[1]
+      elif d[0] == 'Comment':
+        response['comment'] = d[1]
+      elif d[0] == 'Location':
+        response['location'] = d[1]
+      elif d[0] == 'Owner':
+        response['owner_name'] = d[1]
+      elif d[0] == 'Properties':
+        properties = _handle_properties_format(d[1])
+        if properties is not None:
+          response['parameters'] = properties
+
+    return response
+
+
   def _get_standalone_jobs(self, logs):
   def _get_standalone_jobs(self, logs):
     job_ids = set([])
     job_ids = set([])
 
 
@@ -612,6 +695,132 @@ class SparkApi(Api):
     self.user.profile.save()
     self.user.profile.save()
 
 
 
 
+class SparkDescribeTable(Table):
+  def __init__(self, desc_results):
+    self.data = desc_results.get('data', [])
+    self.meta = desc_results.get('meta', [])
+    self.images = desc_results.get('images', [])
+    self.type = desc_results.get('type')
+
+    self.name = ''
+    self.path_location = ''
+    self.comment = ''
+    self.owner = ''
+    self.table_type = ''
+    self.created_time = ''
+    self.serde = ''
+    self.properties = []
+    self.stats = []
+    self.cols = []
+    self.partition_keys = []
+    self.primary_keys = [] # Not implemented
+    self.is_view = False
+    self._details = None
+
+  def handle_describe_format(self):
+    cols_flag = True
+    partition_flag = True
+
+    for d in self.data:
+      if cols_flag:
+        if (d[0] == d[1] == d[2] == '') or d[0].startswith('#'):
+          cols_flag = False
+        else:
+          self.cols.append({
+            'name': d[0],
+            'type': d[1],
+            'comment': d[2]
+          })
+
+      if not cols_flag:
+        if d[0] == '# Partition Information' or partition_flag:
+          if d[0] == d[1] == d[2] == '':
+            partition_flag = False
+
+          if not d[0].startswith('#') and partition_flag:
+            self.partition_keys.append({
+              'name': d[0],
+              'type': d[1]
+            })
+
+        self.properties.append({
+          'col_name': d[0],
+          'data_type': d[1],
+          'comment': d[2]
+        })
+
+      if d[0] == 'Table':
+        self.name = d[1] 
+      elif d[0] == 'Type':
+        if 'view' in d[1].lower():
+          self.is_view = True
+        else:
+          self.table_type = d[1]
+      elif d[0] == 'Location':
+        self.path_location = d[1]
+      elif d[0] == 'Comment':
+        self.comment = d[1]
+      elif d[0] == 'Owner':
+        self.owner = d[1]
+      elif d[0] == 'Created Time':
+        self.created_time = d[1]
+      elif d[0] == 'Serde Library':
+        self.serde = d[1]
+      elif d[0] == 'Table Properties':
+        stat_list = list(map(str.strip, d[1].strip('][').replace('"', '').split(',')))
+        self.stats = [{
+          'data_type': stat.split('=')[0],
+          'col_name': stat.split('=')[1],
+          'comment': ''
+        } for stat in stat_list]
+
+  @property
+  def details(self):
+    if self._details is None:
+      if 'ParquetHiveSerDe' in self.serde:
+        details_format = 'parquet'
+      elif 'LazySimpleSerDe' in self.serde:
+        details_format = 'text'
+      else:
+        details_format = serde.rsplit('.', 1)[-1]
+
+      self._details = {
+        'stats': self.stats,
+        'properties': {
+          'owner': self.owner,
+          'create_time': self.created_time,
+          'table_type': self.table_type,
+          'format': details_format
+        }
+      }
+    return self._details
+
+
+def _handle_properties_format(str_tuple):
+  """
+  Converts tuple string to following format:
+  '((Create-by,hueuser), (Create-date,09/01/2019))' --> "{Create-by=hueuser, Create-date=09/01/2019}"
+  """
+  comma_split = '),('
+  if '), (' in str_tuple:
+    comma_split = '), ('
+
+  try:
+    # Temp conversion to dict
+    prop_dict = dict([x.split(',') for x in str_tuple[2:-2].split(comma_split)])
+
+    prop_str = "{"
+    for k, v in prop_dict.items():
+      prop_str += str(k) + '=' + str(v) + ', '
+    prop_str = prop_str[:-2] + "}"
+
+    return prop_str
+
+  except Exception as e:
+    message = force_unicode(str(e))
+    LOG.error(message)
+
+
 class SparkConfiguration(object):
 class SparkConfiguration(object):
 
 
   APP_NAME = 'spark'
   APP_NAME = 'spark'

+ 220 - 2
desktop/libs/notebook/src/notebook/connectors/spark_shell_tests.py

@@ -210,18 +210,39 @@ class TestSparkApi(object):
     self.api._execute = Mock(
     self.api._execute = Mock(
       return_value='test_value'
       return_value='test_value'
     )
     )
+    self.api.create_session = Mock(
+      return_value={
+        'id': 'test_id'
+      }
+    )
     self.api._check_status_and_fetch_result = Mock(
     self.api._check_status_and_fetch_result = Mock(
       return_value={
       return_value={
         'data': 'test_data',
         'data': 'test_data',
         'meta': 'test_meta'
         'meta': 'test_meta'
       }
       }
     )
     )
-    self.api.create_session = Mock(
+
+    # When table is transactional
+    self.api.describe_table = Mock(
       return_value={
       return_value={
-        'id': 'test_id'
+        'stats': [{
+          'data_type': 'transactional',
+          'col_name': 'true',
+          'comment': ''
+        }]
       }
       }
     )
     )
+    response = self.api.get_sample_data(snippet, 'test_db', 'test_table', 'test_column')
 
 
+    assert_equal(response['rows'], [])
+    assert_equal(response['full_headers'], [])
+
+    # When table is not transactional
+    self.api.describe_table = Mock(
+      return_value={
+        'stats': [] # No details regarding transactionality is present in describe response
+      }
+    )
     response = self.api.get_sample_data(snippet, 'test_db', 'test_table', 'test_column')
     response = self.api.get_sample_data(snippet, 'test_db', 'test_table', 'test_column')
 
 
     assert_equal(response['rows'], 'test_data')
     assert_equal(response['rows'], 'test_data')
@@ -242,6 +263,203 @@ class TestSparkApi(object):
     assert_equal(response, 'SELECT test_column\nFROM test_db.test_table\nLIMIT 100\n')
     assert_equal(response, 'SELECT test_column\nFROM test_db.test_table\nLIMIT 100\n')
 
 
 
 
+  def test_describe_database(self):
+    notebook = Mock()
+    snippet = Mock()
+    self.api.create_session = Mock(
+      return_value={
+        'id': 'test_id'
+      }
+    )
+    self.api._execute = Mock(
+      return_value='test_value'
+    )
+    self.api._check_status_and_fetch_result = Mock(
+      return_value={
+        'data': [
+          ['Namespace Name', 'employees'],
+          ['Comment', 'For software companies'],
+          ['Location', 'hdfs://test_url:8020/warehouse/tablespace/external/hive/employees.db'],
+          ['Owner', 'demo'],
+          ['Properties', '((Create-by,Kevin), (Create-date,09/01/2019))']],
+        'images': [],
+        'meta': [
+          {'comment': '', 'name': 'info_name', 'type': 'string'},
+          {'comment': '', 'name': 'info_value', 'type': 'string'}],
+        'type': 'table'}
+    )
+    response = self.api.describe_database(notebook, snippet, 'employees')
+
+    assert_equal(response, {
+      'comment': 'For software companies',
+      'db_name': 'employees',
+      'location': 'hdfs://test_url:8020/warehouse/tablespace/external/hive/employees.db',
+      'owner_name': 'demo',
+      'parameters': '{Create-by=Kevin, Create-date=09/01/2019}',
+      'status': 0})
+
+
+  def test_describe_table(self):
+    notebook = Mock()
+    snippet = Mock()
+    self.api.create_session = Mock(
+      return_value={
+        'id': 'test_id'
+      }
+    )
+    self.api._execute = Mock(
+      return_value='test_value'
+    )
+    self.api._check_status_and_fetch_result = Mock(
+      return_value={
+        'data': [
+          ['nname', 'string', None],
+          ['# Partition Information', '', ''],
+          ['# col_name', 'data_type', 'comment'],
+          ['state', 'string', 'null'],
+          ['', '', ''],
+          ['# Detailed Table Information', '', ''],
+          ['Database', 'default', ''],
+          ['Table', 'test_nonacid', ''],
+          ['Owner', 'demo', ''],
+          ['Created Time', 'Tue Jun 28 11:35:33 UTC 2022', ''],
+          ['Last Access', 'UNKNOWN', ''],
+          ['Created By', 'Spark 3.3.0.7.2.16.0-94', ''],
+          ['Type', 'EXTERNAL', ''],
+          ['Provider', 'hive', ''],
+          ['Table Properties',
+            '[TRANSLATED_TO_EXTERNAL=TRUE, bucketing_version=2, '
+            'external.table.purge=TRUE, numFilesErasureCoded=0, '
+            'transient_lastDdlTime=1656416152]',
+            ''],
+          ['Statistics', '6 bytes', ''],
+          ['Location',
+            'hdfs://test_url:8020/warehouse/tablespace/external/hive/test_nonacid',
+            ''],
+          ['Serde Library',
+            'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe',
+            ''],
+          ['InputFormat', 'org.apache.hadoop.mapred.TextInputFormat', ''],
+          ['OutputFormat',
+            'org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat',
+            ''],
+          ['Storage Properties', '[serialization.format=1]', ''],
+          ['Partition Provider', 'Catalog', '']],
+        'images': [],
+        'meta': [
+          {'comment': '', 'name': 'col_name', 'type': 'string'},
+          {'comment': '', 'name': 'data_type', 'type': 'string'},
+          {'comment': '', 'name': 'comment', 'type': 'string'}],
+        'type': 'table'
+      }
+    )
+    response = self.api.describe_table(notebook, snippet, 'default', 'test_nonacid')
+
+    assert_equal(response, {
+      'cols': [{'comment': None, 'name': 'nname', 'type': 'string'}],
+      'comment': '',
+      'details': {'properties': {
+        'create_time': 'Tue Jun 28 11:35:33 UTC 2022',
+        'format': 'text',
+        'owner': 'demo',
+        'table_type': 'EXTERNAL'},
+      'stats': [
+        {
+          'col_name': 'TRUE',
+          'comment': '',
+          'data_type': 'TRANSLATED_TO_EXTERNAL'},
+        {
+          'col_name': '2',
+          'comment': '',
+          'data_type': 'bucketing_version'},
+        {
+          'col_name': 'TRUE',
+          'comment': '',
+          'data_type': 'external.table.purge'},
+        {
+          'col_name': '0',
+          'comment': '',
+          'data_type': 'numFilesErasureCoded'},
+        {
+          'col_name': '1656416152',
+          'comment': '',
+          'data_type': 'transient_lastDdlTime'}]},
+      'hdfs_link': '/filebrowser/view=/warehouse/tablespace/external/hive/test_nonacid',
+      'is_view': False,
+      'name': 'test_nonacid',
+      'partition_keys': [{'name': 'state', 'type': 'string'}],
+      'path_location': 'hdfs://test_url:8020/warehouse/tablespace/external/hive/test_nonacid',
+      'primary_keys': [],
+      'properties': [{'col_name': '# Partition Information',
+                      'comment': '',
+                      'data_type': ''},
+                      {'col_name': '# col_name',
+                      'comment': 'comment',
+                      'data_type': 'data_type'},
+                      {'col_name': 'state', 'comment': 'null', 'data_type': 'string'},
+                      {'col_name': '', 'comment': '', 'data_type': ''},
+                      {'col_name': '# Detailed Table Information',
+                      'comment': '',
+                      'data_type': ''},
+                      {'col_name': 'Database', 'comment': '', 'data_type': 'default'},
+                      {'col_name': 'Table',
+                      'comment': '',
+                      'data_type': 'test_nonacid'},
+                      {'col_name': 'Owner', 'comment': '', 'data_type': 'demo'},
+                      {'col_name': 'Created Time',
+                      'comment': '',
+                      'data_type': 'Tue Jun 28 11:35:33 UTC 2022'},
+                      {'col_name': 'Last Access',
+                      'comment': '',
+                      'data_type': 'UNKNOWN'},
+                      {'col_name': 'Created By',
+                      'comment': '',
+                      'data_type': 'Spark 3.3.0.7.2.16.0-94'},
+                      {'col_name': 'Type', 'comment': '', 'data_type': 'EXTERNAL'},
+                      {'col_name': 'Provider', 'comment': '', 'data_type': 'hive'},
+                      {'col_name': 'Table Properties',
+                      'comment': '',
+                      'data_type': '[TRANSLATED_TO_EXTERNAL=TRUE, '
+                                    'bucketing_version=2, external.table.purge=TRUE, '
+                                    'numFilesErasureCoded=0, '
+                                    'transient_lastDdlTime=1656416152]'},
+                      {'col_name': 'Statistics',
+                      'comment': '',
+                      'data_type': '6 bytes'},
+                      {'col_name': 'Location',
+                      'comment': '',
+                      'data_type': 'hdfs://test_url:8020/warehouse/tablespace/external/hive/test_nonacid'},
+                      {'col_name': 'Serde Library',
+                      'comment': '',
+                      'data_type': 'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe'},
+                      {'col_name': 'InputFormat',
+                      'comment': '',
+                      'data_type': 'org.apache.hadoop.mapred.TextInputFormat'},
+                      {'col_name': 'OutputFormat',
+                      'comment': '',
+                      'data_type': 'org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat'},
+                      {'col_name': 'Storage Properties',
+                      'comment': '',
+                      'data_type': '[serialization.format=1]'},
+                      {'col_name': 'Partition Provider',
+                      'comment': '',
+                      'data_type': 'Catalog'}],
+      'stats': [{'col_name': 'TRUE',
+                  'comment': '',
+                  'data_type': 'TRANSLATED_TO_EXTERNAL'},
+                {'col_name': '2', 'comment': '', 'data_type': 'bucketing_version'},
+                {'col_name': 'TRUE',
+                  'comment': '',
+                  'data_type': 'external.table.purge'},
+                {'col_name': '0',
+                  'comment': '',
+                  'data_type': 'numFilesErasureCoded'},
+                {'col_name': '1656416152',
+                  'comment': '',
+                  'data_type': 'transient_lastDdlTime'}],
+      'status': 0})
+
+
   def test_get_jobs(self):
   def test_get_jobs(self):
     local_jobs = [
     local_jobs = [
       {'url': u'http://172.21.1.246:4040/jobs/job/?id=0', 'name': u'0'}
       {'url': u'http://172.21.1.246:4040/jobs/job/?id=0', 'name': u'0'}