QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-07-17 20:25:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(sqllab): mutate per-statement when engine runs statements as one block

Browse Source 
MUTATE_AFTER_SPLIT=True never fired the mutator for engines with
run_multiple_statements_as_one=True (BigQuery, Datastore, Kusto), since
those engines always build a single joined block and the per-block
mutation call always passes is_split=False. Mutate each statement
before joining them into that block so the flag applies consistently.

Verified separately via a full `pre-commit run` (all hooks touching
these files passed: auto-walrus, mypy, ruff-format, ruff, pylint) run
outside the git hook, whose invocation via the system /usr/bin/python3
hits an unrelated pre-existing environment issue installing the zizmor
hook (requires Python >=3.10, system python3 here is older) rather
than anything in this change.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Evan
2026-07-14 04:41:36 -07:00
parent 68c072a524
commit 1f81871fb3
4 changed files with 129 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -389,6 +389,38 @@ def test_prepare_statement_blocks_skips_pre_split_mutation_when_configured(
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: