sadpandajoe commented on code in PR #45232:
URL: https://github.com/apache/superset/pull/45232#discussion_r4276923245
##########
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:
This commits STOPPED and returns True, but the Celery worker's error path
still ends the run as FAILED. `_handle_query_error` (celery_task.py:88-89) only
preserves TIMED_OUT, so when the killed statement raises after this commit
(Trino's `_cancel_requested_query` kill, Hive/Presto after `cursor.cancel()`,
or Postgres/MySQL when the worker commits last), the row flips STOPPED to
FAILED even though `cancel()` reported success. The new docs and the Trino
docstring both say the error handler keeps the stop. Should
`_handle_query_error` refresh and keep STOPPED the way
`sql_lab.handle_query_error` does, with a test that raises after the stop?
##########
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:
The new `not stopped` guard on the mutation/CTAS commit isn't exercised by
any test. The added Celery stop tests return a SELECT with
`select_as_cta=False`, so `has_mutation()` is false and `conn.commit()` is
never reached; removing this guard leaves them green, while a stopped UPDATE or
CTAS would then be committed. Could one case use an UPDATE or CTAS script with
stop intent seen during execution and assert `conn.commit` is not called?
--
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]
