chore: fixes

Author: ovr

packages/cubejs-base-driver/src/queue-driver.interface.ts

@@ -38,10 +38,11 @@ export interface AddToQueueQuery {
 }
 
 export interface AddToQueueOptions {
-  stageQueryKey: string,
+  queueId?: QueueId,
+  stageQueryKey?: any,
   requestId: string,
+  spanId?: string,
   orphanedTimeout?: number,
-  queueId: QueueId,
 }
 
 export interface QueueDriverOptions {

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

@@ -6,7 +6,7 @@ import {
   QueryKeyHash,
   QueueId,
   QueryDef,
-  QueryStageStateResponse
+  QueryStageStateResponse, AddToQueueOptions
 } from '@cubejs-backend/base-driver';
 import { CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver';
 
@@ -23,6 +23,10 @@ export type QueryHandlersMap = Record<string, QueryHandlerFn>;
 export type SendProcessMessageFn = (queryKeyHash: QueryKeyHash, queueId: QueueId | null) => Promise<void> | void;
 export type SendCancelMessageFn = (query: QueryDef, queueId: QueueId | null) => Promise<void> | void;
 
+export type ExecuteInQueueOptions = Omit<AddToQueueOptions, 'queueId'> & {
+  spanId?: string
+};
+
 export type QueryQueueOptions = {
   cacheAndQueueDriver: string;
   logger: (message, event) => void;
@@ -176,9 +180,12 @@ export class QueryQueue {
     queryKey: QueryKey,
     query: QueryDef,
     priority?: number,
-    options?: any,
+    executeOptions?: ExecuteInQueueOptions,
   ) {
-    options.queueId = this.generateQueueId();
+    const options: AddToQueueOptions = {
+      queueId: this.generateQueueId(),
+      ...executeOptions,
+    };
 
     if (this.skipQueue) {
       const queryDef = {
@@ -236,6 +243,7 @@ export class QueryQueue {
       const keyScore = time + (10000 - priority) * 1E14;
 
       options.orphanedTimeout = query.orphanedTimeout;
+
       const orphanedTimeout = 'orphanedTimeout' in query ? query.orphanedTimeout : this.orphanedTimeout;
       const orphanedTime = time + (orphanedTimeout * 1000);
 

packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts

@@ -181,7 +181,9 @@ export function QueryQueueBenchmark(name: string, options: QueryQueueTestOptions
             // eslint-disable-next-line no-bitwise
             payload: 'a'.repeat(benchSettings.queuePayloadSize)
           }, 1, {
-
+            stageQueryKey: 1,
+            requestId: 'request-id',
+            spanId: 'span-id'
           });
 
           counters.queueResolved++;

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

@@ -178,17 +178,29 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
     });
 
     test('stage reporting', async () => {
-      const resultPromise = queue.executeInQueue('delay', '1', { delay: 200, result: '1' }, 0, { stageQueryKey: '1' });
+      const resultPromise = queue.executeInQueue('delay', '1', { delay: 200, result: '1' }, 0, {
+        stageQueryKey: '1',
+        requestId: 'request-id',
+        spanId: 'span-id'
+      });
       await delayFn(null, 50);
       expect((await queue.getQueryStage('1')).stage).toBe('Executing query');
       await resultPromise;
       expect(await queue.getQueryStage('1')).toEqual(undefined);
     });
 
     test('priority stage reporting', async () => {
-      const resultPromise1 = queue.executeInQueue('delay', '31', { delay: 200, result: '1' }, 20, { stageQueryKey: '12' });
+      const resultPromise1 = queue.executeInQueue('delay', '31', { delay: 200, result: '1' }, 20, {
+        stageQueryKey: '12',
+        requestId: 'request-id',
+        spanId: 'span-id'
+      });
       await delayFn(null, 50);
-      const resultPromise2 = queue.executeInQueue('delay', '32', { delay: 200, result: '1' }, 10, { stageQueryKey: '12' });
+      const resultPromise2 = queue.executeInQueue('delay', '32', { delay: 200, result: '1' }, 10, {
+        stageQueryKey: '12',
+        requestId: 'request-id',
+        spanId: 'span-id'
+      });
       await delayFn(null, 50);
 
       expect((await queue.getQueryStage('12', 10)).stage).toBe('#1 in queue');
@@ -275,7 +287,6 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
 
         let orphanedTimeout = 2;
         await connection.addToQueue(keyScore, ['1', []], time + (orphanedTimeout * 1000), 'delay', { isJob: true, orphanedTimeout: time, }, priority, {
-          queueId: 1,
           stageQueryKey: '1',
           requestId: '1',
           orphanedTimeout,
@@ -286,7 +297,6 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
         orphanedTimeout = 60;
 
         await connection.addToQueue(keyScore, ['2', []], time + (orphanedTimeout * 1000), 'delay', { isJob: true, orphanedTimeout: time, }, priority, {
-          queueId: 2,
           stageQueryKey: '2',
           requestId: '2',
           orphanedTimeout,
@@ -423,11 +433,11 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) =>
       await queue.reconcileQueue();
 
       await redisClient.addToQueue(
-        keyScore, 'activated1', time, 'handler', <any>['select'], priority, { stageQueryKey: 'race', requestId: '1', queueId: 1 }
+        keyScore, 'activated1', time, 'handler', <any>['select'], priority, { stageQueryKey: 'race', requestId: '1' }
       );
 
       await redisClient.addToQueue(
-        keyScore + 100, 'activated2', time + 100, 'handler2', <any>['select2'], priority, { stageQueryKey: 'race2', requestId: '1', queueId: 2 }
+        keyScore + 100, 'activated2', time + 100, 'handler2', <any>['select2'], priority, { stageQueryKey: 'race2', requestId: '1' }
       );
 
       const processingId1 = await redisClient.getNextProcessingId();