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;
+  requestTimeout: number;
 };
 
 const FireboltTypeToGeneric: Record<string, string> = {
@@ -91,6 +92,7 @@ export class FireboltDriver extends BaseDriver implements DriverInterface {
 
     this.config = {
       readOnly: true,
+      requestTimeout: 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.requestTimeout },
         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.requestTimeout },
         parameters,
         response: { hydrateRow: this.hydrateRow }
       });

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

@@ -1,6 +1,7 @@
 import { DriverTests } from '@cubejs-backend/testing-shared';
 
 import { FireboltDriver } from '../src';
+import { getEnv } from '@cubejs-backend/shared';
 
 describe('FireboltDriver', () => {
   let tests: DriverTests;
@@ -22,4 +23,22 @@ describe('FireboltDriver', () => {
   test('stream', async () => {
     await tests.testStream();
   });
+
+  test('query should fail on timeout', async () => {
+    const slowQuery = 'SELECT checksum(*) FROM GENERATE_SERIES(1, 1000000000000)';
+    // Set query timeout to 2 seconds
+    process.env.CUBEJS_DB_QUERY_TIMEOUT = '2';
+    const driver = new FireboltDriver({});
+
+    let timedOut = false;
+    try {
+      await driver.query(slowQuery);
+    } catch (e) {
+      if (String(e).toLowerCase().includes('timeout expired (2000 ms)')) {
+        timedOut = true;
+      }
+    }
+
+    expect(timedOut).toBe(true);
+  }, 10000);
 });