aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4280934885
##########
superset/sql/execution/celery_task.py:
##########
@@ -360,27 +431,203 @@ def execute_sql_task(
query_id,
rendered_query,
start_time=start_time,
+ on_commit=mark_committed,
+ on_results_key=track_results_key,
),
can_retry=lambda: not query.progress,
)
except Exception as ex:
logger.exception("Query %d: %s", query_id, ex)
stats_logger = app.config["STATS_LOGGER"]
stats_logger.incr("error_sqllab_unhandled")
- query = _get_query(query_id=query_id)
- return _handle_query_error(ex, query)
+ try:
+ query = _get_query_after_error(query_id)
+ # Failures after the target commit (result serialization,
+ # the results backend, closing the connection) cannot undo
+ # it.
+ return _handle_query_error(
+ ex, query, durable=committed, results_key=results_key
+ )
+ except Exception:
+ # _handle_query_error commits last, so the failure was most
+ # likely not recorded. A commit can still apply and raise.
+ _discard_unreferenced_results(query_id, results_key)
+ raise
+
+
+def _get_query_after_error(query_id: int) -> Query:
+ """Get the query, rolling back a session a failed flush left unusable."""
+ try:
+ return _get_query(query_id=query_id)
+ except Exception: # pylint: disable=broad-except
+ # Rolling back only discards metadata the failed task never committed.
+ db.session.rollback() # pylint: disable=consider-using-transaction
+ return _get_query(query_id=query_id)
+
+
+def _commit_query_unless_stopped(query: Query) -> bool:
+ """Atomically guard terminal metadata writes against a concurrent stop."""
+ with db.session.no_autoflush:
+ # The conditional write acquires the row's write lock until commit.
+ # A separate status read cannot protect the subsequent ORM flush.
+ updated = (
+ db.session.query(Query)
+ .filter(Query.id == query.id, Query.status != QueryStatus.STOPPED)
Review Comment:
Fixed the remaining Trino/Impala stale-metadata overwrite in
`3d14338e31b03a4b5ca191dfa7f21487a510eb2b`: engine preparation only stages
intent, and both the DAO and async-handle callers persist it through the shared
conditional stop guard. A losing stop discards its staged metadata instead of
flushing it over a terminal result.
Regression tests added here:
`tests/unit_tests/daos/test_query_dao.py::test_engine_early_cancel_preserves_concurrently_committed_terminal_metadata`
and `test_engine_early_cancel_commits_with_guarded_stop`. `python3 -m pytest
-q tests/unit_tests/daos/test_query_dao.py -k test_engine_early_cancel -m 'not
manual and not llm_compliance' --tb=short` → 16 failed before at
`840e939ff1a1525f00c986570e8d6104f12f5f8e`, 16 passed after. These exercise the
real Trino/Impala preparation paths; a second session commits
SUCCESS/FAILED/TIMED_OUT with columns/errors/results before the stale stop
proceeds.
`python3 -m pytest -q tests/unit_tests/sql/ tests/unit_tests/daos/
tests/unit_tests/db_engine_specs/test_trino.py
tests/unit_tests/db_engine_specs/test_impala.py -m 'not manual and not
llm_compliance' --tb=short` → 2704 passed, 2 skipped. Staged-file `pre-commit
run` and `pre-commit run --files` covering every file changed by the branch
both passed, including mypy. Concurrency coverage is SQLite only;
PostgreSQL/MySQL concurrency was not run.
Siblings checked: superset/db_engine_specs/trino.py:504 — fixed in
3d14338e31b03a4b5ca191dfa7f21487a510eb2b;
superset/db_engine_specs/impala.py:237 — fixed in
3d14338e31b03a4b5ca191dfa7f21487a510eb2b;
superset/sql/execution/executor.py:1432 — fixed in
3d14338e31b03a4b5ca191dfa7f21487a510eb2b; superset/daos/query.py:65 — not
affected: cancellation and the shared guard already run under no_autoflush in
one transaction; superset/sql_lab.py:971 — not affected: stages early intent
without committing and delegates to the audited engine methods;
superset/sql/execution/executor.py:1449 — not affected: stages intent without
committing and its async caller owns the guarded transaction;
superset/tasks/query_cancel.py:151 — not affected: passes a nonpersistent
cancellation stub and does not write a Query row;
superset/db_engine_specs/base.py:3018 — not affected: default preparation is a
no-op; superset/db_engine_specs/base.py:3117 — not affected: default
cancellation returns
False without writes; superset/db_engine_specs/mysql.py:723 — not affected:
warehouse cancellation only, no Query/session writes;
superset/db_engine_specs/postgres.py:1053 — not affected: warehouse
cancellation only, no Query/session writes;
superset/db_engine_specs/redshift.py:394 — not affected: warehouse cancellation
only, no Query/session writes; superset/db_engine_specs/snowflake.py:545 — not
affected: warehouse cancellation only, no Query/session writes;
superset/db_engine_specs/singlestore.py:571 — not affected: warehouse
cancellation only, no Query/session writes;
superset/db_engine_specs/oracle.py:169 — not affected: warehouse cancellation
only, no Query/session writes; superset/db_engine_specs/ocient.py:392 — not
affected: warehouse cancellation plus in-memory mapping removal, no
metadata-session writes; superset/db_engine_specs/impala.py:151 — not affected:
worker cursor cancellation has no metadata writes;
superset/db_engine_specs/impala.py:220 — not a
ffected: worker progress commit, not called by the stop request;
superset/db_engine_specs/trino.py:345 — not affected: worker handle/progress
metadata writes at 345–404, not called by the stop request;
superset/db_engine_specs/base.py:2032 — not affected: worker cancel-ID
publication, not called by the stop request;
superset/db_engine_specs/base.py:886 — not affected: OAuth setup commit, not a
cancellation-metadata write; superset/db_engine_specs/hive.py:488 — not
affected: worker progress commit and implicit cancellation, no stop-request
Query commit; superset/db_engine_specs/presto.py:1416 — not affected: worker
tracking/progress commits at 1416/1455 and implicit cancellation, not
stop-request writes; superset/db_engine_specs/datastore.py:565 — not affected:
stages worker warnings without committing, outside cancellation;
superset/db_engine_specs/gsheets.py:702 — not affected: spreadsheet catalog
writes a Database, not a Query.
Search: `rg -n 'db.session.(commit|flush)|set_extra_json_key|extra_json|def
.*cancel' superset/db_engine_specs/*.py`, plus every
`cancel_query`/`prepare_cancel_query` call site under `superset/` and AST
inspection of all engine cancellation methods.
Merge coordination: `git merge-tree --write-tree` against
`758e87f60fa44f24e8d7297aa61a4cfb3d4bf28c` conflicts in either order, including
`impala.py`. I recommend landing this PR first, then reconciling
https://github.com/apache/superset/pull/45232 while retaining stage-only engine
preparation and the shared caller guard. This change to `impala.py` is just the
removal of its early commit.
--
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]
