codeant-ai-for-open-source[bot] commented on code in PR #45860:
URL: https://github.com/apache/superset/pull/45860#discussion_r4442977559
##########
superset/tasks/async_queries.py:
##########
@@ -225,48 +229,190 @@ def _get_dependency_cache_key() -> str:
raise SupersetException("Prerequisite task did not publish a cache key")
-@contextmanager
-def _capture_query_cancellation(query_context: "QueryContext") ->
Iterator[None]:
- """Enable engine-level cancellation of this task's warehouse query.
-
- Captures the engine cancel id off the live cursor (for engines that expose
- one before execution) and registers an abort handler that kills the backend
- session over a fresh connection, unblocking the task's ``get_df``. Engines
- without cancel support capture nothing and the task stays non-abortable, so
- an abort/timeout simply frees the task without killing the (uncancellable)
- query โ matching the pre-cancellation behavior for those engines.
+class _QueryCancellation:
+ """Cancel the warehouse operation a chart-data task is running when it
aborts.
+
+ Used as the task's cursor sink (see ``superset.tasks.query_cancel``). It
keeps
+ the cancel id of the operation running on the task's own database. Cursors
on
+ other databases (e.g. an annotation layer's source chart) never replace it,
+ since their id would be cancelled against the wrong database.
+
+ Where the id is read depends on ``has_query_id_before_execute``:
+
+ - True (e.g. PostgreSQL, MySQL): the id names the cursor's session, so it
is
+ read once per cursor, when the cursor opens. Reading it after execution
+ would run SQL on the live cursor (``SELECT pg_backend_pid()``) and
replace
+ the rows the chart is about to fetch.
+ - False (Impala): the id names a single operation, so it is read after each
+ statement is submitted, and retired before the next one. Until then the
+ cursor's last operation is a finished one, such as the ``USE`` that
impyla
+ runs when the URL names a database.
+
+ The task registers a single abort handler. For engines read after
submission
+ it is registered when the cursor opens, so an abort while a statement is
+ being submitted (Impala plans it then) is accepted. Other engines register
+ at their first id, since there is nothing to cancel without one. Abort
+ handlers run once, so the handler records the abort and cancels the
published
+ id; when the worker publishes an id after that, it cancels the id itself.
+ Either way the worker stops before the next fetch or statement. A lock
+ guards the state shared with the handler's thread, which never touches the
+ live cursor.
+
+ The handle is also persisted for the orphan reaper, and replaced or cleared
+ as the running operation changes, so the reaper never cancels an operation
+ the task is done with.
"""
- database = getattr(query_context.datasource, "database", None)
- if database is None:
- yield
- return
- ctx = get_context()
- app = current_app._get_current_object() # noqa: SLF001
- captured = False
-
- def _sink(cursor: Any) -> None:
- nonlocal captured
- if captured:
+ def __init__(self, database: "Database") -> None:
+ self._database = database
+ self._database_id = database.id
+ self._after_submission = not
database.db_engine_spec.has_query_id_before_execute
+ self._ctx = get_context()
+ self._app = current_app._get_current_object() # noqa: SLF001
+ self._lock = threading.Lock()
+ # Shared with the abort handler's thread; guarded by ``_lock``.
+ self._active_id: str | None = None
+ self._abort_requested = False
+ # Worker thread only.
+ self._registered = False
+ self._persisted_id: str | None = None
+
+ def __call__(self, cursor: Any, database: "Database", phase: CursorPhase)
-> None:
+ own = database.id == self._database.id
+ if phase == CursorPhase.OPENED:
+ if own:
+ self._opened(cursor)
+ elif phase == CursorPhase.SUBMITTING:
+ if own and self._after_submission:
+ # The previous statement is over: its id must not be cancelled,
+ # or reaped, in place of the one about to be submitted.
+ self._publish(None)
+ # Checked after the retire, so an abort accepted at any point
before
+ # this check stops the statement from being submitted.
+ self._raise_if_aborted()
+ elif phase == CursorPhase.SUBMITTED:
+ try:
+ if own and self._after_submission:
+ self._submitted(cursor)
+ finally:
+ # Even if publishing or cancelling the id failed, an accepted
+ # abort must stop the worker before it fetches.
+ self._raise_if_aborted()
+ else:
+ if own:
+ self._publish(None)
+ if phase == CursorPhase.FAILED:
+ # Once aborting, the failure is the cancel's doing: report the
+ # abort, not a failed query whose empty result could be cached.
+ self._raise_if_aborted()
+
+ def _opened(self, cursor: Any) -> None:
+ if self._after_submission:
+ self._register()
return
- cancel_id = capture_cancel_query_id(database, cursor)
- if cancel_id is None:
+ cancel_id = self._capture(cursor)
+ if cancel_id is None or self._registered:
+ self._publish(cancel_id)
return
- captured = True
+ with self._lock:
+ self._active_id = cancel_id
+ # Registering the first abort handler flushes the whole property cache,
+ # persisting the handle in that write.
+ self._ctx.set_cancellation(self._database.id, cancel_id)
+ self._persisted_id = cancel_id
+ self._register()
+
+ def _submitted(self, cursor: Any) -> None:
+ cancel_id = self._capture(cursor)
+ if self._publish(cancel_id) and cancel_id is not None:
+ # The abort arrived while the statement was being submitted, when
+ # the handler had nothing to cancel, and the handler does not run
+ # again.
+ cancel_chart_query(self._database, cancel_id, self._app)
+
+ def _capture(self, cursor: Any) -> str | None:
+ try:
+ return capture_cancel_query_id(self._database, cursor)
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ logger.warning("Cancel-id capture failed", exc_info=True)
+ return None
- # Persist the handle so the orphan reaper can cancel the query if this
- # worker dies. Set it before on_abort: registering the first handler
- # flushes the whole property cache, persisting the handle in that
write.
- ctx.set_cancellation(database.id, cancel_id)
+ def _publish(self, cancel_id: str | None) -> bool:
+ """Make ``cancel_id`` the id an abort cancels.
- # Registering the first abort handler marks the task abortable and
starts
- # the abort listener; on abort it cancels the query on a fresh
connection.
- def _cancel() -> None:
- cancel_chart_query(database, cancel_id, app)
+ :returns: whether an abort was requested before it was published
+ """
+ with self._lock:
+ self._active_id = cancel_id
+ abort_requested = self._abort_requested
+ if cancel_id != self._persisted_id:
+ try:
+ self._ctx.persist_cancellation(self._database.id, cancel_id)
+ self._persisted_id = cancel_id
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ # Only the orphan reaper reads the persisted handle; the abort
+ # handler uses the in-memory one.
+ logger.warning("Failed to persist the cancel id",
exc_info=True)
+ return abort_requested
+
+ def _raise_if_aborted(self) -> None:
+ with self._lock:
+ abort_requested = self._abort_requested
+ if abort_requested:
+ raise ChartQueryAbortedError()
+
+ def _register(self) -> None:
+ if self._registered:
+ return
+ self._registered = True
+ # Registering the first abort handler marks the task abortable and
+ # starts the abort listener.
+ self._ctx.on_abort(self._on_abort)
+
+ def _on_abort(self) -> None:
+ """The task's abort handler: runs once, usually on another thread."""
+ with self._lock:
+ self._abort_requested = True
+ cancel_id = self._active_id
+ if cancel_id is not None:
+ cancel_chart_query(self._handler_database(), cancel_id, self._app)
+
+ def _handler_database(self) -> "Database":
+ """The task's database, loaded in the handler thread's own session.
+
+ The worker's metastore writes expire its ``Database`` instance, and
+ reloading that instance from this thread would use the worker's session
+ while the worker uses it. Falls back to the worker's instance when the
+ metastore cannot be read (e.g. a self-fence).
+ """
+ from superset import db
+ from superset.models.core import Database
+
+ try:
+ return db.session.get(Database, self._database_id) or
self._database
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ # Leave this thread's session usable after the failed read.
+ db.session.rollback() # pylint: disable=consider-using-transaction
+ return self._database
Review Comment:
**Suggestion:** When the handler-thread lookup fails during a metastore
outage, this returns the workerโs commit-expired `Database`; cancellation then
needs ORM attributes through that unusable worker session and fails.
**Assessment:** ๐ `Major` ยท ๐ `Occurrence: Rarely` ยท ๐ท๏ธ `Stale reference`
[](https://docs.codeant.ai/cli/resolve-pr-comments-skill)
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=9b15efebb59f4b6690bc5601ae572fa8&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
[](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=9b15efebb59f4b6690bc5601ae572fa8&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
<details>
<summary><b>Prompt for AI Agent ๐ค </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/tasks/async_queries.py
**Line:** 396:396
**Comment:**
*Stale Reference: When the handler-thread lookup fails during a
metastore outage, this returns the workerโs commit-expired `Database`;
cancellation then needs ORM attributes through that unusable worker session and
fails.
Validate the correctness of the flagged issue. If correct, How can I resolve
this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask
user if the user wants to fix the rest of the comments as well. if said yes,
then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F45860&comment_hash=df27b3ae15432cfee94a8611c6c820dd16072664bafe143d049b41c1eeff1fd8&reaction=like'>๐</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F45860&comment_hash=df27b3ae15432cfee94a8611c6c820dd16072664bafe143d049b41c1eeff1fd8&reaction=dislike'>๐</a>
##########
superset/models/core.py:
##########
@@ -1044,20 +1044,25 @@ def _log_query(sql_: str) -> None:
rows = None
description = None
- # Give an active GTF chart-data task a chance to capture an engine
- # cancel id off the live cursor before the (blocking) execute
below,
- # so a concurrent abort/timeout can kill the query. No-op
otherwise.
- from superset.tasks.query_cancel import notify_cursor
-
- notify_cursor(cursor)
-
from superset.sql.execution.cancellation import (
cancellable_cursor,
check_query_deadline,
query_executed,
)
+ from superset.tasks.query_cancel import (
+ CursorPhase,
+ notify_cursor,
+ tracked_cursor,
+ )
- with cancellable_cursor(self, cursor, catalog, schema):
+ # Let an active GTF chart-data task follow the operation this
cursor
+ # runs, so a concurrent abort/timeout can cancel it, and stop
before
+ # the next fetch or statement once the abort is accepted. No-op
+ # otherwise.
+ with (
+ tracked_cursor(cursor, self),
Review Comment:
**Suggestion:** `conn.cursor()` has already run before this context is
entered. Impylaโs cursor-time `USE` can block before abort registration, so a
timeout during it is ignored.
**Assessment:** ๐ `Major` ยท ๐ `Occurrence: Rarely` ยท ๐ท๏ธ `Api mismatch`
[](https://docs.codeant.ai/cli/resolve-pr-comments-skill)
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=67bb335b684144be83289d35270fced0&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
[](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=67bb335b684144be83289d35270fced0&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
<details>
<summary><b>Prompt for AI Agent ๐ค </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/models/core.py
**Line:** 1063:1063
**Comment:**
*Api Mismatch: `conn.cursor()` has already run before this context is
entered. Impylaโs cursor-time `USE` can block before abort registration, so a
timeout during it is ignored.
Validate the correctness of the flagged issue. If correct, How can I resolve
this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask
user if the user wants to fix the rest of the comments as well. if said yes,
then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F45860&comment_hash=2626bfd793277d51b2b408ab68b60d8a00aff690461404896343fd5fdd2ba9e8&reaction=like'>๐</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F45860&comment_hash=2626bfd793277d51b2b408ab68b60d8a00aff690461404896343fd5fdd2ba9e8&reaction=dislike'>๐</a>
##########
superset/tasks/async_queries.py:
##########
@@ -225,48 +229,190 @@ def _get_dependency_cache_key() -> str:
raise SupersetException("Prerequisite task did not publish a cache key")
-@contextmanager
-def _capture_query_cancellation(query_context: "QueryContext") ->
Iterator[None]:
- """Enable engine-level cancellation of this task's warehouse query.
-
- Captures the engine cancel id off the live cursor (for engines that expose
- one before execution) and registers an abort handler that kills the backend
- session over a fresh connection, unblocking the task's ``get_df``. Engines
- without cancel support capture nothing and the task stays non-abortable, so
- an abort/timeout simply frees the task without killing the (uncancellable)
- query โ matching the pre-cancellation behavior for those engines.
+class _QueryCancellation:
+ """Cancel the warehouse operation a chart-data task is running when it
aborts.
+
+ Used as the task's cursor sink (see ``superset.tasks.query_cancel``). It
keeps
+ the cancel id of the operation running on the task's own database. Cursors
on
+ other databases (e.g. an annotation layer's source chart) never replace it,
+ since their id would be cancelled against the wrong database.
+
+ Where the id is read depends on ``has_query_id_before_execute``:
+
+ - True (e.g. PostgreSQL, MySQL): the id names the cursor's session, so it
is
+ read once per cursor, when the cursor opens. Reading it after execution
+ would run SQL on the live cursor (``SELECT pg_backend_pid()``) and
replace
+ the rows the chart is about to fetch.
+ - False (Impala): the id names a single operation, so it is read after each
+ statement is submitted, and retired before the next one. Until then the
+ cursor's last operation is a finished one, such as the ``USE`` that
impyla
+ runs when the URL names a database.
+
+ The task registers a single abort handler. For engines read after
submission
+ it is registered when the cursor opens, so an abort while a statement is
+ being submitted (Impala plans it then) is accepted. Other engines register
+ at their first id, since there is nothing to cancel without one. Abort
+ handlers run once, so the handler records the abort and cancels the
published
+ id; when the worker publishes an id after that, it cancels the id itself.
+ Either way the worker stops before the next fetch or statement. A lock
+ guards the state shared with the handler's thread, which never touches the
+ live cursor.
+
+ The handle is also persisted for the orphan reaper, and replaced or cleared
+ as the running operation changes, so the reaper never cancels an operation
+ the task is done with.
"""
- database = getattr(query_context.datasource, "database", None)
- if database is None:
- yield
- return
- ctx = get_context()
- app = current_app._get_current_object() # noqa: SLF001
- captured = False
-
- def _sink(cursor: Any) -> None:
- nonlocal captured
- if captured:
+ def __init__(self, database: "Database") -> None:
+ self._database = database
+ self._database_id = database.id
+ self._after_submission = not
database.db_engine_spec.has_query_id_before_execute
+ self._ctx = get_context()
+ self._app = current_app._get_current_object() # noqa: SLF001
+ self._lock = threading.Lock()
+ # Shared with the abort handler's thread; guarded by ``_lock``.
+ self._active_id: str | None = None
+ self._abort_requested = False
+ # Worker thread only.
+ self._registered = False
+ self._persisted_id: str | None = None
+
+ def __call__(self, cursor: Any, database: "Database", phase: CursorPhase)
-> None:
+ own = database.id == self._database.id
+ if phase == CursorPhase.OPENED:
+ if own:
+ self._opened(cursor)
+ elif phase == CursorPhase.SUBMITTING:
+ if own and self._after_submission:
+ # The previous statement is over: its id must not be cancelled,
+ # or reaped, in place of the one about to be submitted.
+ self._publish(None)
+ # Checked after the retire, so an abort accepted at any point
before
+ # this check stops the statement from being submitted.
+ self._raise_if_aborted()
+ elif phase == CursorPhase.SUBMITTED:
+ try:
+ if own and self._after_submission:
+ self._submitted(cursor)
+ finally:
+ # Even if publishing or cancelling the id failed, an accepted
+ # abort must stop the worker before it fetches.
+ self._raise_if_aborted()
+ else:
+ if own:
+ self._publish(None)
+ if phase == CursorPhase.FAILED:
+ # Once aborting, the failure is the cancel's doing: report the
+ # abort, not a failed query whose empty result could be cached.
+ self._raise_if_aborted()
+
+ def _opened(self, cursor: Any) -> None:
+ if self._after_submission:
+ self._register()
return
- cancel_id = capture_cancel_query_id(database, cursor)
- if cancel_id is None:
+ cancel_id = self._capture(cursor)
+ if cancel_id is None or self._registered:
+ self._publish(cancel_id)
return
- captured = True
+ with self._lock:
+ self._active_id = cancel_id
+ # Registering the first abort handler flushes the whole property cache,
+ # persisting the handle in that write.
+ self._ctx.set_cancellation(self._database.id, cancel_id)
+ self._persisted_id = cancel_id
+ self._register()
+
+ def _submitted(self, cursor: Any) -> None:
+ cancel_id = self._capture(cursor)
+ if self._publish(cancel_id) and cancel_id is not None:
+ # The abort arrived while the statement was being submitted, when
+ # the handler had nothing to cancel, and the handler does not run
+ # again.
+ cancel_chart_query(self._database, cancel_id, self._app)
+
+ def _capture(self, cursor: Any) -> str | None:
+ try:
+ return capture_cancel_query_id(self._database, cursor)
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ logger.warning("Cancel-id capture failed", exc_info=True)
+ return None
- # Persist the handle so the orphan reaper can cancel the query if this
- # worker dies. Set it before on_abort: registering the first handler
- # flushes the whole property cache, persisting the handle in that
write.
- ctx.set_cancellation(database.id, cancel_id)
+ def _publish(self, cancel_id: str | None) -> bool:
+ """Make ``cancel_id`` the id an abort cancels.
- # Registering the first abort handler marks the task abortable and
starts
- # the abort listener; on abort it cancels the query on a fresh
connection.
- def _cancel() -> None:
- cancel_chart_query(database, cancel_id, app)
+ :returns: whether an abort was requested before it was published
+ """
+ with self._lock:
+ self._active_id = cancel_id
+ abort_requested = self._abort_requested
+ if cancel_id != self._persisted_id:
+ try:
+ self._ctx.persist_cancellation(self._database.id, cancel_id)
+ self._persisted_id = cancel_id
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ # Only the orphan reaper reads the persisted handle; the abort
+ # handler uses the in-memory one.
+ logger.warning("Failed to persist the cancel id",
exc_info=True)
+ return abort_requested
+
+ def _raise_if_aborted(self) -> None:
+ with self._lock:
+ abort_requested = self._abort_requested
+ if abort_requested:
+ raise ChartQueryAbortedError()
+
+ def _register(self) -> None:
+ if self._registered:
+ return
+ self._registered = True
+ # Registering the first abort handler marks the task abortable and
+ # starts the abort listener.
+ self._ctx.on_abort(self._on_abort)
Review Comment:
**Suggestion:** If `on_abort` fails while marking the task abortable or
starting its listener, `_registered` stays true, so later cursor phases skip
registration and aborts are never observed.
**Assessment:** ๐ `Major` ยท ๐ `Occurrence: Rarely` ยท ๐ท๏ธ `Logic error`
[](https://docs.codeant.ai/cli/resolve-pr-comments-skill)
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=c8096d10196248899a3253fa5a61a3c1&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
[](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=c8096d10196248899a3253fa5a61a3c1&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
<details>
<summary><b>Prompt for AI Agent ๐ค </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/tasks/async_queries.py
**Line:** 367:370
**Comment:**
*Logic Error: If `on_abort` fails while marking the task abortable or
starting its listener, `_registered` stays true, so later cursor phases skip
registration and aborts are never observed.
Validate the correctness of the flagged issue. If correct, How can I resolve
this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask
user if the user wants to fix the rest of the comments as well. if said yes,
then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F45860&comment_hash=087e752f5effcf7a9531ddb3a59da62d22e30f9ad1ec7c937d3dff9ccfca4bea&reaction=like'>๐</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F45860&comment_hash=087e752f5effcf7a9531ddb3a59da62d22e30f9ad1ec7c937d3dff9ccfca4bea&reaction=dislike'>๐</a>
--
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]
