codeant-ai-for-open-source[bot] commented on code in PR #44552:
URL: https://github.com/apache/superset/pull/44552#discussion_r4078643235
##########
superset/common/query_context_factory.py:
##########
@@ -111,6 +112,7 @@ def create( # pylint: disable=too-many-arguments
result_format=result_format,
force=force,
force_nonce=force_nonce,
+ client_id=client_id,
Review Comment:
✅ **Customized review instruction saved!**
**Instruction:**
> Do not require propagating client_id through async serialization; async
cancellation is handled by task abort, and the query registry is only for
synchronous Stop requests.
**Applied to:**
- `superset/common/query_context_factory.py`
---
💡 *To manage or update this instruction, visit: [CodeAnt AI
Settings](https://app.codeant.ai/org/settings/learnings)*
##########
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
+ cancel_id = capture_cancel_query_id(target, cursor)
+ if cancel_id is None:
+ return
+ captured = True
+ _publish_cancel_handle(owner_id, query_id, database_id, cancel_id)
+
+ try:
+ with capture_cancel_id(_sink):
+ yield
+ finally:
+ if captured:
+ _discard_cancel_handle(owner_id, query_id)
Review Comment:
✅ **Customized review instruction saved!**
**Instruction:**
> Do not flag the potential stale-handle cleanup race in query cancellation
when each run uses a fresh nanoid and atomic compare-and-delete is unavailable;
a client reusing its own id may miss Stop in the worst case.
**Applied to:**
- `superset/tasks/query_cancel.py`
---
💡 *To manage or update this instruction, visit: [CodeAnt AI
Settings](https://app.codeant.ai/org/settings/learnings)*
--
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]