diff --git a/superset/config.py b/superset/config.py index d9c36955e..984d1b229 100644 --- a/superset/config.py +++ b/superset/config.py @@ -1193,6 +1193,14 @@ def SQL_QUERY_MUTATOR( # pylint: disable=invalid-name,unused-argument return sql +# A variable that chooses whether to apply the SQL_QUERY_MUTATOR before or after splitting the input query +# It allows for using the SQL_QUERY_MUTATOR function for more than comments +# Usage: If you want to apply a change to every statement to a given query, set MUTATE_AFTER_SPLIT = True +# An example use case is if data has role based access controls, and you want to apply +# a SET ROLE statement alongside every user query. Changing this variable maintains +# functionality for both the SQL_Lab and Charts. +MUTATE_AFTER_SPLIT = False + # This allows for a user to add header data to any outgoing emails. For example, # if you need to include metadata in the header or you want to change the specifications # of the email title, header, or sender. diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index fd5942c51..c5fd025f4 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -843,7 +843,8 @@ class SqlaTable(Model, BaseDatasource): # pylint: disable=too-many-public-metho Typically adds comments to the query with context""" sql_query_mutator = config["SQL_QUERY_MUTATOR"] - if sql_query_mutator: + mutate_after_split = config["MUTATE_AFTER_SPLIT"] + if sql_query_mutator and not mutate_after_split: sql = sql_query_mutator( sql, # TODO(john-bodley): Deprecate in 3.0. diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 481ef5796..1aab100c8 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1264,7 +1264,8 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods parsed_query = ParsedQuery(statement) sql = parsed_query.stripped() sql_query_mutator = current_app.config["SQL_QUERY_MUTATOR"] - if sql_query_mutator: + mutate_after_split = current_app.config["MUTATE_AFTER_SPLIT"] + if sql_query_mutator and not mutate_after_split: sql = sql_query_mutator( sql, user_name=get_username(), # TODO(john-bodley): Deprecate in 3.0. diff --git a/superset/models/core.py b/superset/models/core.py index 9e042eeab..403d0f6c1 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -496,6 +496,9 @@ class Database( ) -> pd.DataFrame: sqls = self.db_engine_spec.parse_sql(sql) engine = self._get_sqla_engine(schema) + username = utils.get_username() + mutate_after_split = config["MUTATE_AFTER_SPLIT"] + sql_query_mutator = config["SQL_QUERY_MUTATOR"] def needs_conversion(df_series: pd.Series) -> bool: return ( @@ -518,12 +521,29 @@ class Database( with self.get_raw_connection(schema=schema) as conn: cursor = conn.cursor() for sql_ in sqls[:-1]: + if mutate_after_split: + sql_ = sql_query_mutator( + sql_, + user_name=username, + security_manager=security_manager, + database=None, + ) _log_query(sql_) self.db_engine_spec.execute(cursor, sql_) cursor.fetchall() - _log_query(sqls[-1]) - self.db_engine_spec.execute(cursor, sqls[-1]) + if mutate_after_split: + last_sql = sql_query_mutator( + sqls[-1], + user_name=username, + security_manager=security_manager, + database=None, + ) + _log_query(last_sql) + self.db_engine_spec.execute(cursor, last_sql) + else: + _log_query(sqls[-1]) + self.db_engine_spec.execute(cursor, sqls[-1]) data = self.db_engine_spec.fetch_data(cursor) result_set = SupersetResultSet(