codeant-ai-for-open-source[bot] commented on code in PR #44552:
URL: https://github.com/apache/superset/pull/44552#discussion_r4078643962
##########
superset/tasks/query_cancel.py:
##########
@@ -163,3 +172,140 @@ def cancel_chart_query(
exc_info=True,
)
return False
+
+
+def _registry_key(user_id: int, client_id: str) -> str:
+ """Cache key for a user's in-flight, cancellable chart query.
+
+ The user id is part of the key rather than a field compared after lookup,
so
+ a ``client_id`` is only ever resolvable within the namespace of the user
who
+ registered it. A caller passing somebody else's ``client_id`` gets a miss —
+ it cannot read, cancel, or overwrite another user's entry.
+ """
+ return f"chart-query-cancel:{user_id}:{client_id}"
+
+
+def _registry_ttl() -> int:
+ """How long a cancel handle stays resolvable.
+
+ A synchronous chart query cannot outlive the web request running it, so the
+ webserver timeout is the natural upper bound. Entries are discarded as soon
+ as the query returns; this TTL only bounds the leak when a worker dies
+ mid-query.
+ """
+ from flask import current_app
+
+ return int(current_app.config.get("SUPERSET_WEBSERVER_TIMEOUT", 60))
+
+
+@contextmanager
+def cancellable_chart_query(
+ client_id: "str | None", database: "Database | None"
+) -> Iterator[None]:
+ """Let the requesting user cancel this synchronous chart query by
``client_id``.
+
+ Captures the engine cancel id off the live cursor (for engines that expose
+ one before execution) and publishes it so a concurrent Stop request — which
+ lands on a different worker while this one is blocked on the query — can
kill
+ the backend session. Engines without cancel support capture nothing and the
+ query stays non-cancellable, exactly as before.
+
+ A no-op without a ``client_id``, without a database (e.g. the annotation
+ datasource, which queries Superset's own metadata DB), or for an
+ unauthenticated request — an anonymous viewer of a public dashboard has no
+ user id to scope the handle to, and an unscoped handle would be cancellable
+ by any other anonymous visitor.
+ """
+ from superset.utils.core import get_user_id
+
+ user_id = get_user_id()
+ if not client_id or database is None or user_id is None:
+ yield
+ return
+
+ # Rebound as non-optional locals: mypy does not carry the narrowing above
+ # into the nested function below.
+ owner_id: int = user_id
+ query_id: str = client_id
+ target: "Database" = database
+ database_id = target.id
+ captured = False
+
+ def _sink(cursor: Any) -> None:
+ nonlocal captured
+ if captured:
+ return
Review Comment:
✅ **CodeAnt verified this suggestion was addressed in subsequent commits and
marked this thread resolved** as of `2547ad5`.
The cursor sink now captures and publishes a cancellation handle for every
cursor, rather than using a `captured` guard that only allowed the first handle
to be stored.
<sub>If that's not right, unresolve this thread and CodeAnt will leave it
open.</sub>
<!-- codeant-auto-resolve-reply -->
--
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]
