feat(pinot-driver): add optional oAuth headers (#8953)

Author: KSDaemon

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

@@ -1613,6 +1613,10 @@ const variables: Record<string, (...args: any) => any> = {
     ]
   ),
 
+  /** ***************************************************************
+   * Presto Driver                                                  *
+   **************************************************************** */
+
   /**
    * Presto catalog.
    */
@@ -1626,6 +1630,23 @@ const variables: Record<string, (...args: any) => any> = {
     ]
   ),
 
+  /** ***************************************************************
+   * Pinot Driver                                                  *
+   **************************************************************** */
+
+  /**
+   * Pinot / Startree Auth Token
+   */
+  pinotAuthToken: ({
+    dataSource,
+  }: {
+    dataSource: string,
+  }) => (
+    process.env[
+      keyByDataSource('CUBEJS_DB_PINOT_AUTH_TOKEN', dataSource)
+    ]
+  ),
+
   /** ****************************************************************
    * Cube Store Driver                                               *
    ***************************************************************** */

packages/cubejs-pinot-driver/src/PinotDriver.ts

@@ -27,12 +27,19 @@ export type PinotDriverConfiguration = {
   host?: string;
   port?: string;
   user?: string;
+  database?: string;
   basicAuth?: { user: string, password: string };
+  authToken?: string;
   ssl?: string | TLSConnectionOptions;
   dataSource?: string;
   queryTimeout?: number;
 };
 
+type AuthorizationHeaders = {
+  Authorization: string;
+  database?: string;
+};
+
 type PinotResponse = {
   exceptions: any[],
   minConsumingFreshnessTimeMs: number,
@@ -92,12 +99,14 @@ export class PinotDriver extends BaseDriver implements DriverInterface {
       host: getEnv('dbHost', { dataSource }),
       port: getEnv('dbPort', { dataSource }),
       user: getEnv('dbUser', { dataSource }),
+      database: getEnv('dbName', { dataSource }),
       basicAuth: getEnv('dbPass', { dataSource })
         ? {
           user: getEnv('dbUser', { dataSource }),
           password: getEnv('dbPass', { dataSource }),
         }
         : undefined,
+      authToken: getEnv('pinotAuthToken', { dataSource }),
       ssl: this.getSslOptions(dataSource),
       queryTimeout: getEnv('dbQueryTimeout', { dataSource }),
       ...config
@@ -127,7 +136,17 @@ export class PinotDriver extends BaseDriver implements DriverInterface {
     } : value)));
   }
 
-  public authorizationHeaders(): { Authorization?: string } {
+  public authorizationHeaders(): AuthorizationHeaders | {} {
+    if (this.config.authToken) {
+      const res: AuthorizationHeaders = { Authorization: `Bearer ${this.config.authToken}` };
+
+      if (this.config.database) {
+        res.database = this.config.database;
+      }
+
+      return res;
+    }
+
     if (!this.config.basicAuth) {
       return {};
     }