[phoenixdb] Switch to old user impersonation code for phoenix (#2852)

- Change impersonation to use principal_username broke phoenixdb
- Fix Python code styling issues

Co-authored-by: Richard Antal <[email protected]>
Co-authored-by: Harsh Gupta <[email protected]>

Author: 3 people

desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -89,6 +89,7 @@
 ENGINES = {}
 CONNECTIONS = {}
 ENGINE_KEY = '%(username)s-%(connector_name)s'
+URL_PATTERN = '(?P<driver_name>.+?://)(?P<host>[^:/ ]+):(?P<port>[0-9]*).*'
 
 LOG = logging.getLogger(__name__)
 
@@ -171,6 +172,18 @@ def _create_engine(self):
       s3_staging_dir = url.rsplit('s3_staging_dir=', 1)[1]
       url = url.replace(s3_staging_dir, urllib_quote_plus(s3_staging_dir))
 
+    m = re.search(URL_PATTERN, url)
+    driver_name = m.group('driver_name')
+    if self.options.get('has_impersonation'):
+      if not driver_name:
+        raise QueryError('Driver name of %(url)s could not be found and impersonation is turned on' % {'url': url})
+
+      if driver_name.startswith("phoenix"):
+        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')
@@ -183,7 +196,8 @@ def _create_engine(self):
           self.options.pop('connect_args')
       )
 
-    if self.options.get('has_impersonation'):
+    # phoenixdb does not support impersonation using principal_username parameter
+    if self.options.get('has_impersonation') and not driver_name.startswith("phoenix"):
       self.options.setdefault('connect_args', {}).setdefault('principal_username', self.user.username)
 
     options = self.options.copy()
@@ -258,7 +272,7 @@ def execute(self, notebook, snippet):
     }
     CONNECTIONS[guid] = cache
 
-    response =  {
+    response = {
       'sync': False,
       'has_result_set': result.cursor != None,
       'modified_row_count': 0,
@@ -330,7 +344,7 @@ def check_status(self, notebook, snippet):
   @query_error_handler
   def progress(self, notebook, snippet, logs=''):
     progress = 50
-    if self.options['url'].startswith('presto://') | self.options['url'].startswith('trino://') :
+    if self.options['url'].startswith('presto://') | self.options['url'].startswith('trino://'):
       guid = snippet['result']['handle']['guid']
       handle = CONNECTIONS.get(guid)
       stats = None

desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -244,6 +244,29 @@ def test_create_engine_with_impersonation(self):
                                        connect_args={'principal_username': 'test'},
                                        pool_pre_ping=True)
 
+  def test_create_engine_with_impersonation_phoenix(self):
+    interpreter = {
+      'name': 'phoenix',
+      'options': {
+        'url': 'phoenix://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('phoenix://hue:8080/hue', pool_pre_ping=False)
+
+
+    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('phoenix://test@hue:8080/hue', pool_pre_ping=False)
+
 
   def test_explain(self):