feat: Re-use connection to Cube Store from externalDriver (#5993)

Author: ovr

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

@@ -5,9 +5,20 @@ import { CubeStoreDriver } from './CubeStoreDriver';
 
 export class CubeStoreCacheDriver implements CacheDriverInterface {
   public constructor(
-    protected readonly connection: CubeStoreDriver
+    protected connectionFactory: () => Promise<CubeStoreDriver>,
   ) {}
 
+  protected connection: CubeStoreDriver | null = null;
+
+  protected async getConnection(): Promise<CubeStoreDriver> {
+    if (this.connection) {
+      return this.connection;
+    }
+
+    // eslint-disable-next-line no-return-assign
+    return this.connection = await this.connectionFactory();
+  }
+
   public withLock = (
     key: string,
     cb: () => MaybeCancelablePromise<any>,
@@ -18,11 +29,13 @@ export class CubeStoreCacheDriver implements CacheDriverInterface {
       return false;
     }
 
-    const rows = await this.connection.query('CACHE SET NX TTL ? ? ?', [expiration, key, '1']);
+    const connection = (await this.getConnection());
+
+    const rows = await connection.query('CACHE SET NX TTL ? ? ?', [expiration, key, '1']);
     if (rows && rows.length === 1 && rows[0]?.success === 'true') {
       if (tkn.isCanceled()) {
         if (freeAfter) {
-          await this.connection.query('CACHE REMOVE ?', [
+          await connection.query('CACHE REMOVE ?', [
             key
           ]);
         }
@@ -34,7 +47,7 @@ export class CubeStoreCacheDriver implements CacheDriverInterface {
         await tkn.with(cb());
       } finally {
         if (freeAfter) {
-          await this.connection.query('CACHE REMOVE ?', [
+          await connection.query('CACHE REMOVE ?', [
             key
           ]);
         }
@@ -47,7 +60,7 @@ export class CubeStoreCacheDriver implements CacheDriverInterface {
   });
 
   public async get(key: string) {
-    const rows = await this.connection.query('CACHE GET ?', [
+    const rows = await (await this.getConnection()).query('CACHE GET ?', [
       key
     ]);
     if (rows && rows.length === 1) {
@@ -59,7 +72,7 @@ export class CubeStoreCacheDriver implements CacheDriverInterface {
 
   public async set(key: string, value, expiration) {
     const strValue = JSON.stringify(value);
-    await this.connection.query('CACHE SET TTL ? ? ?', [expiration, key, strValue]);
+    await (await this.getConnection()).query('CACHE SET TTL ? ? ?', [expiration, key, strValue]);
 
     return {
       key,
@@ -68,13 +81,13 @@ export class CubeStoreCacheDriver implements CacheDriverInterface {
   }
 
   public async remove(key: string) {
-    await this.connection.query('CACHE REMOVE ?', [
+    await (await this.getConnection()).query('CACHE REMOVE ?', [
       key
     ]);
   }
 
   public async keysStartingWith(prefix: string) {
-    const rows = await this.connection.query('CACHE KEYS ?', [
+    const rows = await (await this.getConnection()).query('CACHE KEYS ?', [
       prefix
     ]);
     return rows.map((row) => row.key);
@@ -85,6 +98,6 @@ export class CubeStoreCacheDriver implements CacheDriverInterface {
   }
 
   public async testConnection(): Promise<void> {
-    return this.connection.testConnection();
+    return (await this.getConnection()).testConnection();
   }
 }

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

@@ -117,7 +117,7 @@ export interface QueryCacheOptions {
     heartBeatInterval?: number;
   }>;
   redisPool?: any;
-  cubeStoreDriver?: CubeStoreDriver,
+  cubeStoreDriverFactory?: () => Promise<CubeStoreDriver>,
   continueWaitTimeout?: number;
   cacheAndQueueDriver?: CacheAndQueryDriverType;
   maxInMemoryCacheEntries?: number;
@@ -148,7 +148,7 @@ export class QueryCache {
         break;
       case 'cubestore':
         this.cacheDriver = new CubeStoreCacheDriver(
-          options.cubeStoreDriver || new CubeStoreDriver({})
+          options.cubeStoreDriverFactory || (async () => new CubeStoreDriver({}))
         );
         break;
       default:

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

@@ -1,5 +1,6 @@
 import R from 'ramda';
 import { getEnv } from '@cubejs-backend/shared';
+import { CubeStoreDriver } from '@cubejs-backend/cubestore-driver';
 
 import { QueryCache, QueryBody, TempTable } from './QueryCache';
 import { PreAggregations, PreAggregationDescription, getLastUpdatedAtTimestamp } from './PreAggregations';
@@ -59,13 +60,19 @@ export class QueryOrchestrator {
       throw new Error('Only \'redis\', \'memory\' or \'cubestore\' are supported for cacheAndQueueDriver option');
     }
 
+    const { externalDriverFactory, continueWaitTimeout, skipExternalCacheAndQueue } = options;
+
     const redisPool = cacheAndQueueDriver === 'redis' ? new RedisPool(options.redisPoolOptions) : undefined;
     this.redisPool = redisPool;
 
-    // TODO: Re-use connection from external database
-    const cubeStoreDriver = undefined;
+    const cubeStoreDriverFactory = cacheAndQueueDriver === 'cubestore' ? async () => {
+      const externalDriver = await externalDriverFactory();
+      if (externalDriver instanceof CubeStoreDriver) {
+        return externalDriver;
+      }
 
-    const { externalDriverFactory, continueWaitTimeout, skipExternalCacheAndQueue } = options;
+      throw new Error('It`s not possible to use CubeStore as queue & cache driver without using it as external');
+    } : undefined;
 
     this.queryCache = new QueryCache(
       this.redisPrefix,
@@ -75,7 +82,7 @@ export class QueryOrchestrator {
         externalDriverFactory,
         cacheAndQueueDriver,
         redisPool,
-        cubeStoreDriver,
+        cubeStoreDriverFactory,
         continueWaitTimeout,
         skipExternalCacheAndQueue,
         ...options.queryCacheOptions,

packages/cubejs-query-orchestrator/test/integration/cubestore/QueryCacheCubestore.test.ts

@@ -3,7 +3,7 @@ import { QueryCacheTest } from '../../unit/QueryCache.abstract';
 
 let beforeAll;
 let afterAll;
-let cubeStoreDriver = new CubeStoreDriver({});
+let cubeStoreDriverFactory = async () => new CubeStoreDriver({});
 
 if ((process.env.CUBEJS_TESTING_CUBESTORE_AUTO_PROVISIONING || 'true') === 'true') {
   const cubeStoreHandler = new CubeStoreHandler({
@@ -22,14 +22,14 @@ if ((process.env.CUBEJS_TESTING_CUBESTORE_AUTO_PROVISIONING || 'true') === 'true
     await cubeStoreHandler.acquire();
   };
   afterAll = async () => cubeStoreHandler.release(true);
-  cubeStoreDriver = new CubeStoreDevDriver(cubeStoreHandler);
+  cubeStoreDriverFactory = async () => new CubeStoreDevDriver(cubeStoreHandler);
 }
 
 QueryCacheTest(
   'CubeStore Cache Driver',
   {
     cacheAndQueueDriver: 'cubestore',
-    cubeStoreDriver,
+    cubeStoreDriverFactory,
     beforeAll,
     afterAll
   }