From ef1ec529bc35f013ef772e8a93ce7940ad02b6ca Mon Sep 17 00:00:00 2001 From: Evan Date: Sun, 12 Jul 2026 09:21:59 -0700 Subject: [PATCH] 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 --- superset/sql/execution/celery_task.py | 25 +++- superset/sql_lab.py | 17 ++- tests/unit_tests/models/core_test.py | 7 +- tests/unit_tests/sql/execution/conftest.py | 4 + .../sql/execution/test_celery_task.py | 113 +++++++++++++++++- tests/unit_tests/sql_lab_test.py | 43 ++++++- 6 files changed, 201 insertions(+), 8 deletions(-) diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py index 3bcd35fc622..2b3407768b6 100644 --- a/superset/sql/execution/celery_task.py +++ b/superset/sql/execution/celery_task.py @@ -151,11 +151,26 @@ def _prepare_statement_blocks( # `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 = database.mutate_sql_based_on_config( + 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 @@ -179,6 +194,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 ): diff --git a/superset/sql_lab.py b/superset/sql_lab.py index b19eb0e93ae..c8ced0912c8 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -487,11 +487,26 @@ def execute_sql_statements( # noqa: C901 # 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 = database.mutate_sql_based_on_config( + 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 diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py index 6ca9c15b980..27edb8ee93b 100644 --- a/tests/unit_tests/models/core_test.py +++ b/tests/unit_tests/models/core_test.py @@ -262,6 +262,11 @@ 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", [ @@ -290,7 +295,7 @@ def test_mutate_sql_based_on_config_respects_is_split( mocker.patch.dict( current_app.config, { - "SQL_QUERY_MUTATOR": lambda sql, **kwargs: f"-- mutated\n{sql}", + "SQL_QUERY_MUTATOR": _prefixing_sql_query_mutator, "MUTATE_AFTER_SPLIT": mutate_after_split, }, ) diff --git a/tests/unit_tests/sql/execution/conftest.py b/tests/unit_tests/sql/execution/conftest.py index 630f6884529..f8aed71ae1a 100644 --- a/tests/unit_tests/sql/execution/conftest.py +++ b/tests/unit_tests/sql/execution/conftest.py @@ -94,6 +94,10 @@ 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=lambda sql, **kw: sql) return database diff --git a/tests/unit_tests/sql/execution/test_celery_task.py b/tests/unit_tests/sql/execution/test_celery_task.py index 39818dd10fa..c53e2a33ac7 100644 --- a/tests/unit_tests/sql/execution/test_celery_task.py +++ b/tests/unit_tests/sql/execution/test_celery_task.py @@ -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,7 +292,7 @@ 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 = lambda sql, **kw: sql + mock_database.mutate_sql_based_on_config = _passthrough_mutator sql = "SELECT * FROM users" script, blocks = _prepare_statement_blocks( @@ -298,7 +308,7 @@ 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 = lambda sql, **kw: sql + mock_database.mutate_sql_based_on_config = _passthrough_mutator sql = "SELECT * FROM users; SELECT * FROM orders;" script, blocks = _prepare_statement_blocks( @@ -315,7 +325,7 @@ 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 = lambda sql, **kw: sql + mock_database.mutate_sql_based_on_config = _passthrough_mutator sql = "SELECT * FROM users; SELECT * FROM orders;" script, blocks = _prepare_statement_blocks( @@ -339,7 +349,7 @@ def test_prepare_statement_blocks_mutates_before_split_when_configured( mutate_mock = mocker.patch.object( mock_database, "mutate_sql_based_on_config", - side_effect=lambda sql, **kw: f"-- mutated\n{sql}", + side_effect=_prefixing_mutator, ) sql = "SELECT * FROM users; SELECT * FROM orders;" @@ -352,6 +362,55 @@ def test_prepare_statement_blocks_mutates_before_split_when_configured( 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_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) + + # ============================================================================= # Result Finalization Tests # ============================================================================= @@ -404,6 +463,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, diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index 06c3e309463..b766f387ae4 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -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 @@ -194,7 +195,7 @@ def test_execute_sql_statement_within_payload_limit(mocker: MockerFixture, app) def test_execute_sql_statements_mutates_before_split_by_default( - mocker: MockerFixture, app + mocker: MockerFixture, app: SupersetApp ) -> None: """ With the default `MUTATE_AFTER_SPLIT=False`, `execute_sql_statements` should @@ -254,6 +255,46 @@ def test_execute_sql_statements_mutates_before_split_by_default( 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) + + 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={}, + ) + + @freeze_time("2021-04-01T00:00:00Z") def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None: """