You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by vi...@apache.org on 2023/01/12 12:04:12 UTC
[superset] branch master updated: feat(charts): allow query mutator to update queries after splitting original sql (#21645)
This is an automated email from the ASF dual-hosted git repository.
villebro pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new cf00970cde feat(charts): allow query mutator to update queries after splitting original sql (#21645)
cf00970cde is described below
commit cf00970cde573011bb3d8fcdfc43258963f45bdf
Author: solanksh <94...@users.noreply.github.com>
AuthorDate: Thu Jan 12 17:34:03 2023 +0530
feat(charts): allow query mutator to update queries after splitting original sql (#21645)
Co-authored-by: Akash <Ak...@bakerhughes.com>
Co-authored-by: anallani <98...@users.noreply.github.com>
Co-authored-by: Robert Bean <ro...@bakerhughes.com>
Co-authored-by: Akash <an...@umich.edu>
Co-authored-by: AkashN7 <70...@users.noreply.github.com>
Co-authored-by: Ville Brofeldt <vi...@apple.com>
---
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 d9c36955e3..984d1b2295 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 fd5942c51a..c5fd025f4e 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 481ef5796a..1aab100c81 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 9e042eeab5..403d0f6c18 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(