msyavuz commented on code in PR #42305:
URL: https://github.com/apache/superset/pull/42305#discussion_r3665260185


##########
superset-frontend/src/middleware/asyncEvent.test.ts:
##########
@@ -130,6 +130,42 @@ 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' });
+      expect(fetchMock.callHistory.calls(CANCEL_ENDPOINT)).toHaveLength(1);

Review Comment:
   Both cancel tests already await a tick before asserting in the current 
version.



##########
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:
   Relevant, and addressed: the flag write is conditional (`xx=True`) and the 
record is cleared before a terminal event, so a late cancel 404s instead of 
resurrecting the key.



##########
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
+            )
+        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},
+        )

Review Comment:
   API-level tests are in `tests/integration_tests/async_events/api_tests.py` 
(`test_cancel_*`) as of the current version.



##########
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:
   Done, both tasks share `_handle_soft_time_limit` now.



##########
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:
   Added, a non-UUID job id is rejected with 400 before it reaches Redis.



-- 
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]

Reply via email to