mirror of
https://github.com/apache/superset.git
synced 2026-07-18 04:35:40 +00:00
Compare commits
9 Commits
dependabot
...
adopt/sqll
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
464e40619f | ||
|
|
6ab9d7480b | ||
|
|
ed510633ce | ||
|
|
bf85efdbb7 | ||
|
|
d433084b52 | ||
|
|
31ebae7c35 | ||
|
|
408f5112e2 | ||
|
|
bf6d5b1faa | ||
|
|
353250027c |
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user