feat(charts): allow query mutator to update queries after splitting o…

…riginal sql (#21645)

Co-authored-by: Akash <Akash.Nallani@bakerhughes.com>
Co-authored-by: anallani <98122184+anallani@users.noreply.github.com>
Co-authored-by: Robert Bean <robert.bean@bakerhughes.com>
Co-authored-by: Akash <anallani@umich.edu>
Co-authored-by: AkashN7 <70606604+AkashN7@users.noreply.github.com>
Co-authored-by: Ville Brofeldt <ville.brofeldt@apple.com>

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.

superset/connectors/sqla/models.py

@@ -843,7 +843,8 @@ 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:
             sql = sql_query_mutator(
                 sql,
                 # TODO(john-bodley): Deprecate in 3.0.

superset/db_engine_specs/base.py

@@ -1264,7 +1264,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

@@ -496,6 +496,9 @@ def get_df(  # pylint: disable=too-many-locals
     ) -> 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 @@ def _log_query(sql: str) -> None:
         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(