fix(query-orchestrator): Use external_id for queue result retrieval t…

packages/cubejs-backend-shared/src/env.ts

@@ -2281,6 +2281,7 @@ const variables: Record<string, (...args: any) => any> = {
   externalDefault: () => get('CUBEJS_EXTERNAL_DEFAULT')
     .default('true')
     .asBoolStrict(),
+  queueExternalId: () => get('CUBEJS_QUEUE_EXTERNAL_ID').default('false').asBool(),
   scheduledRefreshDefault: () => get(
     'CUBEJS_SCHEDULED_REFRESH_DEFAULT'
   ).default('true').asBoolStrict(),

packages/cubejs-backend-shared/src/logger.ts

@@ -12,6 +12,7 @@ export type LoggerFnParams = {
   error?: Error | string;
   trace?: string,
   warning?: string,
+  level?: LogLevel,
   [key: string]: any,
 };
 

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

@@ -44,6 +44,7 @@ export interface AddToQueueOptions {
   requestId: string,
   spanId?: string,
   orphanedTimeout?: number,
+  externalId?: string,
 }
 
 export interface QueueDriverOptions {
@@ -58,7 +59,7 @@ export interface QueueDriverOptions {
 export interface QueueDriverConnectionInterface {
   redisHash(queryKey: QueryKey): QueryKeyHash;
   getResultBlocking(queryKey: QueryKeyHash, queueId: QueueId): Promise<unknown>;
-  getResult(queryKey: QueryKey): Promise<any>;
+  getResult(queryKey: QueryKey, externalId?: string): Promise<any>;
   /**
    * Adds specified by the queryKey query to the queue, returns tuple
    * with the operation result.

packages/cubejs-cubestore-driver/src/CubeStoreDriver.ts

@@ -21,10 +21,11 @@ import fetch from 'node-fetch';
 import { ConnectionConfig } from './types';
 import { WebSocketConnection } from './WebSocketConnection';
 
-type CubeStoreCapability = 'queueExclusive';
+type CubeStoreCapability = 'queueExclusive' | 'queueExternalId';
 
 const CubeStoreCapabilityMinVersion: Record<CubeStoreCapability, string> = {
   queueExclusive: '1.6.22',
+  queueExternalId: '1.6.26',
 };
 
 const GenericTypeToCubeStore: Record<string, string> = {

packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts

@@ -16,7 +16,7 @@ import {
   GetActiveAndToProcessResponse,
   QueryKeysTuple,
 } from '@cubejs-backend/base-driver';
-import { getProcessUid } from '@cubejs-backend/shared';
+import { getEnv, getProcessUid } from '@cubejs-backend/shared';
 
 import { CubeStoreDriver } from './CubeStoreDriver';
 
@@ -38,11 +38,23 @@ type CubeStoreListResponse = {
   status: string
 };
 
-class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
+export class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
+  private readonly externalIdEnabled: boolean;
+
   public constructor(
     protected readonly driver: CubeStoreDriver,
     protected readonly options: QueueDriverOptions,
-  ) { }
+  ) {
+    this.externalIdEnabled = getEnv('queueExternalId');
+  }
+
+  public async useExternalId(): Promise<boolean> {
+    if (this.externalIdEnabled) {
+      return this.driver.hasCapability('queueExternalId');
+    }
+
+    return false;
+  }
 
   public redisHash(queryKey: QueryKey): QueryKeyHash {
     return hashQueryKey(queryKey, this.options.processUid);
@@ -53,9 +65,9 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
   }
 
   public async addToQueue(
-    keyScore: number,
+    _keyScore: number,
     queryKey: QueryKey,
-    orphanedTime: number,
+    _orphanedTime: number,
     queryHandler: string,
     query: AddToQueueQuery,
     priority: number,
@@ -79,11 +91,16 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
       values.push(options.orphanedTimeout);
     }
 
+    const useExternalId = options.externalId && await this.useExternalId();
+    if (useExternalId) {
+      values.push(options.externalId!);
+    }
+
     values.push(this.prefixKey(this.redisHash(queryKey)));
     values.push(JSON.stringify(data));
 
     const exclusive = queryKey.persistent && await this.driver.hasCapability('queueExclusive');
-    const rows = await this.driver.query(`QUEUE ADD${exclusive ? ' EXCLUSIVE' : ''} PRIORITY ?${options.orphanedTimeout ? ' ORPHANED ?' : ''} ? ?`, values);
+    const rows = await this.driver.query(`QUEUE ADD${exclusive ? ' EXCLUSIVE' : ''} PRIORITY ?${options.orphanedTimeout ? ' ORPHANED ?' : ''}${useExternalId ? ' EXTERNAL_ID ?' : ''} ? ?`, values);
     if (rows && rows.length) {
       return [
         rows[0].added === 'true' ? 1 : 0,
@@ -200,10 +217,17 @@ class CubestoreQueueDriverConnection implements QueueDriverConnectionInterface {
     return [active, toProcess, defs];
   }
 
-  public async getResult(queryKey: QueryKey): Promise<unknown> {
-    const rows = await this.driver.query('QUEUE RESULT ?', [
-      this.prefixKey(this.redisHash(queryKey)),
-    ]);
+  public async getResult(queryKey: QueryKey, externalId?: string): Promise<unknown> {
+    const params: string[] = [];
+
+    const passExternalId = externalId && await this.useExternalId();
+    if (passExternalId) {
+      params.push(externalId);
+    }
+
+    params.push(this.prefixKey(this.redisHash(queryKey)));
+
+    const rows = await this.driver.query(`QUEUE RESULT ${passExternalId ? 'EXTERNAL_ID ? ' : ''}?`, params);
     if (rows && rows.length) {
       return this.decodeQueryDefFromRow(rows[0], 'getResult');
     }

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

@@ -134,7 +134,7 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac
     return res;
   }
 
-  public async getResult(queryKey: QueryKey): Promise<any> {
+  public async getResult(queryKey: QueryKey, _externalId?: string): Promise<any> {
     const resultListKey = this.resultListKey(queryKey);
     if (this.state.resultPromises[resultListKey] && this.state.resultPromises[resultListKey].resolved) {
       return this.getResultBlocking(this.redisHash(queryKey));

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

@@ -193,6 +193,11 @@ export class QueryQueue {
       ...executeOptions,
     };
 
+    if (options.requestId) {
+      const idx = options.requestId.lastIndexOf('-span-');
+      options.externalId = idx !== -1 ? options.requestId.substring(0, idx) : options.requestId;
+    }
+
     if (this.skipQueue) {
       const queryDef = {
         queryHandler,
@@ -233,7 +238,7 @@ export class QueryQueue {
       // Result here won't be fetched for a forced build query and a jobbed build
       // query (initialized by the /cubejs-system/v1/pre-aggregations/jobs
       // endpoint).
-      let result = !query.forceBuild && await queueConnection.getResult(queryKey);
+      let result = !query.forceBuild && await queueConnection.getResult(queryKey, options.externalId);
       if (result && !result.streamResult) {
         return this.parseResult(result);
       }