codeant-ai-for-open-source[bot] commented on code in PR #41127:
URL: https://github.com/apache/superset/pull/41127#discussion_r3565611196
##########
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:
**Suggestion:** Mutating the full SQL and then re-parsing it can change the
number of executable statements (for example, mutators that prepend `SET ROLE`
or other statements). Later, `_finalize_successful_query` zips original
statements with execution results using `strict=True`; if counts diverge, this
raises a runtime error after successful execution and marks the async query as
failed. Keep statement cardinality aligned (eg, track/display using the mutated
script when pre-mutation is applied, or avoid strict pairing against the
unmutated original script). [logic error]
<details>
<summary><b>Severity Level:</b> Critical 🚨</summary>
```mdx
- ❌ Async SQL Lab queries can fail after successful execution.
- ⚠️ Statement/result metadata pairing breaks for some mutated queries.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. Configure a SQL query mutator via `Database.mutate_sql_based_on_config()`
at
`superset/models/core.py:705-24` by setting
`current_app.config["SQL_QUERY_MUTATOR"]` to a
function that prepends an extra statement (e.g. returns `"SET ROLE 'bob';" +
sql_`) and
set `current_app.config["MUTATE_AFTER_SPLIT"] = False` so the mutator runs
on the unsplit
block.
2. Use a database whose engine spec executes statements individually
(`run_multiple_statements_as_one=False`) and submit a multi-statement async
query through
`SQLExecutor.execute_async()` in
`superset/sql/execution/executor.py:137-200`, for example
`"SELECT 1; SELECT 2;"`; `execute_async()` calls `_prepare_sql()`
(executor.py:410-57) and
computes `final_sql = transformed_script.format()` containing the two
statements.
3. In `_submit_query_to_celery()` (executor.py:900-31), the executor submits
the query to
Celery by calling `execute_sql_task.delay(query.id, rendered_sql=final_sql,
...)`, so the
Celery worker receives `rendered_query` equal to this two-statement
`final_sql` while the
`Query.sql` field is also set to `final_sql` when the query record is
created.
4. Inside the Celery worker, `execute_sql_task()`
(superset/sql/execution/celery_task.py:344-369) calls
`_execute_sql_statements()`
(celery_task.py:18-110), which builds `original_script = SQLScript(query.sql,
db_engine_spec.engine)` from the unmutated `final_sql` (two statements) and
then calls
`_prepare_statement_blocks(rendered_query, db_engine_spec, database)`
(celery_task.py:128-164).
5. In `_prepare_statement_blocks()`, because
`db_engine_spec.run_multiple_statements_as_one` is False and
`app.config["MUTATE_AFTER_SPLIT"]` is False, the code at
celery_task.py:146-158 mutates
the full SQL block by calling
`database.mutate_sql_based_on_config(parsed_script.format(...),
is_split=False)`,
producing `mutated_sql` that now contains more statements than the original
(e.g. `"SET
ROLE 'bob'; SELECT 1; SET ROLE 'bob'; SELECT 2;"` = four statements), then
re-parses it as
`parsed_script = SQLScript(mutated_sql, ...)` and builds `blocks` from
`parsed_script.statements`, so `len(blocks)` (and thus
`len(execution_results)`) is now
four while `len(original_script.statements)` remains two.
6. `_execute_sql_statements()` passes these `blocks` to
`execute_sql_with_cursor()`
(executor.py:96-95) with `is_split=True`, which iterates over every block
(executor.py:43-55), mutates each statement again (with `is_split=True`, in
this
configuration the mutator is disabled by the `(is_split ==
MUTATE_AFTER_SPLIT)` guard) and
executes them, returning an `execution_results` list with one tuple per
block (four
tuples).
7. After successful execution, `_execute_sql_statements()` computes
`total_execution_time_ms` and calls `_finalize_successful_query(query,
original_script,
execution_results, payload, total_execution_time_ms)` at
celery_task.py:98-101; inside
`_finalize_successful_query` (celery_task.py:167-229), `original_sqls =
[stmt.format() for
stmt in original_script.statements]` yields two strings, but
`execution_results` has four
entries. The loop `for orig_sql, (exec_sql, result_set, exec_time, rowcount)
in
zip(original_sqls, execution_results, strict=True):` at
celery_task.py:182-184 raises
`ValueError: zip() argument 1 of length 2 does not match length 4` due to
the strict=True
mismatch.
8. That `ValueError` propagates out of `_finalize_successful_query` and
`_execute_sql_statements` back into `execute_sql_task()` where it is caught
by the generic
`except Exception as ex:` handler (celery_task.py:370-375), which logs the
error, marks
the query as failed via `_handle_query_error`, and returns an error payload,
so the async
SQL Lab query appears failed to the user even though all statements were
executed
successfully and the failure is solely due to the post-execution strict
pairing between
the unmutated `original_script` and the mutated/re-parsed statement blocks.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=6959c15ba5614f599a7b78d6cae4e44e&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
[](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=6959c15ba5614f599a7b78d6cae4e44e&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/sql/execution/celery_task.py
**Line:** 146:158
**Comment:**
*Logic Error: Mutating the full SQL and then re-parsing it can change
the number of executable statements (for example, mutators that prepend `SET
ROLE` or other statements). Later, `_finalize_successful_query` zips original
statements with execution results using `strict=True`; if counts diverge, this
raises a runtime error after successful execution and marks the async query as
failed. Keep statement cardinality aligned (eg, track/display using the mutated
script when pre-mutation is applied, or avoid strict pairing against the
unmutated original script).
Validate the correctness of the flagged issue. If correct, How can I resolve
this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask
user if the user wants to fix the rest of the comments as well. if said yes,
then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F41127&comment_hash=d31084a3717f0ece95b520cbfb3c9f271db41245d2689bc5ff09f9b7aed62a42&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F41127&comment_hash=d31084a3717f0ece95b520cbfb3c9f271db41245d2689bc5ff09f9b7aed62a42&reaction=dislike'>👎</a>
--
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]
