aminghadersohi commented on code in PR #45232:
URL: https://github.com/apache/superset/pull/45232#discussion_r4278022823
##########
superset/sql/execution/executor.py:
##########
@@ -1450,33 +1451,60 @@ def _cancel_query(database: Database, query: Any) ->
bool:
:param query: Query model instance to cancel
:returns: True if cancelled successfully, False otherwise
"""
+ from superset.common.db_query_status import FINISHED_QUERY_STATUSES,
QueryStatus
from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query as QueryModel
from superset.utils.core import QuerySource
+ from superset.utils.dates import now_as_float
- # Some engines implicitly handle cancellation
- if database.db_engine_spec.has_implicit_cancel():
- return True
-
- # Some databases may need to make preparations for query cancellation
- database.db_engine_spec.prepare_cancel_query(query)
+ if query.status in FINISHED_QUERY_STATUSES:
+ return False
- # Check early cancellation flag
- if query.extra.get(QUERY_EARLY_CANCEL_KEY):
- return True
+ cancelled = database.db_engine_spec.has_implicit_cancel()
+ if not cancelled:
+ database.db_engine_spec.prepare_cancel_query(query)
+ # The query may have finished while the stop was being prepared.
+ if query.status in FINISHED_QUERY_STATUSES:
+ return False
+ cancelled = bool(query.extra.get(QUERY_EARLY_CANCEL_KEY))
+
+ if not cancelled:
+ cancel_query_id = query.extra.get(QUERY_CANCEL_KEY)
+ if cancel_query_id is None:
+ return False
+
+ with database.get_sqla_engine(
+ catalog=query.catalog,
+ schema=query.schema,
+ source=QuerySource.SQL_LAB,
+ ) as engine:
+ with contextlib.closing(engine.raw_connection()) as conn:
+ with contextlib.closing(conn.cursor()) as cursor:
+ cancelled = database.db_engine_spec.cancel_query(
+ cursor, query, cancel_query_id
+ )
- # Get cancel ID
- cancel_query_id = query.extra.get(QUERY_CANCEL_KEY)
- if cancel_query_id is None:
+ if not cancelled:
return False
+ if query.status == QueryStatus.STOPPED:
+ # Impala commits its stop while preparing cursor-side cancellation.
+ return True
- # Execute cancellation
- with database.get_sqla_engine(
- catalog=query.catalog,
- schema=query.schema,
- source=QuerySource.SQL_LAB,
- ) as engine:
- with contextlib.closing(engine.raw_connection()) as conn:
- with contextlib.closing(conn.cursor()) as cursor:
- return database.db_engine_spec.cancel_query(
- cursor, query, cancel_query_id
- )
+ # Match QueryDAO.stop_query: preserve SUCCESS/TIMED_OUT committed while
+ # cancellation was in flight, but replace a FAILED caused by the kill.
+ stopped = (
+ db.session.query(QueryModel)
+ .filter(
+ QueryModel.id == query.id,
+ QueryModel.status.notin_((QueryStatus.SUCCESS,
QueryStatus.TIMED_OUT)),
Review Comment:
Thanks @sadpandajoe, you're right: a stop committed before the killed
statement raised was rewritten as FAILED by the Celery worker.
Fixed in 758e87f60fa44f24e8d7297aa61a4cfb3d4bf28c. `_handle_query_error` in
`superset/sql/execution/celery_task.py` now does a status-only locking refresh
(no autoflush, so the worker's pending changes aren't written first). If the
row is STOPPED it rolls back and returns STOPPED. A local TIMED_OUT skips the
check, so timeouts behave as before. The payload status is taken from the enum
set in the handler, because after the commit the attribute reloads from a real
session as a `str` and `.value` raised.
Regression test:
`tests/unit_tests/sql/test_async_cancel.py::test_worker_error_after_committed_stop_keeps_stopped`.
It runs `execute_sql_task` against a real SQLite metadata DB. The statement
raises after STOPPED is committed, either by `SQLExecutor` `cancel()` in the
worker's session or by a separate session (the web process). It asserts the
committed row and the payload stay STOPPED, while a timeout stays TIMED_OUT and
an error with no stop stays FAILED. At 8da595275cff9fe901f3e94ce2cfe7e6cf375ad2
with only the enum-payload change applied, the two stop cases fail with `assert
'failed' == 'stopped'`. With the guard, all 6 pass.
Siblings checked:
- superset/sql/execution/celery_task.py:82 — fixed in
758e87f60fa44f24e8d7297aa61a4cfb3d4bf28c
- superset/sql_lab.py:99 — not affected: `handle_query_error` (sync SQL Lab)
already refreshes the status and keeps STOPPED.
- superset/db_engine_specs/trino.py:406 — fixed in
758e87f60fa44f24e8d7297aa61a4cfb3d4bf28c: the `_cancel_requested_query`
docstring claim (error handler keeps the stop) now holds on the Celery path too.
- docs/developer_docs/extensions/development.md:254 — fixed in
758e87f60fa44f24e8d7297aa61a4cfb3d4bf28c: the cancellation docs now hold
whether the worker or the cancel commits first.
- superset/sql/execution/celery_task.py:299 — not affected: the
results-backend write failure runs after the stop re-check. A stop landing
during finalization is the race handled in
https://github.com/apache/superset/pull/45283.
- superset/sqllab/sql_json_executer.py:202 — not affected: this dispatch
failure happens before any worker runs.
Remaining: https://github.com/apache/superset/pull/45283 replaces this
function with a broader version, adding the early-cancel flag, discarding
stored results and making a stop win over TIMED_OUT. Whichever merges second
takes 45283's `_handle_query_error` and drops `_has_committed_stop`. The two
already conflict in `celery_task.py` at 8da59527 (`_make_check_stopped_fn` and
the stop finalization), so that resolution is needed either way.
##########
superset/sql/execution/celery_task.py:
##########
@@ -481,15 +484,23 @@ def _execute_sql_statements(
)
) from ex
- # Check if stopped
- if not execution_results:
- payload.update({"status": QueryStatus.STOPPED.value})
- return payload
+ # Cancelling can clear cursor.description, leaving a nonempty result
+ # list that looks like successful DML. Check intent, not just results.
+ stopped = not execution_results or check_stopped()
# Commit for mutations
- if parsed_script.has_mutation() or query.select_as_cta:
+ if not stopped and (parsed_script.has_mutation() or
query.select_as_cta):
Review Comment:
Thanks @sadpandajoe, good catch: no test exercised that guard.
Fixed in 758e87f60fa44f24e8d7297aa61a4cfb3d4bf28c. Added
`tests/unit_tests/sql/execution/test_celery_task.py::test_stop_seen_during_execution_skips_mutation_commit`,
parametrized over an `UPDATE` script and a CTAS (`select_as_cta=True`). Stop
intent appears during execution through either a STOPPED status or the
early-cancel flag. Execution still returns a rowcount result, which is the case
where the result list alone looks like successful DML. Each case asserts
`conn.commit.assert_not_called()`, a STOPPED payload and status, and that
success finalization never runs.
Regression check: with the guard removed locally (`if
parsed_script.has_mutation() or query.select_as_cta:`), all 4 cases fail with
`Expected 'commit' to not have been called. Called 1 times.` With the guard
restored, all 4 pass.
Siblings checked:
- superset/sql/execution/celery_task.py:514 — fixed in
758e87f60fa44f24e8d7297aa61a4cfb3d4bf28c: the `not stopped` mutation/CTAS
commit guard is now tested.
- superset/sql/execution/celery_task.py:518 — not affected: the
post-connection `if stopped or check_stopped():` re-check is already covered by
`test_stop_with_partial_results_skips_success_finalization`.
- superset/sql_lab.py:632 — not affected: this PR doesn't change it. In the
sync path, a stop seen during a block raises `SqlLabQueryStoppedException` and
returns before this `conn.commit()`.
Remaining: in the sync path, a stop that lands after the last block finishes
is not re-checked before its `conn.commit()`. This PR doesn't touch that path,
and I haven't changed it here.
--
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]
