fix: Re-use external connection for CubeStore in queue (#6028)

cube-js · Jan 18, 2023 · ff724f7 · ff724f7
1 parent 0f4a431
commit ff724f7
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 3 deletions.
diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts
@@ -25,12 +25,14 @@ import {
   StreamOptions,
   UnloadOptions
 } from '@cubejs-backend/base-driver';
+import { CubeStoreDriver } from '@cubejs-backend/cubestore-driver';
 import { PreAggTableToTempTable, Query, QueryBody, QueryCache, QueryTuple, QueryWithParams } from './QueryCache';
 import { ContinueWaitError } from './ContinueWaitError';
 import { DriverFactory, DriverFactoryByDataSource } from './DriverFactory';
 import { QueryQueue } from './QueryQueue';
 import { LargeStreamWarning } from './StreamObjectsCounter';
 import { CacheAndQueryDriverType } from './QueryOrchestrator';
+import { RedisPool } from './RedisPool';
 
 /// Name of the inline table containing the lambda rows.
 export const LAMBDA_TABLE_PREFIX = 'lambda';
@@ -1883,7 +1885,8 @@ type PreAggregationsOptions = {
     orphanedTimeout?: number;
     heartBeatInterval?: number;
   }>;
-  redisPool?: any;
+  redisPool?: RedisPool;
+  cubeStoreDriverFactory?: () => Promise<CubeStoreDriver>;
   continueWaitTimeout?: number;
   cacheAndQueueDriver?: CacheAndQueryDriverType;
   skipExternalCacheAndQueue?: boolean;
@@ -2264,6 +2267,7 @@ export class PreAggregations {
           logger: this.logger,
           cacheAndQueueDriver: this.options.cacheAndQueueDriver,
           redisPool: this.options.redisPool,
+          cubeStoreDriverFactory: this.options.cubeStoreDriverFactory,
           // Centralized continueWaitTimeout that can be overridden in queueOptions
           continueWaitTimeout: this.options.continueWaitTimeout,
           ...(await this.options.queueOptions(dataSource)),
@@ -2310,6 +2314,7 @@ export class PreAggregations {
           logger: this.logger,
           cacheAndQueueDriver: this.options.cacheAndQueueDriver,
           redisPool: this.options.redisPool,
+          cubeStoreDriverFactory: this.options.cubeStoreDriverFactory,
           ...this.options.loadCacheQueueOptions
         }
       );

diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts
@@ -150,8 +150,12 @@ export class QueryCache {
         this.cacheDriver = new LocalCacheDriver();
         break;
       case 'cubestore':
+        if (!options.cubeStoreDriverFactory) {
+          throw new Error('cubeStoreDriverFactory is a required option for Cube Store cache driver');
+        }
+
         this.cacheDriver = new CubeStoreCacheDriver(
-          options.cubeStoreDriverFactory || (async () => new CubeStoreDriver({}))
+          options.cubeStoreDriverFactory
         );
         break;
       default:
@@ -498,6 +502,7 @@ export class QueryCache {
           logger: this.logger,
           cacheAndQueueDriver: this.options.cacheAndQueueDriver,
           redisPool: this.options.redisPool,
+          cubeStoreDriverFactory: this.options.cubeStoreDriverFactory,
           // Centralized continueWaitTimeout that can be overridden in queueOptions
           continueWaitTimeout: this.options.continueWaitTimeout,
           ...(await this.options.queueOptions(dataSource)),
@@ -546,6 +551,7 @@ export class QueryCache {
           logger: this.logger,
           cacheAndQueueDriver: this.options.cacheAndQueueDriver,
           redisPool: this.options.redisPool,
+          cubeStoreDriverFactory: this.options.cubeStoreDriverFactory,
           // Centralized continueWaitTimeout that can be overridden in queueOptions
           continueWaitTimeout: this.options.continueWaitTimeout,
           skipQueue: this.options.skipExternalCacheAndQueue,

diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts
@@ -98,6 +98,7 @@ export class QueryOrchestrator {
         externalDriverFactory,
         cacheAndQueueDriver,
         redisPool,
+        cubeStoreDriverFactory,
         continueWaitTimeout,
         skipExternalCacheAndQueue,
         ...options.preAggregationsOptions,

diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.js b/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.js
@@ -22,8 +22,12 @@ function factoryQueueDriver(cacheAndQueueDriver, queueDriverOptions) {
     case 'memory':
       return new LocalQueueDriver(queueDriverOptions);
     case 'cubestore':
+      if (!queueDriverOptions.cubeStoreDriverFactory) {
+        throw new Error('cubeStoreDriverFactory is a required option for Cube Store queue driver');
+      }
+
       return new CubeStoreQueueDriver(
-        queueDriverOptions.cubeStoreDriverFactory || (async () => new CubeStoreDriver({})),
+        queueDriverOptions.cubeStoreDriverFactory,
         queueDriverOptions
       );
     default: