From f3fb502f0cc2c519bae5f3d965cd93655309466f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?T=C3=B3th=20Andor?= Date: Tue, 27 Mar 2018 16:05:51 +0200 Subject: [PATCH] Enable Impala impersonation by passing cursor configuration. This commit will only work if pull request #298 at cloudera/impyla is accepted, but it may not do no harm if not. It has also touched Hive impersonation, though it should not affect it. --- superset/db_engine_specs.py | 32 +++++++++++++++++++++++++++----- superset/models/core.py | 3 ++- superset/sql_lab.py | 2 +- superset/views/core.py | 5 +++-- 4 files changed, 33 insertions(+), 9 deletions(-) diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py index c9006bc280be9..80d8d1704234c 100644 --- a/superset/db_engine_specs.py +++ b/superset/db_engine_specs.py @@ -280,6 +280,18 @@ def get_configuration_for_impersonation(cls, uri, impersonate_user, username): """ return {} + @classmethod + def get_cursor_configuration_for_impersonation(cls, uri, impersonate_user, username): + """ + Return a cursor configuration dictionary that can be merged with other configs + that can set the correct properties for impersonating users + :param uri: URI string + :param impersonate_user: Bool indicating if impersonation is enabled + :param username: Effective username + :return: Dictionary with configs required for impersonation + """ + return {} + class PostgresBaseEngineSpec(BaseEngineSpec): """ Abstract class for Postgres 'like' databases """ @@ -1046,7 +1058,7 @@ def get_configuration_for_impersonation(cls, uri, impersonate_user, username): if (backend_name == 'hive' and 'auth' in url.query.keys() and impersonate_user is True and username is not None): configuration['hive.server2.proxy.user'] = username - return configuration + return { 'connect_args': {'configuration': configuration} } class MssqlEngineSpec(BaseEngineSpec): @@ -1203,6 +1215,8 @@ class ImpalaEngineSpec(BaseEngineSpec): time_grains = ( Grain('Time Column', _('Time Column'), '{col}'), Grain('minute', _('minute'), "TRUNC({col}, 'MI')"), + Grain('5 minute', _('5 minute'), "CAST(FLOOR(CAST(`{col}` as int)/300)*300 as TIMESTAMP)"), + Grain('10 minute', _('10 minute'), "CAST(FLOOR(CAST(`{col}` as int)/600)*600 as TIMESTAMP)"), Grain('hour', _('hour'), "TRUNC({col}, 'HH')"), Grain('day', _('day'), "TRUNC({col}, 'DD')"), Grain('week', _('week'), "TRUNC({col}, 'WW')"), @@ -1219,10 +1233,18 @@ def convert_dttm(cls, target_type, dttm): return "'{}'".format(dttm.strftime('%Y-%m-%d %H:%M:%S')) @classmethod - def get_schema_names(cls, inspector): - schemas = [row[0] for row in inspector.engine.execute('SHOW SCHEMAS') - if not row[0].startswith('_')] - return schemas + def get_configuration_for_impersonation(cls, uri, impersonate_user, username): + logging.debug('Passing Impala execution_options.cursor_configuration for impersonation') + return { 'execution_options': { 'cursor_configuration': {'impala.doas.user': username } } } + + @classmethod + def get_cursor_configuration_for_impersonation(cls, uri, impersonate_user, username): + logging.debug('Passing Impala cursor configuration for impersonation') + return { 'configuration': {'impala.doas.user': username } } + + @classmethod + def modify_url_for_impersonation(cls, url, impersonate_user, username): + pass class DruidEngineSpec(BaseEngineSpec): diff --git a/superset/models/core.py b/superset/models/core.py index cd7cc44b9eb45..bf9c9fb427dcc 100644 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -665,6 +665,7 @@ def get_sqla_engine(self, schema=None, nullpool=True, user_name=None): logging.info('Database.get_sqla_engine(). Masked URL: {0}'.format(masked_url)) params = extra.get('engine_params', {}) + self.cursor_kwargs = self.db_engine_spec.get_cursor_configuration_for_impersonation(str(url), self.impersonate_user, effective_username) if nullpool: params['poolclass'] = NullPool @@ -676,7 +677,7 @@ def get_sqla_engine(self, schema=None, nullpool=True, user_name=None): self.impersonate_user, effective_username)) if configuration: - params['connect_args'] = {'configuration': configuration} + params.update(configuration) DB_CONNECTION_MUTATOR = config.get('DB_CONNECTION_MUTATOR') if DB_CONNECTION_MUTATOR: diff --git a/superset/sql_lab.py b/superset/sql_lab.py index d28de7c354666..07b5dd7c19e79 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -212,7 +212,7 @@ def handle_error(msg): user_name=user_name, ) conn = engine.raw_connection() - cursor = conn.cursor() + cursor = conn.cursor(**database.cursor_kwargs) logging.info('Running query: \n{}'.format(executed_sql)) logging.info(query.executed_sql) cursor.execute(query.executed_sql, diff --git a/superset/views/core.py b/superset/views/core.py index a080a4b8ebe3a..1e568aab17464 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -1684,10 +1684,11 @@ def testconn(self): .get('engine_params', {}) .get('connect_args', {})) + engine_params = {} if configuration: - connect_args['configuration'] = configuration + engine_params.update(configuration) # impersonation configuration have to include parent keys, like "connect_args" and "configuration" - engine = create_engine(uri, connect_args=connect_args) + engine = create_engine(uri, **engine_params) engine.connect() return json_success(json.dumps(engine.table_names(), indent=4)) except Exception as e: