Skip to content

Commit 9d5b7e8

Browse files
authored
feat(snowflake-driver): OAuth token path support (#8808)
* feat: init * fix: path * fix: set default token path
1 parent 362c32c commit 9d5b7e8

File tree

2 files changed

+46
-3
lines changed

2 files changed

+46
-3
lines changed

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1413,6 +1413,19 @@ const variables: Record<string, (...args: any) => any> = {
14131413
]
14141414
),
14151415

1416+
/**
1417+
* Snowflake OAuth token path.
1418+
*/
1419+
snowflakeOAuthTokenPath: ({
1420+
dataSource
1421+
}: {
1422+
dataSource: string,
1423+
}) => (
1424+
process.env[
1425+
keyByDataSource('CUBEJS_DB_SNOWFLAKE_OAUTH_TOKEN_PATH', dataSource)
1426+
]
1427+
),
1428+
14161429
/**
14171430
* Snowflake private key.
14181431
*/

packages/cubejs-snowflake-driver/src/SnowflakeDriver.ts

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import {
2424
DriverCapabilities,
2525
} from '@cubejs-backend/base-driver';
2626
import { formatToTimeZone } from 'date-fns-timezone';
27+
import fs from 'fs/promises';
2728
import { HydrationMap, HydrationStream } from './HydrationStream';
2829

2930
// eslint-disable-next-line import/order
@@ -163,6 +164,8 @@ interface SnowflakeDriverOptions {
163164
clientSessionKeepAlive?: boolean,
164165
database?: string,
165166
authenticator?: string,
167+
oauthTokenPath?: string,
168+
token?: string,
166169
privateKeyPath?: string,
167170
privateKeyPass?: string,
168171
privateKey?: string,
@@ -207,7 +210,8 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
207210
'CUBEJS_DB_SNOWFLAKE_CLIENT_SESSION_KEEP_ALIVE',
208211
'CUBEJS_DB_SNOWFLAKE_AUTHENTICATOR',
209212
'CUBEJS_DB_SNOWFLAKE_PRIVATE_KEY_PATH',
210-
'CUBEJS_DB_SNOWFLAKE_PRIVATE_KEY_PASS'
213+
'CUBEJS_DB_SNOWFLAKE_PRIVATE_KEY_PASS',
214+
'CUBEJS_DB_SNOWFLAKE_OAUTH_TOKEN_PATH',
211215
];
212216
}
213217

@@ -266,6 +270,7 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
266270
username: getEnv('dbUser', { dataSource }),
267271
password: getEnv('dbPass', { dataSource }),
268272
authenticator: getEnv('snowflakeAuthenticator', { dataSource }),
273+
oauthTokenPath: getEnv('snowflakeOAuthTokenPath', { dataSource }),
269274
privateKeyPath: getEnv('snowflakePrivateKeyPath', { dataSource }),
270275
privateKeyPass: getEnv('snowflakePrivateKeyPass', { dataSource }),
271276
privateKey,
@@ -388,11 +393,35 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
388393
return undefined;
389394
}
390395

396+
private async readOAuthToken() {
397+
const tokenPath = this.config.oauthTokenPath || '/snowflake/session/token';
398+
399+
try {
400+
await fs.access(tokenPath);
401+
} catch (error) {
402+
throw new Error(`File ${tokenPath} provided by CUBEJS_DB_SNOWFLAKE_OAUTH_TOKEN_PATH does not exist.`);
403+
}
404+
405+
const token = await fs.readFile(tokenPath, 'utf8');
406+
return token.trim();
407+
}
408+
409+
private async createConnection() {
410+
if (this.config.authenticator?.toUpperCase() === 'OAUTH') {
411+
this.config.token = await this.readOAuthToken();
412+
}
413+
414+
const connection = snowflake.createConnection(this.config);
415+
416+
return connection;
417+
}
418+
391419
/**
392420
* Test driver's connection.
393421
*/
394422
public async testConnection() {
395-
const connection = snowflake.createConnection(this.config);
423+
const connection = await this.createConnection();
424+
396425
await new Promise(
397426
(resolve, reject) => connection.connect((err, conn) => (err ? reject(err) : resolve(conn)))
398427
);
@@ -411,7 +440,8 @@ export class SnowflakeDriver extends BaseDriver implements DriverInterface {
411440
*/
412441
protected async initConnection() {
413442
try {
414-
const connection = snowflake.createConnection(this.config);
443+
const connection = await this.createConnection();
444+
415445
await new Promise(
416446
(resolve, reject) => connection.connect((err, conn) => (err ? reject(err) : resolve(conn)))
417447
);

0 commit comments

Comments
 (0)