|
@@ -89,6 +89,7 @@ else:
|
|
|
ENGINES = {}
|
|
ENGINES = {}
|
|
|
CONNECTIONS = {}
|
|
CONNECTIONS = {}
|
|
|
ENGINE_KEY = '%(username)s-%(connector_name)s'
|
|
ENGINE_KEY = '%(username)s-%(connector_name)s'
|
|
|
|
|
+URL_PATTERN = '(?P<driver_name>.+?://)(?P<host>[^:/ ]+):(?P<port>[0-9]*).*'
|
|
|
|
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
|
|
|
@@ -171,6 +172,18 @@ class SqlAlchemyApi(Api):
|
|
|
s3_staging_dir = url.rsplit('s3_staging_dir=', 1)[1]
|
|
s3_staging_dir = url.rsplit('s3_staging_dir=', 1)[1]
|
|
|
url = url.replace(s3_staging_dir, urllib_quote_plus(s3_staging_dir))
|
|
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'):
|
|
if self.options.get('credentials_json'):
|
|
|
self.options['credentials_info'] = json.loads(
|
|
self.options['credentials_info'] = json.loads(
|
|
|
self.options.pop('credentials_json')
|
|
self.options.pop('credentials_json')
|
|
@@ -183,7 +196,8 @@ class SqlAlchemyApi(Api):
|
|
|
self.options.pop('connect_args')
|
|
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)
|
|
self.options.setdefault('connect_args', {}).setdefault('principal_username', self.user.username)
|
|
|
|
|
|
|
|
options = self.options.copy()
|
|
options = self.options.copy()
|
|
@@ -258,7 +272,7 @@ class SqlAlchemyApi(Api):
|
|
|
}
|
|
}
|
|
|
CONNECTIONS[guid] = cache
|
|
CONNECTIONS[guid] = cache
|
|
|
|
|
|
|
|
- response = {
|
|
|
|
|
|
|
+ response = {
|
|
|
'sync': False,
|
|
'sync': False,
|
|
|
'has_result_set': result.cursor != None,
|
|
'has_result_set': result.cursor != None,
|
|
|
'modified_row_count': 0,
|
|
'modified_row_count': 0,
|
|
@@ -330,7 +344,7 @@ class SqlAlchemyApi(Api):
|
|
|
@query_error_handler
|
|
@query_error_handler
|
|
|
def progress(self, notebook, snippet, logs=''):
|
|
def progress(self, notebook, snippet, logs=''):
|
|
|
progress = 50
|
|
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']
|
|
guid = snippet['result']['handle']['guid']
|
|
|
handle = CONNECTIONS.get(guid)
|
|
handle = CONNECTIONS.get(guid)
|
|
|
stats = None
|
|
stats = None
|