fix: concurrency lock driver instance & release on error

cube-js · Apr 23, 2021 · dd60f5d · dd60f5d
1 parent 59bc127
commit dd60f5d
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 27 deletions.
diff --git a/packages/cubejs-cli/src/command/generate.ts b/packages/cubejs-cli/src/command/generate.ts
@@ -64,10 +64,7 @@ const generate = async (options) => {
   });
 
   const dbSchema = await driver.tablesSchema();
-
-  if ((<any>driver).release) {
-    await (<any>driver).release();
-  }
+  await driver.release();
 
   logStage('Generating schema files');
   const ScaffoldingTemplate = requireFromPackage<any>(

diff --git a/packages/cubejs-query-orchestrator/src/driver/BaseDriver.js b/packages/cubejs-query-orchestrator/src/driver/BaseDriver.js
@@ -393,6 +393,13 @@ export class BaseDriver {
     }
   }
 
+  /**
+   * @public
+   */
+  async release() {
+    // override, if it's needed
+  }
+
   capabilities() {
     return {};
   }

diff --git a/packages/cubejs-server-core/src/core/server.ts b/packages/cubejs-server-core/src/core/server.ts
@@ -555,36 +555,67 @@ export class CubejsServerCore {
       return this.orchestratorStorage.get(orchestratorId);
     }
 
-    const driverPromise = {};
-    let externalPreAggregationsDriverPromise;
+    const driverPromise: Record<string, Promise<BaseDriver>> = {};
+    let externalPreAggregationsDriverPromise: Promise<BaseDriver>|null = null;
 
     const orchestratorApi = this.createOrchestratorApi({
-      getDriver: async (dataSource) => {
-        if (!driverPromise[dataSource || 'default']) {
-          orchestratorApi.addDataSeenSource(dataSource);
-          const driver = await this.options.driverFactory({ ...context, dataSource });
-          if (driver.setLogger) {
-            driver.setLogger(this.logger);
-          }
-          driverPromise[dataSource || 'default'] = driver.testConnection().then(() => driver).catch(e => {
-            driverPromise[dataSource || 'default'] = null;
-            throw e;
-          });
+      getDriver: async (dataSource = 'default') => {
+        if (driverPromise[dataSource]) {
+          return driverPromise[dataSource];
         }
-        return driverPromise[dataSource || 'default'];
+
+        // eslint-disable-next-line no-return-assign
+        return driverPromise[dataSource] = (async () => {
+          let driver: BaseDriver|null = null;
+
+          try {
+            driver = await this.options.driverFactory({ ...context, dataSource });
+            if (driver.setLogger) {
+              driver.setLogger(this.logger);
+            }
+
+            await driver.testConnection();
+
+            return driver;
+          } catch (e) {
+            driverPromise[dataSource] = null;
+
+            if (driver) {
+              await driver.release();
+            }
+
+            throw e;
+          }
+        })();
       },
       getExternalDriverFactory: this.options.externalDriverFactory && (async () => {
-        if (!externalPreAggregationsDriverPromise) {
-          const driver = await this.options.externalDriverFactory(context);
-          if (driver.setLogger) {
-            driver.setLogger(this.logger);
-          }
-          externalPreAggregationsDriverPromise = driver.testConnection().then(() => driver).catch(e => {
+        if (externalPreAggregationsDriverPromise) {
+          return externalPreAggregationsDriverPromise;
+        }
+
+        // eslint-disable-next-line no-return-assign
+        return externalPreAggregationsDriverPromise = (async () => {
+          let driver: BaseDriver|null = null;
+
+          try {
+            driver = await this.options.externalDriverFactory(context);
+            if (driver.setLogger) {
+              driver.setLogger(this.logger);
+            }
+
+            await driver.testConnection();
+
+            return driver;
+          } catch (e) {
             externalPreAggregationsDriverPromise = null;
+
+            if (driver) {
+              await driver.release();
+            }
+
             throw e;
-          });
-        }
-        return externalPreAggregationsDriverPromise;
+          }
+        })();
       }),
       redisPrefix: orchestratorId,
       orchestratorOptions: this.orchestratorOptions(context)