iop
/
superset
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(charts): allow query mutator to update queries after splitting original 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>
This commit is contained in:
solanksh 2023-01-12 17:34:03 +05:30 committed by GitHub
parent d5ecfbb901
commit cf00970cde
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 34 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
superset/config.py
View File

@ -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.

3
superset/connectors/sqla/models.py
View File

@ -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.

3
superset/db_engine_specs/base.py
View File

@ -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.

24
superset/models/core.py
View File

@ -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(