chore!: update mutator to take kwargs (#19083)

* update mutator to take kwargs

* update updating.md

* lint

* test that the database name is properly passed in to the mutator

UPDATING.md

@@ -37,7 +37,9 @@ assists people when migrating to a new version.
 - [17984](https://github.com/apache/superset/pull/17984): Default Flask SECRET_KEY has changed for security reasons. You should always override with your own secret. Set `PREVIOUS_SECRET_KEY` (ex: PREVIOUS_SECRET_KEY = "\2\1thisismyscretkey\1\2\\e\\y\\y\\h") with your previous key and use `superset re-encrypt-secrets` to rotate you current secrets
 - [15254](https://github.com/apache/superset/pull/15254): Previously `QUERY_COST_FORMATTERS_BY_ENGINE`, `SQL_VALIDATORS_BY_ENGINE` and `SCHEDULED_QUERIES` were expected to be defined in the feature flag dictionary in the `config.py` file. These should now be defined as a top-level config, with the feature flag dictionary being reserved for boolean only values.
 - [17539](https://github.com/apache/superset/pull/17539): all Superset CLI commands (init, load_examples and etc) require setting the FLASK_APP environment variable (which is set by default when `.flaskenv` is loaded)
-- [18970](https://github.com/apache/superset/pull/18970): Changes feature flag for the legacy datasource editor (DISABLE_LEGACY_DATASOURCE_EDITOR) in config.py to True, thus disabling the feature from being shown in the client.
+- [18970](https://github.com/apache/superset/pull/18970): Changes feature
+flag for the legacy datasource editor (DISABLE_LEGACY_DATASOURCE_EDITOR) in config.py to True, thus disabling the feature from being shown in the client.
+- [19083](https://github.com/apache/superset/pull/19083): Updates the mutator function in the config file to take a sql argument and a list of kwargs. Any `SQL_QUERY_MUTATOR` config function overrides will need to be updated to match the new set of params. It is advised regardless of the dictionary args that you list in your function arguments, to keep **kwargs as the last argument to allow for any new kwargs to be passed in.
 - [19017](https://github.com/apache/superset/pull/19017): Removes Python 3.7 support.
 - [19142](https://github.com/apache/superset/pull/19142): Changes feature flag for versioned export(VERSIONED_EXPORT) to be true.
 - [19107](https://github.com/apache/superset/pull/19107): Feature flag `SQLLAB_BACKEND_PERSISTENCE` is now on by default, which enables persisting SQL Lab tabs in the backend instead of the browser's `localStorage`.

superset/config.py

@@ -40,7 +40,6 @@ from flask import Blueprint
 from flask_appbuilder.security.manager import AUTH_DB
 from pandas._libs.parsers import STR_NA_VALUES  # pylint: disable=no-name-in-module
 from typing_extensions import Literal
-from werkzeug.local import LocalProxy
 
 from superset.constants import CHANGE_ME_SECRET_KEY
 from superset.jinja_context import BaseTemplateProcessor
@@ -1048,14 +1047,14 @@ DB_CONNECTION_MUTATOR = None
 # The use case is can be around adding some sort of comment header
 # with information such as the username and worker node information
 #
-#    def SQL_QUERY_MUTATOR(sql, user_name, security_manager, database):
+#    def SQL_QUERY_MUTATOR(sql, user_name=user_name, security_manager=security_manager, database=database):
 #        dttm = datetime.now().isoformat()
 #        return f"-- [SQL LAB] {username} {dttm}\n{sql}"
+# For backward compatibility, you can unpack any of the above arguments in your
+# function definition, but keep the **kwargs as the last argument to allow new args
+# to be added later without any errors.
 def SQL_QUERY_MUTATOR(  # pylint: disable=invalid-name,unused-argument
-    sql: str,
-    user_name: Optional[str],
-    security_manager: LocalProxy,
-    database: "Database",
+    sql: str, **kwargs: Any
 ) -> str:
     return sql
 

superset/connectors/sqla/models.py

@@ -776,7 +776,12 @@ class SqlaTable(Model, BaseDatasource):  # pylint: disable=too-many-public-metho
         sql_query_mutator = config["SQL_QUERY_MUTATOR"]
         if sql_query_mutator:
             username = utils.get_username()
-            sql = sql_query_mutator(sql, username, security_manager, self.database)
+            sql = sql_query_mutator(
+                sql,
+                user_name=username,
+                security_manager=security_manager,
+                database=self.database,
+            )
         return sql
 
     def get_template_processor(self, **kwargs: Any) -> BaseTemplateProcessor:

superset/db_engine_specs/base.py

@@ -1138,7 +1138,12 @@ class BaseEngineSpec:  # pylint: disable=too-many-public-methods
         sql = parsed_query.stripped()
         sql_query_mutator = current_app.config["SQL_QUERY_MUTATOR"]
         if sql_query_mutator:
-            sql = sql_query_mutator(sql, user_name, security_manager, database)
+            sql = sql_query_mutator(
+                sql,
+                user_name=user_name,
+                security_manager=security_manager,
+                database=database,
+            )
 
         return sql
 

superset/sql_lab.py

@@ -225,7 +225,9 @@ def execute_sql_statement(  # pylint: disable=too-many-arguments,too-many-locals
         sql = apply_limit_if_exists(database, increased_limit, query, sql)
 
     # Hook to allow environment-specific mutation (usually comments) to the SQL
-    sql = SQL_QUERY_MUTATOR(sql, user_name, security_manager, database)
+    sql = SQL_QUERY_MUTATOR(
+        sql, user_name=user_name, security_manager=security_manager, database=database
+    )
     try:
         query.executed_sql = sql
         if log_query:

superset/sql_validators/presto_db.py

@@ -55,7 +55,12 @@ class PrestoDBSQLValidator(BaseSQLValidator):
         # Hook to allow environment-specific mutation (usually comments) to the SQL
         sql_query_mutator = config["SQL_QUERY_MUTATOR"]
         if sql_query_mutator:
-            sql = sql_query_mutator(sql, user_name, security_manager, database)
+            sql = sql_query_mutator(
+                sql,
+                user_name=user_name,
+                security_manager=security_manager,
+                database=database,
+            )
 
         # Transform the final statement to an explain call before sending it on
         # to presto to validate

tests/integration_tests/model_tests.py

@@ -506,7 +506,7 @@ class TestSqlaTableModel(SupersetTestCase):
         sql = tbl.get_query_str(query_obj)
         self.assertNotIn("-- COMMENT", sql)
 
-        def mutator(*args):
+        def mutator(*args, **kwargs):
             return "-- COMMENT\n" + args[0]
 
         app.config["SQL_QUERY_MUTATOR"] = mutator
@@ -515,6 +515,33 @@ class TestSqlaTableModel(SupersetTestCase):
 
         app.config["SQL_QUERY_MUTATOR"] = None
 
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    def test_sql_mutator_different_params(self):
+        tbl = self.get_table(name="birth_names")
+        query_obj = dict(
+            groupby=[],
+            metrics=None,
+            filter=[],
+            is_timeseries=False,
+            columns=["name"],
+            granularity=None,
+            from_dttm=None,
+            to_dttm=None,
+            extras={},
+        )
+        sql = tbl.get_query_str(query_obj)
+        self.assertNotIn("-- COMMENT", sql)
+
+        def mutator(sql, database=None, **kwargs):
+            return "-- COMMENT\n--" + "\n" + str(database) + "\n" + sql
+
+        app.config["SQL_QUERY_MUTATOR"] = mutator
+        mutated_sql = tbl.get_query_str(query_obj)
+        self.assertIn("-- COMMENT", mutated_sql)
+        self.assertIn(tbl.database.name, mutated_sql)
+
+        app.config["SQL_QUERY_MUTATOR"] = None
+
     def test_query_with_non_existent_metrics(self):
         tbl = self.get_table(name="birth_names")