Pass CUBEJS_DB_QUERY_TIMEOUT to Firebolt driver

Author: alexanderchiu

packages/cubejs-firebolt-driver/src/FireboltDriver.ts

@@ -30,6 +30,7 @@ export type FireboltDriverConfiguration = {
   readOnly?: boolean;
   apiEndpoint?: string;
   connection: ConnectionOptions;
+  requestTimeoutEnv: number;
 };
 
 const FireboltTypeToGeneric: Record<string, string> = {
@@ -91,6 +92,7 @@ export class FireboltDriver extends BaseDriver implements DriverInterface {
 
     this.config = {
       readOnly: true,
+      requestTimeoutEnv: getEnv('dbQueryTimeout') * 1000,
       apiEndpoint:
         getEnv('fireboltApiEndpoint', { dataSource }) || 'api.app.firebolt.io',
       ...config,
@@ -225,7 +227,7 @@ export class FireboltDriver extends BaseDriver implements DriverInterface {
       const connection = await this.getConnection();
 
       const statement = await connection.execute(query, {
-        settings: { output_format: OutputFormat.JSON },
+        settings: { output_format: OutputFormat.JSON, statement_timeout: this.config.requestTimeoutEnv },
         parameters,
         response: { hydrateRow: this.hydrateRow }
       });
@@ -275,7 +277,7 @@ export class FireboltDriver extends BaseDriver implements DriverInterface {
       const connection = await this.getConnection();
 
       const statement = await connection.execute(query, {
-        settings: { output_format: OutputFormat.JSON },
+        settings: { output_format: OutputFormat.JSON, statement_timeout: this.config.requestTimeoutEnv },
         parameters,
         response: { hydrateRow: this.hydrateRow }
       });

packages/cubejs-firebolt-driver/test/FireboltDriver.test.ts

@@ -22,4 +22,20 @@ describe('FireboltDriver', () => {
   test('stream', async () => {
     await tests.testStream();
   });
+
+  test('query should fail on timeout', async () => {
+    const slowQuery = 'SELECT checksum(*) FROM GENERATE_SERIES(1, 1000000000000)';
+    const driver = new FireboltDriver({});
+
+    let timedOut = false;
+    try {
+      await driver.query(slowQuery);
+    } catch (e) {
+      if (String(e).toLowerCase().includes('Query timeout expired')) {
+        timedOut = true;
+      }
+    }
+
+    expect(timedOut).toBe(true);
+  }, 10000);
 });

packages/cubejs-firebolt-driver/test/test-env.js

@@ -4,6 +4,7 @@ const REQUIRED_ENV_VARS = [
   'CUBEJS_DB_NAME',
   'CUBEJS_FIREBOLT_ENGINE_NAME',
   'CUBEJS_FIREBOLT_ACCOUNT',
+  'CUBEJS_DB_QUERY_TIMEOUT',
 ];
 
 REQUIRED_ENV_VARS.forEach((key) => {