From 6ab9d7480bd84bbde0aa5e491d02344a93e1ceea Mon Sep 17 00:00:00 2001 From: Evan Rusackas Date: Fri, 17 Jul 2026 16:43:29 -0700 Subject: [PATCH] fix(sqllab): guard empty mutated single-block queries, share block prep across paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- superset/sql/execution/celery_task.py | 67 ++----------- superset/sql/execution/executor.py | 95 ++++++++++++++++++- superset/sql_lab.py | 61 ++---------- .../sql/execution/test_celery_task.py | 24 +++++ tests/unit_tests/sql_lab_test.py | 44 +++++++++ 5 files changed, 178 insertions(+), 113 deletions(-) diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py index 03f02f5ddf6..619725563c0 100644 --- a/superset/sql/execution/celery_task.py +++ b/superset/sql/execution/celery_task.py @@ -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 @@ -133,66 +136,12 @@ def _prepare_statement_blocks( """ 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: - if app.config["MUTATE_AFTER_SPLIT"]: - # These engines never actually execute statements individually, so the - # per-block mutation call in `execute_sql_with_cursor` (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. - blocks = [ - ";\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 - ) - ] - 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 later per-statement mutation call in - # `execute_sql_with_cursor`, 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: - # A `SQL_QUERY_MUTATOR` that strips a query down to nothing - # (e.g. only comments/whitespace) would otherwise leave us with - # an empty `blocks` list, skipping the execution loop below and - # surfacing a confusing error instead of a clean one. - raise SupersetErrorException( - SupersetError( - message=__( - "The SQL query mutator removed all executable " - "statements from this query." - ), - error_type=SupersetErrorType.INVALID_SQL_ERROR, - level=ErrorLevel.ERROR, - ) - ) - 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( diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py index 32cb9280922..b000a6edb23 100644 --- a/superset/sql/execution/executor.py +++ b/superset/sql/execution/executor.py @@ -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, diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 259b464b128..bf17daac680 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -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,59 +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 multiple - # statements if they're run separately (especially when using `NullPool`), so we run - # the query as a single block. - if db_engine_spec.run_multiple_statements_as_one: - if app.config["MUTATE_AFTER_SPLIT"]: - # These engines never actually execute statements individually, so the - # per-block mutation call further down (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. - blocks: list[str] = [ - ";\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 - ) - ] - 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 further down, 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: - # A `SQL_QUERY_MUTATOR` that strips a query down to nothing - # (e.g. only comments/whitespace) would otherwise leave us with - # an empty `blocks` list, skipping the execution loop below and - # surfacing a confusing error instead of a clean one. - raise SupersetErrorException( - SupersetError( - message=__( - "The SQL query mutator removed all executable " - "statements from this query." - ), - error_type=SupersetErrorType.INVALID_SQL_ERROR, - level=ErrorLevel.ERROR, - ) - ) - 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, diff --git a/tests/unit_tests/sql/execution/test_celery_task.py b/tests/unit_tests/sql/execution/test_celery_task.py index dc9a1170b13..a69fa138e02 100644 --- a/tests/unit_tests/sql/execution/test_celery_task.py +++ b/tests/unit_tests/sql/execution/test_celery_task.py @@ -443,6 +443,30 @@ def test_prepare_statement_blocks_raises_when_mutator_strips_all_statements( _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 # ============================================================================= diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index d2b1a3635ad..b9bec3cad4c 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -359,6 +359,50 @@ def test_execute_sql_statements_raises_when_mutator_strips_all_statements( ) +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: """