From bdcee60c7c243aafabbc402c988c13fedc1835fb Mon Sep 17 00:00:00 2001
From: Konstantin Burkalev <KSDaemon@gmail.com>
Date: Wed, 13 Nov 2024 16:38:03 +0200
Subject: [PATCH] feat(pinot-driver): add optional oAuth headers

---
 packages/cubejs-backend-shared/src/env.ts     | 21 +++++++++++++++++++
 .../cubejs-pinot-driver/src/PinotDriver.ts    | 21 ++++++++++++++++++-
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts
index a8021d303d9ad..70d64151ee518 100644
--- a/packages/cubejs-backend-shared/src/env.ts
+++ b/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                                               *
    ***************************************************************** */
diff --git a/packages/cubejs-pinot-driver/src/PinotDriver.ts b/packages/cubejs-pinot-driver/src/PinotDriver.ts
index 60fa891c62607..c9f96de3f02b6 100644
--- a/packages/cubejs-pinot-driver/src/PinotDriver.ts
+++ b/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 {};
     }