Эх сурвалжийг харах

Change impersonation to use principal_username

- Change impersonation to use principal_username parameter introduced
  in PyHive 0.6.2
- Fix #1285
byungnam 4 жил өмнө
parent
commit
435d05ff0a

+ 3 - 13
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -89,7 +89,6 @@ 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__)
 
@@ -172,18 +171,6 @@ 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
-      })
-
     if self.options.get('credentials_json'):
       self.options['credentials_info'] = json.loads(
           self.options.pop('credentials_json')
@@ -196,6 +183,9 @@ class SqlAlchemyApi(Api):
           self.options.pop('connect_args')
       )
 
+    if self.options.get('has_impersonation'):
+      self.options.setdefault('connect_args', {}).setdefault('principal_username', self.user.username)
+
     options = self.options.copy()
     options.pop('session', None)
     options.pop('url', None)

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

@@ -240,7 +240,9 @@ class TestApi(object):
     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', pool_pre_ping=True)
+      create_engine.assert_called_with('presto://hue:8080/hue',
+                                       connect_args={'principal_username': 'test'},
+                                       pool_pre_ping=True)
 
 
   def test_explain(self):
@@ -250,7 +252,7 @@ class TestApi(object):
         with patch('notebook.connectors.sql_alchemy.SqlAlchemyApi._get_session') as _get_session:
 
           result = [{"id": 1}, {"select_type": "SIMPLE"}, {"Extra": "No tables used"}]
-          
+
           execute = Mock(return_value=result)
           _create_connection.return_value = Mock(
             execute=execute