feat(charts): allow query mutator to update queries after splitting original sql

superset/config.py

@@ -1049,6 +1049,13 @@ def SQL_QUERY_MUTATOR(  # pylint: disable=invalid-name,unused-argument
 ) -> str:
     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 auth provider is used by background (offline) tasks that need to access
 # protected resources. Can be overridden by end users in order to support

superset/connectors/sqla/models.py

@@ -922,7 +922,9 @@ def mutate_query_from_config(self, sql: str) -> str:
 
         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:
+            username = utils.get_username()
             sql = sql_query_mutator(
                 sql,
                 user_name=get_username(),  # TODO(john-bodley): Deprecate in 3.0.

superset/db_engine_specs/base.py

@@ -1166,7 +1166,8 @@ def process_statement(cls, statement: str, database: "Database") -> str:
         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.

superset/models/core.py

@@ -415,7 +415,11 @@ def get_df(  # pylint: disable=too-many-locals
         mutator: Optional[Callable[[pd.DataFrame], None]] = None,
     ) -> pd.DataFrame:
         sqls = self.db_engine_spec.parse_sql(sql)
+
         engine = self.get_sqla_engine(schema)
+        username = utils.get_username() or username
+        mutate_after_split = config["MUTATE_AFTER_SPLIT"]
+        sql_query_mutator = config["SQL_QUERY_MUTATOR"]
 
         def needs_conversion(df_series: pd.Series) -> bool:
             return (
@@ -438,12 +442,29 @@ def _log_query(sql: str) -> None:
         with closing(engine.raw_connection()) 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(