feat: Introduce single unified CUBEJS_DB_QUERY_TIMEOUT env variable to set all various variables that control database query timeouts (cube-js#3864)

* feat: Introduce single unified CUBEJS_DB_QUERY_TIMEOUT env variable to control all various variables that control database query timeouts

* chore: Introduce single unified CUBEJS_DB_QUERY_TIMEOUT env variable to control all various variables that control database query timeouts -- test fixes

Author: paveltiunov

packages/cubejs-athena-driver/driver/AthenaDriver.js

@@ -18,7 +18,7 @@ class AthenaDriver extends BaseDriver {
       S3OutputLocation: process.env.CUBEJS_AWS_S3_OUTPUT_LOCATION,
       workGroup: process.env.CUBEJS_AWS_ATHENA_WORKGROUP || 'primary',
       ...config,
-      pollTimeout: (config.pollTimeout || getEnv('dbPollTimeout')) * 1000,
+      pollTimeout: (config.pollTimeout || getEnv('dbPollTimeout') || getEnv('dbQueryTimeout')) * 1000,
       pollMaxInterval: (config.pollMaxInterval || getEnv('dbPollMaxInterval')) * 1000,
     };
 

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

@@ -11,7 +11,7 @@ export class InvalidConfiguration extends Error {
 export function convertTimeStrToMs(
   input: string,
   envName: string,
-  description: string = 'Must be number (in seconds) or string in time format (1s, 1m, 1h).',
+  description: string = 'Must be a number in seconds or duration string (1s, 1m, 1h).',
 ) {
   if (/^\d+$/.test(input)) {
     return parseInt(input, 10);
@@ -126,8 +126,16 @@ const variables: Record<string, (...args: any) => any> = {
   preAggregationsSchema: () => get('CUBEJS_PRE_AGGREGATIONS_SCHEMA')
     .asString(),
   dbPollTimeout: () => {
-    const value = process.env.CUBEJS_DB_POLL_TIMEOUT || '15m';
-    return convertTimeStrToMs(value, 'CUBEJS_DB_POLL_TIMEOUT');
+    const value = process.env.CUBEJS_DB_POLL_TIMEOUT;
+    if (value) {
+      return convertTimeStrToMs(value, 'CUBEJS_DB_POLL_TIMEOUT');
+    } else {
+      return null;
+    }
+  },
+  dbQueryTimeout: () => {
+    const value = process.env.CUBEJS_DB_QUERY_TIMEOUT || '10m';
+    return convertTimeStrToMs(value, 'CUBEJS_DB_QUERY_TIMEOUT');
   },
   dbPollMaxInterval: () => {
     const value = process.env.CUBEJS_DB_POLL_MAX_INTERVAL || '5s';

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

@@ -12,7 +12,7 @@ test('convertTimeStrToMs', () => {
 
 test('convertTimeStrToMs(exception)', () => {
   expect(() => convertTimeStrToMs('', 'VARIABLE_ENV')).toThrowError(
-    `Value "" is not valid for VARIABLE_ENV. Must be number (in seconds) or string in time format (1s, 1m, 1h).`
+    `Value "" is not valid for VARIABLE_ENV. Must be a number in seconds or duration string (1s, 1m, 1h).`
   );
 });
 
@@ -68,8 +68,6 @@ describe('getEnv', () => {
   });
 
   test('dbPollTimeout', () => {
-    expect(getEnv('dbPollTimeout')).toBe(15 * 60);
-
     process.env.CUBEJS_DB_POLL_TIMEOUT = '1m';
     expect(getEnv('dbPollTimeout')).toBe(60);
   });

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

@@ -45,7 +45,7 @@ export class BigQueryDriver extends BaseDriver implements DriverInterface {
       exportBucket: getEnv('dbExportBucket') || process.env.CUBEJS_DB_BQ_EXPORT_BUCKET,
       location: getEnv('bigQueryLocation'),
       ...config,
-      pollTimeout: (config.pollTimeout || getEnv('dbPollTimeout')) * 1000,
+      pollTimeout: (config.pollTimeout || getEnv('dbPollTimeout') || getEnv('dbQueryTimeout')) * 1000,
       pollMaxInterval: (config.pollMaxInterval || getEnv('dbPollMaxInterval')) * 1000,
     };
 

packages/cubejs-dremio-driver/driver/DremioDriver.js

@@ -27,7 +27,7 @@ class DremioDriver extends BaseDriver {
       database: config.database || process.env.CUBEJS_DB_NAME,
       ssl: config.ssl || process.env.CUBEJS_DB_SSL,
       ...config,
-      pollTimeout: (config.pollTimeout || getEnv('dbPollTimeout')) * 1000,
+      pollTimeout: (config.pollTimeout || getEnv('dbPollTimeout') || getEnv('dbQueryTimeout')) * 1000,
       pollMaxInterval: (config.pollMaxInterval || getEnv('dbPollMaxInterval')) * 1000,
     };
 

packages/cubejs-postgres-driver/package.json

@@ -28,6 +28,7 @@
   },
   "dependencies": {
     "@cubejs-backend/query-orchestrator": "^0.29.12",
+    "@cubejs-backend/shared": "^0.29.12",
     "@types/pg": "^8.6.0",
     "@types/pg-query-stream": "^1.0.3",
     "moment": "^2.24.0",

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

@@ -1,6 +1,7 @@
 import { types, Pool, PoolConfig, PoolClient, FieldDef } from 'pg';
 // eslint-disable-next-line import/no-extraneous-dependencies
 import { TypeId, TypeFormat } from 'pg-types';
+import { getEnv } from '@cubejs-backend/shared';
 import * as moment from 'moment';
 import {
   BaseDriver,
@@ -80,6 +81,7 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
 
     this.config = {
       ...this.getInitialConfiguration(),
+      executionTimeout: getEnv('dbQueryTimeout'),
       ...config,
     };
   }

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

@@ -1364,7 +1364,7 @@ export class PreAggregations {
 
     this.externalDriverFactory = options.externalDriverFactory;
     this.structureVersionPersistTime = options.structureVersionPersistTime || 60 * 60 * 24 * 30;
-    this.usedTablePersistTime = options.usedTablePersistTime || 600;
+    this.usedTablePersistTime = options.usedTablePersistTime || getEnv('dbQueryTimeout');
     this.externalRefresh = options.externalRefresh;
     this.getQueueEventsBus = options.getQueueEventsBus;
   }

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

@@ -1,4 +1,5 @@
 import R from 'ramda';
+import { getEnv } from '@cubejs-backend/shared';
 
 import { TimeoutError } from './TimeoutError';
 import { ContinueWaitError } from './ContinueWaitError';
@@ -10,7 +11,7 @@ export class QueryQueue {
     this.redisQueuePrefix = redisQueuePrefix;
     this.concurrency = options.concurrency || 2;
     this.continueWaitTimeout = options.continueWaitTimeout || 5;
-    this.executionTimeout = options.executionTimeout || 600;
+    this.executionTimeout = options.executionTimeout || getEnv('dbQueryTimeout');
     this.orphanedTimeout = options.orphanedTimeout || 120;
     this.heartBeatInterval = options.heartBeatInterval || 30;
     this.sendProcessMessageFn = options.sendProcessMessageFn || ((queryKey) => { this.processQuery(queryKey); });

packages/cubejs-redshift-driver/src/RedshiftDriver.ts

@@ -100,7 +100,7 @@ export class RedshiftDriver extends PostgresDriver<RedshiftDriverConfiguration>
         .join(' ');
 
       await this.prepareConnection(conn, {
-        executionTimeout: 600000,
+        executionTimeout: this.config.executionTimeout ? this.config.executionTimeout * 1000 : 600000,
       });
 
       let unloadTotalRows: number | null = null;