Pārlūkot izejas kodu

HUE-9270 [sqlalchemy] Add impersonation flag

With additional impersonation:

    [[[presto]]]
       name = Presto
       interface=sqlalchemy
       options='{"url": "presto://localhost:8080/hive/default", "has_impersonation": true}'
Romain 5 gadi atpakaļ
vecāks
revīzija
e4d046ab60

+ 10 - 6
desktop/conf.dist/hue.ini

@@ -1015,12 +1015,16 @@
     #   interface=oozie
 
     # [[[presto]]]
-    #   name=Presto SQL
-    #   interface=presto
-    #   ## Specific options for connecting to the Presto server.
-    #   ## The JDBC driver presto-jdbc.jar need to be in the CLASSPATH environment variable.
-    #   ## If 'user' and 'password' are omitted, they will be prompted in the UI.
-    #   options='{"url": "jdbc:presto://localhost:8080/catalog/schema", "driver": "io.prestosql.jdbc.PrestoDriver", "user": "root", "password": "root"}'
+    # name=Presto SQL
+    # interface=presto
+    # ## Specific options for connecting to the Presto server.
+    # ## The JDBC driver presto-jdbc.jar need to be in the CLASSPATH environment variable.
+    # ## If 'user' and 'password' are omitted, they will be prompted in the UI.
+    # options='{"url": "jdbc:presto://localhost:8080/catalog/schema", "driver": "io.prestosql.jdbc.PrestoDriver", "user": "root", "password": "root"}'
+    # ## Impersonation on.
+    # # options='{"url": "presto://localhost:8080/hive/default", "has_impersonation": true}'
+    # ## Kerberos.
+    # # options='{"url": "presto://localhost:8080/hive/default?KerberosKeytabPath=/path/to/keytab&KerberosPrincipal=principal&KerberosRemoteServiceName=service&protocol=https"'
 
     # [[[clickhouse]]]
     #   name=ClickHouse

+ 10 - 6
desktop/conf/pseudo-distributed.ini.tmpl

@@ -999,12 +999,16 @@
     #   interface=oozie
 
     # [[[presto]]]
-    #   name=Presto SQL
-    #   interface=presto
-    #   ## Specific options for connecting to the Presto server.
-    #   ## The JDBC driver presto-jdbc.jar need to be in the CLASSPATH environment variable.
-    #   ## If 'user' and 'password' are omitted, they will be prompted in the UI.
-    #   options='{"url": "jdbc:presto://localhost:8080/catalog/schema", "driver": "io.prestosql.jdbc.PrestoDriver", "user": "root", "password": "root"}'
+    # name=Presto SQL
+    # interface=presto
+    # ## Specific options for connecting to the Presto server.
+    # ## The JDBC driver presto-jdbc.jar need to be in the CLASSPATH environment variable.
+    # ## If 'user' and 'password' are omitted, they will be prompted in the UI.
+    # options='{"url": "jdbc:presto://localhost:8080/catalog/schema", "driver": "io.prestosql.jdbc.PrestoDriver", "user": "root", "password": "root"}'
+    # ## Impersonation on.
+    # # options='{"url": "presto://localhost:8080/hive/default", "has_impersonation": true}'
+    # ## Kerberos.
+    # # options='{"url": "presto://localhost:8080/hive/default?KerberosKeytabPath=/path/to/keytab&KerberosPrincipal=principal&KerberosRemoteServiceName=service&protocol=https"'
 
     # [[[clickhouse]]]
     #   name=ClickHouse

+ 15 - 1
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -87,6 +87,7 @@ else:
 ENGINES = {}
 CONNECTIONS = {}
 ENGINE_KEY = '%(username)s-%(connector_name)s'
+URL_PATTERN = '(?P<driver_name>.+?://)(?P<host>[^:/ ]+):(?P<port>[0-9]*).*'
 
 LOG = logging.getLogger(__name__)
 
@@ -116,7 +117,7 @@ def query_error_handler(func):
 class SqlAlchemyApi(Api):
 
   def __init__(self, user, interpreter):
-    super().__init__(user=user, interpreter=interpreter)
+    super(SqlAlchemyApi, self).__init__(user=user, interpreter=interpreter)
     self.options = interpreter['options']
 
     if interpreter.get('dialect_properties'):
@@ -162,10 +163,23 @@ class SqlAlchemyApi(Api):
       s3_staging_dir = url.rsplit('s3_staging_dir=', 1)[1]
       url = url.replace(s3_staging_dir, urllib_quote_plus(s3_staging_dir))
 
+    if self.options.get('has_impersonation'):
+      m = re.search(URL_PATTERN, url)
+      driver_name = m.group('driver_name')
+
+      if not driver_name:
+        raise QueryError('Driver name of %(url)s could not be found and impersonation is turned on' % {'url': url})
+
+      url = url.replace(driver_name, '%(driver_name)s%(username)s@' % {
+        'driver_name': driver_name,
+        'username': self.user.username
+      })
+
     options = self.options.copy()
     options.pop('session', None)
     options.pop('url', None)
     options.pop('has_ssh', None)
+    options.pop('has_impersonation', None)
     options.pop('ssh_server_host', None)
 
     return create_engine(url, **options)

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

@@ -179,6 +179,30 @@ class TestApi(object):
       SqlAlchemyApi(self.user, interpreter)._create_engine()
 
 
+  def test_create_engine_with_impersonation(self):
+    interpreter = {
+      'name': 'hive',
+      'options': {
+        'url': 'presto://hue:8080/hue',
+        'session': {},
+        'has_impersonation': False  # Off
+      }
+    }
+
+    with patch('notebook.connectors.sql_alchemy.create_engine') as create_engine:
+      engine = SqlAlchemyApi(self.user, interpreter)._create_engine()
+
+      create_engine.assert_called_with('presto://hue:8080/hue')
+
+
+    interpreter['options']['has_impersonation'] = True  # On
+
+    with patch('notebook.connectors.sql_alchemy.create_engine') as create_engine:
+      engine = SqlAlchemyApi(self.user, interpreter)._create_engine()
+
+      create_engine.assert_called_with('presto://test@hue:8080/hue')
+
+
   def test_check_status(self):
     notebook = Mock()
 

+ 10 - 2
docs/docs-site/content/administrator/configuration/connectors/_index.md

@@ -137,12 +137,20 @@ The dialect should be added to the Python system or Hue Python virtual environme
 
     ./build/env/bin/pip install pyhive
 
-Then give Hue the information about the database source:
+Then give Hue the information about the database source following the `presto://{presto-coordinator}:{port}/{catalog}/{schema}` format:
 
     [[[presto]]]
        name = Presto
        interface=sqlalchemy
-       options='{"url": "presto://localhost:8080/hive/default"}'
+       options='{"url": "presto://localhost:8080/tpch/default"}'
+
+With impersonation:
+
+        options='{"url": "presto://localhost:8080/tpch/default", "has_impersonation": true}'
+
+With Kerberos:
+
+        options='{"url": "presto://localhost:8080/tpch/default?KerberosKeytabPath=/path/to/keytab&KerberosPrincipal=principal&KerberosRemoteServiceName=service&protocol=https"'
 
 Alternatives.