refactor(query-orchestrator): Migrate QueryQueue to TS

Author: ovr

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

@@ -1,4 +1,4 @@
-export type QueryDef = unknown;
+export type QueryDef = any;
 // Primary key of Queue item
 export type QueueId = string | number | bigint;
 // This was used as a lock for Redis, deprecated.
@@ -64,25 +64,25 @@ export interface QueueDriverConnectionInterface {
    * @param keyScore Redis specific thing
    * @param queryKey
    * @param orphanedTime
-   * @param queryHandler Our queue allows using different handlers. For example, query, cvsQuery, etc.
+   * @param queryHandler Our queue allow to use different handlers. For example query, cvsQuery, etc.
    * @param query
    * @param priority
    * @param options
    */
   addToQueue(keyScore: number, queryKey: QueryKey, orphanedTime: number, queryHandler: string, query: AddToQueueQuery, priority: number, options: AddToQueueOptions): Promise<AddToQueueResponse>;
-  // Return query keys that were sorted by priority and time
+  // Return query keys which was sorted by priority and time
   getToProcessQueries(): Promise<QueryKeysTuple[]>;
   getActiveQueries(): Promise<QueryKeysTuple[]>;
   getQueryDef(hash: QueryKeyHash, queueId: QueueId | null): Promise<QueryDef | null>;
-  // Queries that were added to queue, but was not processed and not needed
+  // Queries which was added to queue, but was not processed and not needed
   getOrphanedQueries(): Promise<QueryKeysTuple[]>;
-  // Queries that were not completed with old heartbeat
+  // Queries which was not completed with old heartbeat
   getStalledQueries(): Promise<QueryKeysTuple[]>;
   getQueryStageState(onlyKeys: boolean): Promise<QueryStageStateResponse>;
   updateHeartBeat(hash: QueryKeyHash, queueId: QueueId | null): Promise<void>;
   getNextProcessingId(): Promise<ProcessingId>;
   // Trying to acquire a lock for processing a queue item, this method can return null when
-  // multiple nodes try to process the same query
+  // multiple nodes tries to process the same query
   retrieveForProcessing(hash: QueryKeyHash, processingId: ProcessingId): Promise<RetrieveForProcessingResponse>;
   freeProcessingLock(hash: QueryKeyHash, processingId: ProcessingId, activated: unknown): Promise<void>;
   optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, processingId: ProcessingId, queueId: QueueId | null): Promise<boolean>;

packages/cubejs-query-orchestrator/CLAUDE.md

@@ -26,7 +26,7 @@ yarn test
 # Run only unit tests
 yarn unit
 
-# Run only integration tests  
+# Run only integration tests
 yarn integration
 
 # Run CubeStore integration tests specifically
@@ -46,16 +46,16 @@ yarn lint:fix
 The Query Orchestrator consists of several interconnected components:
 
 1. **QueryOrchestrator** (`src/orchestrator/QueryOrchestrator.ts`): Main orchestration class that coordinates query execution and manages drivers
-2. **QueryCache** (`src/orchestrator/QueryCache.ts`): Handles query result caching with configurable cache drivers  
-3. **QueryQueue** (`src/orchestrator/QueryQueue.js`): Manages query queuing and background processing
+2. **QueryCache** (`src/orchestrator/QueryCache.ts`): Handles query result caching with configurable cache drivers
+3. **QueryQueue** (`src/orchestrator/QueryQueue.ts`): Manages query queuing and background processing
 4. **PreAggregations** (`src/orchestrator/PreAggregations.ts`): Manages pre-aggregation building and loading
 5. **DriverFactory** (`src/orchestrator/DriverFactory.ts`): Creates and manages database driver instances
 
 ### Cache and Queue Driver Architecture
 
 The orchestrator supports multiple backend drivers:
 - **Memory**: In-memory caching and queuing (development)
-- **CubeStore**: Distributed storage engine (production) 
+- **CubeStore**: Distributed storage engine (production)
 - **Redis**: External Redis-based caching (legacy, being phased out)
 
 Driver selection logic in `QueryOrchestrator.ts:detectQueueAndCacheDriver()`:
@@ -128,4 +128,4 @@ Key configuration options in `QueryOrchestratorOptions`:
 - Inherits linting rules from `@cubejs-backend/linter`
 - Jest configuration extends base repository config
 - Docker Compose setup for integration testing
-- Coverage reports generated in `coverage/` directory
+- Coverage reports generated in `coverage/` directory

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

@@ -418,7 +418,8 @@ export class PreAggregationLoader {
     const queue = await this.preAggregations.getQueue(this.preAggregation.dataSource);
     return queue.executeInQueue(
       'query',
-      this.preAggregationQueryKey(invalidationKeys),
+      // TODO: Sync types
+      this.preAggregationQueryKey(invalidationKeys) as any,
       {
         preAggregation: this.preAggregation,
         preAggregationsTablesToTempTables: this.preAggregationsTablesToTempTables,

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

@@ -376,10 +376,9 @@ export class PreAggregations {
     tables = tables.filter(row => `${schema}.${row.table_name}` === table);
 
     // fetching query result
-    const { queueDriver } = this.queue[dataSource];
-    const conn = await queueDriver.createConnection();
+    const conn = await this.queue[dataSource].getQueueDriver().createConnection();
     const result = await conn.getResult(key);
-    queueDriver.release(conn);
+    this.queue[dataSource].getQueueDriver().release(conn);
 
     // calculating status
     let status: string;

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

@@ -9,10 +9,10 @@ import {
   InlineTables,
   CacheDriverInterface,
   TableStructure,
-  DriverInterface,
+  DriverInterface, QueryKey,
 } from '@cubejs-backend/base-driver';
 
-import { QueryQueue } from './QueryQueue';
+import {QueryQueue, QueryQueueOptions} from './QueryQueue';
 import { ContinueWaitError } from './ContinueWaitError';
 import { LocalCacheDriver } from './LocalCacheDriver';
 import { DriverFactory, DriverFactoryByDataSource } from './DriverFactory';
@@ -455,9 +455,9 @@ export class QueryCache {
     };
 
     if (!persistent) {
-      return queue.executeInQueue('query', cacheKey, _query, priority, opt);
+      return queue.executeInQueue('query', cacheKey as QueryKey, _query, priority, opt);
     } else {
-      return queue.executeInQueue('stream', cacheKey, {
+      return queue.executeInQueue('stream', cacheKey as QueryKey, {
         ..._query,
         aliasNameToMember,
       }, priority, opt);
@@ -563,7 +563,7 @@ export class QueryCache {
     redisPrefix: string,
     clientFactory: DriverFactory,
     executeFn: (client: BaseDriver, req: any) => any,
-    options: Record<string, any> = {}
+    options: Omit<QueryQueueOptions, 'queryHandlers' | 'cancelHandlers'>
   ): QueryQueue {
     const queue: any = new QueryQueue(redisPrefix, {
       queryHandlers: {
@@ -583,57 +583,57 @@ export class QueryCache {
           }
           return result;
         },
-        stream: async (req, target) => {
-          queue.logger('Streaming SQL', { ...req });
-          await (new Promise((resolve, reject) => {
-            let logged = false;
-            Promise
-              .all([clientFactory()])
-              .then(([client]) => (<DriverInterface>client).stream(req.query, req.values, { highWaterMark: getEnv('dbQueryStreamHighWaterMark') }))
-              .then((source) => {
-                const cleanup = async (error) => {
-                  if (source.release) {
-                    const toRelease = source.release;
-                    delete source.release;
-                    await toRelease();
-                  }
-                  if (error && !target.destroyed) {
-                    target.destroy(error);
-                  }
-                  if (!logged && target.destroyed) {
-                    logged = true;
-                    if (error) {
-                      queue.logger('Streaming done with error', {
-                        query: req.query,
-                        query_values: req.values,
-                        error,
-                      });
-                      reject(error);
-                    } else {
-                      queue.logger('Streaming successfully completed', {
-                        requestId: req.requestId,
-                      });
-                      resolve(req.requestId);
-                    }
+      },
+      streamHandler: async (req, target) => {
+        queue.logger('Streaming SQL', { ...req });
+        await (new Promise((resolve, reject) => {
+          let logged = false;
+          Promise
+            .all([clientFactory()])
+            .then(([client]) => (<DriverInterface>client).stream(req.query, req.values, { highWaterMark: getEnv('dbQueryStreamHighWaterMark') }))
+            .then((source) => {
+              const cleanup = async (error) => {
+                if (source.release) {
+                  const toRelease = source.release;
+                  delete source.release;
+                  await toRelease();
+                }
+                if (error && !target.destroyed) {
+                  target.destroy(error);
+                }
+                if (!logged && target.destroyed) {
+                  logged = true;
+                  if (error) {
+                    queue.logger('Streaming done with error', {
+                      query: req.query,
+                      query_values: req.values,
+                      error,
+                    });
+                    reject(error);
+                  } else {
+                    queue.logger('Streaming successfully completed', {
+                      requestId: req.requestId,
+                    });
+                    resolve(req.requestId);
                   }
-                };
-
-                source.rowStream.once('end', () => cleanup(undefined));
-                source.rowStream.once('error', cleanup);
-                source.rowStream.once('close', () => cleanup(undefined));
-
-                target.once('end', () => cleanup(undefined));
-                target.once('error', cleanup);
-                target.once('close', () => cleanup(undefined));
-
-                source.rowStream.pipe(target);
-              })
-              .catch((reason) => {
-                target.emit('error', reason);
-                resolve(reason);
-              });
-          }));
-        },
+                }
+              };
+
+              source.rowStream.once('end', () => cleanup(undefined));
+              source.rowStream.once('error', cleanup);
+              source.rowStream.once('close', () => cleanup(undefined));
+
+              target.once('end', () => cleanup(undefined));
+              target.once('error', cleanup);
+              target.once('close', () => cleanup(undefined));
+
+              source.rowStream.pipe(target);
+            })
+            .catch((reason) => {
+              target.emit('error', reason);
+              resolve(reason);
+            });
+        }));
       },
       cancelHandlers: {
         query: async (req) => {

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

@@ -3,6 +3,7 @@ import R from 'ramda';
 import { getEnv } from '@cubejs-backend/shared';
 import { CubeStoreDriver } from '@cubejs-backend/cubestore-driver';
 
+import { QueryKey } from '@cubejs-backend/base-driver';
 import { QueryCache, QueryBody, TempTable, PreAggTableToTempTable } from './QueryCache';
 import { PreAggregations, PreAggregationDescription, getLastUpdatedAtTimestamp } from './PreAggregations';
 import { DriverFactory, DriverFactoryByDataSource } from './DriverFactory';
@@ -306,7 +307,7 @@ export class QueryOrchestrator {
 
     if (pendingPreAggregationIndex === -1) {
       const qcQueue = await this.queryCache.getQueue(queryBody.dataSource);
-      return qcQueue.getQueryStage(QueryCache.queryCacheKey(queryBody));
+      return qcQueue.getQueryStage(QueryCache.queryCacheKey(queryBody) as QueryKey);
     }
 
     const preAggregation = queryBody.preAggregations[pendingPreAggregationIndex];