From cf00970cde573011bb3d8fcdfc43258963f45bdf Mon Sep 17 00:00:00 2001 From: solanksh <94463220+solanksh@users.noreply.github.com> Date: Thu, 12 Jan 2023 17:34:03 +0530 Subject: [PATCH] feat(charts): allow query mutator to update queries after splitting original sql (#21645) Co-authored-by: Akash Co-authored-by: anallani <98122184+anallani@users.noreply.github.com> Co-authored-by: Robert Bean Co-authored-by: Akash Co-authored-by: AkashN7 <70606604+AkashN7@users.noreply.github.com> Co-authored-by: Ville Brofeldt --- superset/config.py | 8 ++++++++ superset/connectors/sqla/models.py | 3 ++- superset/db_engine_specs/base.py | 3 ++- superset/models/core.py | 24 ++++++++++++++++++++++-- 4 files changed, 34 insertions(+), 4 deletions(-) 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(