chore: fixes

Author: ovr

packages/cubejs-backend-shared/src/type-helpers.ts

@@ -1,5 +1,5 @@
 /**
- * This module export only type helpers for using it across Cube.js project
+ * This module exports only type helpers for using it across the Cube project
  */
 
 export type ResolveAwait<T> = T extends {
@@ -16,3 +16,6 @@ export type Required<T, K extends keyof T> = {
 };
 
 export type Optional<T, K extends keyof T> = Pick<Partial<T>, K> & Omit<T, K>;
+
+// <M extends Method<Class/Interface, M>>
+export type MethodName<T> = { [K in keyof T]: T[K] extends (...args: any[]) => any ? K : never }[keyof T];

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

@@ -12,7 +12,7 @@ import {
   DriverInterface, QueryKey,
 } from '@cubejs-backend/base-driver';
 
-import {QueryQueue, QueryQueueOptions} from './QueryQueue';
+import { QueryQueue, QueryQueueOptions } from './QueryQueue';
 import { ContinueWaitError } from './ContinueWaitError';
 import { LocalCacheDriver } from './LocalCacheDriver';
 import { DriverFactory, DriverFactoryByDataSource } from './DriverFactory';

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

@@ -6,9 +6,9 @@ import {
   QueryKeyHash,
   QueueId,
   QueryDef,
-  QueryStageStateResponse, QueueDriverOptions, QueryKeysTuple
+  QueryStageStateResponse
 } from '@cubejs-backend/base-driver';
-import { CubeStoreDriver, CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver';
+import { CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver';
 
 import { TimeoutError } from './TimeoutError';
 import { ContinueWaitError } from './ContinueWaitError';
@@ -20,11 +20,14 @@ export type QueryHandlerFn = (query: QueryDef, cancelHandler: CancelHandlerFn) =
 export type StreamHandlerFn = (query: QueryDef, stream: QueryStream) => Promise<unknown>;
 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 QueryQueueOptions = {
   cacheAndQueueDriver: string;
   logger: (message, event) => void;
-  sendCancelMessageFn?: (query, queueId) => Promise<void>;
-  sendProcessMessageFn?: (queryKey, queryId) => Promise<void>;
+  sendCancelMessageFn?: SendCancelMessageFn;
+  sendProcessMessageFn?: SendProcessMessageFn;
   cancelHandlers: Record<string, CancelHandlerFn>;
   queryHandlers: QueryHandlersMap;
   streamHandler?: StreamHandlerFn;
@@ -69,9 +72,9 @@ export class QueryQueue {
 
   protected heartBeatInterval: number;
 
-  protected readonly sendProcessMessageFn: any;
+  protected readonly sendProcessMessageFn: SendProcessMessageFn;
 
-  protected readonly sendCancelMessageFn: any;
+  protected readonly sendCancelMessageFn: SendCancelMessageFn;
 
   protected readonly queryHandlers: QueryHandlersMap;
 
@@ -485,7 +488,8 @@ export class QueryQueue {
           preAggregation: query.query?.preAggregation,
           addedToQueueTime: query.addedToQueueTime,
         });
-        await this.sendCancelMessageFn(query);
+
+        await this.sendCancelMessageFn(query, queueId);
       }
 
       return true;
@@ -501,7 +505,7 @@ export class QueryQueue {
    * @private
    * @returns {Promise<void>}
    */
-  async reconcileQueueImpl() {
+  protected async reconcileQueueImpl() {
     const queueConnection = await this.queueDriver.createConnection();
     try {
       const toCancel = await queueConnection.getQueriesToCancel();

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

@@ -1,17 +1,17 @@
-import { CubeStoreDriver, CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver';
+import { CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver';
 import crypto from 'crypto';
-import { createPromiseLock, pausePromise } from '@cubejs-backend/shared';
+import { createPromiseLock, MethodName, pausePromise } from '@cubejs-backend/shared';
 import { QueueDriverConnectionInterface, QueueDriverInterface, } from '@cubejs-backend/base-driver';
-import {LocalQueueDriver, QueryQueue, QueryQueueOptions} from '../../src';
+import { LocalQueueDriver, QueryQueue, QueryQueueOptions } from '../../src';
 
 export type QueryQueueTestOptions = Pick<QueryQueueOptions, 'cacheAndQueueDriver' | 'cubeStoreDriverFactory'> & {
   beforeAll?: () => Promise<void>,
   afterAll?: () => Promise<void>,
 };
 
 function patchQueueDriverConnectionForTrack(connection: QueueDriverConnectionInterface, counters: any): QueueDriverConnectionInterface {
-  function wrapAsyncMethod(methodName: string): any {
-    return async function (...args) {
+  function wrapAsyncMethod<M extends MethodName<QueueDriverConnectionInterface>>(methodName: M): any {
+    return async (...args: Parameters<QueueDriverConnectionInterface[M]>) => {
       if (!(methodName in counters.methods)) {
         counters.methods[methodName] = {
           started: 1,
@@ -21,7 +21,7 @@ function patchQueueDriverConnectionForTrack(connection: QueueDriverConnectionInt
         counters.methods[methodName].started++;
       }
 
-      const result = await connection[methodName](...args);
+      const result = await (connection[methodName] as any)(...args);
       counters.methods[methodName].finished++;
 
       return result;
@@ -122,6 +122,7 @@ export function QueryQueueBenchmark(name: string, options: QueryQueueTestOptions
         },
         cancelHandlers: {
           query: async (_query) => {
+            console.error('Cancel handler was called for query');
           },
         },
         continueWaitTimeout: 60 * 2,

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

@@ -4,7 +4,7 @@ import type { QueryKey, QueueDriverInterface } from '@cubejs-backend/base-driver
 import { pausePromise } from '@cubejs-backend/shared';
 import crypto from 'crypto';
 
-import {QueryQueue, QueryQueueOptions} from '../../src';
+import { QueryQueue, QueryQueueOptions } from '../../src';
 import { processUidRE } from '../../src/orchestrator/utils';
 
 export type QueryQueueTestOptions = Pick<QueryQueueOptions, 'cacheAndQueueDriver' | 'cubeStoreDriverFactory'> & {