aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4280728349
##########
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 in 840e939ff1a1525f00c986570e8d6104f12f5f8e. The stop request uses a
conditional UPDATE excluding SUCCESS, FAILED and TIMED_OUT before any staged
cancellation metadata can autoflush. If publication won, a locking refresh
discards the losing stop's pending fields. The worker's stop-payload writer
uses the same guard and returns the recorded terminal status.
Regression tests added by this PR:
- `test_stop_query_preserves_concurrently_committed_terminal_result`: a
second session commits the terminal status, results key, end time and extra
metadata inside `cancel_query`, after the stop's initial RUNNING read. All six
cases fail at ea6168068bb2473e25fa39e4c9ccaaab143039fa and pass at
840e939ff1a1525f00c986570e8d6104f12f5f8e.
- `test_stopped_payload_preserves_concurrently_committed_terminal_result`:
all three worker-payload sibling cases fail before and pass after.
- `test_stop_query_commits_stop_with_early_cancel`: all five cases fail
before and pass after, checking that an incidental metadata read cannot flush
the early-cancel flag separately from STOPPED.
Evidence: `PYTHONPATH=/tmp/r456-deps:./superset-core/src:. pytest -q
tests/unit_tests/daos/test_query_dao.py
tests/unit_tests/sql/execution/test_utils.py -m 'not manual and not
llm_compliance'` gives `18 failed` in a separate checkout of
ea6168068bb2473e25fa39e4c9ccaaab143039fa and `18 passed` at
840e939ff1a1525f00c986570e8d6104f12f5f8e. The broader SQL/DAO/legacy suites
give `2648 passed, 2 skipped, 1 deselected`; the additional base/SQLite/Trino
engine-spec suite gives `350 passed`. `pre-commit run --files <all PR-changed
paths>` and `pre-commit run mypy` pass. The concurrency tests exercise SQLite;
Postgres/MySQL runtime checks were not run.
Siblings checked:
- superset/daos/query.py:84 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql/execution/celery_task.py:480 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql_lab.py:154 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:161 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:433 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:529 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:647 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:674 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:710 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:712 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:753 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:862 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:914 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/sql_lab.py:923 — not affected: status comparison or payload only;
no STOPPED metadata write
- superset/db_engine_specs/presto.py:1433 — not affected: reads stop intent
to cancel the engine cursor; does not assign STOPPED
- superset/sql/execution/utils.py:125 — not affected: status comparison or
payload only; no STOPPED metadata write
- superset/sql/execution/utils.py:152 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql/execution/utils.py:161 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/db_engine_specs/trino.py:375 — not affected: reads stop intent to
cancel the engine cursor; does not assign STOPPED
- superset/sql/execution/celery_task.py:469 — not affected: conditional
worker terminal write excludes STOPPED; does not write STOPPED
- superset/sql/execution/celery_task.py:499 — not affected: status
comparison or payload only; no STOPPED metadata write
- superset/db_engine_specs/impala.py:199 — not affected: reads stop intent
to cancel the engine cursor; does not assign STOPPED
- superset/db_engine_specs/hive.py:447 — not affected: reads stop intent to
cancel the engine cursor; does not assign STOPPED
- superset/sql/execution/executor.py:1339 — not affected: maps a stored
status into the public status enum; no metadata write
- superset/common/db_query_status.py:23 — not affected: enum definition only
- superset/charts/schemas.py:1865 — not affected: schema validation choice
only
- superset/utils/core.py:422 — not affected: enum definition only
- superset/mcp_service/chart/query_result.py:55 — not affected:
result-status classification only
Searches: `rg -n 'QueryStatus.STOPPED|status\s*=\s*["\x27]stopped'
superset/` and `rg -n "[\'\"]stopped[\'\"]" superset/ --glob "*.py"`.
##########
superset/sql_lab.py:
##########
@@ -62,6 +63,7 @@
from superset.extensions import celery_app, event_logger
from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
+from superset.sql.execution.celery_task import stop_seen_during_execution
Review Comment:
Addressed in 840e939ff1a1525f00c986570e8d6104f12f5f8e. Both paths import
`_extract_errors`, `_discard_unreferenced_results`, the delete helper, and
`stop_seen_during_execution` from `superset/sql/execution/utils.py`; the legacy
path does not import the Celery task module. The cleanup rule is unchanged,
with the caller's results backend passed explicitly.
`test_execution_helpers_are_shared` checks all four helpers are the
identical function objects from the neutral module (including the legacy
delete-helper alias). All four cases fail at
ea6168068bb2473e25fa39e4c9ccaaab143039fa and pass at
840e939ff1a1525f00c986570e8d6104f12f5f8e.
Evidence: `PYTHONPATH=/tmp/r456-deps:./superset-core/src:. pytest -q
tests/unit_tests/daos/test_query_dao.py
tests/unit_tests/sql/execution/test_utils.py -m 'not manual and not
llm_compliance'` gives `18 failed` before / `18 passed` after. Existing error
extraction, cleanup, referenced-blob preservation, and stop-observation tests
remain green in the SQL/DAO/legacy suites: `2648 passed, 2 skipped, 1
deselected`. Additional engine-spec tests: `350 passed`. `pre-commit run
--files <all PR-changed paths>` and `pre-commit run mypy` pass.
Siblings checked:
- superset/sql_lab.py:66 — fixed in 840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql/execution/celery_task.py:54 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql/execution/utils.py:40 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql/execution/utils.py:51 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql/execution/utils.py:81 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql/execution/utils.py:114 — fixed in
840e939ff1a1525f00c986570e8d6104f12f5f8e
- superset/sql/execution/executor.py:1230 — not affected: imports the Celery
task to queue asynchronous work, not to reuse these helpers
`rg -n '^def
(_extract_errors|_discard_stored_results|_discard_unpublished_results|_discard_unreferenced_results|stop_seen_during_execution)'
superset/ --glob '*.py'` finds only the four neutral definitions.
--
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]
