aminghadersohi commented on code in PR #45232:
URL: https://github.com/apache/superset/pull/45232#discussion_r4288033686
##########
superset/daos/query.py:
##########
@@ -70,26 +70,50 @@ def stop_query(client_id: str) -> None:
if not query:
raise QueryNotFoundException(f"Query with client_id {client_id}
not found")
- if query.status in [
- QueryStatus.FAILED,
- QueryStatus.SUCCESS,
- QueryStatus.TIMED_OUT,
- ]:
+ if query.status in FINISHED_QUERY_STATUSES:
logger.warning(
"Query with client_id could not be stopped: query already
complete",
)
return
- if not sql_lab.cancel_query(query):
+ cancelled = sql_lab.cancel_query(query)
+ if query.status in FINISHED_QUERY_STATUSES:
+ # Preparing the stop can observe the worker's terminal commit.
+ logger.warning(
+ "Query with client_id could not be stopped: query already
complete",
+ )
+ return
+ if not cancelled:
raise SupersetCancelQueryException("Could not cancel query")
+ if query.status == QueryStatus.STOPPED:
+ # The engine spec already committed the stop while preparing it
+ # (see ImpalaEngineSpec.prepare_cancel_query).
+ return
+
# cancel_query() may have staged an early-cancel flag on query.extra
# without committing it (see its docstring/comments); the
# @transaction decorator commits it together with status=STOPPED
# below in one transaction, closing the window where another
# request could observe the flag set but the status still RUNNING.
- query.status = QueryStatus.STOPPED
- query.end_time = now_as_float()
+ #
+ # The write is conditional: the worker can commit a terminal status
+ # after the check above, and that result must not be overwritten.
+ stopped = (
+ db.session.query(Query)
+ .filter(
+ Query.id == query.id,
+ Query.status.notin_(FINISHED_QUERY_STATUSES),
Review Comment:
Remaining: generic async cancellation persistence and Celery error/stop
finalization are outside this PR and belong to
https://github.com/apache/superset/pull/45283 . My earlier async-persistence
and error-handler follow-ups above are superseded by the scope reduction
requested in
https://github.com/apache/superset/pull/45232#pullrequestreview-5537157969 .
The synchronous stop claim and Trino follow-up remain in this PR. Current head:
daf2d551660983edb8a8f38f150fdb3619dbc4a4.
##########
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:
Remaining: generic async cancellation persistence and Celery error/stop
finalization are outside this PR and belong to
https://github.com/apache/superset/pull/45283 . My earlier async-persistence
and error-handler follow-ups above are superseded by the scope reduction
requested in
https://github.com/apache/superset/pull/45232#pullrequestreview-5537157969 .
The synchronous stop claim and Trino follow-up remain in this PR. Current head:
daf2d551660983edb8a8f38f150fdb3619dbc4a4.
--
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]
