Compare commits

...

9 Commits

Author SHA1 Message Date
Evan Rusackas
464e40619f test: cover the unparseable-mutated-SQL validation fallback
The superset/sql/ tree is held to 100% unit-test coverage in CI; the
SupersetParseError branch of the joined-block emptiness validation was
the only uncovered path after the shared build_statement_blocks refactor.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-17 17:42:18 -07:00
Evan Rusackas
6ab9d7480b fix(sqllab): guard empty mutated single-block queries, share block prep across paths
The empty-statement guard added for split engines never covered the
run_multiple_statements_as_one branch: with MUTATE_AFTER_SPLIT=True the
per-statement mutator outputs were joined into a single block with no
check that any executable SQL remained, so an empty/comment-only mutator
result bypassed the clean INVALID_SQL_ERROR and reached execution as an
empty block. (An earlier review reply claimed this was fixed, but the
commit never landed.)

Per review feedback, the whole run_multiple_statements_as_one ×
MUTATE_AFTER_SPLIT matrix now lives in one shared helper,
build_statement_blocks() in superset/sql/execution/executor.py, used by
both the sync (sql_lab.py) and async (celery_task.py) paths, so the two
can't drift again. The new guard validates the joined block via SQLScript
and raises the same INVALID_SQL_ERROR as the split-engine branch; if the
mutator emits engine-specific SQL our parser can't handle, validation is
skipped and the database stays the authority on validity.

Regression tests cover the previously-unguarded matrix cell in both the
sync and async suites.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-17 17:42:18 -07:00
Evan
ed510633ce fix(sql-lab): address minor type-hint nits from codeant review
Add explicit type annotations for the `blocks` list in sql_lab.py and
the mock/cursor test variables, and replace an untyped lambda mutator
with a typed helper to keep MUTATE_AFTER_SPLIT type-hint coverage
consistent with the rest of the PR.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-17 17:42:18 -07:00
Evan
bf85efdbb7 fix(sqllab): mutate per-statement when engine runs statements as one block
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>
2026-07-17 17:42:18 -07:00
Evan
d433084b52 fix(tests): mock results_backend in mutator-strips-statements test
Without the mock, allow_run_async=True triggers SupersetResultsBackendNotConfigureException
(a SupersetErrorException subclass) before reaching the empty-statement guard, so the test
passed for the wrong reason.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-17 17:42:18 -07:00
Evan
31ebae7c35 fix(tests): mock mutate_sql_based_on_config passthrough in sqllab_tests
test_execute_sql_statements and test_execute_sql_statements_ctas mocked
`query.database` as a bare MagicMock, so the new pre-split mutation call
in execute_sql_statements (which re-parses the mutator's return value as
SQL) got back an unconfigured MagicMock instead of a string, leaving zero
parsed statements and tripping the "mutator removed all executable
statements" guard. Configure mutate_sql_based_on_config to pass the SQL
through unchanged, matching the real Database method's default behavior
and the pattern already used in the new unit tests.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-17 17:42:18 -07:00
Evan
408f5112e2 fix: cover MUTATE_AFTER_SPLIT branches, guard empty/mismatched mutated statements
Add missing coverage for the MUTATE_AFTER_SPLIT=True branch, add type
hints to satisfy dev-standard lint, and guard against a SQL_QUERY_MUTATOR
that strips all statements or changes statement count when mutating
before the split.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-17 17:42:17 -07:00
Evan
bf6d5b1faa fix(sqllab): mutate the whole query before splitting when MUTATE_AFTER_SPLIT=False
The prior fix derived is_split purely from run_multiple_statements_as_one, so
for the common case (engines that execute statements individually) with the
default MUTATE_AFTER_SPLIT=False, is_split was always True and the mutator
never fired at all -- a regression from the pre-fix behavior where it always
ran. Mutate the whole, un-split query once before splitting into per-statement
blocks when MUTATE_AFTER_SPLIT is False, in both the sync (sql_lab.py) and
async (celery_task.py) SQL Lab paths, and add regression tests covering both
flag values across both engine modes.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-17 17:42:17 -07:00
Evan
353250027c 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 <lucas.wolkersdorfer@rise-world.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-17 17:42:17 -07:00
10 changed files with 697 additions and 35 deletions

View File

@@ -2221,6 +2221,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

View File

@@ -27,7 +27,7 @@ import dataclasses
import logging
import traceback
import uuid
from typing import Any
from typing import Any, TYPE_CHECKING
import msgpack
from celery.exceptions import SoftTimeLimitExceeded
@@ -49,7 +49,10 @@ from superset.exceptions import (
from superset.extensions import celery_app
from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
from superset.sql.execution.executor import execute_sql_with_cursor
from superset.sql.execution.executor import (
build_statement_blocks,
execute_sql_with_cursor,
)
from superset.sql.parse import SQLScript
from superset.sqllab.utils import write_ipc_buffer
from superset.utils import json
@@ -57,6 +60,9 @@ from superset.utils.core import override_user, zlib_compress
from superset.utils.dates import now_as_float
from superset.utils.decorators import stats_timing
if TYPE_CHECKING:
from superset.models.core import Database
logger = logging.getLogger(__name__)
BYTES_IN_MB = 1024 * 1024
@@ -125,26 +131,17 @@ def _serialize_payload(payload: dict[Any, Any]) -> bytes:
def _prepare_statement_blocks(
rendered_query: str,
db_engine_spec: Any,
database: Database,
) -> tuple[SQLScript, list[str]]:
"""
Parse SQL and build statement blocks for execution.
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 when the database engine spec requires it.
Delegates to the shared ``build_statement_blocks`` so the sync
(``sql_lab``) and async (this module) SQL Lab paths apply
``SQL_QUERY_MUTATOR``/``MUTATE_AFTER_SPLIT`` identically.
"""
parsed_script = SQLScript(rendered_query, engine=db_engine_spec.engine)
# Build statement blocks for execution
if db_engine_spec.run_multiple_statements_as_one:
blocks = [parsed_script.format(comments=db_engine_spec.allows_sql_comments)]
else:
blocks = [
statement.format(comments=db_engine_spec.allows_sql_comments)
for statement in parsed_script.statements
]
return parsed_script, blocks
return build_statement_blocks(parsed_script, db_engine_spec, database)
def _finalize_successful_query(
@@ -162,6 +159,14 @@ def _finalize_successful_query(
# Get original statement strings
original_sqls = [stmt.format() for stmt in original_script.statements]
if len(original_sqls) != len(execution_results):
# A `SQL_QUERY_MUTATOR` that changes the number of statements (e.g. by
# prepending a `SET ROLE` statement when run on the whole, un-split
# query) can leave the un-mutated `original_script` no longer aligned
# 1:1 with `execution_results`. Fall back to labeling each result with
# its own executed SQL rather than crash a query that ran successfully.
original_sqls = [exec_sql for exec_sql, *_ in execution_results]
for orig_sql, (exec_sql, result_set, exec_time, rowcount) in zip(
original_sqls, execution_results, strict=True
):
@@ -423,7 +428,9 @@ def _execute_sql_statements(
original_script = SQLScript(query.sql, engine=db_engine_spec.engine)
# Parse transformed SQL (with RLS, limits, etc.)
parsed_script, blocks = _prepare_statement_blocks(rendered_query, db_engine_spec)
parsed_script, blocks = _prepare_statement_blocks(
rendered_query, db_engine_spec, database
)
with database.get_raw_connection(
catalog=query.catalog,
@@ -445,6 +452,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

View File

@@ -62,15 +62,18 @@ import logging
import time
import uuid
from datetime import datetime
from typing import Any, TYPE_CHECKING
from typing import Any, NoReturn, TYPE_CHECKING
from flask import current_app as app, g, has_app_context
from flask_babel import gettext as __
from superset import db
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import (
OAuth2Error,
OAuth2RedirectError,
SupersetErrorException,
SupersetParseError,
SupersetSecurityException,
SupersetTimeoutException,
)
@@ -87,12 +90,102 @@ if TYPE_CHECKING:
StatementResult,
)
from superset.db_engine_specs.base import BaseEngineSpec
from superset.models.core import Database
from superset.result_set import SupersetResultSet
logger = logging.getLogger(__name__)
def _raise_all_statements_stripped() -> NoReturn:
"""Raise a clean error for a mutator that stripped a query down to nothing."""
raise SupersetErrorException(
SupersetError(
message=__(
"The SQL query mutator removed all executable "
"statements from this query."
),
error_type=SupersetErrorType.INVALID_SQL_ERROR,
level=ErrorLevel.ERROR,
)
)
def _has_executable_statements(sql: str, engine: str) -> bool:
"""Best-effort check that mutated SQL still contains executable statements."""
try:
return bool(SQLScript(sql, engine=engine).statements)
except SupersetParseError:
# A mutator may emit engine-specific SQL our parser can't handle; the
# database itself is the authority on validity in that case.
return True
def build_statement_blocks(
parsed_script: SQLScript,
db_engine_spec: type[BaseEngineSpec],
database: Database,
) -> tuple[SQLScript, list[str]]:
"""
Build the SQL blocks to execute from a parsed script, applying
``SQL_QUERY_MUTATOR`` according to ``MUTATE_AFTER_SPLIT``.
Some databases (like BigQuery and Kusto) do not persist state across multiple
statements if they're run separately (especially when using `NullPool`), so the
query runs as a single joined block when the engine spec requires it; otherwise
each statement becomes its own block. Shared by the sync (``sql_lab``) and
async (``celery_task``) SQL Lab paths so the
``run_multiple_statements_as_one`` × ``MUTATE_AFTER_SPLIT`` matrix behaves
identically in both.
Returns the (possibly re-parsed) script and the blocks to execute.
:raises SupersetErrorException: if the mutator strips the query down to
nothing executable (e.g. only comments/whitespace)
"""
blocks: list[str]
if db_engine_spec.run_multiple_statements_as_one:
if app.config["MUTATE_AFTER_SPLIT"]:
# These engines never actually execute statements individually, so
# the per-block mutation call at execution time (whose `is_split` is
# always `False` here) would never fire. Mutate each statement here,
# before joining them into the single block this engine requires, so
# `MUTATE_AFTER_SPLIT=True` still applies the mutator per statement.
joined_block = ";\n".join(
database.mutate_sql_based_on_config(
statement.format(comments=db_engine_spec.allows_sql_comments),
is_split=True,
)
for statement in parsed_script.statements
)
if not _has_executable_statements(joined_block, db_engine_spec.engine):
_raise_all_statements_stripped()
blocks = [joined_block]
else:
blocks = [parsed_script.format(comments=db_engine_spec.allows_sql_comments)]
else:
if not app.config["MUTATE_AFTER_SPLIT"]:
# `MUTATE_AFTER_SPLIT=False` means the mutator should see the whole,
# un-split query, but this engine executes statements individually.
# Mutate the whole block up front and re-parse it, so the
# per-statement split below (and the per-block mutation call at
# execution time, which is a no-op here since its `is_split=True` no
# longer matches the config) operate on the already-mutated SQL.
mutated_sql: str = database.mutate_sql_based_on_config(
parsed_script.format(comments=db_engine_spec.allows_sql_comments),
is_split=False,
)
parsed_script = SQLScript(mutated_sql, engine=db_engine_spec.engine)
if not parsed_script.statements:
_raise_all_statements_stripped()
blocks = [
statement.format(comments=db_engine_spec.allows_sql_comments)
for statement in parsed_script.statements
]
return parsed_script, blocks
def execute_sql_with_cursor(
database: Database,
cursor: Any,
@@ -101,6 +194,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 +213,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 +238,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

View File

@@ -57,6 +57,7 @@ from superset.exceptions import (
from superset.extensions import celery_app, event_logger
from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
from superset.sql.execution.executor import build_statement_blocks
from superset.sql.parse import BaseSQLStatement, CTASMethod, SQLScript, Table
from superset.sqllab.limiting_factor import LimitingFactor
from superset.sqllab.utils import write_ipc_buffer
@@ -485,16 +486,13 @@ 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
# 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:
blocks = [parsed_script.format(comments=db_engine_spec.allows_sql_comments)]
else:
blocks = [
statement.format(comments=db_engine_spec.allows_sql_comments)
for statement in parsed_script.statements
]
# Build the execution blocks, applying `SQL_QUERY_MUTATOR` per
# `MUTATE_AFTER_SPLIT` (shared with the async path in `celery_task` so the
# `run_multiple_statements_as_one` × `MUTATE_AFTER_SPLIT` matrix behaves
# identically in both).
parsed_script, blocks = build_statement_blocks(
parsed_script, db_engine_spec, database
)
with database.get_raw_connection(
catalog=query.catalog,
@@ -528,8 +526,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)

View File

@@ -673,6 +673,9 @@ class TestSqlLab(SupersetTestCase):
mock_cursor
)
mock_query.database.db_engine_spec.run_multiple_statements_as_one = False
mock_query.database.mutate_sql_based_on_config.side_effect = (
lambda sql_, **kwargs: sql_
)
mock_get_query.return_value = mock_query
execute_sql_statements(
@@ -773,6 +776,9 @@ class TestSqlLab(SupersetTestCase):
mock_cursor
)
mock_query.database.db_engine_spec.run_multiple_statements_as_one = False
mock_query.database.mutate_sql_based_on_config.side_effect = (
lambda sql_, **kwargs: sql_
)
mock_get_query.return_value = mock_query
# set the query to CTAS

View File

@@ -264,6 +264,49 @@ def test_table_column_database() -> None:
assert TableColumn(database=database).database is database
def _prefixing_sql_query_mutator(sql: str, **kwargs: Any) -> str:
"""`SQL_QUERY_MUTATOR` stand-in that prepends a marker comment."""
return f"-- mutated\n{sql}"
@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": _prefixing_sql_query_mutator,
"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.

View File

@@ -77,6 +77,11 @@ def mock_query() -> MagicMock:
return query
def _passthrough_mutate_sql_based_on_config(sql: str, **kwargs: Any) -> str:
"""Mirror the real `Database.mutate_sql_based_on_config` no-op default."""
return sql
@pytest.fixture
def mock_database() -> MagicMock:
"""Create a mock Database."""
@@ -94,6 +99,12 @@ def mock_database() -> MagicMock:
database.db_engine_spec.get_cancel_query_id = MagicMock(return_value=None)
database.db_engine_spec.patch = MagicMock()
database.db_engine_spec.fetch_data = MagicMock(return_value=[])
# Mirrors the real `Database.mutate_sql_based_on_config` default (no-op
# when no `SQL_QUERY_MUTATOR` is configured), so SQL parsed from its
# return value stays valid instead of an un-parseable `MagicMock`.
database.mutate_sql_based_on_config = MagicMock(
side_effect=_passthrough_mutate_sql_based_on_config
)
return database

View File

@@ -34,6 +34,16 @@ from superset.exceptions import SupersetErrorException, SupersetErrorsException
# fixtures are imported from conftest.py
def _passthrough_mutator(sql: str, **kwargs: Any) -> str:
"""SQL mutator stand-in that returns the SQL unchanged."""
return sql
def _prefixing_mutator(sql: str, **kwargs: Any) -> str:
"""SQL mutator stand-in that prepends a marker comment."""
return f"-- mutated\n{sql}"
# =============================================================================
# Query Retrieval Tests
# =============================================================================
@@ -282,9 +292,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_mutator
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 +308,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_mutator
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 +325,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_mutator
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 +519,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,

View File

@@ -47,7 +47,11 @@ 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 (
_passthrough_mutate_sql_based_on_config,
create_mock_cursor,
mock_query_execution,
)
# =============================================================================
# Basic Execution Tests
@@ -871,6 +875,47 @@ 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: MagicMock = mocker.patch.object(
database,
"mutate_sql_based_on_config",
side_effect=_passthrough_mutate_sql_based_on_config,
)
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: MagicMock = 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
# =============================================================================
@@ -2340,3 +2385,37 @@ def test_cached_async_result_get_result_returns_cached(
assert retrieved_result.status == QueryStatus.SUCCESS
assert sum(s.row_count for s in retrieved_result.statements) == 3
assert retrieved_result is cached_result
def test_build_statement_blocks_skips_validation_for_unparseable_mutated_sql(
mocker: MockerFixture, mock_database: MagicMock, app_context: None
) -> None:
"""
A mutator may emit engine-specific SQL the parser can't handle. The
empty-statement validation on the joined block is skipped in that case
instead of blocking execution, leaving the database as the authority
on validity.
"""
from superset.exceptions import SupersetParseError
from superset.sql.execution.executor import build_statement_blocks
from superset.sql.parse import SQLScript
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: f"ENGINE SPECIFIC {sql}",
)
parsed_script = SQLScript("SELECT 1; SELECT 2;", engine="bigquery")
mocker.patch(
"superset.sql.execution.executor.SQLScript",
side_effect=SupersetParseError("ENGINE SPECIFIC SQL"),
)
_, blocks = build_statement_blocks(
parsed_script, mock_database.db_engine_spec, mock_database
)
assert len(blocks) == 1
assert blocks[0].count("ENGINE SPECIFIC") == 2

View File

@@ -25,6 +25,7 @@ import pytest
from freezegun import freeze_time
from pytest_mock import MockerFixture
from superset.app import SupersetApp
from superset.common.db_query_status import QueryStatus
from superset.db_engine_specs.postgres import PostgresEngineSpec
from superset.errors import ErrorLevel, SupersetErrorType
@@ -193,6 +194,215 @@ def test_execute_sql_statement_within_payload_limit(mocker: MockerFixture, app)
)
def test_execute_sql_statements_mutates_before_split_by_default(
mocker: MockerFixture, app: SupersetApp
) -> None:
"""
With the default `MUTATE_AFTER_SPLIT=False`, `execute_sql_statements` should
mutate the whole, un-split query once before splitting it into individual
statement blocks, for engines that execute statements individually rather
than as one. Regression guard for issue #30169.
"""
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
query.database.db_engine_spec.engine = "sqlite"
query.database.db_engine_spec.run_multiple_statements_as_one = False
query.database.db_engine_spec.allows_sql_comments = True
mutate_mock = mocker.patch.object(
query.database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: sql,
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("sys.getsizeof", return_value=10000000)
mocker.patch(
"superset.sql_lab._serialize_payload",
side_effect=lambda payload, use_msgpack: "serialized_payload",
)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
execute_sql_statements(
query_id=1,
rendered_query="SELECT 1; SELECT 2;",
return_results=True,
store_results=True,
start_time=None,
expand_data=False,
log_params={},
)
is_split_values = [
call.kwargs.get("is_split") for call in mutate_mock.call_args_list
]
# The mutator is called once on the whole, un-split query before splitting...
assert is_split_values[0] is False
first_call_sql = mutate_mock.call_args_list[0].args[0]
assert "1" in first_call_sql
assert "2" in first_call_sql
# Both statements are present in a single, un-split call.
assert first_call_sql.count("SELECT") == 2
# ...and once again per already-split statement (a no-op when
# `MUTATE_AFTER_SPLIT=False`, since `is_split=True` won't match the config).
assert all(value is True for value in is_split_values[1:])
assert len(is_split_values) == 3
def test_execute_sql_statements_mutates_per_statement_when_run_as_one(
mocker: MockerFixture, app: SupersetApp
) -> 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 further down, so with
`MUTATE_AFTER_SPLIT=True` the mutator must instead be applied to each
statement up front, before they're joined into that single block.
"""
mocker.patch.dict(app.config, {"MUTATE_AFTER_SPLIT": True})
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
query.database.db_engine_spec.engine = "bigquery"
query.database.db_engine_spec.run_multiple_statements_as_one = True
query.database.db_engine_spec.allows_sql_comments = True
mutate_mock = mocker.patch.object(
query.database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: sql,
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("sys.getsizeof", return_value=10000000)
mocker.patch(
"superset.sql_lab._serialize_payload",
side_effect=lambda payload, use_msgpack: "serialized_payload",
)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
execute_sql_statements(
query_id=1,
rendered_query="SELECT 1; SELECT 2;",
return_results=True,
store_results=True,
start_time=None,
expand_data=False,
log_params={},
)
is_split_values = [
call.kwargs.get("is_split") for call in mutate_mock.call_args_list
]
# Mutated once per statement before joining into the single block...
assert is_split_values[0] is True
assert is_split_values[1] is True
first_call_sql = mutate_mock.call_args_list[0].args[0]
second_call_sql = mutate_mock.call_args_list[1].args[0]
assert "1" in first_call_sql
assert "2" in second_call_sql
# ...and the later per-block call is a no-op (`is_split=False` never matches
# `MUTATE_AFTER_SPLIT=True`), so the mutator isn't applied a second time.
assert is_split_values[2] is False
assert len(is_split_values) == 3
def test_execute_sql_statements_raises_when_mutator_strips_all_statements(
mocker: MockerFixture, app: SupersetApp
) -> 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.
"""
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
query.database.db_engine_spec.engine = "sqlite"
query.database.db_engine_spec.run_multiple_statements_as_one = False
query.database.db_engine_spec.allows_sql_comments = True
mocker.patch.object(
query.database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: "-- just a comment",
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
with pytest.raises(SupersetErrorException):
execute_sql_statements(
query_id=1,
rendered_query="SELECT 1;",
return_results=True,
store_results=True,
start_time=None,
expand_data=False,
log_params={},
)
def test_execute_sql_statements_raises_when_mutator_strips_single_block(
mocker: MockerFixture, app: SupersetApp
) -> 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.
"""
mocker.patch.dict(app.config, {"MUTATE_AFTER_SPLIT": True})
query = mocker.MagicMock()
query.limit = 1
query.database = mocker.MagicMock()
query.database.cache_timeout = 100
query.status = "RUNNING"
query.select_as_cta = False
query.database.allow_run_async = True
query.database.db_engine_spec.engine = "bigquery"
query.database.db_engine_spec.run_multiple_statements_as_one = True
query.database.db_engine_spec.allows_sql_comments = True
mocker.patch.object(
query.database,
"mutate_sql_based_on_config",
side_effect=lambda sql, **kw: "-- just a comment",
)
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
with pytest.raises(SupersetErrorException):
execute_sql_statements(
query_id=1,
rendered_query="SELECT 1; SELECT 2;",
return_results=True,
store_results=True,
start_time=None,
expand_data=False,
log_params={},
)
@freeze_time("2021-04-01T00:00:00Z")
def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None:
"""