You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ru...@apache.org on 2026/06/16 18:48:12 UTC
(superset) 01/01: fix(sqllab): apply SQL_QUERY_MUTATOR in SQL Lab when MUTATE_AFTER_SPLIT is set
This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch adopt/sqllab-mutator-after-split
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 6a2cd30ac48ce11291d4a51d1492e13f5b479d2a
Author: Evan <ev...@preset.io>
AuthorDate: Tue Jun 16 11:47:41 2026 -0700
fix(sqllab): apply SQL_QUERY_MUTATOR in SQL Lab when MUTATE_AFTER_SPLIT is set
SQL Lab passed the default is_split=False to mutate_sql_based_on_config for
every block, so SQL_QUERY_MUTATOR fired only when MUTATE_AFTER_SPLIT=False and
never when True. Pass is_split based on whether the engine runs statements as
one block, matching the canonical usage in db_engine_specs and the new
execution engine, and thread the same flag through execute_sql_with_cursor so
the async (Celery) path is consistent. Adds unit tests for both the guard
semantics and the executor contract, plus a config note.
Closes #30169
Supersedes #34111
Co-authored-by: Lucas Wolkersdorfer <lu...@rise-world.com>
Co-Authored-By: Claude Opus 4.8 <no...@anthropic.com>
---
superset/config.py | 3 ++
superset/sql/execution/celery_task.py | 3 ++
superset/sql/execution/executor.py | 7 ++++-
superset/sql_lab.py | 13 ++++++--
tests/unit_tests/models/core_test.py | 38 +++++++++++++++++++++++
tests/unit_tests/sql/execution/test_executor.py | 41 ++++++++++++++++++++++++-
6 files changed, 100 insertions(+), 5 deletions(-)
diff --git a/superset/config.py b/superset/config.py
index 0bb0cacb9f5..80356e99d1d 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -2005,6 +2005,9 @@ def SQL_QUERY_MUTATOR( # pylint: disable=invalid-name,unused-argument # noqa:
# 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.
+# This applies consistently in SQL Lab: with MUTATE_AFTER_SPLIT = True the mutator runs
+# on each individual statement, and with MUTATE_AFTER_SPLIT = False it runs once on the
+# un-split query block.
MUTATE_AFTER_SPLIT = False
diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py
index dabb5f0915b..4310991edee 100644
--- a/superset/sql/execution/celery_task.py
+++ b/superset/sql/execution/celery_task.py
@@ -445,6 +445,9 @@ def _execute_sql_statements(
log_query_fn=_make_log_query_fn(database),
check_stopped_fn=_make_check_stopped_fn(query),
execute_fn=_make_execute_fn(query, db_engine_spec),
+ # `blocks` is a single un-split block when the engine runs multiple
+ # statements as one; otherwise each block is an individual statement.
+ is_split=not db_engine_spec.run_multiple_statements_as_one,
)
except SoftTimeLimitExceeded as ex:
query.status = QueryStatus.TIMED_OUT
diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py
index 13e12fc1a47..7a69510abfd 100644
--- a/superset/sql/execution/executor.py
+++ b/superset/sql/execution/executor.py
@@ -101,6 +101,7 @@ def execute_sql_with_cursor(
log_query_fn: Any | None = None,
check_stopped_fn: Any | None = None,
execute_fn: Any | None = None,
+ is_split: bool = True,
) -> list[tuple[str, SupersetResultSet | None, float, int]]:
"""
Execute SQL statements with a cursor and return all result sets.
@@ -119,6 +120,10 @@ def execute_sql_with_cursor(
:param execute_fn: Optional custom execute function. If not provided, uses
database.db_engine_spec.execute(cursor, sql, database). Custom function
should accept (cursor, sql) and handle execution.
+ :param is_split: Whether `statements` are individual split-out statements (True)
+ or a single un-split block (False, e.g. when the engine spec runs multiple
+ statements as one). Passed to the SQL mutator so `MUTATE_AFTER_SPLIT` can
+ decide whether to fire.
:returns: List of (statement_sql, result_set, execution_time_ms, rowcount) tuples
Returns empty list if stopped. Raises exception on error (fail-fast).
"""
@@ -140,7 +145,7 @@ def execute_sql_with_cursor(
# Apply SQL mutation
stmt_sql = database.mutate_sql_based_on_config(
statement,
- is_split=True,
+ is_split=is_split,
)
# Log query
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 34611d38a33..f17a1cea5d1 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -474,7 +474,7 @@ def execute_sql_statements( # noqa: C901
for statement in parsed_script.statements:
apply_limit(query, statement)
- # some databases (like BigQuery and Kusto) do not persist state across mmultiple
+ # some databases (like BigQuery and Kusto) do not persist state across multiple
# statements if they're run separately (especially when using `NullPool`), so we run
# the query as a single block.
if db_engine_spec.run_multiple_statements_as_one:
@@ -517,8 +517,15 @@ def execute_sql_statements( # noqa: C901
query.set_extra_json_key("progress", msg)
db.session.commit()
- # Hook to allow environment-specific mutation (usually comments) to the SQL
- query.executed_sql = database.mutate_sql_based_on_config(block)
+ # Hook to allow environment-specific mutation (usually comments) to the SQL.
+ # `is_split` reflects whether this block is an individual statement: when
+ # the engine runs everything as one block the SQL is not split, otherwise
+ # each block is a single split-out statement. This lets `MUTATE_AFTER_SPLIT`
+ # decide correctly whether the mutator fires here.
+ query.executed_sql = database.mutate_sql_based_on_config(
+ block,
+ is_split=not db_engine_spec.run_multiple_statements_as_one,
+ )
try:
result_set = execute_query(query, cursor, log_params)
diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py
index 8d15b23077b..6ca9c15b980 100644
--- a/tests/unit_tests/models/core_test.py
+++ b/tests/unit_tests/models/core_test.py
@@ -262,6 +262,44 @@ def test_table_column_database() -> None:
assert TableColumn(database=database).database is database
+@pytest.mark.parametrize(
+ "is_split,mutate_after_split,expect_mutated",
+ [
+ # A split-out statement is mutated only when the mutator is meant to run
+ # after the split, and an un-split block only when it runs before.
+ (True, True, True),
+ (True, False, False),
+ (False, False, True),
+ (False, True, False),
+ ],
+)
+def test_mutate_sql_based_on_config_respects_is_split(
+ app_context: None,
+ mocker: MockerFixture,
+ is_split: bool,
+ mutate_after_split: bool,
+ expect_mutated: bool,
+) -> None:
+ """
+ `mutate_sql_based_on_config` fires `SQL_QUERY_MUTATOR` only when the call
+ site's `is_split` matches the `MUTATE_AFTER_SPLIT` config. Regression guard
+ for issue #30169, where SQL Lab always passed the default `is_split=False`
+ and so never mutated when `MUTATE_AFTER_SPLIT=True`.
+ """
+ database = Database(database_name="db", sqlalchemy_uri="sqlite://")
+ mocker.patch.dict(
+ current_app.config,
+ {
+ "SQL_QUERY_MUTATOR": lambda sql, **kwargs: f"-- mutated\n{sql}",
+ "MUTATE_AFTER_SPLIT": mutate_after_split,
+ },
+ )
+
+ result = database.mutate_sql_based_on_config("SELECT 1", is_split=is_split)
+
+ assert result == ("-- mutated\nSELECT 1" if expect_mutated else "SELECT 1")
+
+
def test_catalog_cache() -> None:
"""
Test the catalog cache.
diff --git a/tests/unit_tests/sql/execution/test_executor.py b/tests/unit_tests/sql/execution/test_executor.py
index fdaf52f8a79..0e12de4b034 100644
--- a/tests/unit_tests/sql/execution/test_executor.py
+++ b/tests/unit_tests/sql/execution/test_executor.py
@@ -47,7 +47,7 @@ from superset.models.core import Database
# Note: database, database_with_dml, mock_db_session fixtures and
# mock_query_execution helper are imported from conftest.py
-from .conftest import mock_query_execution
+from .conftest import create_mock_cursor, mock_query_execution
# =============================================================================
# Basic Execution Tests
@@ -871,6 +871,45 @@ def test_execute_applies_sql_mutator(
mutate_mock.assert_called()
+@pytest.mark.parametrize("is_split", [True, False])
+def test_execute_sql_with_cursor_forwards_is_split(
+ mocker: MockerFixture,
+ database: Database,
+ app_context: None,
+ mock_db_session: MagicMock,
+ mock_query: MagicMock,
+ is_split: bool,
+) -> None:
+ """
+ `execute_sql_with_cursor` must forward `is_split` to the SQL mutator.
+
+ `Database.mutate_sql_based_on_config` only fires `SQL_QUERY_MUTATOR` when
+ `is_split == MUTATE_AFTER_SPLIT`, so passing the wrong value silently skips
+ mutation (the SQL Lab bug behind issue #30169). This guards the contract.
+ """
+ from superset.sql.execution.executor import execute_sql_with_cursor
+
+ mutate_mock = mocker.patch.object(
+ database, "mutate_sql_based_on_config", side_effect=lambda sql, **kw: sql
+ )
+ mocker.patch.object(database.db_engine_spec, "execute")
+ mocker.patch.object(database.db_engine_spec, "fetch_data", return_value=[(1,)])
+ mocker.patch("superset.result_set.SupersetResultSet", return_value=MagicMock())
+
+ cursor = create_mock_cursor(["id"], data=[(1,)])
+
+ execute_sql_with_cursor(
+ database=database,
+ cursor=cursor,
+ statements=["SELECT id FROM t"],
+ query=mock_query,
+ is_split=is_split,
+ )
+
+ mutate_mock.assert_called_once()
+ assert mutate_mock.call_args.kwargs["is_split"] is is_split
+
+
# =============================================================================
# Progress Tracking Tests
# =============================================================================