fix(query-orchestrator): don't report a cancelled query as a query error

packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts

@@ -9,6 +9,7 @@ import {
   QueryStageStateResponse,
   AddToQueueOptions,
   QueuePriority,
+  QueueDriverConnectionInterface,
   RetrieveForProcessingSuccess
 } from '@cubejs-backend/base-driver';
 import { CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver';
@@ -854,6 +855,9 @@ export class QueryQueue {
     try {
       let executionResult;
       let queryExecutionFinished = false;
+      // Set once the queue item is known to be gone, so that the rejection the cancellation
+      // causes isn't reported as a query error.
+      let queryCancelled = false;
       // Set by the query handler's setCancelHandler callback once execution begins.
       // Not available on the original query def from retrieveForProcessing.
       let localCancelHandler: unknown = null;
@@ -902,6 +906,7 @@ export class QueryQueue {
             try {
               const currentDef = await queueConnection.getQueryDef(queryKeyHashed, queueId);
               if (!currentDef && !queryExecutionFinished) {
+                queryCancelled = true;
                 this.logger('Cancelling query due to external cancellation', {
                   queueId,
                   queryKey: query.queryKey,
@@ -996,7 +1001,17 @@ export class QueryQueue {
         executionResult = {
           error: (e.message || e).toString() // TODO error handling
         };
-        this.logger('Error while querying', {
+
+        // The queue removes an item when it cancels it - orphaned, stalled or explicit - and the
+        // driver then rejects the execution still in flight. Reporting that rejection as a query
+        // failure surfaces the cancellation in query history, which orphaned queries are kept out
+        // of. The lookup happens after the rejection, so a query which genuinely fails as its item
+        // is orphaned is reclassified too: harmless, nothing is waiting for either result. A
+        // timeout is excluded - it cancels a query which is still queued, so it stays an error.
+        const cancelled = !(e instanceof TimeoutError) &&
+          (queryCancelled || await this.isQueryRemovedFromQueue(queueConnection, query, queryKeyHashed, queueId));
+
+        const logEvent = {
           queueId,
           queueSize,
           duration: ((new Date()).getTime() - startQueryTime),
@@ -1009,8 +1024,24 @@ export class QueryQueue {
           newVersionEntry: query.query?.newVersionEntry,
           preAggregation: query.query?.preAggregation,
           addedToQueueTime: query.addedToQueueTime,
-          error: (e.stack || e).toString()
-        });
+        };
+
+        if (cancelled) {
+          // `warning` rather than `error` on purpose: an `error` field is what marks a query as
+          // failed downstream, and a cancellation is not a failure. It still has to be one of the
+          // two, otherwise the default logger drops the event at the default `info` level.
+          this.logger('Cancelled query execution', {
+            ...logEvent,
+            warning: 'Query execution was rejected because the query had been cancelled',
+            cancellationError: (e.stack || e).toString()
+          });
+        } else {
+          this.logger('Error while querying', {
+            ...logEvent,
+            error: (e.stack || e).toString()
+          });
+        }
+
         if (e instanceof TimeoutError) {
           const queryWithCancelHandle = await queueConnection.getQueryDef(queryKeyHashed, queueId);
           if (queryWithCancelHandle) {
@@ -1063,6 +1094,32 @@ export class QueryQueue {
     }
   }
 
+  /**
+   * Whether the queue item is gone, which means a cancellation - orphaned, stalled or explicit -
+   * has removed it while its execution was still in flight. A queue lookup failure is reported as
+   * `false` so that a genuine query error is never hidden by a queue storage issue.
+   */
+  protected async isQueryRemovedFromQueue(
+    queueConnection: QueueDriverConnectionInterface,
+    query: QueryDef,
+    queryKeyHashed: QueryKeyHash,
+    queueId: QueueId
+  ): Promise<boolean> {
+    try {
+      return !(await queueConnection.getQueryDef(queryKeyHashed, queueId));
+    } catch (e: any) {
+      this.logger('Error while checking query cancellation', {
+        queueId,
+        queryKey: query.queryKey,
+        queuePrefix: this.redisQueuePrefix,
+        requestId: query.requestId,
+        error: (e.stack || e).toString()
+      });
+
+      return false;
+    }
+  }
+
   /**
    * Processing cancel query flow.
    */

packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts

@@ -42,6 +42,9 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
     const processMessagePromises: Promise<any>[] = [];
     const processCancelPromises: Promise<any>[] = [];
     let cancelledQuery;
+    // Rejects of the in-flight `cancelable` queries, so that a cancellation can reject the
+    // running handler the way a driver rejects a query it has stopped.
+    let cancelableRejects: ((error: Error) => void)[] = [];
     let streamCallOrder: string[] = [];
 
     const tenantPrefix = crypto.randomBytes(6).toString('hex');
@@ -55,6 +58,17 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
           await setCancelHandler(result);
           return delayFn(result, query.delay);
         },
+        cancelable: async (query, setCancelHandler) => {
+          await setCancelHandler(query.result);
+
+          return new Promise((resolve, reject) => {
+            const timer = setTimeout(() => resolve(query.result), query.delay);
+            cancelableRejects.push((error) => {
+              clearTimeout(timer);
+              reject(error);
+            });
+          });
+        },
       },
       streamHandler: async (query, stream) => {
         streamCount++;
@@ -82,6 +96,10 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
         delay: async (query) => {
           console.log(`cancel call: ${JSON.stringify(query)}`);
           cancelledQuery = query.queryKey;
+        },
+        cancelable: async (query) => {
+          cancelledQuery = query.queryKey;
+          cancelableRejects.splice(0).forEach((reject) => reject(new Error('Query was cancelled')));
         }
       },
       continueWaitTimeout: 1,
@@ -111,6 +129,7 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
       streamCount = 0;
       streamHandlerDelay = 250;
       streamCallOrder = [];
+      cancelableRejects = [];
     });
 
     afterAll(async () => {
@@ -187,6 +206,38 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
       expect(logger.mock.calls[3][0]).toEqual('Error while querying');
     });
 
+    test('cancelled query is not reported as an error', async () => {
+      cancelledQuery = null;
+
+      const queryKey: QueryKey = ['select * from cancelled', []];
+      // The client gives up on ContinueWaitError long before the handler would resolve
+      const pending = queue
+        .executeInQueue('cancelable', queryKey, { delay: 60 * 1000, result: '1' }, QueuePriority.Background)
+        .catch(e => e);
+
+      // executionTimeout is 2s, so the cancellation has to reach a handler which is already
+      // running, otherwise the query fails with a timeout instead. The deadline stays well under
+      // it so that a worker which never picks the query up fails here rather than as a confusing
+      // assertion on the log events below.
+      const deadline = Date.now() + 750;
+      while (cancelableRejects.length === 0 && Date.now() < deadline) {
+        await pausePromise(10);
+      }
+      expect(cancelableRejects.length).toEqual(1);
+
+      await queue.cancelQuery(queue.redisHash(queryKey), null);
+      expect(cancelledQuery).toEqual(queryKey);
+      expect(await pending).toBeInstanceOf(ContinueWaitError);
+      await awaitProcessing();
+
+      // The rejection the cancellation causes is a cancellation, not a query failure: reporting
+      // it as one would surface it in query history
+      const events = logger.mock.calls.map(([message]) => message);
+      expect(events).toContain('Cancelling query manual');
+      expect(events).toContain('Cancelled query execution');
+      expect(events).not.toContain('Error while querying');
+    });
+
     test('stage reporting', async () => {
       const resultPromise = queue.executeInQueue('delay', '1', { delay: 200, result: '1' }, QueuePriority.Background, {
         stageQueryKey: '1',