refactor(query-orchestrator): Migrate QueryQueue to TS

Author: ovr

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

@@ -1,7 +1,14 @@
-export type QueryDef = unknown;
-export type QueryKey = (string | [string, any[]]) & {
+export type QueryDef = any;
+export type QueryKeyOptions = any;
+export type QueryKeyTuple = [
+  sql: string,
+  params: unknown[],
+  options?: QueryKeyOptions
+];
+export type QueryKey = (string | QueryKeyTuple) & {
   persistent?: true,
 };
+export type QueryMessageId = string;
 
 export type AddToQueueResponse = [added: number, _b: any, _c: any, queueSize: number, addedToQueueTime: number];
 export type QueryStageStateResponse = [active: string[], toProcess: string[]] | [active: string[], toProcess: string[], defs: Record<string, QueryDef>];
@@ -27,7 +34,7 @@ export interface QueueDriverOptions {
 }
 
 export interface QueueDriverConnectionInterface {
-  redisHash(queryKey: QueryKey): string;
+  redisHash(queryKey: QueryKey): QueryMessageId;
   getResultBlocking(queryKey: QueryKey): Promise<unknown>;
   getResult(queryKey: QueryKey): Promise<any>;
   addToQueue(keyScore: number, queryKey: QueryKey, orphanedTime: any, queryHandler: any, query: AddToQueueQuery, priority: number, options: AddToQueueOptions): Promise<AddToQueueResponse>;
@@ -45,10 +52,10 @@ export interface QueueDriverConnectionInterface {
   // Trying to acquire a lock for processing a queue item, this method can return null when
   // multiple nodes tries to process the same query
   retrieveForProcessing(queryKey: QueryKey, processingId: number | string): Promise<RetrieveForProcessingResponse>;
-  freeProcessingLock(queryKey: QueryKey, processingId: string | number, activated: unknown): Promise<void>;
+  freeProcessingLock(queryKey: QueryKey, processingId: string | number, activated: unknown): Promise<void | string>;
   optimisticQueryUpdate(queryKey: QueryKey, toUpdate, processingId): Promise<boolean>;
   cancelQuery(queryKey: QueryKey): Promise<QueryDef | null>;
-  getQueryAndRemove(queryKey: QueryKey): Promise<[QueryDef]>;
+  getQueryAndRemove(queryKey: QueryMessageId): Promise<[QueryDef]>;
   setResultAndRemoveQuery(queryKey: QueryKey, executionResult: any, processingId: any): Promise<unknown>;
   release(): void;
   //

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

@@ -1,8 +1,13 @@
-import { QueueDriverConnectionInterface, QueueDriverInterface } from '@cubejs-backend/base-driver';
+import {
+  QueryKey,
+  QueryMessageId,
+  QueueDriverConnectionInterface,
+  QueueDriverInterface,
+} from '@cubejs-backend/base-driver';
 import { getCacheHash } from './utils';
 
 export abstract class BaseQueueDriver implements QueueDriverInterface {
-  public redisHash(queryKey: string) {
+  public redisHash(queryKey: QueryKey): QueryMessageId {
     return getCacheHash(queryKey);
   }
 

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

@@ -845,7 +845,7 @@ export class PreAggregationLoader {
     const queue = await this.preAggregations.getQueue(this.preAggregation.dataSource);
     return queue.executeInQueue(
       'query',
-      this.preAggregationQueryKey(invalidationKeys),
+      <any> this.preAggregationQueryKey(invalidationKeys),
       {
         preAggregation: this.preAggregation,
         preAggregationsTablesToTempTables: this.preAggregationsTablesToTempTables,
@@ -2015,7 +2015,7 @@ export class PreAggregations {
     tables = tables.filter(row => `${schema}.${row.table_name}` === table);
 
     // fetching query result
-    const { queueDriver } = this.queue[dataSource];
+    const queueDriver = this.queue[dataSource].getQueueDriver();
     const conn = await queueDriver.createConnection();
     const result = await conn.getResult(key);
     queueDriver.release(conn);

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

@@ -475,12 +475,12 @@ export class QueryCache {
     };
 
     if (!persistent) {
-      return queue.executeInQueue('query', cacheKey, _query, priority, opt);
+      return queue.executeInQueue('query', cacheKey as any, _query, priority, opt);
     } else {
-      const stream = queue.setQueryStream(cacheKey, aliasNameToMember);
+      const stream = queue.setQueryStream(cacheKey as any, aliasNameToMember);
       // we don't want to handle error here as we want it to bubble up
       // to the api gateway
-      queue.executeInQueue('stream', cacheKey, _query, priority, opt);
+      queue.executeInQueue('stream', cacheKey as any, _query, priority, opt);
       return stream;
     }
   }