feat(dremio-driver): Add Dremio Cloud Support (#8956)

* feat: Add Dremio Cloud Support

- Adds support to cubejs-dremio-driver for Dremio Cloud
- Adds support for personal access tokens for Dremio Server

* feat: Introduce dremioAuthToken config

- Adds the dremioAuthToken config to support the dremio cloud driver

---------

Co-authored-by: Konstantin Burkalev <[email protected]>

Author: maladroitthief

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

@@ -1647,6 +1647,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                                               *
    ***************************************************************** */

packages/cubejs-dremio-driver/README.md

@@ -9,11 +9,26 @@
 
 Pure Javascript Dremio driver.
 
+## Dremio Cloud
+
+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_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
 
-This package is **community supported** and should be used at your own risk. 
+This package is **community supported** and should be used at your own risk.
 
-While the Cube Dev team is happy to review and accept future community contributions, we don't have active plans for further development. This includes bug fixes unless they affect different parts of Cube.js. **We're looking for maintainers for this package.** If you'd like to become a maintainer, please contact us in Cube.js Slack. 
+While the Cube Dev team is happy to review and accept future community contributions, we don't have active plans for further development. This includes bug fixes unless they affect different parts of Cube.js. **We're looking for maintainers for this package.** If you'd like to become a maintainer, please contact us in Cube.js Slack.
 
 ## License
 

packages/cubejs-dremio-driver/driver/DremioDriver.js

@@ -49,6 +49,14 @@ class DremioDriver extends BaseDriver {
       assertDataSource('default');
 
     this.config = {
+      dbUrl:
+        config.dbUrl ||
+        getEnv('dbUrl', { dataSource }) ||
+        '',
+      dremioAuthToken:
+        config.dremioAuthToken ||
+        getEnv('dremioAuthToken', { dataSource }) ||
+        '',
       host:
         config.host ||
         getEnv('dbHost', { dataSource }) ||
@@ -80,11 +88,20 @@ class DremioDriver extends BaseDriver {
         getEnv('dbPollMaxInterval', { dataSource })
       ) * 1000,
     };
-    const protocol = (this.config.ssl === true || this.config.ssl === 'true')
-      ? 'https'
-      : 'http';
-    this.config.url =
-      `${protocol}://${this.config.host}:${this.config.port}`;
+
+    if (this.config.dbUrl) {
+      this.config.url = this.config.dbUrl;
+      this.config.apiVersion = '';
+      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';
+    }
   }
 
   /**
@@ -103,6 +120,20 @@ class DremioDriver extends BaseDriver {
    * @protected
    */
   async getToken() {
+    if (this.config.dremioAuthToken) {
+      const bearerToken = `Bearer ${this.config.dremioAuthToken}`;
+      await axios.get(
+        `${this.config.url}${this.config.apiVersion}/catalog`,
+        {
+          headers: {
+            Authorization: bearerToken
+          },
+        },
+      );
+
+      return bearerToken;
+    }
+
     if (this.authToken && this.authToken.expires > new Date().getTime()) {
       return `_dremio${this.authToken.token}`;
     }
@@ -129,7 +160,7 @@ class DremioDriver extends BaseDriver {
 
     return axios.request({
       method,
-      url: `${this.config.url}${url}`,
+      url: `${this.config.url}${this.config.apiVersion}${url}`,
       headers: {
         Authorization: token
       },
@@ -141,7 +172,7 @@ class DremioDriver extends BaseDriver {
    * @protected
    */
   async getJobStatus(jobId) {
-    const { data } = await this.restDremioQuery('get', `/api/v3/job/${jobId}`);
+    const { data } = await this.restDremioQuery('get', `/job/${jobId}`);
 
     if (data.jobState === 'FAILED') {
       throw new Error(data.errorMessage);
@@ -162,7 +193,7 @@ class DremioDriver extends BaseDriver {
    * @protected
    */
   async getJobResults(jobId, limit = 500, offset = 0) {
-    return this.restDremioQuery('get', `/api/v3/job/${jobId}/results?offset=${offset}&limit=${limit}`);
+    return this.restDremioQuery('get', `/job/${jobId}/results?offset=${offset}&limit=${limit}`);
   }
 
   /**
@@ -171,7 +202,7 @@ class DremioDriver extends BaseDriver {
    * @return {Promise<*>}
    */
   async executeQuery(sql) {
-    const { data } = await this.restDremioQuery('post', '/api/v3/sql', { sql });
+    const { data } = await this.restDremioQuery('post', '/sql', { sql });
     return data.id;
   }
 
@@ -216,7 +247,7 @@ class DremioDriver extends BaseDriver {
   }
 
   async refreshTablesSchema(path) {
-    const { data } = await this.restDremioQuery('get', `/api/v3/catalog/by-path/${path}`);
+    const { data } = await this.restDremioQuery('get', `/catalog/by-path/${path}`);
     if (!data || !data.children) {
       return true;
     }