Pārlūkot izejas kodu

HUE-8740 [sqlalchemy] Support complex types in the autocomplete

e.g. ARRAY, JSON... columns in left assist or autocomplete
Romain 5 gadi atpakaļ
vecāks
revīzija
04db0b7677

+ 1 - 1
desktop/core/requirements.txt

@@ -4,7 +4,7 @@ avro-python3==1.8.2
 Babel==2.5.1
 BabelDjango==0.2.2
 boto==2.46.1
-celery==4.2.1
+celery[redis]==4.4.5  # For Python 3.8
 cffi==1.13.2
 channels==2.1.6
 channels-redis==2.3.2

+ 12 - 4
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -66,8 +66,7 @@ from string import Template
 from django.core.cache import caches
 from django.utils.translation import ugettext as _
 from sqlalchemy import create_engine, inspect, Table, MetaData
-from sqlalchemy.exc import OperationalError
-from sqlalchemy.types import NullType
+from sqlalchemy.exc import OperationalError, UnsupportedCompilationError, CompileError
 
 from desktop.lib import export_csvxls
 from desktop.lib.i18n import force_unicode
@@ -362,7 +361,7 @@ class SqlAlchemyApi(Api):
           'default': col.get('default'),
           'name': col.get('name'),
           'nullable': col.get('nullable'),
-          'type': str(col.get('type')) if not isinstance(col.get('type'), NullType) else 'Null',
+          'type': self._get_column_type_name(col),
         }
         for col in columns
       ]
@@ -393,7 +392,7 @@ class SqlAlchemyApi(Api):
       columns = assist.get_columns(database, table)
       response['full_headers'] = [{
           'name': col.get('name'),
-          'type': str(col.get('type')) if not isinstance(col.get('type'), NullType) else 'Null',
+          'type': self._get_column_type_name(col),
           'comment': ''
         } for col in columns
       ]
@@ -420,6 +419,15 @@ class SqlAlchemyApi(Api):
     })
 
 
+  def _get_column_type_name(self, col):
+    try:
+      name = str(col.get('type'))
+    except (UnsupportedCompilationError, CompileError):
+      name = col.get('type').__visit_name__.lower()
+
+    return name
+
+
   def _fix_phoenix_empty_database(self, database):
     return None if self.options['url'].startswith('phoenix://') and database == 'NULL' else database
 

+ 15 - 2
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -21,7 +21,8 @@ import logging
 import sys
 
 from nose.tools import assert_equal, assert_not_equal, assert_true, assert_false, raises
-from sqlalchemy.types import NullType
+from sqlalchemy.exc import UnsupportedCompilationError
+from sqlalchemy.types import NullType, ARRAY, JSON, VARCHAR
 
 from desktop.auth.backend import rewrite_user
 from desktop.lib.django_test_util import make_logged_in_client
@@ -369,7 +370,7 @@ class TestAutocomplete(object):
           data = SqlAlchemyApi(self.user, interpreter).autocomplete(snippet, database='database', table='table')
 
           assert_equal(data['columns'], ['col1', 'col2'])
-          assert_equal([col['type'] for col in data['extended_columns']], ['string', 'Null'])
+          assert_equal([col['type'] for col in data['extended_columns']], ['string', 'null'])
 
   def test_get_keys(self):
 
@@ -406,6 +407,18 @@ class TestUtils():
       },
     }
 
+  def test_get_column_type_name_complex(self):
+    api = SqlAlchemyApi(self.user, self.interpreter)
+
+    with patch('notebook.connectors.sql_alchemy.str') as str:
+      str.side_effect = UnsupportedCompilationError(None, None)
+
+      assert_equal(api._get_column_type_name({'type': VARCHAR}), 'varchar')  # Not complex but not testable otherwise
+      assert_equal(api._get_column_type_name({'type': NullType}), 'null')
+      assert_equal(api._get_column_type_name({'type': ARRAY}), 'array')
+      assert_equal(api._get_column_type_name({'type': JSON}), 'json')
+
+
   def test_fix_bigquery_db_prefixes(self):
     interpreter = {
       'name': 'bigquery',