diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py index 4310991edee..3bcd35fc622 100644 --- a/superset/sql/execution/celery_task.py +++ b/superset/sql/execution/celery_task.py @@ -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 @@ -57,6 +57,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,6 +128,7 @@ 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. @@ -139,6 +143,19 @@ def _prepare_statement_blocks( if db_engine_spec.run_multiple_statements_as_one: 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 = 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) blocks = [ statement.format(comments=db_engine_spec.allows_sql_comments) for statement in parsed_script.statements @@ -423,7 +440,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, diff --git a/superset/sql_lab.py b/superset/sql_lab.py index f17a1cea5d1..b19eb0e93ae 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -480,6 +480,18 @@ def execute_sql_statements( # noqa: C901 if db_engine_spec.run_multiple_statements_as_one: 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 = 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) blocks = [ statement.format(comments=db_engine_spec.allows_sql_comments) for statement in parsed_script.statements diff --git a/tests/unit_tests/sql/execution/test_celery_task.py b/tests/unit_tests/sql/execution/test_celery_task.py index fa3ae8f3d51..39818dd10fa 100644 --- a/tests/unit_tests/sql/execution/test_celery_task.py +++ b/tests/unit_tests/sql/execution/test_celery_task.py @@ -282,9 +282,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 = lambda sql, **kw: sql 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 +298,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 = lambda sql, **kw: sql 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 +315,43 @@ 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 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=lambda sql, **kw: f"-- mutated\n{sql}", + ) + 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) + + # ============================================================================= # Result Finalization Tests # ============================================================================= diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index 0f715b11a18..06c3e309463 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -193,6 +193,67 @@ def test_execute_sql_statement_within_payload_limit(mocker: MockerFixture, app) ) +def test_execute_sql_statements_mutates_before_split_by_default( + mocker: MockerFixture, app +) -> 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 + + @freeze_time("2021-04-01T00:00:00Z") def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None: """