rusackas commented on code in PR #44552:
URL: https://github.com/apache/superset/pull/44552#discussion_r4078640776
##########
superset-frontend/src/explore/components/ExploreViewContainer/index.tsx:
##########
@@ -584,9 +585,26 @@ function ExploreViewContainer(props:
ExploreViewContainerProps) {
);
function onStop() {
- if (props.chart && props.chart.queryController) {
+ // Abort the in-flight HTTP request so the UI stops waiting immediately.
+ if (props.chart?.queryController) {
props.chart.queryController.abort();
}
+
+ // Aborting only drops the response; the database keeps executing the
query.
+ // Ask the backend to cancel it too. The backend resolves `client_id`
within
+ // the requesting user's own in-flight queries, so this can only ever
cancel
+ // our own query. A falsy `stopped` just means there was nothing running to
+ // cancel (or the engine has no cancel support), which is not an error.
+ const clientId = props.chart?.latestQueryId;
+ if (clientId) {
+ SupersetClient.post({
+ endpoint: '/api/v1/chart/data/stop',
+ body: JSON.stringify({ client_id: clientId }),
+ headers: { 'Content-Type': 'application/json' },
+ }).catch(() => {
+ props.addDangerToast(t('Failed to stop query.'));
+ });
Review Comment:
The gap is between the request landing and the cursor opening. A miss just
returns `stopped: false`, which is exactly what `master` does today. Closing it
means the execution path checking a cancel flag before it runs anything, and I
would rather not grow this PR that way.
##########
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:
That is on purpose. The async path already has its own cancellation via the
task abort (`_capture_query_cancellation` in `async_queries.py`), so this
registry only backs the synchronous Stop, and an async request just gets
`stopped: false` from the endpoint.
##########
superset-frontend/src/explore/exploreUtils/index.ts:
##########
@@ -329,6 +333,9 @@ export const buildV1ChartDataPayload = async ({
}
});
}
+ if (clientId) {
+ payload.client_id = clientId;
+ }
Review Comment:
The queries in a payload run one after another, and each `get_df_payload`
publishes its own handle and drops it when it returns, so the key always names
whichever query is executing at that moment. Nothing gets overwritten
mid-flight.
##########
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:
Good catch, the sink now republishes for every cursor it sees instead of
stopping at the first, with a test for the two-statement case.
##########
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:
The ids are a fresh `nanoid` per run, so the only way to collide is a client
reusing its own id, and the worst case is that user's own Stop missing. Not
worth a compare-and-delete the cache cannot do atomically anyway.
##########
superset-frontend/src/explore/components/ExploreViewContainer/index.tsx:
##########
@@ -584,9 +585,26 @@ function ExploreViewContainer(props:
ExploreViewContainerProps) {
);
function onStop() {
- if (props.chart && props.chart.queryController) {
+ // Abort the in-flight HTTP request so the UI stops waiting immediately.
+ if (props.chart?.queryController) {
props.chart.queryController.abort();
}
+
+ // Aborting only drops the response; the database keeps executing the
query.
+ // Ask the backend to cancel it too. The backend resolves `client_id`
within
+ // the requesting user's own in-flight queries, so this can only ever
cancel
+ // our own query. A falsy `stopped` just means there was nothing running to
+ // cancel (or the engine has no cancel support), which is not an error.
+ const clientId = props.chart?.latestQueryId;
+ if (clientId) {
+ SupersetClient.post({
+ endpoint: '/api/v1/chart/data/stop',
+ body: JSON.stringify({ client_id: clientId }),
+ headers: { 'Content-Type': 'application/json' },
+ }).catch(() => {
+ props.addDangerToast(t('Failed to stop query.'));
+ });
+ }
Review Comment:
`SupersetClient.post` already sends the CSRF token: `request()` merges
`this.headers`, which carries `X-CSRFToken`, into every call, same as every
other post in the codebase. I did add a `logging.error` before the toast so a
failed stop leaves a trace.
##########
tests/unit_tests/tasks/test_query_cancel.py:
##########
@@ -205,3 +206,208 @@ 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)
Review Comment:
There is no `BITO.md` in this repo, and we do not annotate inferable locals
in tests anywhere else, so leaving these as they are.
##########
tests/unit_tests/tasks/test_query_cancel.py:
##########
@@ -205,3 +206,208 @@ 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:
Review Comment:
Added docstrings to the new tests.
--
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]
