From 0de1978f86977483c9f7c0d1d74cfbd2d0a412be Mon Sep 17 00:00:00 2001 From: johancube Date: Fri, 11 Oct 2024 15:33:53 +0300 Subject: [PATCH 1/3] feat: init --- packages/cubejs-backend-shared/src/env.ts | 13 +++++++++ .../src/SnowflakeDriver.ts | 27 ++++++++++++++++--- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts index 570719c558ad9..ecc15837b31d0 100644 --- a/packages/cubejs-backend-shared/src/env.ts +++ b/packages/cubejs-backend-shared/src/env.ts @@ -1413,6 +1413,19 @@ const variables: Record any> = { ] ), + /** + * Snowflake OAuth token path. + */ + snowflakeOAuthTokenPath: ({ + dataSource + }: { + dataSource: string, + }) => ( + process.env[ + keyByDataSource('CUBEJS_DB_SNOWFLAKE_OAUTH_TOKEN_PATH', dataSource) + ] + ), + /** * Snowflake private key. */ diff --git a/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts b/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts index 76dcc5a873248..4e73839ddf4c5 100644 --- a/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts +++ b/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts @@ -24,6 +24,7 @@ import { DriverCapabilities, } from '@cubejs-backend/base-driver'; import { formatToTimeZone } from 'date-fns-timezone'; +import fs from 'fs/promises'; import { HydrationMap, HydrationStream } from './HydrationStream'; // eslint-disable-next-line import/order @@ -163,6 +164,7 @@ interface SnowflakeDriverOptions { clientSessionKeepAlive?: boolean, database?: string, authenticator?: string, + token?: string, privateKeyPath?: string, privateKeyPass?: string, privateKey?: string, @@ -207,7 +209,8 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface { 'CUBEJS_DB_SNOWFLAKE_CLIENT_SESSION_KEEP_ALIVE', 'CUBEJS_DB_SNOWFLAKE_AUTHENTICATOR', 'CUBEJS_DB_SNOWFLAKE_PRIVATE_KEY_PATH', - 'CUBEJS_DB_SNOWFLAKE_PRIVATE_KEY_PASS' + 'CUBEJS_DB_SNOWFLAKE_PRIVATE_KEY_PASS', + 'CUBEJS_DB_SNOWFLAKE_OAUTH_TOKEN_PATH', ]; } @@ -388,11 +391,28 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface { return undefined; } + private async readOAuthToken(dataSource: string) { + const tokenPath = getEnv('snowflakeOAuthTokenPath', { dataSource }); + const token = await fs.readFile(tokenPath, 'utf8'); + return token.trim(); + } + + private async createConnection() { + if (this.config.authenticator?.toUpperCase() === 'OAUTH') { + this.config.token = await this.readOAuthToken(this.config.dataSource); + } + + const connection = snowflake.createConnection(this.config); + + return connection; + } + /** * Test driver's connection. */ public async testConnection() { - const connection = snowflake.createConnection(this.config); + const connection = await this.createConnection(); + await new Promise( (resolve, reject) => connection.connect((err, conn) => (err ? reject(err) : resolve(conn))) ); @@ -411,7 +431,8 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface { */ protected async initConnection() { try { - const connection = snowflake.createConnection(this.config); + const connection = await this.createConnection(); + await new Promise( (resolve, reject) => connection.connect((err, conn) => (err ? reject(err) : resolve(conn))) ); From a2b277ba424f50d307022ea27f97a80af4676210 Mon Sep 17 00:00:00 2001 From: johancube Date: Sun, 13 Oct 2024 19:13:06 +0300 Subject: [PATCH 2/3] fix: path --- .../src/SnowflakeDriver.ts | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts b/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts index 4e73839ddf4c5..8cb06199c89a2 100644 --- a/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts +++ b/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts @@ -164,6 +164,7 @@ interface SnowflakeDriverOptions { clientSessionKeepAlive?: boolean, database?: string, authenticator?: string, + oauthTokenPath?: string, token?: string, privateKeyPath?: string, privateKeyPass?: string, @@ -269,6 +270,7 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface { username: getEnv('dbUser', { dataSource }), password: getEnv('dbPass', { dataSource }), authenticator: getEnv('snowflakeAuthenticator', { dataSource }), + oauthTokenPath: getEnv('snowflakeOAuthTokenPath', { dataSource }), privateKeyPath: getEnv('snowflakePrivateKeyPath', { dataSource }), privateKeyPass: getEnv('snowflakePrivateKeyPass', { dataSource }), privateKey, @@ -391,15 +393,26 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface { return undefined; } - private async readOAuthToken(dataSource: string) { - const tokenPath = getEnv('snowflakeOAuthTokenPath', { dataSource }); + private async readOAuthToken() { + const tokenPath = this.config.oauthTokenPath; + + if (!tokenPath) { + return undefined; + } + + try { + await fs.access(tokenPath); + } catch (error) { + throw new Error(`File ${tokenPath} provided by CUBEJS_DB_SNOWFLAKE_OAUTH_TOKEN_PATH does not exist.`); + } + const token = await fs.readFile(tokenPath, 'utf8'); return token.trim(); } private async createConnection() { if (this.config.authenticator?.toUpperCase() === 'OAUTH') { - this.config.token = await this.readOAuthToken(this.config.dataSource); + this.config.token = await this.readOAuthToken(); } const connection = snowflake.createConnection(this.config); From d5af1cb65fb9e83e9aa0594a9ee797a543dd6261 Mon Sep 17 00:00:00 2001 From: johancube Date: Mon, 14 Oct 2024 14:08:19 +0300 Subject: [PATCH 3/3] fix: set default token path --- packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts b/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts index 8cb06199c89a2..8a4b893ef37dc 100644 --- a/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts +++ b/packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts @@ -394,11 +394,7 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface { } private async readOAuthToken() { - const tokenPath = this.config.oauthTokenPath; - - if (!tokenPath) { - return undefined; - } + const tokenPath = this.config.oauthTokenPath || '/snowflake/session/token'; try { await fs.access(tokenPath);