databricks connector

Author: mbostock

package.json

@@ -63,6 +63,7 @@
     "vite": "^7.0.0"
   },
   "devDependencies": {
+    "@databricks/sql": "^1.11.0",
     "@duckdb/node-api": "^1.3.2-alpha.26",
     "@eslint/js": "^9.29.0",
     "@google-cloud/bigquery": "^8.1.1",
@@ -80,12 +81,16 @@
     "vitest": "^3.2.4"
   },
   "peerDependencies": {
+    "@databricks/sql": "^1.11.0",
     "@duckdb/node-api": "^1.3.2-alpha.26",
     "@google-cloud/bigquery": "^8.1.1",
     "postgres": "^3.4.7",
     "snowflake-sdk": "^2.1.3"
   },
   "peerDependenciesMeta": {
+    "@databricks/sql": {
+      "optional": true
+    },
     "@duckdb/node-api": {
       "optional": true
     },

src/databases/databricks.ts

@@ -0,0 +1,74 @@
+import {DBSQLClient, DBSQLLogger, LogLevel, thrift} from "@databricks/sql";
+import type {DBSQLSession} from "@databricks/sql";
+import type {DatabricksConfig, QueryTemplateFunction} from "./index.js";
+import {ColumnSchema} from "../runtime/index.js";
+
+type IOperation = Awaited<ReturnType<DBSQLSession["executeStatement"]>>;
+type TTableSchema = NonNullable<Awaited<ReturnType<IOperation["getSchema"]>>>;
+type TColumnDesc = TTableSchema["columns"][0];
+type TTypeDesc = TColumnDesc["typeDesc"];
+const TTypeId = thrift.TCLIService_types.TTypeId;
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export default function databricks({type, ...options}: DatabricksConfig): QueryTemplateFunction {
+  return async (strings, ...params) => {
+    const logger = new DBSQLLogger({level: LogLevel.error});
+    const client = new DBSQLClient({logger});
+    await client.connect(options);
+    try {
+      const session = await client.openSession();
+      try {
+        const date = new Date();
+        const operation = await session.executeStatement(strings.join("?"), {
+          runAsync: true,
+          ordinalParameters: params,
+          maxRows: 10000
+        });
+        try {
+          const rows = (await operation.fetchAll()) as Record<string, unknown>[];
+          const schema = await operation.getSchema();
+          return {rows, schema: getTableSchema(schema!), duration: Date.now() - +date, date};
+        } finally {
+          await operation.close();
+        }
+      } finally {
+        await session.close();
+      }
+    } finally {
+      await client.close();
+    }
+  };
+}
+
+function getTableSchema({columns}: TTableSchema): ColumnSchema[] {
+  return columns.map(getColumnSchema);
+}
+
+function getColumnSchema(column: TColumnDesc): ColumnSchema {
+  return {name: column.columnName, type: getColumnType(column.typeDesc)};
+}
+
+function getColumnType({types: [type]}: TTypeDesc): ColumnSchema["type"] {
+  switch (type.primitiveEntry?.type) {
+    case TTypeId.BINARY_TYPE:
+      return "buffer";
+    case TTypeId.BOOLEAN_TYPE:
+      return "boolean";
+    case TTypeId.BIGINT_TYPE:
+    case TTypeId.TINYINT_TYPE:
+    case TTypeId.SMALLINT_TYPE:
+    case TTypeId.INT_TYPE:
+    case TTypeId.DECIMAL_TYPE:
+      return "integer";
+    case TTypeId.DOUBLE_TYPE:
+    case TTypeId.FLOAT_TYPE:
+      return "number";
+    case TTypeId.DATE_TYPE:
+    case TTypeId.TIMESTAMP_TYPE:
+    case TTypeId.INTERVAL_DAY_TIME_TYPE:
+    case TTypeId.INTERVAL_YEAR_MONTH_TYPE:
+      return "date";
+    default:
+      return "string";
+  }
+}

src/databases/index.ts

@@ -7,6 +7,7 @@ import type {ColumnSchema, QueryParam} from "../runtime/index.js";
 
 export type DatabaseConfig =
   | BigQueryConfig
+  | DatabricksConfig
   | DuckDBConfig
   | SQLiteConfig
   | SnowflakeConfig
@@ -20,6 +21,15 @@ export type BigQueryConfig = {
   projectId?: string;
 };
 
+export type DatabricksConfig = {
+  type: "databricks";
+  host: string;
+  path: string;
+} & (
+  | {authType?: "access-token"; token: string}
+  | {authType: "databricks-oauth"; oauthClientId?: string; oauthClientSecret?: string}
+);
+
 export type DuckDBConfig = {
   type: "duckdb";
   path?: string;
@@ -83,7 +93,7 @@ export async function getDatabaseConfig(
     else if (databaseName === "duckdb") config = {type: "duckdb"};
     else if (databaseName === "sqlite") config = {type: "sqlite"};
     else if (/\.duckdb$/i.test(databaseName)) config = {type: "duckdb", path: databaseName};
-    else if (/\.db$/i.test(databaseName)) config = {type: "sqlite", path: databaseName}; // TODO disambiguate
+    else if (/\.db$/i.test(databaseName)) config = {type: "sqlite", path: databaseName};
     else throw new Error(`database not found: ${databaseName}`);
   }
   return config;
@@ -93,10 +103,12 @@ export async function getDatabase(config: DatabaseConfig): Promise<QueryTemplate
   switch (config.type) {
     case "bigquery":
       return (await import("./bigquery.js")).default(config);
+    case "databricks":
+      return (await import("./databricks.js")).default(config);
     case "duckdb":
       return (await import("./duckdb.js")).default(config);
     case "sqlite":
-      return (await import(process.versions.bun ? "./sqlite-bun.js" : "./sqlite-node.js")).default(config);
+      return (await import(process.versions.bun ? "./sqlite-bun.js" : "./sqlite-node.js")).default(config); // prettier-ignore
     case "snowflake":
       return (await import("./snowflake.js")).default(config);
     case "postgres":