From a5fe556bbd986ebc54fdfaa9f11dc8df72ee83dc Mon Sep 17 00:00:00 2001 From: Konstantin Burkalev Date: Fri, 6 Jun 2025 11:41:33 +0300 Subject: [PATCH] feat(prestodb-driver, trino-driver): Support dbUseSelectTestConnection flag With this flag set to true - use `SELECT 1` query for testing connection --- .github/actions/integration/trino.sh | 15 ++++ .github/workflows/push.yml | 2 +- packages/cubejs-backend-shared/src/env.ts | 27 ++++++ .../src/PrestoDriver.ts | 13 +++ .../test/presto-driver.test.ts | 4 +- .../cubejs-trino-driver/docker-compose.yml | 13 +++ packages/cubejs-trino-driver/package.json | 8 +- .../cubejs-trino-driver/src/TrinoDriver.ts | 5 ++ .../test/trino-driver.test.ts | 85 +++++++++++++++++++ 9 files changed, 168 insertions(+), 4 deletions(-) create mode 100755 .github/actions/integration/trino.sh create mode 100644 packages/cubejs-trino-driver/docker-compose.yml create mode 100644 packages/cubejs-trino-driver/test/trino-driver.test.ts diff --git a/.github/actions/integration/trino.sh b/.github/actions/integration/trino.sh new file mode 100755 index 0000000000000..4371338dbace8 --- /dev/null +++ b/.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::" diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index b1d622e22c8a3..1b25207fd35e0 100644 --- a/.github/workflows/push.yml +++ b/.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' ] diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts index 94bbdea86ad3f..411240c1e247c 100644 --- a/packages/cubejs-backend-shared/src/env.ts +++ b/packages/cubejs-backend-shared/src/env.ts @@ -337,6 +337,33 @@ const variables: Record 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 */ diff --git a/packages/cubejs-prestodb-driver/src/PrestoDriver.ts b/packages/cubejs-prestodb-driver/src/PrestoDriver.ts index fe50e631413d8..fc600e5dd45e8 100644 --- a/packages/cubejs-prestodb-driver/src/PrestoDriver.ts +++ b/packages/cubejs-prestodb-driver/src/PrestoDriver.ts @@ -68,6 +68,8 @@ export class PrestoDriver extends BaseDriver implements DriverInterface { protected client: any; + protected useSelectTestConnection: boolean; + /** * Class constructor. */ @@ -86,6 +88,8 @@ export class PrestoDriver extends BaseDriver implements DriverInterface { 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 }), @@ -109,6 +113,10 @@ export class PrestoDriver extends BaseDriver implements DriverInterface { } public async testConnection(): Promise { + 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 @@ -122,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 { return > this.queryPromised(this.prepareQueryWithParams(query, values), false); } diff --git a/packages/cubejs-prestodb-driver/test/presto-driver.test.ts b/packages/cubejs-prestodb-driver/test/presto-driver.test.ts index 1c32b49fc77ca..77646035f35bc 100644 --- a/packages/cubejs-prestodb-driver/test/presto-driver.test.ts +++ b/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\''); }); }); }); diff --git a/packages/cubejs-trino-driver/docker-compose.yml b/packages/cubejs-trino-driver/docker-compose.yml new file mode 100644 index 0000000000000..5a8497cc5ee32 --- /dev/null +++ b/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 diff --git a/packages/cubejs-trino-driver/package.json b/packages/cubejs-trino-driver/package.json index c87c1c3e2b188..0d6990eef1e09 100644 --- a/packages/cubejs-trino-driver/package.json +++ b/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.21" + "@cubejs-backend/linter": "1.3.21", + "@types/jest": "^29", + "jest": "^29", + "testcontainers": "^10.13.0", + "typescript": "~5.2.2" }, "eslintConfig": { "extends": "../cubejs-linter" diff --git a/packages/cubejs-trino-driver/src/TrinoDriver.ts b/packages/cubejs-trino-driver/src/TrinoDriver.ts index 45c880bf88e96..800673f8dcf77 100644 --- a/packages/cubejs-trino-driver/src/TrinoDriver.ts +++ b/packages/cubejs-trino-driver/src/TrinoDriver.ts @@ -11,7 +11,12 @@ export class TrinoDriver extends PrestoDriver { return PrestodbQuery; } + // eslint-disable-next-line consistent-return public override async testConnection(): Promise { + 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`; diff --git a/packages/cubejs-trino-driver/test/trino-driver.test.ts b/packages/cubejs-trino-driver/test/trino-driver.test.ts new file mode 100644 index 0000000000000..9c54df14dce4b --- /dev/null +++ b/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\''); + }); + }); +});