codeant-ai-for-open-source[bot] commented on code in PR #42305:
URL: https://github.com/apache/superset/pull/42305#discussion_r3656347154
##########
superset/async_events/async_query_manager.py:
##########
@@ -380,3 +412,64 @@ def update_job(
self._cache.xadd(scoped_stream_name, event_data, "*",
self._stream_limit)
self._cache.xadd(full_stream_name, event_data, "*",
self._stream_limit_firehose)
+
+ # Once a job reaches a terminal state its cancel record is dead weight
+ # (and a stale record would let a caller "cancel" a finished job), so
+ # drop it. TTL is the backstop if this is never reached.
+ if status in (self.STATUS_DONE, self.STATUS_ERROR,
self.STATUS_CANCELLED):
+ if job_id := job_metadata.get("job_id"):
+ self._cache.delete(self._job_registry_key(job_id))
+
+ def is_job_cancelled(self, job_id: str) -> bool:
+ """Whether ``cancel_job`` has flagged this job for cancellation."""
+ if not self._cache:
+ return False
+ raw = self._cache.get(self._job_registry_key(job_id))
+ if raw is None:
+ return False
+ try:
+ return bool(json.loads(raw).get("cancelled"))
+ except (ValueError, AttributeError):
+ 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. On success the running Celery task is revoked; the SIGUSR1 it
+ receives surfaces as ``SoftTimeLimitExceeded`` in the worker, which
+ reads the cancelled flag set here and emits the terminal
+ ``STATUS_CANCELLED`` event — so exactly one terminal event is written.
+
+ :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.
+ self._cache.set(
+ key,
+ json.dumps({**record, "cancelled": True}),
+ ex=self._jwt_expiration_seconds or None,
+ )
Review Comment:
Yes, the comment is relevant. There is a real read–write race between
fetching the registry record and setting the cancellation flag: `update_job()`
can delete the key after the `get()` but before the `set()`. A normal
conditional `xx=True` write prevents recreating a key that no longer exists,
which is the right minimal fix:
```python
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")
```
However, this only makes the write conditional; it does not make the
authorization read and flag update fully atomic. Between authorization and the
conditional write, another caller could still read the same record. If strict
single-winner cancellation semantics are required, use a Redis transaction or
Lua script that verifies the record and updates it atomically. For the stated
stale-key issue, the highlighted `xx=True` change is relevant and addresses the
primary vulnerability, with the existing test covering the
key-deleted-before-write case.
--
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]
