QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-07-26 08:32:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(sqllab): apply SQL_QUERY_MUTATOR in SQL Lab when MUTATE_AFTER_SPLIT is set (#41127)

Browse Source 
Co-authored-by: Lucas Wolkersdorfer <lucas.wolkersdorfer@rise-world.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Evan Rusackas
2026-07-20 15:36:45 -07:00
committed by GitHub
parent 158e321992
commit b3757870cc
11 changed files with 746 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

199
tests/unit_tests/sql/execution/test_celery_task.py
View File

@@ -32,6 +32,12 @@ from superset.exceptions import SupersetErrorException, SupersetErrorsException
# Note: mock_query, mock_database, mock_result_set, and mock_db_session
# fixtures are imported from conftest.py
from .conftest import _passthrough_mutate_sql_based_on_config
def _prefixing_mutator(sql: str, **kwargs: Any) -> str:
"""SQL mutator stand-in that prepends a marker comment."""
return f"-- mutated\n{sql}"
# =============================================================================
@@ -282,9 +288,12 @@ def test_prepare_statement_blocks_single_statement(
"""Test statement block preparation for single statement."""
from superset.sql.execution.celery_task import _prepare_statement_blocks
mock_database.mutate_sql_based_on_config = _passthrough_mutate_sql_based_on_config
sql = "SELECT * FROM users"
script, blocks = _prepare_statement_blocks(sql, mock_database.db_engine_spec)
script, blocks = _prepare_statement_blocks(
sql, mock_database.db_engine_spec, mock_database
)
assert len(blocks) == 1
@@ -295,9 +304,12 @@ def test_prepare_statement_blocks_multiple_statements(
"""Test statement block preparation for multiple statements."""
from superset.sql.execution.celery_task import _prepare_statement_blocks
mock_database.mutate_sql_based_on_config = _passthrough_mutate_sql_based_on_config
sql = "SELECT * FROM users; SELECT * FROM orders;"
script, blocks = _prepare_statement_blocks(sql, mock_database.db_engine_spec)
script, blocks = _prepare_statement_blocks(
sql, mock_database.db_engine_spec, mock_database
)
assert len(blocks) == 2
@@ -309,13 +321,148 @@ def test_prepare_statement_blocks_run_as_one(
from superset.sql.execution.celery_task import _prepare_statement_blocks
mock_database.db_engine_spec.run_multiple_statements_as_one = True
mock_database.mutate_sql_based_on_config = _passthrough_mutate_sql_based_on_config
sql = "SELECT * FROM users; SELECT * FROM orders;"
script, blocks = _prepare_statement_blocks(sql, mock_database.db_engine_spec)
script, blocks = _prepare_statement_blocks(
sql, mock_database.db_engine_spec, mock_database
)
assert len(blocks) == 1
def test_prepare_statement_blocks_mutates_before_split_when_configured(
app_context: None, mock_database: MagicMock, mocker: MockerFixture
) -> None:
"""
`MUTATE_AFTER_SPLIT=False` should mutate the whole, un-split query before
it gets broken into per-statement blocks, for engines that execute
statements individually. Regression guard for issue #30169.
"""
from superset.sql.execution.celery_task import _prepare_statement_blocks
mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": False})
mutate_mock = mocker.patch.object(
mock_database,
"mutate_sql_based_on_config",
side_effect=_prefixing_mutator,
)
sql = "SELECT * FROM users; SELECT * FROM orders;"
_, blocks = _prepare_statement_blocks(
sql, mock_database.db_engine_spec, mock_database
)
assert len(blocks) == 2
assert "mutated" in blocks[0]
mutate_mock.assert_called_once_with(mocker.ANY, is_split=False)
def test_prepare_statement_blocks_skips_pre_split_mutation_when_configured(
app_context: None, mock_database: MagicMock, mocker: MockerFixture
) -> None:
"""
`MUTATE_AFTER_SPLIT=True` means the mutator should see each already-split
statement instead, so `_prepare_statement_blocks` must not mutate the
whole, un-split query up front.
"""
from superset.sql.execution.celery_task import _prepare_statement_blocks
mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": True})
mutate_mock = mocker.patch.object(
mock_database,
"mutate_sql_based_on_config",
side_effect=_prefixing_mutator,
)
sql = "SELECT * FROM users; SELECT * FROM orders;"
_, blocks = _prepare_statement_blocks(
sql, mock_database.db_engine_spec, mock_database
)
assert len(blocks) == 2
assert "mutated" not in blocks[0]
mutate_mock.assert_not_called()
def test_prepare_statement_blocks_mutates_per_statement_when_run_as_one(
app_context: None, mock_database: MagicMock, mocker: MockerFixture
) -> None:
"""
Engines that always run statements as a single block (e.g. BigQuery, Kusto)
never see `is_split=True` in the per-block mutation call, so with
`MUTATE_AFTER_SPLIT=True` the mutator must instead be applied to each
statement here, before they're joined into that single block.
"""
from superset.sql.execution.celery_task import _prepare_statement_blocks
mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": True})
mock_database.db_engine_spec.run_multiple_statements_as_one = True
mutate_mock = mocker.patch.object(
mock_database,
"mutate_sql_based_on_config",
side_effect=_prefixing_mutator,
)
sql = "SELECT * FROM users; SELECT * FROM orders;"
_, blocks = _prepare_statement_blocks(
sql, mock_database.db_engine_spec, mock_database
)
assert len(blocks) == 1
assert blocks[0].count("mutated") == 2
is_split_values = [
call.kwargs.get("is_split") for call in mutate_mock.call_args_list
]
assert is_split_values == [True, True]
def test_prepare_statement_blocks_raises_when_mutator_strips_all_statements(
app_context: None, mock_database: MagicMock, mocker: MockerFixture
) -> None:
"""
A `SQL_QUERY_MUTATOR` that strips a query down to nothing (e.g. only
comments/whitespace) must raise a clean error instead of silently
producing an empty block list.
"""
from superset.sql.execution.celery_task import _prepare_statement_blocks
mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": False})
mocker.patch.object(
mock_database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: "-- just a comment",
)
sql = "SELECT * FROM users"
with pytest.raises(SupersetErrorException):
_prepare_statement_blocks(sql, mock_database.db_engine_spec, mock_database)
def test_prepare_statement_blocks_raises_when_mutator_strips_single_block(
app_context: None, mock_database: MagicMock, mocker: MockerFixture
) -> None:
"""
The empty-statement guard must also cover engines that run all statements
as one block: with `MUTATE_AFTER_SPLIT=True` the per-statement mutator
outputs are joined into a single block, and a comment-only/empty result
must raise a clean error instead of reaching execution as an empty block.
"""
from superset.sql.execution.celery_task import _prepare_statement_blocks
mocker.patch.dict(current_app.config, {"MUTATE_AFTER_SPLIT": True})
mock_database.db_engine_spec.run_multiple_statements_as_one = True
mocker.patch.object(
mock_database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: "-- just a comment",
)
sql = "SELECT * FROM users; SELECT * FROM orders;"
with pytest.raises(SupersetErrorException):
_prepare_statement_blocks(sql, mock_database.db_engine_spec, mock_database)
# =============================================================================
# Result Finalization Tests
# =============================================================================
@@ -368,6 +515,52 @@ def test_finalize_successful_query(
)
def test_finalize_successful_query_statement_count_mismatch(
mocker: MockerFixture,
app_context: None,
mock_query: MagicMock,
mock_result_set: MagicMock,
mock_database: MagicMock,
) -> None:
"""
A `SQL_QUERY_MUTATOR` that changes the number of statements (e.g. by
prepending a statement when run on the whole, un-split query) must not
crash finalization when `original_script`'s statement count no longer
matches `execution_results`.
"""
from superset.sql.execution.celery_task import _finalize_successful_query
from superset.sql.parse import SQLScript
mocker.patch("superset.results_backend_use_msgpack", False)
mocker.patch("superset.dataframe.df_to_records", return_value=[{"id": 1}])
payload: dict[str, Any] = {}
# Only one statement in the un-mutated original...
original_script = SQLScript(
"SELECT * FROM users", mock_database.db_engine_spec.engine
)
# ...but the mutator turned it into two.
execution_results = [
("SET ROLE 'bob'", None, 1.0, 0),
("SELECT * FROM users", mock_result_set, 10.5, 2),
]
_finalize_successful_query(
mock_query,
original_script,
execution_results, # type: ignore[arg-type]
payload,
11.5,
)
assert mock_query.rows == 2
assert payload["status"] == QueryStatusEnum.SUCCESS
assert len(payload["statements"]) == 2
# Falls back to labeling each result with its own executed SQL.
assert payload["statements"][0]["original_sql"] == "SET ROLE 'bob'"
assert payload["statements"][1]["original_sql"] == "SELECT * FROM users"
def test_finalize_successful_query_with_msgpack(
mocker: MockerFixture,
app_context: None,