This is an automated email from the ASF dual-hosted git repository.
msyavuz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new b983161eeff feat(async): cancel running chart queries under
GLOBAL_ASYNC_QUERIES (#42305)
b983161eeff is described below
commit b983161eeff4a28f5c6088f4f69065e2bf1aa686
Author: Mehmet Salih Yavuz <[email protected]>
AuthorDate: Tue Jul 28 18:02:16 2026 +0300
feat(async): cancel running chart queries under GLOBAL_ASYNC_QUERIES
(#42305)
---
.../src/components/Chart/chartAction.ts | 29 +++-
.../src/middleware/asyncEvent.test.ts | 38 +++++
superset-frontend/src/middleware/asyncEvent.ts | 18 ++-
superset/async_events/api.py | 80 +++++++++-
superset/async_events/async_query_manager.py | 115 ++++++++++++++
superset/async_events/cache_backend.py | 18 +++
superset/tasks/async_queries.py | 28 +++-
tests/integration_tests/async_events/api_tests.py | 72 +++++++++
.../async_events/async_query_manager_tests.py | 170 +++++++++++++++++++++
tests/unit_tests/tasks/test_async_queries.py | 125 +++++++++++++++
10 files changed, 678 insertions(+), 15 deletions(-)
diff --git a/superset-frontend/src/components/Chart/chartAction.ts
b/superset-frontend/src/components/Chart/chartAction.ts
index 2e6854e0791..4a0b86db022 100644
--- a/superset-frontend/src/components/Chart/chartAction.ts
+++ b/superset-frontend/src/components/Chart/chartAction.ts
@@ -713,8 +713,8 @@ export function handleChartDataResponse(
// Query is running asynchronously and we must await the results.
// When status is 202, result contains async event data (job_id,
channel_id, etc.)
// which differs from QueryData. We cast through unknown to handle
this safely.
- // The optional signal lets a caller (e.g. StatefulChart) cancel the
wait
- // when its chart is superseded or unmounted, avoiding leaked
listeners.
+ // The optional signal lets a caller abort the wait (Stop pressed,
chart
+ // superseded or unmounted), cancelling the job and avoiding leaked
listeners.
if (useLegacyApi) {
return waitForAsyncData(
result[0] as unknown as Parameters<typeof waitForAsyncData>[0],
@@ -786,15 +786,25 @@ export function exploreJSON(
const [useLegacyApi] = getQuerySettings(formData);
const chartDataRequestCaught = chartDataRequest
.then(({ response, json }) =>
- handleChartDataResponse(response, json, useLegacyApi),
+ handleChartDataResponse(
+ response,
+ json,
+ useLegacyApi,
+ controller.signal,
+ ),
)
.then(queriesResponse => {
- // Drop stale responses: if a newer query has started for this chart,
- // its controller will have replaced ours in state, so ignore this
- // response to avoid clobbering newer data with older results.
+ // Drop stale responses: if this request was aborted (Stop, or a newer
+ // query that aborted ours), or a newer query has since replaced our
+ // controller in state, ignore the result so we don't clobber newer
+ // data or a 'stopped' status. Checking the signal is authoritative
+ // because the reducer nulls out queryController when a query stops.
if (key != null) {
const currentController = getState().charts?.[key]?.queryController;
- if (currentController && currentController !== controller) {
+ if (
+ controller.signal.aborted ||
+ (currentController != null && currentController !== controller)
+ ) {
return undefined;
}
}
@@ -853,7 +863,10 @@ export function exploreJSON(
// so a slow earlier request can't mark a newer one as failed.
if (key != null) {
const currentController =
getState().charts?.[key]?.queryController;
- if (currentController && currentController !== controller) {
+ if (
+ controller.signal.aborted ||
+ (currentController != null && currentController !== controller)
+ ) {
return undefined;
}
}
diff --git a/superset-frontend/src/middleware/asyncEvent.test.ts
b/superset-frontend/src/middleware/asyncEvent.test.ts
index 909f1887076..ab7dde9b968 100644
--- a/superset-frontend/src/middleware/asyncEvent.test.ts
+++ b/superset-frontend/src/middleware/asyncEvent.test.ts
@@ -134,6 +134,44 @@ describe('asyncEvent middleware', () => {
expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(1);
});
+ test('rejects with an AbortError and cancels the job when the signal
aborts', async () => {
+ const CANCEL_ENDPOINT = 'glob:*/api/v1/async_event/*/cancel';
+ fetchMock.post(CANCEL_ENDPOINT, { status: 200, body: {} });
+
+ const controller = new AbortController();
+ const promise = asyncEvent.waitForAsyncData(
+ asyncPendingEvent,
+ controller.signal,
+ );
+ controller.abort();
+
+ let error: any = null;
+ try {
+ await promise;
+ } catch (err) {
+ error = err;
+ }
+ expect(error?.name).toBe('AbortError');
+ // The cancel POST is fire-and-forget; let its microtask flush.
+ await new Promise(resolve => setTimeout(resolve, 0));
+ expect(fetchMock.callHistory.calls(CANCEL_ENDPOINT)).toHaveLength(1);
+ });
+
+ test('rejects immediately when given an already-aborted signal', async ()
=> {
+ const CANCEL_ENDPOINT = 'glob:*/api/v1/async_event/*/cancel';
+ fetchMock.post(CANCEL_ENDPOINT, { status: 200, body: {} });
+
+ const controller = new AbortController();
+ controller.abort();
+
+ await expect(
+ asyncEvent.waitForAsyncData(asyncPendingEvent, controller.signal),
+ ).rejects.toMatchObject({ name: 'AbortError' });
+ // The cancel POST is fire-and-forget; let its microtask flush.
+ await new Promise(resolve => setTimeout(resolve, 0));
+ expect(fetchMock.callHistory.calls(CANCEL_ENDPOINT)).toHaveLength(1);
+ });
+
test('rejects on event error status', async () => {
fetchMock.clearHistory().removeRoutes();
fetchMock.get(EVENTS_ENDPOINT, {
diff --git a/superset-frontend/src/middleware/asyncEvent.ts
b/superset-frontend/src/middleware/asyncEvent.ts
index 622926a49fa..d82d613ac3a 100644
--- a/superset-frontend/src/middleware/asyncEvent.ts
+++ b/superset-frontend/src/middleware/asyncEvent.ts
@@ -104,6 +104,17 @@ const fetchCachedData = async (
return { status, data };
};
+const cancelAsyncJob = (jobId: string) => {
+ // Best-effort server-side cancel; the request stops the running Celery task
+ // so it no longer consumes warehouse resources. Failures are non-fatal: the
+ // client has already stopped waiting on the job.
+ SupersetClient.post({
+ endpoint: `/api/v1/async_event/${jobId}/cancel`,
+ }).catch(error => {
+ logging.warn('Failed to cancel async job', jobId, error);
+ });
+};
+
export const waitForAsyncData = async (
asyncResponse: AsyncEvent,
signal?: AbortSignal,
@@ -122,6 +133,7 @@ export const waitForAsyncData = async (
// Bail immediately if the caller has already aborted (e.g. the chart was
// unmounted before the job started), avoiding a leaked listener.
if (signal?.aborted) {
+ cancelAsyncJob(jobId);
reject(new DOMException('Aborted', 'AbortError'));
return;
}
@@ -161,11 +173,13 @@ export const waitForAsyncData = async (
}
};
- // When the caller aborts (chart superseded/unmounted), stop listening so
the
- // listener and its retained closure don't leak and keep the poller busy.
+ // When the caller aborts (Stop pressed, chart superseded/unmounted), stop
+ // listening so the listener and its retained closure don't leak, and ask
the
+ // server to cancel the job so it stops consuming warehouse resources.
if (signal) {
onAbort = () => {
cleanup();
+ cancelAsyncJob(jobId);
reject(new DOMException('Aborted', 'AbortError'));
};
signal.addEventListener('abort', onAbort, { once: true });
diff --git a/superset/async_events/api.py b/superset/async_events/api.py
index 55539d86aa4..7968f956b54 100644
--- a/superset/async_events/api.py
+++ b/superset/async_events/api.py
@@ -15,14 +15,19 @@
# specific language governing permissions and limitations
# under the License.
import logging
+import uuid
from flask import request, Response
from flask_appbuilder import expose
from flask_appbuilder.api import safe
from flask_appbuilder.security.decorators import permission_name, protect
-from superset.async_events.async_query_manager import AsyncQueryTokenException
+from superset.async_events.async_query_manager import (
+ AsyncQueryJobException,
+ AsyncQueryTokenException,
+)
from superset.extensions import async_query_manager, event_logger
+from superset.utils.core import get_user_id
from superset.views.base_api import BaseSupersetApi, statsd_metrics
logger = logging.getLogger(__name__)
@@ -99,3 +104,76 @@ class AsyncEventsRestApi(BaseSupersetApi):
return self.response_401()
return self.response(200, result=events)
+
+ @expose("/<job_id>/cancel", methods=("POST",))
+ @event_logger.log_this
+ @protect()
+ @safe
+ @statsd_metrics
+ @permission_name("cancel")
+ def cancel(self, job_id: str) -> Response:
+ """Cancel a running async query job.
+ ---
+ post:
+ summary: Cancel a running async query job
+ description: >-
+ Revokes the Celery task backing an in-flight async query. The
+ caller is authorized against the job's original owner (channel and
+ user), both resolved server-side from the request, so a client
+ cannot cancel a job it did not submit.
+ parameters:
+ - in: path
+ name: job_id
+ required: true
+ description: The job ID returned when the async query was submitted
+ schema:
+ type: string
+ responses:
+ 200:
+ description: Job cancelled
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ result:
+ type: object
+ properties:
+ job_id:
+ type: string
+ status:
+ type: string
+ 400:
+ $ref: '#/components/responses/400'
+ 401:
+ $ref: '#/components/responses/401'
+ 403:
+ $ref: '#/components/responses/403'
+ 404:
+ $ref: '#/components/responses/404'
+ 500:
+ $ref: '#/components/responses/500'
+ """
+ try:
+ uuid.UUID(job_id)
+ except ValueError:
+ return self.response_400(message="Invalid job ID")
+
+ try:
+ async_channel_id =
async_query_manager.parse_channel_id_from_request(
+ request
+ )
+ except AsyncQueryTokenException:
+ return self.response_401()
+
+ try:
+ async_query_manager.cancel_job(job_id, async_channel_id,
get_user_id())
+ except AsyncQueryTokenException:
+ return self.response_403()
+ except AsyncQueryJobException:
+ return self.response_404()
+
+ return self.response(
+ 200,
+ result={"job_id": job_id, "status":
async_query_manager.STATUS_CANCELLED},
+ )
diff --git a/superset/async_events/async_query_manager.py
b/superset/async_events/async_query_manager.py
index cda66bfde4a..4ae08fcfb43 100644
--- a/superset/async_events/async_query_manager.py
+++ b/superset/async_events/async_query_manager.py
@@ -106,6 +106,10 @@ class AsyncQueryManager:
STATUS_RUNNING = "running"
STATUS_ERROR = "error"
STATUS_DONE = "done"
+ STATUS_CANCELLED = "cancelled"
+ # Redis key prefix (within the GAQ stream namespace) for the per-job record
+ # that authorizes cancellation and flags a job as cancelled for the worker.
+ _JOB_REGISTRY_PREFIX = "job-cancel:"
def __init__(self) -> None:
super().__init__()
@@ -281,10 +285,32 @@ class AsyncQueryManager:
def init_job(self, channel_id: str, user_id: Optional[int]) -> dict[str,
Any]:
job_id = str(uuid.uuid4())
+ self._register_cancellable_job(job_id, channel_id, user_id)
return build_job_metadata(
channel_id, job_id, user_id, status=self.STATUS_PENDING
)
+ def _job_registry_key(self, job_id: str) -> str:
+ return f"{self._stream_prefix}{self._JOB_REGISTRY_PREFIX}{job_id}"
+
+ def _register_cancellable_job(
+ self, job_id: str, channel_id: str, user_id: Optional[int]
+ ) -> None:
+ """
+ Persist the identity a later cancel request must match. Keyed by
+ ``job_id`` (also the Celery task id — see ``submit_chart_data_job``) so
+ the cancel endpoint can authorize the caller against the job's original
+ owner without trusting the client-supplied id. Expires with the JWT so
+ it never outlives the job it guards.
+ """
+ if not self._cache:
+ return
+ self._cache.set(
+ self._job_registry_key(job_id),
+ json.dumps({"channel_id": channel_id, "user_id": user_id}),
+ ex=self._jwt_expiration_seconds or None,
+ )
+
# pylint: disable=too-many-arguments
def submit_explore_json_job(
self,
@@ -307,6 +333,9 @@ class AsyncQueryManager:
response_type,
force,
],
+ # Use job_id as the Celery task id so the cancel endpoint can
revoke
+ # the running task by the id the client already holds.
+ task_id=job_metadata["job_id"],
expires=self._jwt_expiration_seconds,
)
return job_metadata
@@ -332,6 +361,9 @@ class AsyncQueryManager:
else job_metadata,
form_data,
],
+ # Use job_id as the Celery task id so the cancel endpoint can
revoke
+ # the running task by the id the client already holds.
+ task_id=job_metadata["job_id"],
expires=self._jwt_expiration_seconds,
)
return job_metadata
@@ -383,5 +415,88 @@ class AsyncQueryManager:
logger.debug("********** logging event data to stream %s",
scoped_stream_name)
logger.debug(event_data)
+ # Drop the cancel record before announcing the result, so a cancel that
+ # arrives once the job is done finds nothing to flag and is reported as
+ # such instead of contradicting the event below. A cancelled job keeps
+ # its record until the TTL: the worker still has to recognize the
+ # SIGUSR1 it is about to receive as a cancellation.
+ if status in (self.STATUS_DONE, self.STATUS_ERROR):
+ if job_id := job_metadata.get("job_id"):
+ self._cache.delete(self._job_registry_key(job_id))
+
self._cache.xadd(scoped_stream_name, event_data, "*",
self._stream_limit)
self._cache.xadd(full_stream_name, event_data, "*",
self._stream_limit_firehose)
+
+ def is_job_cancelled(self, job_id: str) -> bool:
+ """
+ Whether ``cancel_job`` has flagged this job for cancellation.
+
+ Called from the worker's exception handler, so any cache failure is
+ swallowed and treated as "not cancelled" — a Redis blip must never mask
+ the original error (e.g. a genuine timeout) with a connection error.
+ """
+ if not self._cache:
+ return False
+ try:
+ raw = self._cache.get(self._job_registry_key(job_id))
+ if raw is None:
+ return False
+ return bool(json.loads(raw).get("cancelled"))
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Failed to read cancellation flag for job %s", job_id,
exc_info=True
+ )
+ return False
+
+ def cancel_job(self, job_id: str, channel_id: str, user_id: Optional[int])
-> None:
+ """
+ Authorize and cancel a running async job.
+
+ The caller's ``channel_id`` and ``user_id`` (resolved server-side from
+ the request, never taken from the client) must match the job's original
+ owner. The terminal ``STATUS_CANCELLED`` event is emitted here rather
+ than by the worker, which never runs for a task revoked while it was
+ still queued; the worker only logs the cancellation it is told about
+ through the flag, so a job still gets exactly one terminal event.
+
+ :raises AsyncQueryJobException: the job is unknown or already terminal
+ :raises AsyncQueryTokenException: the caller does not own the job
+ """
+ if not self._cache:
+ raise CacheBackendNotInitialized("Cache backend not initialized")
+
+ key = self._job_registry_key(job_id)
+ raw = self._cache.get(key)
+ if raw is None:
+ raise AsyncQueryJobException("Job not found or already completed")
+
+ record = json.loads(raw)
+ if record.get("channel_id") != channel_id or record.get("user_id") !=
user_id:
+ raise AsyncQueryTokenException("Not authorized to cancel this job")
+
+ # Flag before revoking so the worker's timeout handler, which may fire
+ # almost immediately, reliably sees the cancellation. Write only if the
+ # key still exists (``xx``): if the job finished and cleared its record
+ # between the read above and here, don't recreate a stale record or
+ # revoke a task that is already gone — report it as not found instead.
+ flagged = self._cache.set(
+ key,
+ json.dumps({**record, "cancelled": True}),
+ ex=self._jwt_expiration_seconds or None,
+ xx=True,
+ )
+ if not flagged:
+ raise AsyncQueryJobException("Job not found or already completed")
+
+ # pylint: disable=import-outside-toplevel
+ from superset.extensions import celery_app
+
+ # SIGUSR1 raises SoftTimeLimitExceeded inside the running task rather
+ # than hard-killing the process, so it unwinds through the task's
+ # exception handling instead of dying mid-query.
+ celery_app.control.revoke(job_id, terminate=True, signal="SIGUSR1")
+
+ self.update_job(
+ build_job_metadata(channel_id, job_id, user_id),
+ self.STATUS_CANCELLED,
+ )
diff --git a/superset/async_events/cache_backend.py
b/superset/async_events/cache_backend.py
index 55c22715165..4e528b82694 100644
--- a/superset/async_events/cache_backend.py
+++ b/superset/async_events/cache_backend.py
@@ -99,6 +99,15 @@ class RedisCacheBackend(RedisCache):
"""
return self._cache.set(name, value, ex=ex, px=px, nx=nx, xx=xx)
+ def get(self, name: str) -> Any:
+ """
+ Get the raw value at key ``name``.
+
+ :param name: Key name
+ :returns: The stored value (bytes), or None if the key is absent
+ """
+ return self._cache.get(name)
+
def delete(self, *names: str) -> int:
"""
Delete one or more keys.
@@ -283,6 +292,15 @@ class RedisSentinelCacheBackend(RedisSentinelCache):
"""
return self._cache.set(name, value, ex=ex, px=px, nx=nx, xx=xx)
+ def get(self, name: str) -> Any:
+ """
+ Get the raw value at key ``name``.
+
+ :param name: Key name
+ :returns: The stored value (bytes), or None if the key is absent
+ """
+ return self._cache.get(name)
+
def delete(self, *names: str) -> int:
"""
Delete one or more keys.
diff --git a/superset/tasks/async_queries.py b/superset/tasks/async_queries.py
index f9aa9efc2fe..85830120bce 100644
--- a/superset/tasks/async_queries.py
+++ b/superset/tasks/async_queries.py
@@ -84,6 +84,28 @@ def _load_user_from_job_metadata(job_metadata: dict[str,
Any]) -> User:
return user
+def _handle_soft_time_limit(
+ job_metadata: dict[str, Any], ex: Exception, activity: str
+) -> None:
+ """
+ SoftTimeLimitExceeded is raised both by a genuine timeout and by a
+ user-initiated cancel (revoke sends SIGUSR1). The cancel endpoint has
+ already emitted the terminal event for the latter - it has to, since a task
+ revoked while still queued never reaches this handler - so only a timeout
+ is reported here, and without one the client would wait forever.
+ """
+ if async_query_manager.is_job_cancelled(job_metadata["job_id"]):
+ logger.info("Cancelled by the user while %s", activity)
+ return
+
+ logger.warning("A timeout occurred while %s, error: %s", activity, ex)
+ async_query_manager.update_job(
+ job_metadata,
+ async_query_manager.STATUS_ERROR,
+ errors=[{"message": f"A timeout occurred while {activity}"}],
+ )
+
+
@celery_app.task(name="load_chart_data_into_cache",
soft_time_limit=query_timeout)
def load_chart_data_into_cache(
job_metadata: dict[str, Any],
@@ -106,7 +128,7 @@ def load_chart_data_into_cache(
result_url=result_url,
)
except SoftTimeLimitExceeded as ex:
- logger.warning("A timeout occurred while loading chart data,
error: %s", ex)
+ _handle_soft_time_limit(job_metadata, ex, "loading chart data")
raise
except Exception as ex:
# Extract SIP-40 style errors when available
@@ -176,9 +198,7 @@ def load_explore_json_into_cache( # pylint:
disable=too-many-locals
result_url=result_url,
)
except SoftTimeLimitExceeded as ex:
- logger.warning(
- "A timeout occurred while loading explore json, error: %s", ex
- )
+ _handle_soft_time_limit(job_metadata, ex, "loading explore json")
raise
except Exception as ex:
if isinstance(ex, SupersetVizException):
diff --git a/tests/integration_tests/async_events/api_tests.py
b/tests/integration_tests/async_events/api_tests.py
index 44580c35bca..5be8022d5b7 100644
--- a/tests/integration_tests/async_events/api_tests.py
+++ b/tests/integration_tests/async_events/api_tests.py
@@ -17,6 +17,10 @@
from typing import Any, Optional, Type
from unittest import mock
+from superset.async_events.async_query_manager import (
+ AsyncQueryJobException,
+ AsyncQueryTokenException,
+)
from superset.async_events.cache_backend import (
RedisCacheBackend,
RedisSentinelCacheBackend,
@@ -30,12 +34,16 @@ from tests.integration_tests.test_app import app
class TestAsyncEventApi(SupersetTestCase):
UUID = "943c920-32a5-412a-977d-b8e47d36f5a4"
+ JOB_ID = "10a0bd9a-03c8-4737-9345-f4234ba86512"
def fetch_events(self, last_id: Optional[str] = None):
base_uri = "api/v1/async_event/"
uri = f"{base_uri}?last_id={last_id}" if last_id else base_uri
return self.client.get(uri)
+ def cancel_event(self, job_id: str):
+ return self.client.post(f"api/v1/async_event/{job_id}/cancel")
+
def run_test_with_cache_backend(self, cache_backend_cls: Type[Any],
test_func):
app._got_first_request = False
async_query_manager_factory.init_app(app)
@@ -136,3 +144,67 @@ class TestAsyncEventApi(SupersetTestCase):
self.client.set_cookie(app.config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_NAME"], "")
rv = self.fetch_events()
assert rv.status_code == 401
+
+ def _test_cancel_logic(self, mock_cache):
+ with mock.patch.object(async_query_manager, "cancel_job") as
mock_cancel:
+ rv = self.cancel_event(self.JOB_ID)
+
+ assert rv.status_code == 200
+ mock_cancel.assert_called_once()
+ assert mock_cancel.call_args.args[0] == self.JOB_ID
+ response = json.loads(rv.data.decode("utf-8"))
+ assert response["result"] == {"job_id": self.JOB_ID, "status":
"cancelled"}
+
+ def _test_cancel_invalid_job_id_logic(self, mock_cache):
+ # job_id lands in a Redis key, so reject anything that isn't a UUID
+ # before it reaches the cache backend.
+ with mock.patch.object(async_query_manager, "cancel_job") as
mock_cancel:
+ rv = self.cancel_event("not-a-uuid")
+ assert rv.status_code == 400
+ mock_cancel.assert_not_called()
+
+ def _test_cancel_forbidden_logic(self, mock_cache):
+ with mock.patch.object(
+ async_query_manager,
+ "cancel_job",
+ side_effect=AsyncQueryTokenException("nope"),
+ ):
+ rv = self.cancel_event(self.JOB_ID)
+ assert rv.status_code == 403
+
+ def _test_cancel_not_found_logic(self, mock_cache):
+ with mock.patch.object(
+ async_query_manager,
+ "cancel_job",
+ side_effect=AsyncQueryJobException("gone"),
+ ):
+ rv = self.cancel_event(self.JOB_ID)
+ assert rv.status_code == 404
+
+ @mock.patch("uuid.uuid4", return_value=UUID)
+ def test_cancel_redis_cache_backend(self, mock_uuid4):
+ self.run_test_with_cache_backend(RedisCacheBackend,
self._test_cancel_logic)
+
+ @mock.patch("uuid.uuid4", return_value=UUID)
+ def test_cancel_invalid_job_id(self, mock_uuid4):
+ self.run_test_with_cache_backend(
+ RedisCacheBackend, self._test_cancel_invalid_job_id_logic
+ )
+
+ @mock.patch("uuid.uuid4", return_value=UUID)
+ def test_cancel_forbidden(self, mock_uuid4):
+ self.run_test_with_cache_backend(
+ RedisCacheBackend, self._test_cancel_forbidden_logic
+ )
+
+ @mock.patch("uuid.uuid4", return_value=UUID)
+ def test_cancel_not_found(self, mock_uuid4):
+ self.run_test_with_cache_backend(
+ RedisCacheBackend, self._test_cancel_not_found_logic
+ )
+
+ def test_cancel_no_login(self):
+ app._got_first_request = False
+ async_query_manager_factory.init_app(app)
+ rv = self.cancel_event(self.JOB_ID)
+ assert rv.status_code == 401
diff --git a/tests/unit_tests/async_events/async_query_manager_tests.py
b/tests/unit_tests/async_events/async_query_manager_tests.py
index f5b5c97c8a2..d0879306da8 100644
--- a/tests/unit_tests/async_events/async_query_manager_tests.py
+++ b/tests/unit_tests/async_events/async_query_manager_tests.py
@@ -24,6 +24,7 @@ from pytest import fixture, mark, raises # noqa: PT013
from superset import security_manager
from superset.async_events.async_query_manager import (
+ AsyncQueryJobException,
AsyncQueryManager,
AsyncQueryTokenException,
)
@@ -31,6 +32,7 @@ from superset.async_events.cache_backend import (
RedisCacheBackend,
RedisSentinelCacheBackend,
)
+from superset.utils import json
JWT_TOKEN_SECRET = "some_secret" # noqa: S105
JWT_TOKEN_COOKIE_NAME = "superset_async_jwt" # noqa: S105
@@ -306,6 +308,7 @@ def test_submit_chart_data_job_as_guest_user(
},
{},
],
+ task_id=ANY,
expires=3600,
)
@@ -414,7 +417,174 @@ def test_submit_explore_json_job_as_guest_user(
"json",
False,
],
+ task_id=ANY,
expires=3600,
)
assert "guest_token" not in job_meta
+
+
+@fixture
+def cancellable_manager():
+ """A manager wired to a mock Redis backend for cancellation tests."""
+ manager = AsyncQueryManager()
+ manager._jwt_expiration_seconds = 3600
+ manager._stream_prefix = "async-events-"
+ manager._cache = mock.Mock(spec=RedisCacheBackend)
+ return manager
+
+
+def test_init_job_registers_cancellable_record(cancellable_manager):
+ """init_job persists the owner identity a later cancel must match."""
+ cancellable_manager.init_job("chan-1", 7)
+
+ cancellable_manager._cache.set.assert_called_once()
+ key, value = cancellable_manager._cache.set.call_args.args
+ assert key.startswith("async-events-job-cancel:")
+ assert json.loads(value) == {"channel_id": "chan-1", "user_id": 7}
+
+
+def test_cancel_job_authorized_revokes_task(cancellable_manager):
+ cancellable_manager._stream_limit = 100
+ cancellable_manager._stream_limit_firehose = 1000
+ cancellable_manager._cache.get.return_value = json.dumps(
+ {"channel_id": "chan-1", "user_id": 7}
+ )
+ cancellable_manager._cache.set.return_value = True
+
+ with mock.patch("superset.extensions.celery_app") as celery_app:
+ cancellable_manager.cancel_job("job-1", "chan-1", 7)
+
+ celery_app.control.revoke.assert_called_once_with(
+ "job-1", terminate=True, signal="SIGUSR1"
+ )
+ # The job is flagged cancelled (conditionally, xx=True) so the worker knows
+ # what the signal it is about to receive means.
+ assert cancellable_manager._cache.set.call_args.kwargs["xx"] is True
+ flagged = json.loads(cancellable_manager._cache.set.call_args.args[1])
+ assert flagged["cancelled"] is True
+
+
+def test_cancel_job_emits_the_terminal_event(cancellable_manager):
+ """A task revoked before a worker picks it up never reports on itself."""
+ cancellable_manager._stream_limit = 100
+ cancellable_manager._stream_limit_firehose = 1000
+ cancellable_manager._cache.get.return_value = json.dumps(
+ {"channel_id": "chan-1", "user_id": 7}
+ )
+ cancellable_manager._cache.set.return_value = True
+
+ with mock.patch("superset.extensions.celery_app"):
+ cancellable_manager.cancel_job("job-1", "chan-1", 7)
+
+ scoped_stream, event_data =
cancellable_manager._cache.xadd.call_args_list[0].args[
+ :2
+ ]
+ assert scoped_stream == "async-events-chan-1"
+ assert json.loads(event_data["data"]) == {
+ "channel_id": "chan-1",
+ "job_id": "job-1",
+ "user_id": 7,
+ "status": AsyncQueryManager.STATUS_CANCELLED,
+ "errors": [],
+ "result_url": None,
+ }
+ # the record has to outlive the event: the worker still needs to recognize
+ # the signal on its way as a cancellation
+ cancellable_manager._cache.delete.assert_not_called()
+
+
+def test_cancel_job_completed_between_read_and_flag(cancellable_manager):
+ """If the job's record is cleared after the auth read, don't revoke."""
+ cancellable_manager._cache.get.return_value = json.dumps(
+ {"channel_id": "chan-1", "user_id": 7}
+ )
+ # Conditional (xx) write finds no key: the job finished and cleaned up.
+ cancellable_manager._cache.set.return_value = None
+
+ with (
+ mock.patch("superset.extensions.celery_app") as celery_app,
+ raises(AsyncQueryJobException),
+ ):
+ cancellable_manager.cancel_job("job-1", "chan-1", 7)
+
+ celery_app.control.revoke.assert_not_called()
+
+
+def test_cancel_job_wrong_user_is_rejected(cancellable_manager):
+ cancellable_manager._cache.get.return_value = json.dumps(
+ {"channel_id": "chan-1", "user_id": 7}
+ )
+
+ with (
+ mock.patch("superset.extensions.celery_app") as celery_app,
+ raises(AsyncQueryTokenException),
+ ):
+ cancellable_manager.cancel_job("job-1", "chan-1", 999)
+
+ celery_app.control.revoke.assert_not_called()
+
+
+def test_cancel_job_wrong_channel_is_rejected(cancellable_manager):
+ """A matching user on a different channel still cannot cancel the job."""
+ cancellable_manager._cache.get.return_value = json.dumps(
+ {"channel_id": "chan-1", "user_id": 7}
+ )
+
+ with (
+ mock.patch("superset.extensions.celery_app") as celery_app,
+ raises(AsyncQueryTokenException),
+ ):
+ cancellable_manager.cancel_job("job-1", "other-chan", 7)
+
+ celery_app.control.revoke.assert_not_called()
+
+
+def test_cancel_job_unknown_raises(cancellable_manager):
+ cancellable_manager._cache.get.return_value = None
+
+ with (
+ mock.patch("superset.extensions.celery_app") as celery_app,
+ raises(AsyncQueryJobException),
+ ):
+ cancellable_manager.cancel_job("job-1", "chan-1", 7)
+
+ celery_app.control.revoke.assert_not_called()
+
+
+def test_is_job_cancelled(cancellable_manager):
+ cancellable_manager._cache.get.return_value = json.dumps(
+ {"channel_id": "chan-1", "user_id": 7, "cancelled": True}
+ )
+ assert cancellable_manager.is_job_cancelled("job-1") is True
+
+ cancellable_manager._cache.get.return_value = json.dumps(
+ {"channel_id": "chan-1", "user_id": 7}
+ )
+ assert cancellable_manager.is_job_cancelled("job-1") is False
+
+ cancellable_manager._cache.get.return_value = None
+ assert cancellable_manager.is_job_cancelled("job-1") is False
+
+
+def test_is_job_cancelled_swallows_cache_errors(cancellable_manager):
+ """A cache failure must not escape and mask the worker's original error."""
+ cancellable_manager._cache.get.side_effect = RuntimeError("redis down")
+ assert cancellable_manager.is_job_cancelled("job-1") is False
+
+
+def test_update_job_clears_registry_before_terminal_event(cancellable_manager):
+ """Clearing first is what makes a cancel that lost the race a 404."""
+ calls = []
+ cancellable_manager._stream_limit = 100
+ cancellable_manager._stream_limit_firehose = 1000
+ cancellable_manager._cache.delete.side_effect = lambda *_:
calls.append("delete")
+ cancellable_manager._cache.xadd.side_effect = lambda *_:
calls.append("xadd")
+ job_metadata = {"channel_id": "chan-1", "job_id": "job-1", "user_id": 7}
+
+ cancellable_manager.update_job(job_metadata, AsyncQueryManager.STATUS_DONE)
+
+ cancellable_manager._cache.delete.assert_called_once_with(
+ "async-events-job-cancel:job-1"
+ )
+ assert calls == ["delete", "xadd", "xadd"]
diff --git a/tests/unit_tests/tasks/test_async_queries.py
b/tests/unit_tests/tasks/test_async_queries.py
index c442f309549..8ddbeb73013 100644
--- a/tests/unit_tests/tasks/test_async_queries.py
+++ b/tests/unit_tests/tasks/test_async_queries.py
@@ -18,6 +18,7 @@ from typing import Any
from unittest import mock
import pytest
+from celery.exceptions import SoftTimeLimitExceeded
from flask_babel import lazy_gettext as _
from superset.commands.chart.exceptions import ChartDataQueryFailedError
@@ -63,6 +64,62 @@ def test_load_chart_data_into_cache_with_error(
)
[email protected]("superset.tasks.async_queries.security_manager")
[email protected]("superset.tasks.async_queries.async_query_manager")
[email protected]("superset.tasks.async_queries.ChartDataQueryContextSchema")
+def test_load_chart_data_into_cache_cancelled_emits_no_event(
+ mock_query_context_schema_cls, mock_async_query_manager,
mock_security_manager
+):
+ """A revoke leaves the terminal event to the cancel request that sent
it."""
+ from superset.tasks.async_queries import load_chart_data_into_cache
+
+ job_metadata = {"user_id": 1, "job_id": "job-1"}
+ form_data: dict[str, Any] = {}
+
+ mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+ # Sync Mock: is_job_cancelled is a plain method, but patching the manager
+ # yields an AsyncMock whose calls would otherwise return truthy coroutines.
+ mock_async_query_manager.is_job_cancelled = mock.Mock(return_value=True)
+ mock_query_context_schema_cls.return_value.load.side_effect = (
+ SoftTimeLimitExceeded()
+ )
+
+ with pytest.raises(SoftTimeLimitExceeded):
+ load_chart_data_into_cache(job_metadata, form_data)
+
+ mock_async_query_manager.is_job_cancelled.assert_called_once_with("job-1")
+ mock_async_query_manager.update_job.assert_not_called()
+
+
[email protected]("superset.tasks.async_queries.security_manager")
[email protected]("superset.tasks.async_queries.async_query_manager")
[email protected]("superset.tasks.async_queries.ChartDataQueryContextSchema")
+def test_load_chart_data_into_cache_timeout_emits_error(
+ mock_query_context_schema_cls, mock_async_query_manager,
mock_security_manager
+):
+ """A genuine timeout reports an error, or the client waits forever."""
+ from superset.tasks.async_queries import load_chart_data_into_cache
+
+ job_metadata = {"user_id": 1, "job_id": "job-1"}
+ form_data: dict[str, Any] = {}
+
+ mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+ mock_async_query_manager.STATUS_ERROR = "error"
+ mock_async_query_manager.is_job_cancelled = mock.Mock(return_value=False)
+ mock_query_context_schema_cls.return_value.load.side_effect = (
+ SoftTimeLimitExceeded()
+ )
+
+ with pytest.raises(SoftTimeLimitExceeded):
+ load_chart_data_into_cache(job_metadata, form_data)
+
+ mock_async_query_manager.update_job.assert_called_once_with(
+ job_metadata,
+ "error",
+ errors=[{"message": "A timeout occurred while loading chart data"}],
+ )
+
+
@mock.patch("superset.tasks.async_queries.security_manager")
@mock.patch("superset.tasks.async_queries.async_query_manager")
@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema")
@@ -373,3 +430,71 @@ def
test_load_explore_json_into_cache_falls_back_to_string_for_generic_exception
errors = mock_async_query_manager.update_job.call_args[1]["errors"]
assert errors == ["boom"]
+
+
[email protected]("superset.tasks.async_queries.security_manager")
[email protected]("superset.tasks.async_queries.async_query_manager")
[email protected]("superset.tasks.async_queries.get_viz")
[email protected]("superset.tasks.async_queries.get_datasource_info")
+def test_load_explore_json_into_cache_cancelled_emits_no_event(
+ mock_get_datasource_info,
+ mock_get_viz,
+ mock_async_query_manager,
+ mock_security_manager,
+):
+ """A revoke leaves the terminal event to the cancel request that sent
it."""
+ from superset.tasks.async_queries import load_explore_json_into_cache
+
+ job_metadata = {"user_id": 1, "job_id": "job-1"}
+ form_data: dict = {}
+
+ mock_get_datasource_info.return_value = (1, "table")
+ mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+ # Sync Mock: is_job_cancelled is a plain method, but patching the manager
+ # yields an AsyncMock whose calls would otherwise return truthy coroutines.
+ mock_async_query_manager.is_job_cancelled = mock.Mock(return_value=True)
+
+ viz_obj = mock.MagicMock()
+ viz_obj.get_payload.side_effect = SoftTimeLimitExceeded()
+ mock_get_viz.return_value = viz_obj
+
+ with pytest.raises(SoftTimeLimitExceeded):
+ load_explore_json_into_cache(job_metadata, form_data)
+
+ mock_async_query_manager.is_job_cancelled.assert_called_once_with("job-1")
+ mock_async_query_manager.update_job.assert_not_called()
+
+
[email protected]("superset.tasks.async_queries.security_manager")
[email protected]("superset.tasks.async_queries.async_query_manager")
[email protected]("superset.tasks.async_queries.get_viz")
[email protected]("superset.tasks.async_queries.get_datasource_info")
+def test_load_explore_json_into_cache_timeout_emits_error(
+ mock_get_datasource_info,
+ mock_get_viz,
+ mock_async_query_manager,
+ mock_security_manager,
+):
+ """A genuine timeout reports an error, or the client waits forever."""
+ from superset.tasks.async_queries import load_explore_json_into_cache
+
+ job_metadata = {"user_id": 1, "job_id": "job-1"}
+ form_data: dict = {}
+
+ mock_get_datasource_info.return_value = (1, "table")
+ mock_security_manager.get_user_by_id.return_value = mock.MagicMock()
+ mock_async_query_manager.STATUS_ERROR = "error"
+ mock_async_query_manager.is_job_cancelled = mock.Mock(return_value=False)
+
+ viz_obj = mock.MagicMock()
+ viz_obj.get_payload.side_effect = SoftTimeLimitExceeded()
+ mock_get_viz.return_value = viz_obj
+
+ with pytest.raises(SoftTimeLimitExceeded):
+ load_explore_json_into_cache(job_metadata, form_data)
+
+ mock_async_query_manager.update_job.assert_called_once_with(
+ job_metadata,
+ "error",
+ errors=[{"message": "A timeout occurred while loading explore json"}],
+ )
