feat(prestodb-driver, trino-driver): Support dbUseSelectTestConnection flag

.github/actions/integration/trino.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+set -eo pipefail
+
+# Debug log for test containers
+export DEBUG=testcontainers
+
+export TEST_PRESTO_VERSION=341-SNAPSHOT
+export TEST_PGSQL_VERSION=12.4
+
+echo "::group::Trino ${TEST_PRESTO_VERSION} with PostgreSQL ${TEST_PGSQL_VERSION}"
+docker pull lewuathe/presto-coordinator:${TEST_PRESTO_VERSION}
+docker pull lewuathe/presto-worker:${TEST_PRESTO_VERSION}
+docker pull postgres:${TEST_PGSQL_VERSION}
+yarn lerna run --concurrency 1 --stream --no-prefix integration:trino
+echo "::endgroup::"

.github/workflows/push.yml

@@ -412,7 +412,7 @@ jobs:
       matrix:
         node-version: [22.x]
         db: [
-            'athena', 'bigquery', 'snowflake',
+            'athena', 'bigquery', 'snowflake', 'trino',
             'clickhouse', 'druid', 'elasticsearch', 'mssql', 'mysql', 'postgres', 'prestodb',
             'mysql-aurora-serverless', 'crate', 'mongobi', 'firebolt', 'dremio', 'vertica'
         ]

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

@@ -337,6 +337,33 @@ const variables: Record<string, (...args: any) => any> = {
     ]
   ),
 
+  /**
+   * Use `SELECT 1` query for testConnection.
+   * It might be used in any driver where there is a specific testConnection
+   * like a REST call, but for some reason it's not possible to use it in
+   * deployment environment.
+   */
+  dbUseSelectTestConnection: ({
+    dataSource,
+  }: {
+    dataSource: string,
+  }) => {
+    const val = process.env[
+      keyByDataSource('CUBEJS_DB_USE_SELECT_TEST_CONNECTION', dataSource)
+    ] || 'false';
+    if (val.toLocaleLowerCase() === 'true') {
+      return true;
+    } else if (val.toLowerCase() === 'false') {
+      return false;
+    } else {
+      throw new TypeError(
+        `The ${
+          keyByDataSource('CUBEJS_DB_USE_SELECT_TEST_CONNECTION', dataSource)
+        } must be either 'true' or 'false'.`
+      );
+    }
+  },
+
   /**
    * Kafka host for direct downloads from ksqlDb
    */
@@ -1798,7 +1825,7 @@ const variables: Record<string, (...args: any) => any> = {
     return [];
   },
   /** ***************************************************************
-   * Presto Driver                                                  *
+   * Presto/Trino Driver                                                  *
    **************************************************************** */
 
   /**
@@ -1814,12 +1841,25 @@ const variables: Record<string, (...args: any) => any> = {
     ]
   ),
 
+  /**
+   * Presto/Trino Auth Token
+   */
+  prestoAuthToken: ({
+    dataSource,
+  }: {
+    dataSource: string,
+  }) => (
+    process.env[
+      keyByDataSource('CUBEJS_DB_PRESTO_AUTH_TOKEN', dataSource)
+    ]
+  ),
+
   /** ***************************************************************
    * Pinot Driver                                                  *
    **************************************************************** */
 
   /**
-   * Pinot / Startree Auth Token
+   * Pinot/Startree Auth Token
    */
   pinotAuthToken: ({
     dataSource,

packages/cubejs-prestodb-driver/src/PrestoDriver.ts

@@ -42,6 +42,8 @@ export type PrestoDriverConfiguration = PrestoDriverExportBucket & {
   schema?: string;
   user?: string;
   // eslint-disable-next-line camelcase
+  custom_auth?: string;
+  // eslint-disable-next-line camelcase
   basic_auth?: { user: string, password: string };
   ssl?: string | TLSConnectionOptions;
   dataSource?: string;
@@ -66,6 +68,8 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
 
   protected client: any;
 
+  protected useSelectTestConnection: boolean;
+
   /**
    * Class constructor.
    */
@@ -76,6 +80,16 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
       config.dataSource ||
       assertDataSource('default');
 
+    const dbUser = getEnv('dbUser', { dataSource });
+    const dbPassword = getEnv('dbPass', { dataSource });
+    const authToken = getEnv('prestoAuthToken', { dataSource });
+
+    if (authToken && dbPassword) {
+      throw new Error('Both user/password and auth token are set. Please remove password or token.');
+    }
+
+    this.useSelectTestConnection = getEnv('dbUseSelectTestConnection', { dataSource });
+
     this.config = {
       host: getEnv('dbHost', { dataSource }),
       port: getEnv('dbPort', { dataSource }),
@@ -85,13 +99,9 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
       schema:
         getEnv('dbName', { dataSource }) ||
         getEnv('dbSchema', { dataSource }),
-      user: getEnv('dbUser', { dataSource }),
-      basic_auth: getEnv('dbPass', { dataSource })
-        ? {
-          user: getEnv('dbUser', { dataSource }),
-          password: getEnv('dbPass', { dataSource }),
-        }
-        : undefined,
+      user: dbUser,
+      ...(authToken ? { custom_auth: `Bearer ${authToken}` } : {}),
+      ...(dbPassword ? { basic_auth: { user: dbUser, password: dbPassword } } : {}),
       ssl: this.getSslOptions(dataSource),
       bucketType: getEnv('dbExportBucketType', { supported: ['gcs'], dataSource }),
       exportBucket: getEnv('dbExportBucket', { dataSource }),
@@ -103,6 +113,10 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
   }
 
   public async testConnection(): Promise<void> {
+    if (this.useSelectTestConnection) {
+      return this.testConnectionViaSelect();
+    }
+
     return new Promise((resolve, reject) => {
       // Get node list of presto cluster and return it.
       // @see https://prestodb.io/docs/current/rest/node.html
@@ -116,6 +130,11 @@ export class PrestoDriver extends BaseDriver implements DriverInterface {
     });
   }
 
+  protected async testConnectionViaSelect() {
+    const query = SqlString.format('SELECT 1', []);
+    await this.queryPromised(query, false);
+  }
+
   public query(query: string, values: unknown[]): Promise<any[]> {
     return <Promise<any[]>> this.queryPromised(this.prepareQueryWithParams(query, values), false);
   }

packages/cubejs-prestodb-driver/test/presto-driver.test.ts

@@ -78,8 +78,8 @@ describe('PrestoHouseDriver', () => {
   // eslint-disable-next-line func-names
   it('should test informationSchemaQuery', async () => {
     await doWithDriver(async (driver: any) => {
-      const informationSchemaQuery=driver.informationSchemaQuery();
-      expect(informationSchemaQuery).toContain("columns.table_schema = 'sf1'");
+      const informationSchemaQuery = driver.informationSchemaQuery();
+      expect(informationSchemaQuery).toContain('columns.table_schema = \'sf1\'');
     });
   });
 });

packages/cubejs-trino-driver/docker-compose.yml

@@ -0,0 +1,13 @@
+version: '2.2'
+
+services:
+  coordinator:
+    image: trinodb/trino
+    ports:
+      - "8080:8080"
+    container_name: "coordinator"
+    healthcheck:
+      test: "trino --execute 'SELECT 1' || exit 1"
+      interval: 10s
+      timeout: 5s
+      retries: 5

packages/cubejs-trino-driver/package.json

@@ -21,6 +21,8 @@
     "build": "rm -rf dist && npm run tsc",
     "tsc": "tsc",
     "watch": "tsc -w",
+    "integration": "jest dist/test",
+    "integration:trino": "jest dist/test",
     "lint": "eslint src/* --ext .ts",
     "lint:fix": "eslint --fix src/* --ext .ts"
   },
@@ -38,7 +40,11 @@
     "access": "public"
   },
   "devDependencies": {
-    "@cubejs-backend/linter": "1.3.19"
+    "@cubejs-backend/linter": "1.3.19",
+    "@types/jest": "^29",
+    "jest": "^29",
+    "testcontainers": "^10.13.0",
+    "typescript": "~5.2.2"
   },
   "eslintConfig": {
     "extends": "../cubejs-linter"

packages/cubejs-trino-driver/src/TrinoDriver.ts

@@ -11,13 +11,20 @@ export class TrinoDriver extends PrestoDriver {
     return PrestodbQuery;
   }
 
+  // eslint-disable-next-line consistent-return
   public override async testConnection(): Promise<void> {
-    const { host, port, ssl, basic_auth: basicAuth } = this.config;
+    if (this.useSelectTestConnection) {
+      return this.testConnectionViaSelect();
+    }
+
+    const { host, port, ssl, basic_auth: basicAuth, custom_auth: customAuth } = this.config;
     const protocol = ssl ? 'https' : 'http';
     const url = `${protocol}://${host}:${port}/v1/info`;
     const headers: Record<string, string> = {};
 
-    if (basicAuth) {
+    if (customAuth) {
+      headers.Authorization = customAuth;
+    } else if (basicAuth) {
       const { user, password } = basicAuth;
       const encoded = Buffer.from(`${user}:${password}`).toString('base64');
       headers.Authorization = `Basic ${encoded}`;

packages/cubejs-trino-driver/test/trino-driver.test.ts

@@ -0,0 +1,85 @@
+import { TrinoDriver } from '../src/TrinoDriver';
+
+const path = require('path');
+const { DockerComposeEnvironment, Wait } = require('testcontainers');
+
+describe('TrinoDriver', () => {
+  jest.setTimeout(6 * 60 * 1000);
+
+  let env: any;
+  let config: any;
+
+  const doWithDriver = async (callback: any) => {
+    const driver = new TrinoDriver(config);
+
+    await callback(driver);
+  };
+
+  // eslint-disable-next-line consistent-return,func-names
+  beforeAll(async () => {
+    const authOpts = {
+      basic_auth: {
+        user: 'presto',
+        password: ''
+      }
+    };
+
+    if (process.env.TEST_PRESTO_HOST) {
+      config = {
+        host: process.env.TEST_PRESTO_HOST || 'localhost',
+        port: process.env.TEST_PRESTO_PORT || '8080',
+        catalog: process.env.TEST_PRESTO_CATALOG || 'tpch',
+        schema: 'sf1',
+        ...authOpts
+      };
+
+      return;
+    }
+
+    const dc = new DockerComposeEnvironment(
+      path.resolve(path.dirname(__filename), '../../'),
+      'docker-compose.yml'
+    );
+
+    env = await dc
+      .withStartupTimeout(240 * 1000)
+      .withWaitStrategy('coordinator', Wait.forHealthCheck())
+      .up();
+
+    config = {
+      host: env.getContainer('coordinator').getHost(),
+      port: env.getContainer('coordinator').getMappedPort(8080),
+      catalog: 'tpch',
+      schema: 'sf1',
+      ...authOpts
+    };
+  });
+
+  // eslint-disable-next-line consistent-return,func-names
+  afterAll(async () => {
+    if (env) {
+      await env.down();
+    }
+  });
+
+  it('should construct', async () => {
+    await doWithDriver(() => {
+      //
+    });
+  });
+
+  // eslint-disable-next-line func-names
+  it('should test connection', async () => {
+    await doWithDriver(async (driver: any) => {
+      await driver.testConnection();
+    });
+  });
+
+  // eslint-disable-next-line func-names
+  it('should test informationSchemaQuery', async () => {
+    await doWithDriver(async (driver: any) => {
+      const informationSchemaQuery = driver.informationSchemaQuery();
+      expect(informationSchemaQuery).toContain('columns.table_schema = \'sf1\'');
+    });
+  });
+});