fix(databricks-jdbc-driver): ensure default UID property is set (#7168)

cube-js · Sep 28, 2023 · 0f357f8 · 0f357f8
1 parent e58af39
commit 0f357f8
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 32 deletions.
diff --git a/packages/cubejs-databricks-jdbc-driver/src/DatabricksDriver.ts b/packages/cubejs-databricks-jdbc-driver/src/DatabricksDriver.ts
@@ -8,8 +8,6 @@ import {
   getEnv,
   assertDataSource,
 } from '@cubejs-backend/shared';
-import fs from 'fs';
-import path from 'path';
 import { S3, GetObjectCommand } from '@aws-sdk/client-s3';
 import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
 import {
@@ -25,7 +23,7 @@ import {
   JDBCDriverConfiguration,
 } from '@cubejs-backend/jdbc-driver';
 import { DatabricksQuery } from './DatabricksQuery';
-import { downloadJDBCDriver } from './installer';
+import { resolveJDBCDriver, extractUidFromJdbcUrl } from './helpers';
 
 export type DatabricksDriverConfiguration = JDBCDriverConfiguration &
   {
@@ -91,16 +89,6 @@ export type DatabricksDriverConfiguration = JDBCDriverConfiguration &
     token?: string,
   };
 
-async function fileExistsOr(
-  fsPath: string,
-  fn: () => Promise<string>,
-): Promise<string> {
-  if (fs.existsSync(fsPath)) {
-    return fsPath;
-  }
-  return fn();
-}
-
 type ShowTableRow = {
   database: string,
   tableName: string,
@@ -115,25 +103,6 @@ const DatabricksToGenericType: Record<string, string> = {
   'decimal(10,0)': 'bigint',
 };
 
-async function resolveJDBCDriver(): Promise<string> {
-  return fileExistsOr(
-    path.join(process.cwd(), 'DatabricksJDBC42.jar'),
-    async () => fileExistsOr(
-      path.join(__dirname, '..', 'download', 'DatabricksJDBC42.jar'),
-      async () => {
-        const pathOrNull = await downloadJDBCDriver();
-        if (pathOrNull) {
-          return pathOrNull;
-        }
-        throw new Error(
-          'Please download and place DatabricksJDBC42.jar inside your ' +
-          'project directory'
-        );
-      }
-    )
-  );
-}
-
 /**
  * Databricks driver class.
  */
@@ -202,6 +171,7 @@ export class DatabricksDriver extends JDBCDriver {
       drivername: 'com.databricks.client.jdbc.Driver',
       customClassPath: undefined,
       properties: {
+        UID: extractUidFromJdbcUrl(url),
         // PWD-parameter passed to the connection string has higher priority,
         // so we can set this one to an empty string to avoid a Java error.
         PWD:

diff --git a/packages/cubejs-databricks-jdbc-driver/src/helpers.ts b/packages/cubejs-databricks-jdbc-driver/src/helpers.ts
@@ -0,0 +1,40 @@
+import fs from 'fs';
+import path from 'path';
+
+import { downloadJDBCDriver } from './installer';
+
+async function fileExistsOr(
+  fsPath: string,
+  fn: () => Promise<string>,
+): Promise<string> {
+  if (fs.existsSync(fsPath)) {
+    return fsPath;
+  }
+  return fn();
+}
+
+export async function resolveJDBCDriver(): Promise<string> {
+  return fileExistsOr(
+    path.join(process.cwd(), 'DatabricksJDBC42.jar'),
+    async () => fileExistsOr(
+      path.join(__dirname, '..', 'download', 'DatabricksJDBC42.jar'),
+      async () => {
+        const pathOrNull = await downloadJDBCDriver();
+        if (pathOrNull) {
+          return pathOrNull;
+        }
+        throw new Error(
+          'Please download and place DatabricksJDBC42.jar inside your ' +
+          'project directory'
+        );
+      }
+    )
+  );
+}
+
+export function extractUidFromJdbcUrl(jdbcUrl: string): string {
+  const { pathname } = new URL(jdbcUrl);
+  const [_, ...params] = pathname.split(';');
+  const searchParams = new URLSearchParams(params.join('&'));
+  return searchParams.get('UID') || 'token';
+}