feat: Introduce dremioAuthToken config

- Adds the dremioAuthToken config to support the dremio cloud driver
cube-js · Dec 2, 2024 · 068560b · 068560b
1 parent 11437c2
commit 068560b
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 18 deletions.
diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts
@@ -1626,6 +1626,23 @@ const variables: Record<string, (...args: any) => any> = {
     ]
   ),
 
+  /** ****************************************************************
+   * Dremio Driver                                                   *
+   ***************************************************************** */
+
+  /**
+   * Dremio Auth Token
+   */
+  dremioAuthToken: ({
+    dataSource,
+  }: {
+    dataSource: string,
+  }) => (
+    process.env[
+      keyByDataSource('CUBEJS_DB_DREMIO_AUTH_TOKEN', dataSource)
+    ]
+  ),
+
   /** ****************************************************************
    * Cube Store Driver                                               *
    ***************************************************************** */

diff --git a/packages/cubejs-dremio-driver/README.md b/packages/cubejs-dremio-driver/README.md
@@ -13,14 +13,16 @@ Pure Javascript Dremio driver.
 
 To use this driver with [Dremio Cloud](https://docs.dremio.com/cloud/reference/api/), use the following setup:
 
-| Environment Variable | Value                                              |
-| -------------------- | -------------------------------------------------- |
-| CUBEJS_DB_TYPE       | dremio                                             |
-| CUBEJS_DB_URL        | https://api.dremio.cloud/v0/projects/${PROJECT_ID} |
-| CUBEJS_DB_NAME       | ${DB_NAME}                                         |
-| CUBEJS_DB_PASS       | Bearer ${PERSONAL_ACCESS_TOKEN}                    |
-
-> It's important to note that "Bearer" is required for using personal access tokens.
+| Environment Variable        | Value                                              |
+| --------------------------- | -------------------------------------------------- |
+| CUBEJS_DB_TYPE              | dremio                                             |
+| CUBEJS_DB_URL               | https://api.dremio.cloud/v0/projects/${PROJECT_ID} |
+| CUBEJS_DB_NAME              | ${DB_NAME}                                         |
+| CUBEJS_DB_DREMIO_AUTH_TOKEN | ${PERSONAL_ACCESS_TOKEN}                           |
+
+> [!NOTE]
+> When `CUBEJS_DB_URL` is set it takes precedence over `CUBEJS_DB_HOST` and it
+> is assumed that the driver is connecting to the Dremio Cloud API.
 
 ## Support
 

diff --git a/packages/cubejs-dremio-driver/driver/DremioDriver.js b/packages/cubejs-dremio-driver/driver/DremioDriver.js
@@ -53,6 +53,10 @@ class DremioDriver extends BaseDriver {
         config.dbUrl ||
         getEnv('dbUrl', { dataSource }) ||
         '',
+      dremioAuthToken:
+        config.dremioAuthToken ||
+        getEnv('dremioAuthToken', { dataSource }) ||
+        '',
       host:
         config.host ||
         getEnv('dbHost', { dataSource }) ||
@@ -88,20 +92,16 @@ class DremioDriver extends BaseDriver {
     if (this.config.dbUrl) {
       this.config.url = this.config.dbUrl;
       this.config.apiVersion = '';
-      this.config.auth = 'PAT';
+      if (this.config.dremioAuthToken === '') {
+        throw new Error('dremioAuthToken is blank');
+      }
     } else {
       const protocol = (this.config.ssl === true || this.config.ssl === 'true')
         ? 'https'
         : 'http';
       this.config.url = `${protocol}://${this.config.host}:${this.config.port}`;
       this.config.apiVersion = '/api/v3';
     }
-
-    if (this.config.auth === 'PAT' || this.config.password.startsWith('Bearer ')) {
-      this.config.auth = 'PAT';
-    } else {
-      this.config.auth = 'BASIC';
-    }
   }
 
   /**
@@ -120,17 +120,18 @@ class DremioDriver extends BaseDriver {
    * @protected
    */
   async getToken() {
-    if (this.config.auth === 'PAT') {
+    if (this.config.dremioAuthToken) {
+      const bearerToken = `Bearer ${this.config.dremioAuthToken}`;
       await axios.get(
         `${this.config.url}${this.config.apiVersion}/catalog`,
         {
           headers: {
-            Authorization: this.config.password
+            Authorization: bearerToken
           },
         },
       );
 
-      return this.config.password;
+      return bearerToken;
     }
 
     if (this.authToken && this.authToken.expires > new Date().getTime()) {