From d645579cdd64b7fe7f9dde4a8da000dd7db51ce9 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Fri, 18 Mar 2022 16:01:27 -0700 Subject: [PATCH] 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 | 4 +++- superset/config.py | 11 +++++----- superset/connectors/sqla/models.py | 7 ++++++- superset/db_engine_specs/base.py | 7 ++++++- superset/sql_lab.py | 4 +++- superset/sql_validators/presto_db.py | 7 ++++++- tests/integration_tests/model_tests.py | 29 +++++++++++++++++++++++++- 7 files changed, 57 insertions(+), 12 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index 4272ecc56..c28975387 100644 --- a/UPDATING.md +++ b/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`. diff --git a/superset/config.py b/superset/config.py index dc93ad7f5..86e35be71 100644 --- a/superset/config.py +++ b/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 diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index 9261f97ae..62ae8c9eb 100644 --- a/superset/connectors/sqla/models.py +++ b/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: diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index d7e457baa..e867f200d 100644 --- a/superset/db_engine_specs/base.py +++ b/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 diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 8fac419cf..613db963e 100644 --- a/superset/sql_lab.py +++ b/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: diff --git a/superset/sql_validators/presto_db.py b/superset/sql_validators/presto_db.py index bf77f474a..6d0311f12 100644 --- a/superset/sql_validators/presto_db.py +++ b/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 diff --git a/tests/integration_tests/model_tests.py b/tests/integration_tests/model_tests.py index 5ffa65e58..c63886013 100644 --- a/tests/integration_tests/model_tests.py +++ b/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")