feat(databricks-jdbc-driver): Implement connection checking without waking up SQL warehouse

Author: KSDaemon

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

@@ -20,6 +20,7 @@ import { JDBCDriver, JDBCDriverConfiguration, } from '@cubejs-backend/jdbc-drive
 import { DatabricksQuery } from './DatabricksQuery';
 import {
   extractAndRemoveUidPwdFromJdbcUrl,
+  parseDatabricksJdbcUrl,
   resolveJDBCDriver
 } from './helpers';
 
@@ -124,6 +125,11 @@ type ColumnInfo = {
   type: GenericDataBaseType;
 };
 
+export type ParsedConnectionProperties = {
+  host: string,
+  warehouseId: string,
+};
+
 const DatabricksToGenericType: Record<string, string> = {
   binary: 'hll_datasketches',
   'decimal(10,0)': 'bigint',
@@ -143,6 +149,8 @@ export class DatabricksDriver extends JDBCDriver {
    */
   protected readonly config: DatabricksDriverConfiguration;
 
+  private readonly parsedConnectionProperties: ParsedConnectionProperties;
+
   public static dialectClass() {
     return DatabricksQuery;
   }
@@ -262,38 +270,46 @@ export class DatabricksDriver extends JDBCDriver {
 
     super(config);
     this.config = config;
+    this.parsedConnectionProperties = parseDatabricksJdbcUrl(url);
     this.showSparkProtocolWarn = showSparkProtocolWarn;
   }
 
-  /**
-   * @override
-   */
-  public readOnly() {
+  public override readOnly() {
     return !!this.config.readOnly;
   }
 
-  /**
-   * @override
-   */
-  public capabilities(): DriverCapabilities {
+  public override capabilities(): DriverCapabilities {
     return {
       unloadWithoutTempTable: true,
       incrementalSchemaLoading: true
     };
   }
 
-  /**
-   * @override
-   */
-  public setLogger(logger: any) {
+  public override setLogger(logger: any) {
     super.setLogger(logger);
     this.showDeprecations();
   }
 
-  /**
-   * @override
-   */
-  public async loadPreAggregationIntoTable(
+  public override async testConnection() {
+    const 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 },
+    });
+
+    if (!res.ok) {
+      throw new Error(`Databricks API error: ${res.statusText}`);
+    }
+
+    const data = await res.json();
+
+    // We don't need to check the actual status, because SQL WH will auto resume on real connection
+    // if (data.state !== 'RUNNING') {
+    //   throw new Error(`Warehouse not running (current state: ${data.state})`);
+    // }
+  }
+
+  public override async loadPreAggregationIntoTable(
     preAggregationTableName: string,
     loadSql: string,
     params: unknown[],
@@ -320,10 +336,7 @@ export class DatabricksDriver extends JDBCDriver {
     }
   }
 
-  /**
-   * @override
-   */
-  public async query<R = unknown>(
+  public override async query<R = unknown>(
     query: string,
     values: unknown[],
   ): Promise<R[]> {
@@ -357,10 +370,7 @@ export class DatabricksDriver extends JDBCDriver {
     }
   }
 
-  /**
-   * @override
-   */
-  public dropTable(tableName: string, options?: QueryOptions): Promise<unknown> {
+  public override dropTable(tableName: string, options?: QueryOptions): Promise<unknown> {
     const tableFullName = `${
       this.config?.catalog ? `${this.config.catalog}.` : ''
     }${tableName}`;
@@ -392,10 +402,7 @@ export class DatabricksDriver extends JDBCDriver {
     }
   }
 
-  /**
-   * @override
-   */
-  protected async getCustomClassPath() {
+  protected override async getCustomClassPath() {
     return resolveJDBCDriver();
   }
 

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

@@ -2,6 +2,7 @@ import fs from 'fs';
 import path from 'path';
 
 import { downloadJDBCDriver, OSS_DRIVER_VERSION } from './installer';
+import type { ParsedConnectionProperties } from './DatabricksDriver';
 
 async function fileExistsOr(
   fsPath: string,
@@ -51,3 +52,33 @@ export function extractAndRemoveUidPwdFromJdbcUrl(jdbcUrl: string): [uid: string
 
   return [uid, pwd, cleanedUrl];
 }
+
+export function parseDatabricksJdbcUrl(jdbcUrl: string): ParsedConnectionProperties {
+  const jdbcPrefix = 'jdbc:databricks://';
+  const urlWithoutPrefix = jdbcUrl.slice(jdbcPrefix.length);
+
+  const [hostPortAndPath, ...params] = urlWithoutPrefix.split(';');
+  const [host] = hostPortAndPath.split(':');
+
+  const paramMap = new Map<string, string>();
+  for (const param of params) {
+    const [key, value] = param.split('=');
+    if (key && value) {
+      paramMap.set(key, value);
+    }
+  }
+
+  const httpPath = paramMap.get('httpPath');
+  if (!httpPath) {
+    throw new Error('Missing httpPath in JDBC URL');
+  }
+
+  const warehouseMatch = httpPath.match(/\/warehouses\/([a-zA-Z0-9]+)/);
+  if (!warehouseMatch) {
+    throw new Error('Could not extract warehouseId from httpPath');
+  }
+
+  const warehouseId = warehouseMatch[1];
+
+  return { host, warehouseId };
+}