diff --git a/superset/config.py b/superset/config.py index 0bb0cacb9f5..80356e99d1d 100644 --- a/superset/config.py +++ b/superset/config.py @@ -2005,6 +2005,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 diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py index dabb5f0915b..4310991edee 100644 --- a/superset/sql/execution/celery_task.py +++ b/superset/sql/execution/celery_task.py @@ -445,6 +445,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 diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py index 13e12fc1a47..7a69510abfd 100644 --- a/superset/sql/execution/executor.py +++ b/superset/sql/execution/executor.py @@ -101,6 +101,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 +120,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 +145,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 diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 34611d38a33..f17a1cea5d1 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -474,7 +474,7 @@ 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 + # 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: @@ -517,8 +517,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) diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py index 8d15b23077b..6ca9c15b980 100644 --- a/tests/unit_tests/models/core_test.py +++ b/tests/unit_tests/models/core_test.py @@ -262,6 +262,44 @@ def test_table_column_database() -> None: assert TableColumn(database=database).database is database +@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": lambda sql, **kwargs: f"-- mutated\n{sql}", + "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. diff --git a/tests/unit_tests/sql/execution/test_executor.py b/tests/unit_tests/sql/execution/test_executor.py index fdaf52f8a79..0e12de4b034 100644 --- a/tests/unit_tests/sql/execution/test_executor.py +++ b/tests/unit_tests/sql/execution/test_executor.py @@ -47,7 +47,7 @@ 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 create_mock_cursor, mock_query_execution # ============================================================================= # Basic Execution Tests @@ -871,6 +871,45 @@ 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 = mocker.patch.object( + database, "mutate_sql_based_on_config", side_effect=lambda sql, **kw: sql + ) + 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 = 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 # =============================================================================