aminghadersohi commented on code in PR #44552:
URL: https://github.com/apache/superset/pull/44552#discussion_r4336410058
##########
superset/tasks/query_cancel.py:
##########
@@ -168,3 +186,275 @@ 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 _cancel_requested_key(user_id: int, client_id: str) -> str:
+ """Cache key recording that a user asked to stop a chart query run.
+
+ Scoped by user for the same reason as ``_registry_key``: a ``client_id``
+ can only ever flag the requester's own run.
+ """
+ return f"chart-query-cancel-requested:{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.
+ """
+ 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.
+ """
+ # Inline (also below, and in _publish_cancel_handle/_discard_cancel_handle/
+ # cancel_chart_query_for_user): the unit tests patch get_user_id and
+ # cache_manager at their defining modules (superset.utils.core,
+ # superset.extensions), not here. A module-level `from ... import X` binds
+ # X once at import time, before any test patch runs, so the patch would
+ # silently never take effect; re-importing on every call picks up the
+ # patched object instead.
+ 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:
+ # Republish for every cursor rather than only the first: one query
+ # object can execute several statements in turn (e.g. the grouping-sets
+ # fallback), each on whatever connection the pool hands out, and the
+ # handle must always name the session that is executing right now.
+ nonlocal captured
+ # A Stop recorded before this statement had a handle to cancel is
+ # honoured by not running the statement at all.
+ if chart_query_cancel_requested(query_id):
+ raise ChartQueryStoppedError(_("Query was stopped"))
+ 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)
+ # A Stop landing between the check above and the publish saw no handle
+ # to kill; recheck now that the handle is visible.
+ if chart_query_cancel_requested(query_id):
+ _discard_cancel_handle(owner_id, query_id, cancel_id)
+ raise ChartQueryStoppedError(_("Query was stopped"))
+
+ try:
+ with capture_cancel_id(_sink):
+ yield
+ finally:
+ if captured:
+ _discard_cancel_handle(owner_id, query_id)
+
+
+def _publish_cancel_handle(
+ user_id: int, client_id: str, database_id: int, cancel_query_id: str
+) -> None:
+ """Publish a cancel handle for the owning user. Best-effort."""
+ # Inline for test-patchability; see cancellable_chart_query's comment.
+ from superset.extensions import cache_manager
+
+ try:
+ cache_manager.chart_query_cancel_cache.set(
+ _registry_key(user_id, client_id),
+ {"database_id": database_id, "cancel_query_id": cancel_query_id},
+ timeout=_registry_ttl(),
+ )
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ # Forfeits cancellability for this query; never breaks its execution.
+ logger.warning("Could not publish chart query cancel handle",
exc_info=True)
+
+
+def _discard_cancel_handle(
+ user_id: int, client_id: str, expected_cancel_query_id: "str | None" = None
+) -> None:
+ """Drop a cancel handle once its query is no longer running. Best-effort.
+
+ When ``expected_cancel_query_id`` is given, the entry is only deleted if it
+ still names that id. Sequential cursors within one query (e.g. the
+ grouping-sets fallback) republish the same key as each statement starts;
+ without this check, discarding a handle resolved just before a newer
+ cursor's republish would delete that newer, still-live entry instead of
+ the stale one that was actually acted on.
+ """
+ # Inline for test-patchability; see cancellable_chart_query's comment.
+ from superset.extensions import cache_manager
+
+ key = _registry_key(user_id, client_id)
+ try:
+ if expected_cancel_query_id is not None:
+ current = cache_manager.chart_query_cancel_cache.get(key)
+ if (
+ not current
+ or current.get("cancel_query_id") != expected_cancel_query_id
+ ):
+ return
+ cache_manager.chart_query_cancel_cache.delete(key)
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ logger.warning("Could not discard chart query cancel handle",
exc_info=True)
+
+
+def _request_chart_query_cancel(user_id: int, client_id: str) -> None:
+ """Record that the user asked to stop this run. Best-effort.
+
+ Killing the statement that is executing right now is only half of a Stop:
+ a chart-data run can be several statements (contribution totals, the
+ grouping-sets fallback's one-query-per-level, one query object after
+ another), and a Stop that lands before the first starts, or in the gap
+ between two, finds no handle to cancel. The intent recorded here is what
+ the execution loop checks before it starts each statement (see
+ ``chart_query_cancel_requested``), so the rest of the run is skipped
+ instead of carrying on as if nothing happened.
+ """
+ # Inline for test-patchability; see cancellable_chart_query's comment.
+ from superset.extensions import cache_manager
+
+ try:
+ cache_manager.chart_query_cancel_cache.set(
+ _cancel_requested_key(user_id, client_id),
+ True,
+ timeout=_registry_ttl(),
+ )
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ logger.warning("Could not record chart query cancel request",
exc_info=True)
+
+
+def chart_query_cancel_requested(client_id: "str | None") -> bool:
+ """Whether the requesting user has asked to stop the run named by
``client_id``.
+
+ Checked by the execution loop before each statement of a chart-data run.
+ False without a ``client_id`` or a user (nothing could have been
+ requested), and on a cache failure (never block a query on the registry).
+ """
+ # Inline for test-patchability; see cancellable_chart_query's comment.
+ from superset.extensions import cache_manager
+ from superset.utils.core import get_user_id
+
+ user_id = get_user_id()
+ if not client_id or user_id is None:
+ return False
+
+ try:
+ return bool(
+ cache_manager.chart_query_cancel_cache.get(
+ _cancel_requested_key(user_id, client_id)
+ )
+ )
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ logger.warning("Could not read chart query cancel request",
exc_info=True)
+ return False
+
+
+@contextmanager
+def chart_query_run(client_id: "str | None") -> Iterator[None]:
+ """Scope one chart-data run so its cancel intent is cleared when it ends.
+
+ The intent is deliberately *not* cleared on entry: a Stop that arrives
+ before the run reaches its first statement must still take effect. Each
+ run carries a fresh ``client_id``, so an intent recorded for a run that is
+ already over can never leak into a later one; clearing it here just keeps
+ the registry tidy rather than leaving it to the TTL.
+ """
+ # Inline for test-patchability; see cancellable_chart_query's comment.
+ from superset.extensions import cache_manager
+ from superset.utils.core import get_user_id
+
+ user_id = get_user_id()
+ if not client_id or user_id is None:
+ yield
+ return
+
+ try:
+ yield
+ finally:
+ try:
+ cache_manager.chart_query_cancel_cache.delete(
+ _cancel_requested_key(user_id, client_id)
+ )
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ logger.warning("Could not clear chart query cancel request",
exc_info=True)
+
+
+def cancel_chart_query_for_user(client_id: str) -> bool:
+ """Cancel the requesting user's running chart query, if it is cancellable.
+
+ Always records the cancel intent for the run first, so whatever the run
+ has not started yet is skipped at its next checkpoint even when there is
+ no statement executing right now to kill.
+
+ :returns: True if the engine reported the query cancelled. False covers
both
+ "no such in-flight query for this user" and "the engine declined" — the
+ caller cannot distinguish them, which is deliberate: it keeps the
+ endpoint from confirming whether a given ``client_id`` exists.
+ """
+ # Inline for test-patchability; see cancellable_chart_query's comment.
+ from superset.daos.database import DatabaseDAO
+ from superset.extensions import cache_manager
+ from superset.utils.core import get_user_id
+
+ user_id = get_user_id()
+ if user_id is None:
+ return False
+
+ _request_chart_query_cancel(user_id, client_id)
+
+ try:
+ handle = cache_manager.chart_query_cancel_cache.get(
+ _registry_key(user_id, client_id)
+ )
+ except Exception: # noqa: BLE001 pylint: disable=broad-except
+ logger.warning("Could not read chart query cancel handle",
exc_info=True)
+ return False
+
+ if not handle:
+ return False
+
+ database = DatabaseDAO.find_by_id(handle["database_id"])
Review Comment:
With `all_datasource_access` but no per-database grants, chart queries can
run but this lookup filters out their database, so Stop never calls engine
cancellation. The user-scoped handle already identifies the database of the
caller's authorized query.
```suggestion
database = DatabaseDAO.find_by_id(handle["database_id"],
skip_base_filter=True)
```
--
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]
