rusackas commented on code in PR #41127:
URL: https://github.com/apache/superset/pull/41127#discussion_r3566667744
##########
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
Review Comment:
Fixed — replaced the untyped lambda with a typed `_passthrough_mutator`
helper.
##########
tests/unit_tests/sql/execution/test_celery_task.py:
##########
@@ -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}",
Review Comment:
Fixed — replaced with a typed `_prefixing_mutator` helper.
##########
superset/sql/execution/celery_task.py:
##########
@@ -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(
Review Comment:
Fixed — `mutated_sql` now has an explicit `str` annotation.
##########
tests/unit_tests/models/core_test.py:
##########
@@ -262,6 +262,44 @@ def test_table_column_database() -> None:
assert TableColumn(database=database).database is database
[email protected](
+ "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}",
Review Comment:
Fixed — replaced with a typed `_prefixing_sql_query_mutator` helper.
##########
tests/unit_tests/sql_lab_test.py:
##########
@@ -193,6 +193,67 @@ def mock_serialize_payload(payload, use_msgpack):
)
+def test_execute_sql_statements_mutates_before_split_by_default(
+ mocker: MockerFixture, app
+) -> None:
Review Comment:
Fixed — `app` is now typed as `SupersetApp`.
##########
superset/sql_lab.py:
##########
@@ -474,12 +474,24 @@ 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:
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)
Review Comment:
Fixed — `_prepare_statement_blocks` now raises `SupersetErrorException` when
the mutator strips a query down to zero statements, in both `celery_task.py`
and `sql_lab.py`. Added regression tests for both.
##########
superset/sql/execution/celery_task.py:
##########
@@ -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)
Review Comment:
Fixed — a mutator that changes statement count now falls back to labeling
each result with its own executed SQL instead of crashing the strict zip in
`_finalize_successful_query`. Added a regression test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
