aminghadersohi commented on code in PR #44552:
URL: https://github.com/apache/superset/pull/44552#discussion_r4304741714
##########
tests/unit_tests/tasks/test_query_cancel.py:
##########
@@ -205,3 +206,249 @@ def test_capture_is_noop_without_a_database() -> None:
with _capture_query_cancellation(qc):
assert _cancel_id_sink.get() is None
get_context.assert_not_called()
+
+
+# --- synchronous Explore cancellation registry ----------------------------
+
+
+class _FakeCache:
+ """Minimal dict-backed stand-in for ``cache_manager.cache``.
+
+ Real storage (rather than a MagicMock) so the user-scoping of the registry
+ keys is genuinely exercised end to end.
+ """
+
+ def __init__(self) -> None:
+ self.store: dict[str, object] = {}
+
+ def set(self, key: str, value: object, timeout: int | None = None) -> None:
+ self.store[key] = value
+
+ def get(self, key: str) -> object | None:
+ return self.store.get(key)
+
+ def delete(self, key: str) -> None:
+ self.store.pop(key, None)
+
+
+@contextmanager
+def _registry_env(user_id: int | None, cache: "_FakeCache"):
+ """Run with a given current user and a shared fake cache backend."""
+ # ``cache_manager.cache`` is a read-only property; swap the backing
attribute.
+ with (
+ patch("superset.utils.core.get_user_id", return_value=user_id),
+ patch("superset.extensions.cache_manager._cache", cache),
+ patch("superset.tasks.query_cancel._registry_ttl", return_value=60),
+ ):
+ yield
+
+
+def _run_cancellable(client_id, database, cache, user_id,
cancel_id="engine-1"):
+ """Drive one cancellable query, returning whether a handle was live inside
it."""
+ from superset.tasks.query_cancel import cancellable_chart_query
+
+ with _registry_env(user_id, cache):
+ with patch(
+ "superset.tasks.query_cancel.capture_cancel_query_id",
+ return_value=cancel_id,
+ ):
+ with cancellable_chart_query(client_id, database):
+ # Stands in for Database._execute_sql_with_mutation_and_logging
+ # handing the live cursor to the active sink before executing.
+ notify_cursor(MagicMock())
+ inside = dict(cache.store)
+ return inside
+
+
+def test_registry_key_is_scoped_by_user() -> None:
+ """The same ``client_id`` under two users maps to two distinct keys."""
+ from superset.tasks.query_cancel import _registry_key
+
+ # Same client_id under two users must not collide.
+ assert _registry_key(1, "abc") != _registry_key(2, "abc")
+ assert "abc" in _registry_key(1, "abc")
+
+
+def test_cancellable_chart_query_publishes_then_discards_the_handle() -> None:
+ """A captured handle is live for the query's duration, then removed."""
+ cache = _FakeCache()
+ database = MagicMock()
+ database.id = 5
+
+ inside = _run_cancellable("client-1", database, cache, user_id=1)
+
+ # Live for the duration of the query...
+ assert inside == {
+ "chart-query-cancel:1:client-1": {
+ "database_id": 5,
+ "cancel_query_id": "engine-1",
+ }
+ }
+ # ...and cleaned up once it returns.
+ assert cache.store == {}
+
+
+def test_cancellable_chart_query_republishes_for_each_cursor() -> None:
+ """Every cursor the query opens refreshes the handle to its own session.
+
+ One query object can run several statements in turn (the grouping-sets
+ fallback does), each on whatever pooled connection it gets, so the handle
+ must follow the statement that is executing now rather than the first one.
+ """
+ from superset.tasks.query_cancel import cancellable_chart_query
+
+ cache = _FakeCache()
+ database = MagicMock()
+ database.id = 5
+ seen: list[object] = []
+
+ with _registry_env(1, cache):
+ with patch(
+ "superset.tasks.query_cancel.capture_cancel_query_id",
+ side_effect=["engine-1", "engine-2"],
+ ):
+ with cancellable_chart_query("client-1", database):
+ notify_cursor(MagicMock())
+ seen.append(cache.store["chart-query-cancel:1:client-1"])
+ notify_cursor(MagicMock())
+ seen.append(cache.store["chart-query-cancel:1:client-1"])
+
+ assert seen == [
+ {"database_id": 5, "cancel_query_id": "engine-1"},
+ {"database_id": 5, "cancel_query_id": "engine-2"},
+ ]
+ assert cache.store == {}
+
+
+def test_cancellable_chart_query_is_noop_without_a_client_id() -> None:
+ """Requests that carry no ``client_id`` register nothing."""
+ cache = _FakeCache()
+ database = MagicMock()
+ database.id = 5
+
+ assert _run_cancellable(None, database, cache, user_id=1) == {}
+
+
+def test_cancellable_chart_query_is_noop_without_a_database() -> None:
+ """Datasources without a database (e.g. annotations) register nothing."""
+ cache = _FakeCache()
+
+ assert _run_cancellable("client-1", None, cache, user_id=1) == {}
+
+
+def test_cancellable_chart_query_is_noop_for_an_anonymous_request() -> None:
+ """An anonymous request has no user id to scope a handle to."""
+ # An anonymous viewer has no user id to scope the handle to, and an
unscoped
+ # handle would be cancellable by any other anonymous visitor.
+ cache = _FakeCache()
+ database = MagicMock()
+ database.id = 5
+
+ assert _run_cancellable("client-1", database, cache, user_id=None) == {}
+
+
+def
test_cancellable_chart_query_publishes_nothing_when_engine_has_no_cancel_id()
-> (
+ None
+):
+ """Engines with no pre-execution cancel id leave the query
non-cancellable."""
+ cache = _FakeCache()
+ database = MagicMock()
+ database.id = 5
+
+ assert _run_cancellable("client-1", database, cache, 1, cancel_id=None) ==
{}
+
+
+def test_cancel_chart_query_for_user_cancels_its_own_query() -> None:
+ """Cancelling one's own registered query hits the engine and cleans up."""
+ from superset.tasks.query_cancel import cancel_chart_query_for_user
+
+ cache = _FakeCache()
+ cache.store["chart-query-cancel:1:client-1"] = {
+ "database_id": 5,
+ "cancel_query_id": "engine-1",
+ }
+ database = MagicMock()
+
+ with _registry_env(1, cache):
+ with (
+ patch(
+ "superset.daos.database.DatabaseDAO.find_by_id",
return_value=database
+ ),
+ patch(
+ "superset.tasks.query_cancel.cancel_chart_query",
return_value=True
+ ) as cancel,
+ ):
+ assert cancel_chart_query_for_user("client-1") is True
+
+ cancel.assert_called_once_with(database, "engine-1")
+ # A cancelled query's handle is dropped rather than left to expire.
+ assert cache.store == {}
+
+
+def test_cancel_chart_query_for_user_cannot_reach_another_users_query() ->
None:
+ """Regression test: ``client_id`` is untrusted, client-supplied input.
+
+ User 1 has a running, cancellable query. User 2 asks to cancel that exact
+ ``client_id``. Because the registry key embeds the requesting user, the
+ lookup misses: no engine cancellation is attempted and user 1's handle is
+ left untouched.
+ """
+ from superset.tasks.query_cancel import cancel_chart_query_for_user
+
+ cache = _FakeCache()
+ victim_handle = {"database_id": 5, "cancel_query_id": "engine-1"}
+ cache.store["chart-query-cancel:1:client-1"] = victim_handle
Review Comment:
The literal victim key makes this test pass even if `_registry_key` drops
user scoping: the mutated lookup simply misses it. Seed through the production
key builder so that regression reaches the cancellation assertions.
```suggestion
from superset.tasks.query_cancel import (
_registry_key,
cancel_chart_query_for_user,
)
cache = _FakeCache()
victim_handle = {"database_id": 5, "cancel_query_id": "engine-1"}
cache.store[_registry_key(1, "client-1")] = victim_handle
```
##########
superset/tasks/query_cancel.py:
##########
@@ -163,3 +172,150 @@ 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.
+ """
+ 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)
+
+ 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.cache.set(
Review Comment:
At `fb560033f051`, default `NullCache` still discards every handle, so Stop
never reaches cancellation. A safe fix spans `cache_manager.py`, `config.py`
and registry helpers, so no single-span suggestion is possible.
##########
superset/tasks/query_cancel.py:
##########
@@ -163,3 +172,150 @@ 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.
+ """
+ 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)
+
+ try:
+ with capture_cancel_id(_sink):
+ yield
+ finally:
+ if captured:
+ _discard_cancel_handle(owner_id, query_id)
Review Comment:
At `fb560033f051`, if a cursor republishes during a successful cancel, the
unconditional delete still erases its handle and Stop cannot reach the live
statement. Safe publication/deletion coordination spans multiple sites, so no
single-span suggestion is possible.
--
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]
