feat: add support for sparksql dialect (acryldata#4)

Author: skrydal

dev_requirements.txt

@@ -9,6 +9,7 @@ pytest-random==0.2
 pytest-timeout==1.2.0
 
 # actual dependencies: let things break if a package changes
+sqlalchemy==1.3.24
 requests>=1.0.0
 requests_kerberos>=0.12.0
 sasl>=0.2.1

pyhive/sqlalchemy_hive.py

@@ -15,7 +15,6 @@
 from sqlalchemy import exc
 from sqlalchemy import processors
 from sqlalchemy import types
-from sqlalchemy import sql
 from sqlalchemy.sql import sqltypes
 from sqlalchemy import util
 # TODO shouldn't use mysql type
@@ -249,6 +248,8 @@ class HiveDialect(default.DefaultDialect):
     supports_multivalues_insert = True
     type_compiler = HiveTypeCompiler
     supports_sane_rowcount = False
+    info_rows_delimiter = ('# Detailed Table Information', None, None)
+    partition_columns_names = ['# Partition Information']
 
     @classmethod
     def dbapi(cls):
@@ -316,7 +317,7 @@ def get_columns(self, connection, table_name, schema=None, **kw):
         rows = [row for row in rows if row[0] and row[0] != '# col_name']
         result = []
         for (col_name, full_col_type, comment) in rows:
-            if col_name == '# Partition Information':
+            if col_name in self.partition_columns_names:
                 break
             # Take out the more detailed type information
             # e.g. 'map<int,int>' -> 'map'
@@ -353,7 +354,7 @@ def get_indexes(self, connection, table_name, schema=None, **kw):
         # Filter out empty rows and comment
         rows = [row for row in rows if row[0] and row[0] != '# col_name']
         for i, (col_name, _col_type, _comment) in enumerate(rows):
-            if col_name == '# Partition Information':
+            if col_name in self.partition_columns_names:
                 break
         # Handle partition columns
         col_names = []
@@ -369,12 +370,12 @@ def get_table_names(self, connection, schema=None, **kw):
         if schema:
             query += ' IN ' + self.identifier_preparer.quote_identifier(schema)
         return [row[0] for row in connection.execute(query)]
-    
+
     def get_table_comment(self, connection, table_name, schema=None, **kw):
         rows = self._get_table_columns(connection, table_name, schema, extended=True)
 
         # Remove the column type specs.
-        start_detailed_info_index = rows.index(('# Detailed Table Information', None, None))
+        start_detailed_info_index = rows.index(self.info_rows_delimiter)
         assert start_detailed_info_index >= 0
         rows = rows[start_detailed_info_index:]
 
@@ -396,7 +397,7 @@ def get_table_comment(self, connection, table_name, schema=None, **kw):
                 # col_name == "", data_type is not None
                 prop_name = "{} {}".format(active_heading, data_type.rstrip())
                 properties[prop_name] = value.rstrip()
-        
+
         return {'text': properties.get('Table Parameters: comment', None), 'properties': properties}
 
     def do_rollback(self, dbapi_connection):

pyhive/sqlalchemy_sparksql.py

@@ -0,0 +1,57 @@
+import re
+
+from . import sqlalchemy_hive
+
+
+from sqlalchemy import exc
+from sqlalchemy.engine import default
+
+
+class SparkSqlDialect(sqlalchemy_hive.HiveDialect):
+    name = b'sparksql'
+    execution_ctx_cls = default.DefaultExecutionContext
+    info_rows_delimiter = ('# Detailed Table Information', '', '')
+    partition_columns_names = ['# Partition Information', '# Partitioning']
+
+    def _get_table_columns(self, connection, table_name, schema, extended=False):
+        full_table = table_name
+        if schema:
+            full_table = schema + '.' + table_name
+        # TODO using TGetColumnsReq hangs after sending TFetchResultsReq.
+        # Using DESCRIBE works but is uglier.
+        try:
+            # we need to set this to avoid sparksql truncating long column types (i.e. structs), arbitrarily chose 1kk
+            connection.execute("SET spark.sql.debug.maxToStringFields=1000000")
+            extended = " FORMATTED" if extended else ""
+            rows = connection.execute('DESCRIBE{} {}'.format(extended, full_table)).fetchall()
+        except exc.OperationalError as e:
+            regex_fmt = r'TExecuteStatementResp.*AnalysisException.*Table or view not found:.*{}'
+            hive_regex_fmt = r'org.apache.spark.SparkException: Cannot recognize hive type ' \
+                             r'string'
+            if re.search(regex_fmt.format(re.escape(table_name)), e.args[0]):
+                raise exc.NoSuchTableError(full_table)
+            elif re.search(hive_regex_fmt, e.args[0]):
+                raise exc.UnreflectableTableError
+            else:
+                raise
+        else:
+            return rows
+
+    def get_table_names(self, connection, schema=None, **kw):
+        # SHOW TABLES will show tables and views, SHOW VIEWS only views, if it is needed we could potentially
+        # subtract set of views from set of tables to get a proper list of only tables
+        # since hive dialect implementation does not support views extraction get_view_names need to be reimplemented
+        # too
+        query = 'SHOW TABLES'
+        if schema:
+            query += ' IN ' + self.identifier_preparer.quote_identifier(schema)
+        # returns tuples ('database', 'tableName', 'isTemporary')
+        result = connection.execute(query)
+        a = [row[1] for row in result if not row[-1]]
+        return a
+
+    def has_table(self, connection, table_name, schema=None):
+        try:
+            return super().has_table(connection, table_name, schema)
+        except exc.UnreflectableTableError:
+            return False

pyhive/tests/sqlalchemy_test_case.py

@@ -34,6 +34,9 @@ def wrapped_fn(self, *args, **kwargs):
 
 
 class SqlAlchemyTestCase(with_metaclass(abc.ABCMeta, object)):
+    complex_table = "one_row_complex"
+    complex_null_table = "one_row_complex_null"
+
     @with_engine_connection
     def test_basic_query(self, engine, connection):
         rows = connection.execute('SELECT * FROM one_row').fetchall()
@@ -43,7 +46,7 @@ def test_basic_query(self, engine, connection):
 
     @with_engine_connection
     def test_one_row_complex_null(self, engine, connection):
-        one_row_complex_null = Table('one_row_complex_null', MetaData(bind=engine), autoload=True)
+        one_row_complex_null = Table(self.complex_null_table, MetaData(bind=engine), autoload=True)
         rows = one_row_complex_null.select().execute().fetchall()
         self.assertEqual(len(rows), 1)
         self.assertEqual(list(rows[0]), [None] * len(rows[0]))
@@ -62,7 +65,7 @@ def test_reflect_no_such_table(self, engine, connection):
     @with_engine_connection
     def test_reflect_include_columns(self, engine, connection):
         """When passed include_columns, reflecttable should filter out other columns"""
-        one_row_complex = Table('one_row_complex', MetaData(bind=engine))
+        one_row_complex = Table(self.complex_table, MetaData(bind=engine))
         engine.dialect.reflecttable(
             connection, one_row_complex, include_columns=['int'],
             exclude_columns=[], resolve_fks=True)
@@ -123,7 +126,7 @@ def test_get_table_names(self, engine, connection):
         meta = MetaData()
         meta.reflect(bind=engine)
         self.assertIn('one_row', meta.tables)
-        self.assertIn('one_row_complex', meta.tables)
+        self.assertIn(self.complex_table, meta.tables)
 
         insp = sqlalchemy.inspect(engine)
         self.assertIn(
@@ -138,7 +141,7 @@ def test_has_table(self, engine, connection):
 
     @with_engine_connection
     def test_char_length(self, engine, connection):
-        one_row_complex = Table('one_row_complex', MetaData(bind=engine), autoload=True)
+        one_row_complex = Table(self.complex_table, MetaData(bind=engine), autoload=True)
         result = sqlalchemy.select([
             sqlalchemy.func.char_length(one_row_complex.c.string)
         ]).execute().scalar()