fix(sqllab): mutate per-statement when engine runs statements as one block

MUTATE_AFTER_SPLIT=True never fired the mutator for engines with
run_multiple_statements_as_one=True (BigQuery, Datastore, Kusto), since
those engines always build a single joined block and the per-block
mutation call always passes is_split=False. Mutate each statement
before joining them into that block so the flag applies consistently.

Verified separately via a full `pre-commit run` (all hooks touching
these files passed: auto-walrus, mypy, ruff-format, ruff, pylint) run
outside the git hook, whose invocation via the system /usr/bin/python3
hits an unrelated pre-existing environment issue installing the zizmor
hook (requires Python >=3.10, system python3 here is older) rather
than anything in this change.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Evan
2026-07-14 04:41:36 -07:00
parent 68c072a524
commit 1f81871fb3
4 changed files with 129 additions and 2 deletions

View File

@@ -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,