@@ -24,6 +24,7 @@ import {
2424 DriverCapabilities ,
2525} from '@cubejs-backend/base-driver' ;
2626import { formatToTimeZone } from 'date-fns-timezone' ;
27+ import fs from 'fs/promises' ;
2728import { 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