feat(base-driver): add an option to skip connection test probes

feat(bigquery-driver): add an option to skip connection test probes
feat(aurora-serverless-driver): add an option to skip connection test probes
feat(postgres-driver): add an option to skip connection test probes

Author: KSDaemon

packages/cubejs-backend-shared/src/env.ts

@@ -479,6 +479,17 @@ const variables: Record<string, (...args: any) => any> = {
     return val;
   },
 
+  /**
+   * Flag to disable driver's test connection probes
+   */
+  dbDisableTestConnection: ({
+    dataSource
+  }: {
+    dataSource: string,
+  }) => (
+    !!process.env[keyByDataSource('CUBEJS_DB_DISABLE_TEST_CONNECTION', dataSource)]
+  ),
+
   /**
    * Database max pool size.
    */
@@ -500,7 +511,7 @@ const variables: Record<string, (...args: any) => any> = {
   ),
 
   /**
-   * Max polling interval. Currenly used in BigQuery and Databricks.
+   * Max polling interval. Currently used in BigQuery and Databricks.
    * TODO: clarify this env.
    */
   dbPollMaxInterval: ({
@@ -514,7 +525,7 @@ const variables: Record<string, (...args: any) => any> = {
   },
 
   /**
-   * Polling timeout. Currenly used in BigQuery, Dremio and Athena.
+   * Polling timeout. Currently used in BigQuery, Dremio and Athena.
    * TODO: clarify this env.
    */
   dbPollTimeout: ({
@@ -532,7 +543,7 @@ const variables: Record<string, (...args: any) => any> = {
   },
 
   /**
-   * Query timeout. Currenly used in BigQuery, Dremio, Postgres, Snowflake
+   * Query timeout. Currently used in BigQuery, Dremio, Postgres, Snowflake
    * and Athena drivers and the orchestrator (queues, pre-aggs). For the
    * orchestrator this variable did not split by the datasource.
    *

packages/cubejs-base-driver/src/BaseDriver.ts

@@ -171,10 +171,12 @@ const DbTypeValueMatcher: Record<string, ((v: any) => boolean)> = {
  * Base driver class.
  */
 export abstract class BaseDriver implements DriverInterface {
-  private testConnectionTimeoutValue = 10000;
+  private readonly testConnectionTimeoutValue: number = 10000;
 
   protected logger: any;
 
+  protected isTestConnectionDisabledFlag: boolean = false;
+
   /**
    * Class constructor.
    */
@@ -184,8 +186,14 @@ export abstract class BaseDriver implements DriverInterface {
      * request before determining it as not valid. Default - 10000 ms.
      */
     testConnectionTimeout?: number,
+
+    /**
+     * Flag for serverless DWHs to omit test connection probes.
+     */
+    isTestConnectionDisabled?: boolean,
   } = {}) {
     this.testConnectionTimeoutValue = _options.testConnectionTimeout || 10000;
+    this.isTestConnectionDisabledFlag = _options.isTestConnectionDisabled || false;
   }
 
   protected informationSchemaQuery() {
@@ -680,6 +688,10 @@ export abstract class BaseDriver implements DriverInterface {
     query.query = `SELECT * FROM (${query.query}) AS t LIMIT ${query.limit}`;
   }
 
+  public isTestConnectionDisabled(): boolean {
+    return !!this.isTestConnectionDisabledFlag;
+  }
+
   /**
    * Returns an array of signed AWS S3 URLs of the unloaded csv files.
    */

packages/cubejs-bigquery-driver/src/BigQueryDriver.ts

@@ -93,14 +93,15 @@ export class BigQueryDriver extends BaseDriver implements DriverInterface {
       testConnectionTimeout?: number,
     } = {}
   ) {
-    super({
-      testConnectionTimeout: config.testConnectionTimeout,
-    });
-
     const dataSource =
       config.dataSource ||
       assertDataSource('default');
 
+    super({
+      testConnectionTimeout: config.testConnectionTimeout,
+      isTestConnectionDisabled: getEnv('dbDisableTestConnection', { dataSource }),
+    });
+
     this.options = {
       scopes: [
         'https://www.googleapis.com/auth/bigquery',
@@ -155,6 +156,10 @@ export class BigQueryDriver extends BaseDriver implements DriverInterface {
   }
 
   public async testConnection() {
+    if (this.isTestConnectionDisabled()) {
+      return;
+    }
+
     await this.bigquery.query({
       query: 'SELECT ? AS number',
       params: ['1'],

packages/cubejs-mysql-aurora-serverless-driver/driver/AuroraServerlessMySqlDriver.js

@@ -33,14 +33,15 @@ class AuroraServerlessMySqlDriver extends BaseDriver {
    * Class constructor.
    */
   constructor(config = {}) {
-    super({
-      testConnectionTimeout: config.testConnectionTimeout,
-    });
-    
     const dataSource =
       config.dataSource ||
       assertDataSource('default');
-    
+
+    super({
+      testConnectionTimeout: config.testConnectionTimeout,
+      isTestConnectionDisabled: getEnv('dbDisableTestConnection', { dataSource }),
+    });
+
     this.config = {
       secretArn:
         config.secretArn ||
@@ -66,7 +67,11 @@ class AuroraServerlessMySqlDriver extends BaseDriver {
   }
 
   async testConnection() {
-    return this.dataApi.query('SELECT 1');
+    if (this.isTestConnectionDisabled()) {
+      return;
+    }
+
+    await this.dataApi.query('SELECT 1');
   }
 
   positionBindings(sql) {

packages/cubejs-postgres-driver/src/PostgresDriver.ts

@@ -111,14 +111,15 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
       testConnectionTimeout?: number,
     } = {}
   ) {
+    const dataSource =
+      config.dataSource ||
+      assertDataSource('default');
+
     super({
       testConnectionTimeout: config.testConnectionTimeout,
+      isTestConnectionDisabled: getEnv('dbDisableTestConnection', { dataSource }),
     });
 
-    const dataSource =
-      config.dataSource ||
-      assertDataSource('default');
-    
     this.pool = new Pool({
       idleTimeoutMillis: 30000,
       max:
@@ -146,9 +147,9 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
   }
 
   protected primaryKeysQuery(conditionString?: string): string | null {
-    return `SELECT 
+    return `SELECT
       columns.table_schema as ${this.quoteIdentifier('table_schema')},
-      columns.table_name as ${this.quoteIdentifier('table_name')}, 
+      columns.table_name as ${this.quoteIdentifier('table_name')},
       columns.column_name as ${this.quoteIdentifier('column_name')}
     FROM information_schema.table_constraints tc
     JOIN information_schema.constraint_column_usage AS ccu USING (constraint_schema, constraint_name)
@@ -225,6 +226,10 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
   }
 
   public async testConnection(): Promise<void> {
+    if (this.isTestConnectionDisabled()) {
+      return;
+    }
+
     try {
       await this.pool.query('SELECT $1::int AS number', ['1']);
     } catch (e) {