Bläddra i källkod

[sqlalchemy] Add vars 'work_group' and 'catalog_name' from python lib PyAthena for AWS Athena connector (#3327)

Co-authored-by: Aleksandr Goltsov <agoltsov@evolution.com>
pid998877 2 år sedan
förälder
incheckning
25de430c66

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

@@ -76,16 +76,15 @@ from notebook.connectors.base import Api, QueryError, QueryExpired, _get_snippet
 from notebook.models import escape_rows
 
 if sys.version_info[0] > 2:
-  from urllib.parse import quote_plus as urllib_quote_plus
+  from urllib.parse import quote_plus as urllib_quote_plus, urlparse as urllib_urlparse, parse_qs as urllib_parse_qs
   from past.builtins import long
   from io import StringIO
   from django.utils.translation import gettext as _
 else:
   from django.utils.translation import ugettext as _
-  from urllib import quote_plus as urllib_quote_plus
+  from urllib import quote_plus as urllib_quote_plus, urlparse as urllib_urlparse, parse_qs as urllib_parse_qs
   from cStringIO import StringIO
 
-
 ENGINES = {}
 CONNECTIONS = {}
 ENGINE_KEY = '%(username)s-%(connector_name)s'
@@ -169,8 +168,13 @@ class SqlAlchemyApi(Api):
     if url.startswith('awsathena+rest://'):
       url = url.replace(url[17:37], urllib_quote_plus(url[17:37]))
       url = url.replace(url[38:50], urllib_quote_plus(url[38:50]))
-      s3_staging_dir = url.rsplit('s3_staging_dir=', 1)[1]
+      parsed = urllib_urlparse(url)
+      s3_staging_dir = urllib_parse_qs(parsed.query)['s3_staging_dir'][0]
       url = url.replace(s3_staging_dir, urllib_quote_plus(s3_staging_dir))
+      work_group = urllib_parse_qs(parsed.query)['work_group'][0]
+      url = url.replace(work_group, urllib_quote_plus(work_group))
+      catalog_name = urllib_parse_qs(parsed.query)['catalog_name'][0]
+      url = url.replace(catalog_name, urllib_quote_plus(catalog_name))
 
     m = re.search(URL_PATTERN, url)
     driver_name = m.group('driver_name')

+ 1 - 1
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -78,7 +78,7 @@ class TestApi(object):
       'name': 'hive',
       'options': {
         'url': 'awsathena+rest://XXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?'
-            's3_staging_dir=s3://gethue-athena/scratch'
+            's3_staging_dir=s3://gethue-athena/scratch&work_group=demo_group&catalog_name=AwsDataCatalog'
       }
     }
 

+ 4 - 1
tools/docker/hue/Dockerfile

@@ -85,7 +85,10 @@ RUN ./build/env/bin/pip install --no-cache-dir \
   # install after sqlalchemy-clickhouse and version == 1.0.4 \
   # otherwise Code: 516, Authentication failed will display \
   infi.clickhouse_orm==1.0.4 \
-  mysqlclient
+  mysqlclient \
+  PyAthena==2.25.2
+  # PyAthena == 3.x.x is latest, but not working with current configuration \
+  # otherwise, 'VisitableType' object is not subscriptable in Hue UI
 
 USER hue