aminghadersohi commented on code in PR #44552:
URL: https://github.com/apache/superset/pull/44552#discussion_r4333474718
##########
superset/common/query_context_processor.py:
##########
@@ -560,8 +574,36 @@ def get_query_result(self, query_object: QueryObject) ->
QueryResult:
"""
if query_object.grouping_sets and not self._supports_grouping_sets():
return self._grouping_sets_fallback(query_object)
+ return self._execute_query_object(query_object)
+
+ def _execute_query_object(self, query_object: QueryObject) -> QueryResult:
+ """Run one statement on the datasource, unless its owner asked to stop.
+
+ Every statement of a chart-data run passes through here, so a Stop
+ that lands before the run starts, or between two of its statements,
+ is honoured at the next one instead of being lost. Cancelling the
+ statement executing right now is ``cancellable_chart_query``'s job.
+ """
+ if chart_query_cancel_requested(self._query_context.client_id):
+ return self._stopped_query_result()
return self._qc_datasource.get_query_result(query_object)
Review Comment:
A Stop that kills a time-offset subquery leaves an empty comparison frame,
but the parent still reports SUCCESS and caches it. Recheck the stop intent
after the datasource call so cancelled results cannot become cache hits.
```suggestion
if chart_query_cancel_requested(self._query_context.client_id):
return self._stopped_query_result()
result = self._qc_datasource.get_query_result(query_object)
if chart_query_cancel_requested(self._query_context.client_id):
return self._stopped_query_result()
return result
```
##########
superset/tasks/query_cancel.py:
##########
@@ -168,3 +177,266 @@ 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
+ 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)
Review Comment:
A Stop during connection setup records intent and misses the unpublished
handle. This sink then publishes without checking that intent, so the statement
still executes; internal time-offset statements bypass the only check too.
--
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]
