feat(databricks-jdbc-driver): Support M2M OAuth Authentication

Author: KSDaemon

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

@@ -997,6 +997,32 @@ const variables: Record<string, (...args: any) => any> = {
     keyByDataSource('CUBEJS_DB_DATABRICKS_CATALOG', dataSource)
   ],
 
+  /**
+   * Databricks OAuth Client ID (Same as the service principal UUID)
+   */
+  databricksOAuthClientId: ({
+    dataSource,
+  }: {
+    dataSource: string,
+  }) => (
+    process.env[
+      keyByDataSource('CUBEJS_DB_DATABRICKS_OAUTH_CLIENT_ID', dataSource)
+    ]
+  ),
+
+  /**
+   * Databricks OAuth Client Secret.
+   */
+  databricksOAuthClientSecret: ({
+    dataSource,
+  }: {
+    dataSource: string,
+  }) => (
+    process.env[
+      keyByDataSource('CUBEJS_DB_DATABRICKS_OAUTH_CLIENT_SECRET', dataSource)
+    ]
+  ),
+
   /** ****************************************************************
    * Athena Driver                                                   *
    ***************************************************************** */

packages/cubejs-databricks-jdbc-driver/src/DatabricksDriver.ts

@@ -90,6 +90,16 @@ export type DatabricksDriverConfiguration = JDBCDriverConfiguration &
      */
     token?: string,
 
+    /**
+     * Databricks OAuth Client ID.
+     */
+    oauthClientId?: string,
+
+    /**
+     * Databricks OAuth Client Secret.
+     */
+    oauthClientSecret?: string,
+
     /**
      * Azure tenant Id
      */
@@ -200,6 +210,41 @@ export class DatabricksDriver extends JDBCDriver {
     }
 
     const [uid, pwd, cleanedUrl] = extractAndRemoveUidPwdFromJdbcUrl(url);
+    const passwd = conf?.token ||
+          getEnv('databricksToken', { dataSource }) ||
+          pwd;
+    const oauthClientId = conf?.oauthClientId || getEnv('databricksOAuthClientId', { dataSource });
+    const oauthClientSecret = conf?.oauthClientSecret || getEnv('databricksOAuthClientSecret', { dataSource });
+
+    if (oauthClientId && !oauthClientSecret) {
+      throw new Error('Invalid credentials: No OAuth Client Secret provided');
+    } else if (!oauthClientId && oauthClientSecret) {
+      throw new Error('Invalid credentials: No OAuth Client ID provided');
+    } else if (!oauthClientId && !oauthClientSecret && !passwd) {
+      throw new Error('No credentials provided');
+    }
+
+    const user = uid || 'token';
+
+    let authProps = {};
+
+    // OAuth has an advantage over UID+PWD
+    // For magic numbers below - see Databricks docs:
+    // https://docs.databricks.com/aws/en/integrations/jdbc-oss/configure#authenticate-the-driver
+    if (oauthClientId) {
+      authProps = {
+        OAuth2ClientID: oauthClientId,
+        OAuth2Secret: oauthClientSecret,
+        AuthMech: 11,
+        Auth_Flow: 1,
+      };
+    } else {
+      authProps = {
+        UID: user,
+        PWD: passwd,
+        AuthMech: 3,
+      };
+    }
 
     const config: DatabricksDriverConfiguration = {
       ...conf,
@@ -208,11 +253,7 @@ export class DatabricksDriver extends JDBCDriver {
       drivername: 'com.databricks.client.jdbc.Driver',
       customClassPath: undefined,
       properties: {
-        UID: uid,
-        PWD:
-          conf?.token ||
-          getEnv('databricksToken', { dataSource }) ||
-          pwd,
+        ...authProps,
         UserAgentEntry: 'CubeDev_Cube',
       },
       catalog:
@@ -292,7 +333,31 @@ export class DatabricksDriver extends JDBCDriver {
   }
 
   public override async testConnection() {
-    const token = `Bearer ${this.config.properties.PWD}`;
+    let token: string;
+
+    // Databricks docs on accessing REST API
+    // https://docs.databricks.com/aws/en/dev-tools/auth/oauth-m2m
+    if (this.config.properties.OAuth2Secret) {
+      // Need to exchange client ID + Secret => Access token
+
+      const basicAuth = Buffer.from(`${this.config.properties.OAuth2ClientID}:${this.config.properties.OAuth2Secret}`).toString('base64');
+
+      const res = await fetch(`https://${this.parsedConnectionProperties.host}/oidc/v1/token`, {
+        method: 'POST',
+        headers: {
+          Authorization: `Basic ${basicAuth}`,
+          'Content-Type': 'application/x-www-form-urlencoded',
+        },
+        body: new URLSearchParams({
+          grant_type: 'client_credentials',
+          scope: 'all-apis',
+        }),
+      });
+      const resp = await res.json();
+      token = resp.access_token;
+    } else {
+      token = `Bearer ${this.config.properties.PWD}`;
+    }
 
     const res = await fetch(`https://${this.parsedConnectionProperties.host}/api/2.0/sql/warehouses/${this.parsedConnectionProperties.warehouseId}`, {
       headers: { Authorization: token },

packages/cubejs-databricks-jdbc-driver/src/helpers.ts

@@ -35,20 +35,21 @@ export async function resolveJDBCDriver(): Promise<string> {
 
 /**
  * Extract if exist UID and PWD from URL and return UID, PWD and URL without these params.
- * New Databricks OSS driver throws an error if UID and PWD are provided in the URL and as a separate params
+ * New Databricks OSS driver throws an error if any parameter is provided in the URL and as a separate param
  * passed to the driver instance. That's why we strip them out from the URL if they exist there.
  * @param jdbcUrl
  */
 export function extractAndRemoveUidPwdFromJdbcUrl(jdbcUrl: string): [uid: string, pwd: string, cleanedUrl: string] {
   const uidMatch = jdbcUrl.match(/UID=([^;]*)/i);
   const pwdMatch = jdbcUrl.match(/PWD=([^;]*)/i);
 
-  const uid = uidMatch?.[1] || 'token';
+  const uid = uidMatch?.[1] || '';
   const pwd = pwdMatch?.[1] || '';
 
   const cleanedUrl = jdbcUrl
     .replace(/;?UID=[^;]*/i, '')
-    .replace(/;?PWD=[^;]*/i, '');
+    .replace(/;?PWD=[^;]*/i, '')
+    .replace(/;?AuthMech=[^;]*/i, '');
 
   return [uid, pwd, cleanedUrl];
 }