diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py index 2b3407768b6..03f02f5ddf6 100644 --- a/superset/sql/execution/celery_task.py +++ b/superset/sql/execution/celery_task.py @@ -141,7 +141,23 @@ def _prepare_statement_blocks( # 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)] + 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, diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 4d6824c0c3f..527be41773a 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -489,7 +489,23 @@ def execute_sql_statements( # noqa: C901 # 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)] + 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 = [ + ";\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, diff --git a/tests/unit_tests/sql/execution/test_celery_task.py b/tests/unit_tests/sql/execution/test_celery_task.py index c53e2a33ac7..dc9a1170b13 100644 --- a/tests/unit_tests/sql/execution/test_celery_task.py +++ b/tests/unit_tests/sql/execution/test_celery_task.py @@ -389,6 +389,38 @@ def test_prepare_statement_blocks_skips_pre_split_mutation_when_configured( 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: diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index b4ce388b617..d2b1a3635ad 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -255,6 +255,69 @@ def test_execute_sql_statements_mutates_before_split_by_default( 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: