alexandrusoare commented on code in PR #42305:
URL: https://github.com/apache/superset/pull/42305#discussion_r3656319854
##########
superset/tasks/async_queries.py:
##########
@@ -106,7 +106,18 @@ 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)
+ # SoftTimeLimitExceeded is raised both by a genuine timeout and by
a
+ # user-initiated cancel (revoke sends SIGUSR1). Emit the matching
+ # terminal event so a cancelled job is reported as cancelled.
+ if async_query_manager.is_job_cancelled(job_metadata["job_id"]):
Review Comment:
Both `submit_explore_json_job` and `submit_chart_data_job` are now
cancellable, but only `load_chart_data_into_cache` handles the cancellation.
`load_explore_json_into_cache` needs the same treatment.
##########
superset/async_events/api.py:
##########
@@ -99,3 +103,69 @@ def events(self) -> Response:
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("list")
+ 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
+ 401:
+ $ref: '#/components/responses/401'
+ 403:
+ $ref: '#/components/responses/403'
+ 404:
+ $ref: '#/components/responses/404'
+ 500:
+ $ref: '#/components/responses/500'
+ """
+ try:
+ async_channel_id =
async_query_manager.parse_channel_id_from_request(
+ request
+ )
Review Comment:
Nit: job_id comes from the URL and is used directly in a Redis key. Worth
validating it's a UUID first?
##########
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:
Do you think this comment is relevant?
--
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]
