feat(prestodb-driver, trino-driver): Support custom auth headers (JWT)

Author: KSDaemon

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

@@ -1798,7 +1798,7 @@ const variables: Record<string, (...args: any) => any> = {
     return [];
   },
   /** ***************************************************************
-   * Presto Driver                                                  *
+   * Presto/Trino Driver                                                  *
    **************************************************************** */
 
   /**
@@ -1814,12 +1814,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;
@@ -76,6 +78,14 @@ 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.config = {
       host: getEnv('dbHost', { dataSource }),
       port: getEnv('dbPort', { dataSource }),
@@ -85,13 +95,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 }),

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

@@ -12,12 +12,14 @@ export class TrinoDriver extends PrestoDriver {
   }
 
   public override async testConnection(): Promise<void> {
-    const { host, port, ssl, basic_auth: basicAuth } = this.config;
+    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}`;