fix(query-orchestrator): Queue - clear timer on streaming (#7501)

Author: ovr

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

@@ -310,22 +310,32 @@ export class QueryQueue {
       if (queryHandler === 'stream') {
         const self = this;
         result = await new Promise((resolve) => {
+          let timeoutTimerId = null;
+
           const onStreamStarted = (streamStartedHash) => {
             if (streamStartedHash === queryKeyHash) {
+              if (timeoutTimerId) {
+                clearTimeout(timeoutTimerId);
+              }
+
               resolve(self.getQueryStream(queryKeyHash));
             }
           };
 
-          setTimeout(() => {
-            self.streamEvents.removeListener('streamStarted', onStreamStarted);
-            resolve(null);
-          }, this.continueWaitTimeout * 1000);
-
           self.streamEvents.addListener('streamStarted', onStreamStarted);
-          const stream = this.getQueryStream(this.redisHash(queryKey));
+
+          const stream = this.getQueryStream(queryKeyHash);
           if (stream) {
             self.streamEvents.removeListener('streamStarted', onStreamStarted);
             resolve(stream);
+          } else {
+            timeoutTimerId = setTimeout(
+              () => {
+                self.streamEvents.removeListener('streamStarted', onStreamStarted);
+                resolve(null);
+              },
+              this.continueWaitTimeout * 10000
+            );
           }
         });
       } else {

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

@@ -2,7 +2,7 @@ import * as stream from 'stream';
 import { getEnv } from '@cubejs-backend/shared';
 
 export class QueryStream extends stream.Transform {
-  private timeout = 5 * 60000 || getEnv('dbQueryTimeout');
+  private timeout = 5 * 60 * 1000;
 
   private timer = null;
 

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

@@ -155,7 +155,8 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions = {}
     test('timeout', async () => {
       const query = ['select * from 3'];
 
-      await queue.executeInQueue('delay', query, { delay: 60 * 60 * 1000, result: '1', isJob: true });
+      // executionTimeout is 2s, 5s is enough
+      await queue.executeInQueue('delay', query, { delay: 5 * 1000, result: '1', isJob: true });
       await awaitProcessing();
 
       expect(logger.mock.calls.length).toEqual(5);
@@ -324,9 +325,15 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions = {}
     test('stream handler', async () => {
       const key: QueryKey = ['select * from table', []];
       key.persistent = true;
-      await queue.executeInQueue('stream', key, { aliasNameToMember: {} }, 0);
+      const stream = await queue.executeInQueue('stream', key, { aliasNameToMember: {} }, 0);
       await awaitProcessing();
 
+      // QueryStream has a debounce timer to destroy stream
+      // without reading it, timer will block exit for jest
+      for await (const chunk of stream) {
+        console.log('streaming chunk: ', chunk);
+      }
+
       expect(streamCount).toEqual(1);
       expect(logger.mock.calls[logger.mock.calls.length - 1][0]).toEqual('Performing query completed');
     });