Browse Source

HUE-8815 [connectors] Automatically quote some SqlAlchemy engine URLs

Romain 6 years ago
parent
commit
afbe1d3b0a

+ 9 - 0
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -41,6 +41,8 @@ Note: this is currently supporting concurrent querying by one users as engine is
 session at some point.
 Note: using the task server would not leverage any caching.
 '''
+from future import standard_library
+standard_library.install_aliases()
 
 import datetime
 import json
@@ -49,6 +51,7 @@ import uuid
 import sys
 
 from string import Template
+from urllib.parse import quote_plus
 
 from django.utils.translation import ugettext as _
 from sqlalchemy import create_engine, inspect
@@ -106,6 +109,12 @@ class SqlAlchemyApi(Api):
     else:
       url = self.options['url']
 
+    if url.startswith('awsathena+rest://'):
+      url = url.replace(url[17:37], quote_plus(url[17:37]))
+      url = url.replace(url[38:50], quote_plus(url[38:50]))
+      s3_staging_dir = url.rsplit('s3_staging_dir=', 1)[1]
+      url = url.replace(s3_staging_dir, quote_plus(s3_staging_dir))
+
     options = self.options.copy()
     options.pop('session', None)
     options.pop('url', None)

+ 9 - 0
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -42,6 +42,15 @@ class TestApi():
     grant_access("test", "default", "notebook")
 
 
+  def test_create_athena_engine(self):
+    interpreter = {
+      'options': {"url": "awsathena+rest://XXXXXXXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://gethue-athena/scratch"}
+    }
+
+    with patch('notebook.connectors.sql_alchemy.create_engine') as create_engine:
+      SqlAlchemyApi(self.user, interpreter)._create_engine()
+
+
   def test_fetch_result_empty(self):
     interpreter = {
       'options': {}

+ 7 - 1
docs/docs-site/content/administrator/configuration/editor/_index.md

@@ -319,7 +319,13 @@ Then give Hue the information about the database source:
     [[[athena]]]
        name = AWS Athena
        interface=sqlalchemy
-       options='{"url": "awsathena+rest://..."}'
+        options='{"url": "awsathena+rest://${AWS_ACCESS_KEY_ID}:${AWS_SECRET_ACCESS_KEY}@athena.${REGION}.amazonaws.com:443/${SCHEMA}?s3_staging_dir=${S3_BUCKET_DIRECTORY}"}'
+
+e.g.
+
+    options='{"url": "awsathena+rest://XXXXXXXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://gethue-athena/scratch"}'
+
+Note: Keys and S3 buckets need to be URL quoted but Hue does it automatically for you.
 
 ### Apache Druid