@databricks/appkit 0.1.4 → 0.1.5

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (82) hide show
  1. package/AGENTS.md +37 -12
  2. package/CLAUDE.md +37 -12
  3. package/NOTICE.md +2 -0
  4. package/dist/analytics/analytics.d.ts +33 -8
  5. package/dist/analytics/analytics.d.ts.map +1 -1
  6. package/dist/analytics/analytics.js +51 -24
  7. package/dist/analytics/analytics.js.map +1 -1
  8. package/dist/analytics/defaults.js.map +1 -1
  9. package/dist/analytics/query.js +4 -4
  10. package/dist/analytics/query.js.map +1 -1
  11. package/dist/appkit/package.js +1 -1
  12. package/dist/cache/defaults.js.map +1 -1
  13. package/dist/cache/index.d.ts +1 -0
  14. package/dist/cache/index.d.ts.map +1 -1
  15. package/dist/cache/index.js +1 -2
  16. package/dist/cache/index.js.map +1 -1
  17. package/dist/cache/storage/memory.js.map +1 -1
  18. package/dist/connectors/lakebase/client.js +7 -8
  19. package/dist/connectors/lakebase/client.js.map +1 -1
  20. package/dist/connectors/lakebase/defaults.js.map +1 -1
  21. package/dist/connectors/sql-warehouse/client.js.map +1 -1
  22. package/dist/connectors/sql-warehouse/defaults.js.map +1 -1
  23. package/dist/context/execution-context.js +75 -0
  24. package/dist/context/execution-context.js.map +1 -0
  25. package/dist/context/index.js +27 -0
  26. package/dist/context/index.js.map +1 -0
  27. package/dist/context/service-context.js +149 -0
  28. package/dist/context/service-context.js.map +1 -0
  29. package/dist/context/user-context.js +15 -0
  30. package/dist/context/user-context.js.map +1 -0
  31. package/dist/core/appkit.d.ts +3 -0
  32. package/dist/core/appkit.d.ts.map +1 -1
  33. package/dist/core/appkit.js +7 -0
  34. package/dist/core/appkit.js.map +1 -1
  35. package/dist/index.d.ts +5 -6
  36. package/dist/index.js +3 -10
  37. package/dist/plugin/interceptors/cache.js.map +1 -1
  38. package/dist/plugin/interceptors/retry.js.map +1 -1
  39. package/dist/plugin/interceptors/telemetry.js.map +1 -1
  40. package/dist/plugin/interceptors/timeout.js.map +1 -1
  41. package/dist/plugin/plugin.d.ts +39 -5
  42. package/dist/plugin/plugin.d.ts.map +1 -1
  43. package/dist/plugin/plugin.js +82 -6
  44. package/dist/plugin/plugin.js.map +1 -1
  45. package/dist/plugin/to-plugin.d.ts +4 -0
  46. package/dist/plugin/to-plugin.d.ts.map +1 -1
  47. package/dist/plugin/to-plugin.js +3 -0
  48. package/dist/plugin/to-plugin.js.map +1 -1
  49. package/dist/server/index.d.ts +3 -0
  50. package/dist/server/index.d.ts.map +1 -1
  51. package/dist/server/index.js +3 -4
  52. package/dist/server/index.js.map +1 -1
  53. package/dist/server/remote-tunnel/remote-tunnel-manager.js.map +1 -1
  54. package/dist/server/utils.js.map +1 -1
  55. package/dist/server/vite-dev-server.js +0 -2
  56. package/dist/server/vite-dev-server.js.map +1 -1
  57. package/dist/shared/src/sql/helpers.js.map +1 -1
  58. package/dist/stream/arrow-stream-processor.js.map +1 -1
  59. package/dist/stream/buffers.js.map +1 -1
  60. package/dist/stream/sse-writer.js.map +1 -1
  61. package/dist/stream/stream-manager.js.map +1 -1
  62. package/dist/stream/stream-registry.js.map +1 -1
  63. package/dist/telemetry/instrumentations.js.map +1 -1
  64. package/dist/type-generator/index.js.map +1 -1
  65. package/dist/type-generator/query-registry.js.map +1 -1
  66. package/dist/type-generator/types.js.map +1 -1
  67. package/dist/type-generator/vite-plugin.js.map +1 -1
  68. package/dist/utils/env-validator.js +1 -5
  69. package/dist/utils/env-validator.js.map +1 -1
  70. package/dist/utils/merge.js +1 -5
  71. package/dist/utils/merge.js.map +1 -1
  72. package/dist/utils/vite-config-merge.js +1 -5
  73. package/dist/utils/vite-config-merge.js.map +1 -1
  74. package/llms.txt +37 -12
  75. package/package.json +3 -1
  76. package/dist/index.js.map +0 -1
  77. package/dist/utils/databricks-client-middleware.d.ts +0 -17
  78. package/dist/utils/databricks-client-middleware.d.ts.map +0 -1
  79. package/dist/utils/databricks-client-middleware.js +0 -117
  80. package/dist/utils/databricks-client-middleware.js.map +0 -1
  81. package/dist/utils/index.js +0 -26
  82. package/dist/utils/index.js.map +0 -1
@@ -1 +1 @@
1
- {"version":3,"file":"memory.js","names":["oldestKey: string | null"],"sources":["../../../src/cache/storage/memory.ts"],"sourcesContent":["import type { CacheConfig, CacheEntry, CacheStorage } from \"shared\";\nimport { inMemoryStorageDefaults } from \"./defaults\";\n\n/**\n * In-memory cache storage implementation. Uses a least recently used (LRU) eviction policy\n * to manage memory usage and ensure efficient cache operations.\n */\nexport class InMemoryStorage implements CacheStorage {\n private cache: Map<string, CacheEntry> = new Map();\n private accessOrder: Map<string, number> = new Map();\n private accessCounter: number;\n private maxSize: number;\n\n constructor(config: CacheConfig) {\n this.cache = new Map();\n this.accessOrder = new Map();\n this.maxSize = config.maxSize ?? inMemoryStorageDefaults.maxSize;\n this.accessCounter = 0;\n }\n\n /** Get an entry from the cache */\n async get<T>(key: string): Promise<CacheEntry<T> | null> {\n const entry = this.cache.get(key);\n if (!entry) return null;\n\n this.accessOrder.set(key, ++this.accessCounter);\n return entry as CacheEntry<T>;\n }\n\n /** Set an entry in the cache */\n async set<T>(key: string, entry: CacheEntry<T>): Promise<void> {\n if (this.cache.size >= this.maxSize && !this.cache.has(key)) {\n this.evictLRU();\n }\n\n this.cache.set(key, entry);\n this.accessOrder.set(key, ++this.accessCounter);\n }\n\n /** Delete an entry from the cache */\n async delete(key: string): Promise<void> {\n this.cache.delete(key);\n this.accessOrder.delete(key);\n }\n\n /** Clean in-memory cache */\n async clear(): Promise<void> {\n this.cache.clear();\n this.accessOrder.clear();\n this.accessCounter = 0;\n }\n\n /** Check if the cache has an entry */\n async has(key: string): Promise<boolean> {\n const entry = this.cache.get(key);\n if (!entry) return false;\n\n if (Date.now() > entry.expiry) {\n this.cache.delete(key);\n this.accessOrder.delete(key);\n return false;\n }\n\n return true;\n }\n\n /** Get the size of the cache */\n async size(): Promise<number> {\n return this.cache.size;\n }\n\n /** Check if the cache is persistent */\n isPersistent(): boolean {\n return false;\n }\n\n /** Check the health of the cache */\n async healthCheck(): Promise<boolean> {\n return true;\n }\n\n /** Close the cache */\n async close(): Promise<void> {\n this.cache.clear();\n this.accessOrder.clear();\n this.accessCounter = 0;\n }\n\n /** Evict the least recently used entry (LRU) */\n private evictLRU(): void {\n let oldestKey: string | null = null;\n let oldestAccess = Infinity;\n\n for (const [key, accessTime] of this.accessOrder) {\n if (accessTime < oldestAccess) {\n oldestAccess = accessTime;\n oldestKey = key;\n }\n }\n\n if (oldestKey) {\n this.cache.delete(oldestKey);\n this.accessOrder.delete(oldestKey);\n }\n }\n}\n"],"mappings":";;;;;;;AAOA,IAAa,kBAAb,MAAqD;CAMnD,YAAY,QAAqB;+BALQ,IAAI,KAAK;qCACP,IAAI,KAAK;AAKlD,OAAK,wBAAQ,IAAI,KAAK;AACtB,OAAK,8BAAc,IAAI,KAAK;AAC5B,OAAK,UAAU,OAAO,WAAW,wBAAwB;AACzD,OAAK,gBAAgB;;;CAIvB,MAAM,IAAO,KAA4C;EACvD,MAAM,QAAQ,KAAK,MAAM,IAAI,IAAI;AACjC,MAAI,CAAC,MAAO,QAAO;AAEnB,OAAK,YAAY,IAAI,KAAK,EAAE,KAAK,cAAc;AAC/C,SAAO;;;CAIT,MAAM,IAAO,KAAa,OAAqC;AAC7D,MAAI,KAAK,MAAM,QAAQ,KAAK,WAAW,CAAC,KAAK,MAAM,IAAI,IAAI,CACzD,MAAK,UAAU;AAGjB,OAAK,MAAM,IAAI,KAAK,MAAM;AAC1B,OAAK,YAAY,IAAI,KAAK,EAAE,KAAK,cAAc;;;CAIjD,MAAM,OAAO,KAA4B;AACvC,OAAK,MAAM,OAAO,IAAI;AACtB,OAAK,YAAY,OAAO,IAAI;;;CAI9B,MAAM,QAAuB;AAC3B,OAAK,MAAM,OAAO;AAClB,OAAK,YAAY,OAAO;AACxB,OAAK,gBAAgB;;;CAIvB,MAAM,IAAI,KAA+B;EACvC,MAAM,QAAQ,KAAK,MAAM,IAAI,IAAI;AACjC,MAAI,CAAC,MAAO,QAAO;AAEnB,MAAI,KAAK,KAAK,GAAG,MAAM,QAAQ;AAC7B,QAAK,MAAM,OAAO,IAAI;AACtB,QAAK,YAAY,OAAO,IAAI;AAC5B,UAAO;;AAGT,SAAO;;;CAIT,MAAM,OAAwB;AAC5B,SAAO,KAAK,MAAM;;;CAIpB,eAAwB;AACtB,SAAO;;;CAIT,MAAM,cAAgC;AACpC,SAAO;;;CAIT,MAAM,QAAuB;AAC3B,OAAK,MAAM,OAAO;AAClB,OAAK,YAAY,OAAO;AACxB,OAAK,gBAAgB;;;CAIvB,AAAQ,WAAiB;EACvB,IAAIA,YAA2B;EAC/B,IAAI,eAAe;AAEnB,OAAK,MAAM,CAAC,KAAK,eAAe,KAAK,YACnC,KAAI,aAAa,cAAc;AAC7B,kBAAe;AACf,eAAY;;AAIhB,MAAI,WAAW;AACb,QAAK,MAAM,OAAO,UAAU;AAC5B,QAAK,YAAY,OAAO,UAAU"}
1
+ {"version":3,"file":"memory.js","names":[],"sources":["../../../src/cache/storage/memory.ts"],"sourcesContent":["import type { CacheConfig, CacheEntry, CacheStorage } from \"shared\";\nimport { inMemoryStorageDefaults } from \"./defaults\";\n\n/**\n * In-memory cache storage implementation. Uses a least recently used (LRU) eviction policy\n * to manage memory usage and ensure efficient cache operations.\n */\nexport class InMemoryStorage implements CacheStorage {\n private cache: Map<string, CacheEntry> = new Map();\n private accessOrder: Map<string, number> = new Map();\n private accessCounter: number;\n private maxSize: number;\n\n constructor(config: CacheConfig) {\n this.cache = new Map();\n this.accessOrder = new Map();\n this.maxSize = config.maxSize ?? inMemoryStorageDefaults.maxSize;\n this.accessCounter = 0;\n }\n\n /** Get an entry from the cache */\n async get<T>(key: string): Promise<CacheEntry<T> | null> {\n const entry = this.cache.get(key);\n if (!entry) return null;\n\n this.accessOrder.set(key, ++this.accessCounter);\n return entry as CacheEntry<T>;\n }\n\n /** Set an entry in the cache */\n async set<T>(key: string, entry: CacheEntry<T>): Promise<void> {\n if (this.cache.size >= this.maxSize && !this.cache.has(key)) {\n this.evictLRU();\n }\n\n this.cache.set(key, entry);\n this.accessOrder.set(key, ++this.accessCounter);\n }\n\n /** Delete an entry from the cache */\n async delete(key: string): Promise<void> {\n this.cache.delete(key);\n this.accessOrder.delete(key);\n }\n\n /** Clean in-memory cache */\n async clear(): Promise<void> {\n this.cache.clear();\n this.accessOrder.clear();\n this.accessCounter = 0;\n }\n\n /** Check if the cache has an entry */\n async has(key: string): Promise<boolean> {\n const entry = this.cache.get(key);\n if (!entry) return false;\n\n if (Date.now() > entry.expiry) {\n this.cache.delete(key);\n this.accessOrder.delete(key);\n return false;\n }\n\n return true;\n }\n\n /** Get the size of the cache */\n async size(): Promise<number> {\n return this.cache.size;\n }\n\n /** Check if the cache is persistent */\n isPersistent(): boolean {\n return false;\n }\n\n /** Check the health of the cache */\n async healthCheck(): Promise<boolean> {\n return true;\n }\n\n /** Close the cache */\n async close(): Promise<void> {\n this.cache.clear();\n this.accessOrder.clear();\n this.accessCounter = 0;\n }\n\n /** Evict the least recently used entry (LRU) */\n private evictLRU(): void {\n let oldestKey: string | null = null;\n let oldestAccess = Infinity;\n\n for (const [key, accessTime] of this.accessOrder) {\n if (accessTime < oldestAccess) {\n oldestAccess = accessTime;\n oldestKey = key;\n }\n }\n\n if (oldestKey) {\n this.cache.delete(oldestKey);\n this.accessOrder.delete(oldestKey);\n }\n }\n}\n"],"mappings":";;;;;;;AAOA,IAAa,kBAAb,MAAqD;CAMnD,YAAY,QAAqB;+BALQ,IAAI,KAAK;qCACP,IAAI,KAAK;AAKlD,OAAK,wBAAQ,IAAI,KAAK;AACtB,OAAK,8BAAc,IAAI,KAAK;AAC5B,OAAK,UAAU,OAAO,WAAW,wBAAwB;AACzD,OAAK,gBAAgB;;;CAIvB,MAAM,IAAO,KAA4C;EACvD,MAAM,QAAQ,KAAK,MAAM,IAAI,IAAI;AACjC,MAAI,CAAC,MAAO,QAAO;AAEnB,OAAK,YAAY,IAAI,KAAK,EAAE,KAAK,cAAc;AAC/C,SAAO;;;CAIT,MAAM,IAAO,KAAa,OAAqC;AAC7D,MAAI,KAAK,MAAM,QAAQ,KAAK,WAAW,CAAC,KAAK,MAAM,IAAI,IAAI,CACzD,MAAK,UAAU;AAGjB,OAAK,MAAM,IAAI,KAAK,MAAM;AAC1B,OAAK,YAAY,IAAI,KAAK,EAAE,KAAK,cAAc;;;CAIjD,MAAM,OAAO,KAA4B;AACvC,OAAK,MAAM,OAAO,IAAI;AACtB,OAAK,YAAY,OAAO,IAAI;;;CAI9B,MAAM,QAAuB;AAC3B,OAAK,MAAM,OAAO;AAClB,OAAK,YAAY,OAAO;AACxB,OAAK,gBAAgB;;;CAIvB,MAAM,IAAI,KAA+B;EACvC,MAAM,QAAQ,KAAK,MAAM,IAAI,IAAI;AACjC,MAAI,CAAC,MAAO,QAAO;AAEnB,MAAI,KAAK,KAAK,GAAG,MAAM,QAAQ;AAC7B,QAAK,MAAM,OAAO,IAAI;AACtB,QAAK,YAAY,OAAO,IAAI;AAC5B,UAAO;;AAGT,SAAO;;;CAIT,MAAM,OAAwB;AAC5B,SAAO,KAAK,MAAM;;;CAIpB,eAAwB;AACtB,SAAO;;;CAIT,MAAM,cAAgC;AACpC,SAAO;;;CAIT,MAAM,QAAuB;AAC3B,OAAK,MAAM,OAAO;AAClB,OAAK,YAAY,OAAO;AACxB,OAAK,gBAAgB;;;CAIvB,AAAQ,WAAiB;EACvB,IAAI,YAA2B;EAC/B,IAAI,eAAe;AAEnB,OAAK,MAAM,CAAC,KAAK,eAAe,KAAK,YACnC,KAAI,aAAa,cAAc;AAC7B,kBAAe;AACf,eAAY;;AAIhB,MAAI,WAAW;AACb,QAAK,MAAM,OAAO,UAAU;AAC5B,QAAK,YAAY,OAAO,UAAU"}
@@ -2,14 +2,13 @@ import { __toCommonJS } from "../../_virtual/rolldown_runtime.js";
2
2
  import { TelemetryManager } from "../../telemetry/telemetry-manager.js";
3
3
  import { SpanStatusCode } from "../../telemetry/index.js";
4
4
  import { deepMerge } from "../../utils/merge.js";
5
- import { init_utils, utils_exports } from "../../utils/index.js";
6
5
  import { lakebaseDefaults } from "./defaults.js";
6
+ import { context_exports, init_context } from "../../context/index.js";
7
7
  import { randomUUID } from "node:crypto";
8
8
  import { ApiClient, Config } from "@databricks/sdk-experimental";
9
9
  import pg from "pg";
10
10
 
11
11
  //#region src/connectors/lakebase/client.ts
12
- init_utils();
13
12
  /**
14
13
  * Enterprise-grade connector for Databricks Lakebase
15
14
  * @example Simplest - everything from env/context
@@ -194,16 +193,16 @@ var LakebaseConnector = class {
194
193
  process.on("SIGINT", () => this.close());
195
194
  this.close();
196
195
  }
197
- /** Get Databricks workspace client - from config or request context */
196
+ /** Get Databricks workspace client - from config or execution context */
198
197
  getWorkspaceClient() {
199
198
  if (this.config.workspaceClient) return this.config.workspaceClient;
200
199
  try {
201
- const { getRequestContext } = (init_utils(), __toCommonJS(utils_exports));
202
- const { serviceDatabricksClient } = getRequestContext();
203
- this.config.workspaceClient = serviceDatabricksClient;
204
- return serviceDatabricksClient;
200
+ const { getWorkspaceClient: getClient } = (init_context(), __toCommonJS(context_exports));
201
+ const client = getClient();
202
+ this.config.workspaceClient = client;
203
+ return client;
205
204
  } catch (_error) {
206
- throw new Error("Databricks workspace client not available. Either pass it in config or use within AppKit request context.");
205
+ throw new Error("Databricks workspace client not available. Either pass it in config or ensure ServiceContext is initialized.");
207
206
  }
208
207
  }
209
208
  /** Get or create connection pool */
@@ -1 +1 @@
1
- {"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/lakebase/client.ts"],"sourcesContent":["import { randomUUID } from \"node:crypto\";\nimport type { WorkspaceClient } from \"@databricks/sdk-experimental\";\nimport { ApiClient, Config } from \"@databricks/sdk-experimental\";\nimport pg from \"pg\";\nimport {\n type Counter,\n type Histogram,\n SpanStatusCode,\n TelemetryManager,\n type TelemetryProvider,\n} from \"@/telemetry\";\nimport { deepMerge } from \"../../utils\";\nimport { lakebaseDefaults } from \"./defaults\";\nimport type {\n LakebaseConfig,\n LakebaseConnectionConfig,\n LakebaseCredentials,\n} from \"./types\";\n\n/**\n * Enterprise-grade connector for Databricks Lakebase\n * @example Simplest - everything from env/context\n * ```typescript\n * const connector = new LakebaseConnector();\n * await connector.query('SELECT * FROM users');\n * ```\n *\n * @example With explicit connection string\n * ```typescript\n * const connector = new LakebaseConnector({\n * connectionString: 'postgresql://...'\n * });\n * ```\n */\nexport class LakebaseConnector {\n private readonly name: string = \"lakebase\";\n private readonly CACHE_BUFFER_MS = 2 * 60 * 1000;\n private readonly config: LakebaseConfig;\n private readonly connectionConfig: LakebaseConnectionConfig;\n private pool: pg.Pool | null = null;\n private credentials: LakebaseCredentials | null = null;\n\n // telemetry\n private readonly telemetry: TelemetryProvider;\n private readonly telemetryMetrics: {\n queryCount: Counter;\n queryDuration: Histogram;\n };\n\n constructor(userConfig?: Partial<LakebaseConfig>) {\n this.config = deepMerge(lakebaseDefaults, userConfig);\n this.connectionConfig = this.parseConnectionConfig();\n\n this.telemetry = TelemetryManager.getProvider(\n this.name,\n this.config.telemetry,\n );\n this.telemetryMetrics = {\n queryCount: this.telemetry\n .getMeter()\n .createCounter(\"lakebase.query.count\", {\n description: \"Total number of queries executed\",\n unit: \"1\",\n }),\n queryDuration: this.telemetry\n .getMeter()\n .createHistogram(\"lakebase.query.duration\", {\n description: \"Duration of queries executed\",\n unit: \"ms\",\n }),\n };\n\n // validate configuration\n if (this.config.maxPoolSize < 1) {\n throw new Error(\"maxPoolSize must be at least 1\");\n }\n }\n\n /**\n * Execute a SQL query\n *\n * @example\n * ```typescript\n * const users = await connector.query('SELECT * FROM users');\n * const user = await connector.query('SELECT * FROM users WHERE id = $1', [123]);\n * ```\n */\n async query<T extends pg.QueryResultRow>(\n sql: string,\n params?: any[],\n retryCount: number = 0,\n ): Promise<pg.QueryResult<T>> {\n const startTime = Date.now();\n\n return this.telemetry.startActiveSpan(\n \"lakebase.query\",\n {\n attributes: {\n \"db.system\": \"lakebase\",\n \"db.statement\": sql.substring(0, 500),\n \"db.retry_count\": retryCount,\n },\n },\n async (span) => {\n try {\n const pool = await this.getPool();\n const result = await pool.query<T>(sql, params);\n span.setAttribute(\"db.rows_affected\", result.rowCount ?? 0);\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n } catch (error) {\n // retry on auth failure\n if (this.isAuthError(error)) {\n span.addEvent(\"auth_error_retry\");\n await this.rotateCredentials();\n const newPool = await this.getPool();\n const result = await newPool.query<T>(sql, params);\n span.setAttribute(\"db.rows_affected\", result.rowCount ?? 0);\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n }\n\n // retry on transient errors, but only once\n if (this.isTransientError(error) && retryCount < 1) {\n span.addEvent(\"transient_error_retry\");\n await new Promise((resolve) => setTimeout(resolve, 100));\n return await this.query<T>(sql, params, retryCount + 1);\n }\n\n span.recordException(error as Error);\n span.setStatus({ code: SpanStatusCode.ERROR });\n\n throw error;\n } finally {\n const duration = Date.now() - startTime;\n this.telemetryMetrics.queryCount.add(1);\n this.telemetryMetrics.queryDuration.record(duration);\n span.end();\n }\n },\n );\n }\n\n /**\n * Execute a transaction\n *\n * COMMIT and ROLLBACK are automatically managed by the transaction function.\n *\n * @param callback - Callback function to execute within the transaction context\n * @example\n * ```typescript\n * await connector.transaction(async (client) => {\n * await client.query('INSERT INTO accounts (name) VALUES ($1)', ['Alice']);\n * await client.query('INSERT INTO logs (action) VALUES ($1)', ['Created Alice']);\n * });\n * ```\n */\n async transaction<T>(\n callback: (client: pg.PoolClient) => Promise<T>,\n retryCount: number = 0,\n ): Promise<T> {\n const startTime = Date.now();\n return this.telemetry.startActiveSpan(\n \"lakebase.transaction\",\n {\n attributes: {\n \"db.system\": \"lakebase\",\n \"db.retry_count\": retryCount,\n },\n },\n async (span) => {\n const pool = await this.getPool();\n const client = await pool.connect();\n try {\n await client.query(\"BEGIN\");\n const result = await callback(client);\n await client.query(\"COMMIT\");\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n } catch (error) {\n try {\n await client.query(\"ROLLBACK\");\n } catch {}\n // retry on auth failure\n if (this.isAuthError(error)) {\n span.addEvent(\"auth_error_retry\");\n client.release();\n await this.rotateCredentials();\n const newPool = await this.getPool();\n const retryClient = await newPool.connect();\n try {\n await client.query(\"BEGIN\");\n const result = await callback(retryClient);\n await client.query(\"COMMIT\");\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n } catch (retryError) {\n try {\n await retryClient.query(\"ROLLBACK\");\n } catch {}\n throw retryError;\n } finally {\n retryClient.release();\n }\n }\n\n // retry on transient errors, but only once\n if (this.isTransientError(error) && retryCount < 1) {\n span.addEvent(\"transaction_error_retry\");\n client.release();\n await new Promise((resolve) => setTimeout(resolve, 100));\n return await this.transaction<T>(callback, retryCount + 1);\n }\n span.recordException(error as Error);\n span.setStatus({ code: SpanStatusCode.ERROR });\n throw error;\n } finally {\n client.release();\n const duration = Date.now() - startTime;\n this.telemetryMetrics.queryCount.add(1);\n this.telemetryMetrics.queryDuration.record(duration);\n span.end();\n }\n },\n );\n }\n\n /** Check if database connection is healthy */\n async healthCheck(): Promise<boolean> {\n return this.telemetry.startActiveSpan(\n \"lakebase.healthCheck\",\n {},\n async (span) => {\n try {\n const result = await this.query<{ result: number }>(\n \"SELECT 1 as result\",\n );\n const healthy = result.rows[0]?.result === 1;\n span.setAttribute(\"db.healthy\", healthy);\n span.setStatus({ code: SpanStatusCode.OK });\n return healthy;\n } catch {\n span.setAttribute(\"db.healthy\", false);\n span.setStatus({ code: SpanStatusCode.ERROR });\n return false;\n } finally {\n span.end();\n }\n },\n );\n }\n\n /** Close connection pool (call on shutdown) */\n async close(): Promise<void> {\n if (this.pool) {\n await this.pool.end().catch((error: unknown) => {\n console.error(\"Error closing connection pool:\", error);\n });\n this.pool = null;\n }\n this.credentials = null;\n }\n\n /** Setup graceful shutdown to close connection pools */\n shutdown(): void {\n process.on(\"SIGTERM\", () => this.close());\n process.on(\"SIGINT\", () => this.close());\n this.close();\n }\n\n /** Get Databricks workspace client - from config or request context */\n private getWorkspaceClient(): WorkspaceClient {\n if (this.config.workspaceClient) {\n return this.config.workspaceClient;\n }\n\n try {\n const { getRequestContext } = require(\"../../utils\");\n const { serviceDatabricksClient } = getRequestContext();\n\n // cache it for subsequent calls\n this.config.workspaceClient = serviceDatabricksClient;\n return serviceDatabricksClient;\n } catch (_error) {\n throw new Error(\n \"Databricks workspace client not available. Either pass it in config or use within AppKit request context.\",\n );\n }\n }\n\n /** Get or create connection pool */\n private async getPool(): Promise<pg.Pool> {\n if (!this.connectionConfig) {\n throw new Error(\n \"Lakebase connection not configured. \" +\n \"Set PGHOST, PGDATABASE, PGAPPNAME env vars, provide a connectionString, or pass explicit config.\",\n );\n }\n\n if (!this.pool) {\n const creds = await this.getCredentials();\n this.pool = this.createPool(creds);\n }\n return this.pool;\n }\n\n /** Create PostgreSQL pool */\n private createPool(credentials: {\n username: string;\n password: string;\n }): pg.Pool {\n const { host, database, port, sslMode } = this.connectionConfig;\n\n const pool = new pg.Pool({\n host,\n port,\n database,\n user: credentials.username,\n password: credentials.password,\n max: this.config.maxPoolSize,\n idleTimeoutMillis: this.config.idleTimeoutMs,\n connectionTimeoutMillis: this.config.connectionTimeoutMs,\n ssl: sslMode === \"require\" ? { rejectUnauthorized: true } : false,\n });\n\n pool.on(\"error\", (error: Error & { code?: string }) => {\n console.error(\"Connection pool error:\", error.message, {\n code: error.code,\n });\n });\n\n return pool;\n }\n\n /** Get or fetch credentials with caching */\n private async getCredentials(): Promise<{\n username: string;\n password: string;\n }> {\n const now = Date.now();\n\n // return cached if still valid\n if (\n this.credentials &&\n now < this.credentials.expiresAt - this.CACHE_BUFFER_MS\n ) {\n return this.credentials;\n }\n\n // fetch new credentials\n const username = await this.fetchUsername();\n const { token, expiresAt } = await this.fetchPassword();\n\n this.credentials = {\n username,\n password: token,\n expiresAt,\n };\n\n return { username, password: token };\n }\n\n /** Rotate credentials and recreate pool */\n private async rotateCredentials(): Promise<void> {\n // clear cached credentials\n this.credentials = null;\n\n if (this.pool) {\n const oldPool = this.pool;\n this.pool = null;\n oldPool.end().catch((error: unknown) => {\n console.error(\n \"Error closing old connection pool during rotation:\",\n error,\n );\n });\n }\n }\n\n /** Fetch username from Databricks */\n private async fetchUsername(): Promise<string> {\n const workspaceClient = this.getWorkspaceClient();\n const user = await workspaceClient.currentUser.me();\n if (!user.userName) {\n throw new Error(\"Failed to get current user from Databricks workspace\");\n }\n return user.userName;\n }\n\n /** Fetch password (OAuth token) from Databricks */\n private async fetchPassword(): Promise<{ token: string; expiresAt: number }> {\n const workspaceClient = this.getWorkspaceClient();\n const config = new Config({ host: workspaceClient.config.host });\n const apiClient = new ApiClient(config);\n\n if (!this.connectionConfig.appName) {\n throw new Error(`Database app name not found in connection config`);\n }\n\n const credentials = await apiClient.request({\n path: `/api/2.0/database/credentials`,\n method: \"POST\",\n headers: new Headers(),\n raw: false,\n payload: {\n instance_names: [this.connectionConfig.appName],\n request_id: randomUUID(),\n },\n });\n\n if (!this.validateCredentials(credentials)) {\n throw new Error(\n `Failed to generate credentials for instance: ${this.connectionConfig.appName}`,\n );\n }\n\n const expiresAt = new Date(credentials.expiration_time).getTime();\n\n return { token: credentials.token, expiresAt };\n }\n\n /** Check if error is auth failure */\n private isAuthError(error: unknown): boolean {\n return (\n typeof error === \"object\" &&\n error !== null &&\n \"code\" in error &&\n (error as any).code === \"28P01\"\n );\n }\n\n /** Check if error is transient */\n private isTransientError(error: unknown): boolean {\n if (typeof error !== \"object\" || error === null || !(\"code\" in error)) {\n return false;\n }\n\n const code = (error as any).code;\n return (\n code === \"ECONNRESET\" ||\n code === \"ECONNREFUSED\" ||\n code === \"ETIMEDOUT\" ||\n code === \"57P01\" || // admin_shutdown\n code === \"57P03\" || // cannot_connect_now\n code === \"08006\" || // connection_failure\n code === \"08003\" || // connection_does_not_exist\n code === \"08000\" // connection_exception\n );\n }\n\n /** Type guard for credentials */\n private validateCredentials(\n value: unknown,\n ): value is { token: string; expiration_time: string } {\n if (typeof value !== \"object\" || value === null) {\n return false;\n }\n\n const credentials = value as { token: string; expiration_time: string };\n return (\n \"token\" in credentials &&\n typeof credentials.token === \"string\" &&\n \"expiration_time\" in credentials &&\n typeof credentials.expiration_time === \"string\" &&\n new Date(credentials.expiration_time).getTime() > Date.now()\n );\n }\n\n /** Parse connection configuration from config or environment */\n private parseConnectionConfig(): LakebaseConnectionConfig {\n if (this.config.connectionString) {\n return this.parseConnectionString(this.config.connectionString);\n }\n\n // get connection from config\n if (this.config.host && this.config.database && this.config.appName) {\n return {\n host: this.config.host,\n database: this.config.database,\n port: this.config.port ?? 5432,\n sslMode: this.config.sslMode ?? \"require\",\n appName: this.config.appName,\n };\n }\n\n // get connection from environment variables\n const pgHost = process.env.PGHOST;\n const pgDatabase = process.env.PGDATABASE;\n const pgAppName = process.env.PGAPPNAME;\n if (!pgHost || !pgDatabase || !pgAppName) {\n throw new Error(\n \"Lakebase connection not configured. Required env vars: PGHOST, PGDATABASE, PGAPPNAME. \" +\n \"Optional: PGPORT (default: 5432), PGSSLMODE (default: require).\",\n );\n }\n const pgPort = process.env.PGPORT;\n const port = pgPort ? parseInt(pgPort, 10) : 5432;\n\n if (Number.isNaN(port)) {\n throw new Error(`Invalid port: ${pgPort}. Must be a number.`);\n }\n\n const pgSSLMode = process.env.PGSSLMODE;\n const sslMode =\n (pgSSLMode as \"require\" | \"disable\" | \"prefer\") || \"require\";\n\n return {\n host: pgHost,\n database: pgDatabase,\n port,\n sslMode,\n appName: pgAppName,\n };\n }\n\n private parseConnectionString(\n connectionString: string,\n ): LakebaseConnectionConfig {\n const url = new URL(connectionString);\n const appName = url.searchParams.get(\"appName\");\n if (!appName) {\n throw new Error(\"Connection string must include appName parameter\");\n }\n\n return {\n host: url.hostname,\n database: url.pathname.slice(1), // remove leading slash\n port: url.port ? parseInt(url.port, 10) : 5432,\n sslMode:\n (url.searchParams.get(\"sslmode\") as \"require\" | \"disable\" | \"prefer\") ??\n \"require\",\n appName: appName,\n };\n }\n}\n"],"mappings":";;;;;;;;;;;YAWwC;;;;;;;;;;;;;;;;AAuBxC,IAAa,oBAAb,MAA+B;CAe7B,YAAY,YAAsC;cAdlB;yBACG,MAAS;cAGb;qBACmB;AAUhD,OAAK,SAAS,UAAU,kBAAkB,WAAW;AACrD,OAAK,mBAAmB,KAAK,uBAAuB;AAEpD,OAAK,YAAY,iBAAiB,YAChC,KAAK,MACL,KAAK,OAAO,UACb;AACD,OAAK,mBAAmB;GACtB,YAAY,KAAK,UACd,UAAU,CACV,cAAc,wBAAwB;IACrC,aAAa;IACb,MAAM;IACP,CAAC;GACJ,eAAe,KAAK,UACjB,UAAU,CACV,gBAAgB,2BAA2B;IAC1C,aAAa;IACb,MAAM;IACP,CAAC;GACL;AAGD,MAAI,KAAK,OAAO,cAAc,EAC5B,OAAM,IAAI,MAAM,iCAAiC;;;;;;;;;;;CAarD,MAAM,MACJ,KACA,QACA,aAAqB,GACO;EAC5B,MAAM,YAAY,KAAK,KAAK;AAE5B,SAAO,KAAK,UAAU,gBACpB,kBACA,EACE,YAAY;GACV,aAAa;GACb,gBAAgB,IAAI,UAAU,GAAG,IAAI;GACrC,kBAAkB;GACnB,EACF,EACD,OAAO,SAAS;AACd,OAAI;IAEF,MAAM,SAAS,OADF,MAAM,KAAK,SAAS,EACP,MAAS,KAAK,OAAO;AAC/C,SAAK,aAAa,oBAAoB,OAAO,YAAY,EAAE;AAC3D,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AAEd,QAAI,KAAK,YAAY,MAAM,EAAE;AAC3B,UAAK,SAAS,mBAAmB;AACjC,WAAM,KAAK,mBAAmB;KAE9B,MAAM,SAAS,OADC,MAAM,KAAK,SAAS,EACP,MAAS,KAAK,OAAO;AAClD,UAAK,aAAa,oBAAoB,OAAO,YAAY,EAAE;AAC3D,UAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,YAAO;;AAIT,QAAI,KAAK,iBAAiB,MAAM,IAAI,aAAa,GAAG;AAClD,UAAK,SAAS,wBAAwB;AACtC,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,IAAI,CAAC;AACxD,YAAO,MAAM,KAAK,MAAS,KAAK,QAAQ,aAAa,EAAE;;AAGzD,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAE9C,UAAM;aACE;IACR,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,WAAW,IAAI,EAAE;AACvC,SAAK,iBAAiB,cAAc,OAAO,SAAS;AACpD,SAAK,KAAK;;IAGf;;;;;;;;;;;;;;;;CAiBH,MAAM,YACJ,UACA,aAAqB,GACT;EACZ,MAAM,YAAY,KAAK,KAAK;AAC5B,SAAO,KAAK,UAAU,gBACpB,wBACA,EACE,YAAY;GACV,aAAa;GACb,kBAAkB;GACnB,EACF,EACD,OAAO,SAAS;GAEd,MAAM,SAAS,OADF,MAAM,KAAK,SAAS,EACP,SAAS;AACnC,OAAI;AACF,UAAM,OAAO,MAAM,QAAQ;IAC3B,MAAM,SAAS,MAAM,SAAS,OAAO;AACrC,UAAM,OAAO,MAAM,SAAS;AAC5B,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AACd,QAAI;AACF,WAAM,OAAO,MAAM,WAAW;YACxB;AAER,QAAI,KAAK,YAAY,MAAM,EAAE;AAC3B,UAAK,SAAS,mBAAmB;AACjC,YAAO,SAAS;AAChB,WAAM,KAAK,mBAAmB;KAE9B,MAAM,cAAc,OADJ,MAAM,KAAK,SAAS,EACF,SAAS;AAC3C,SAAI;AACF,YAAM,OAAO,MAAM,QAAQ;MAC3B,MAAM,SAAS,MAAM,SAAS,YAAY;AAC1C,YAAM,OAAO,MAAM,SAAS;AAC5B,WAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,aAAO;cACA,YAAY;AACnB,UAAI;AACF,aAAM,YAAY,MAAM,WAAW;cAC7B;AACR,YAAM;eACE;AACR,kBAAY,SAAS;;;AAKzB,QAAI,KAAK,iBAAiB,MAAM,IAAI,aAAa,GAAG;AAClD,UAAK,SAAS,0BAA0B;AACxC,YAAO,SAAS;AAChB,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,IAAI,CAAC;AACxD,YAAO,MAAM,KAAK,YAAe,UAAU,aAAa,EAAE;;AAE5D,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,UAAM;aACE;AACR,WAAO,SAAS;IAChB,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,WAAW,IAAI,EAAE;AACvC,SAAK,iBAAiB,cAAc,OAAO,SAAS;AACpD,SAAK,KAAK;;IAGf;;;CAIH,MAAM,cAAgC;AACpC,SAAO,KAAK,UAAU,gBACpB,wBACA,EAAE,EACF,OAAO,SAAS;AACd,OAAI;IAIF,MAAM,WAHS,MAAM,KAAK,MACxB,qBACD,EACsB,KAAK,IAAI,WAAW;AAC3C,SAAK,aAAa,cAAc,QAAQ;AACxC,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;WACD;AACN,SAAK,aAAa,cAAc,MAAM;AACtC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,WAAO;aACC;AACR,SAAK,KAAK;;IAGf;;;CAIH,MAAM,QAAuB;AAC3B,MAAI,KAAK,MAAM;AACb,SAAM,KAAK,KAAK,KAAK,CAAC,OAAO,UAAmB;AAC9C,YAAQ,MAAM,kCAAkC,MAAM;KACtD;AACF,QAAK,OAAO;;AAEd,OAAK,cAAc;;;CAIrB,WAAiB;AACf,UAAQ,GAAG,iBAAiB,KAAK,OAAO,CAAC;AACzC,UAAQ,GAAG,gBAAgB,KAAK,OAAO,CAAC;AACxC,OAAK,OAAO;;;CAId,AAAQ,qBAAsC;AAC5C,MAAI,KAAK,OAAO,gBACd,QAAO,KAAK,OAAO;AAGrB,MAAI;GACF,MAAM,EAAE;GACR,MAAM,EAAE,4BAA4B,mBAAmB;AAGvD,QAAK,OAAO,kBAAkB;AAC9B,UAAO;WACA,QAAQ;AACf,SAAM,IAAI,MACR,4GACD;;;;CAKL,MAAc,UAA4B;AACxC,MAAI,CAAC,KAAK,iBACR,OAAM,IAAI,MACR,uIAED;AAGH,MAAI,CAAC,KAAK,MAAM;GACd,MAAM,QAAQ,MAAM,KAAK,gBAAgB;AACzC,QAAK,OAAO,KAAK,WAAW,MAAM;;AAEpC,SAAO,KAAK;;;CAId,AAAQ,WAAW,aAGP;EACV,MAAM,EAAE,MAAM,UAAU,MAAM,YAAY,KAAK;EAE/C,MAAM,OAAO,IAAI,GAAG,KAAK;GACvB;GACA;GACA;GACA,MAAM,YAAY;GAClB,UAAU,YAAY;GACtB,KAAK,KAAK,OAAO;GACjB,mBAAmB,KAAK,OAAO;GAC/B,yBAAyB,KAAK,OAAO;GACrC,KAAK,YAAY,YAAY,EAAE,oBAAoB,MAAM,GAAG;GAC7D,CAAC;AAEF,OAAK,GAAG,UAAU,UAAqC;AACrD,WAAQ,MAAM,0BAA0B,MAAM,SAAS,EACrD,MAAM,MAAM,MACb,CAAC;IACF;AAEF,SAAO;;;CAIT,MAAc,iBAGX;EACD,MAAM,MAAM,KAAK,KAAK;AAGtB,MACE,KAAK,eACL,MAAM,KAAK,YAAY,YAAY,KAAK,gBAExC,QAAO,KAAK;EAId,MAAM,WAAW,MAAM,KAAK,eAAe;EAC3C,MAAM,EAAE,OAAO,cAAc,MAAM,KAAK,eAAe;AAEvD,OAAK,cAAc;GACjB;GACA,UAAU;GACV;GACD;AAED,SAAO;GAAE;GAAU,UAAU;GAAO;;;CAItC,MAAc,oBAAmC;AAE/C,OAAK,cAAc;AAEnB,MAAI,KAAK,MAAM;GACb,MAAM,UAAU,KAAK;AACrB,QAAK,OAAO;AACZ,WAAQ,KAAK,CAAC,OAAO,UAAmB;AACtC,YAAQ,MACN,sDACA,MACD;KACD;;;;CAKN,MAAc,gBAAiC;EAE7C,MAAM,OAAO,MADW,KAAK,oBAAoB,CACd,YAAY,IAAI;AACnD,MAAI,CAAC,KAAK,SACR,OAAM,IAAI,MAAM,uDAAuD;AAEzE,SAAO,KAAK;;;CAId,MAAc,gBAA+D;EAG3E,MAAM,YAAY,IAAI,UADP,IAAI,OAAO,EAAE,MADJ,KAAK,oBAAoB,CACC,OAAO,MAAM,CAAC,CACzB;AAEvC,MAAI,CAAC,KAAK,iBAAiB,QACzB,OAAM,IAAI,MAAM,mDAAmD;EAGrE,MAAM,cAAc,MAAM,UAAU,QAAQ;GAC1C,MAAM;GACN,QAAQ;GACR,SAAS,IAAI,SAAS;GACtB,KAAK;GACL,SAAS;IACP,gBAAgB,CAAC,KAAK,iBAAiB,QAAQ;IAC/C,YAAY,YAAY;IACzB;GACF,CAAC;AAEF,MAAI,CAAC,KAAK,oBAAoB,YAAY,CACxC,OAAM,IAAI,MACR,gDAAgD,KAAK,iBAAiB,UACvE;EAGH,MAAM,YAAY,IAAI,KAAK,YAAY,gBAAgB,CAAC,SAAS;AAEjE,SAAO;GAAE,OAAO,YAAY;GAAO;GAAW;;;CAIhD,AAAQ,YAAY,OAAyB;AAC3C,SACE,OAAO,UAAU,YACjB,UAAU,QACV,UAAU,SACT,MAAc,SAAS;;;CAK5B,AAAQ,iBAAiB,OAAyB;AAChD,MAAI,OAAO,UAAU,YAAY,UAAU,QAAQ,EAAE,UAAU,OAC7D,QAAO;EAGT,MAAM,OAAQ,MAAc;AAC5B,SACE,SAAS,gBACT,SAAS,kBACT,SAAS,eACT,SAAS,WACT,SAAS,WACT,SAAS,WACT,SAAS,WACT,SAAS;;;CAKb,AAAQ,oBACN,OACqD;AACrD,MAAI,OAAO,UAAU,YAAY,UAAU,KACzC,QAAO;EAGT,MAAM,cAAc;AACpB,SACE,WAAW,eACX,OAAO,YAAY,UAAU,YAC7B,qBAAqB,eACrB,OAAO,YAAY,oBAAoB,YACvC,IAAI,KAAK,YAAY,gBAAgB,CAAC,SAAS,GAAG,KAAK,KAAK;;;CAKhE,AAAQ,wBAAkD;AACxD,MAAI,KAAK,OAAO,iBACd,QAAO,KAAK,sBAAsB,KAAK,OAAO,iBAAiB;AAIjE,MAAI,KAAK,OAAO,QAAQ,KAAK,OAAO,YAAY,KAAK,OAAO,QAC1D,QAAO;GACL,MAAM,KAAK,OAAO;GAClB,UAAU,KAAK,OAAO;GACtB,MAAM,KAAK,OAAO,QAAQ;GAC1B,SAAS,KAAK,OAAO,WAAW;GAChC,SAAS,KAAK,OAAO;GACtB;EAIH,MAAM,SAAS,QAAQ,IAAI;EAC3B,MAAM,aAAa,QAAQ,IAAI;EAC/B,MAAM,YAAY,QAAQ,IAAI;AAC9B,MAAI,CAAC,UAAU,CAAC,cAAc,CAAC,UAC7B,OAAM,IAAI,MACR,wJAED;EAEH,MAAM,SAAS,QAAQ,IAAI;EAC3B,MAAM,OAAO,SAAS,SAAS,QAAQ,GAAG,GAAG;AAE7C,MAAI,OAAO,MAAM,KAAK,CACpB,OAAM,IAAI,MAAM,iBAAiB,OAAO,qBAAqB;AAO/D,SAAO;GACL,MAAM;GACN,UAAU;GACV;GACA,SARgB,QAAQ,IAAI,aAEuB;GAOnD,SAAS;GACV;;CAGH,AAAQ,sBACN,kBAC0B;EAC1B,MAAM,MAAM,IAAI,IAAI,iBAAiB;EACrC,MAAM,UAAU,IAAI,aAAa,IAAI,UAAU;AAC/C,MAAI,CAAC,QACH,OAAM,IAAI,MAAM,mDAAmD;AAGrE,SAAO;GACL,MAAM,IAAI;GACV,UAAU,IAAI,SAAS,MAAM,EAAE;GAC/B,MAAM,IAAI,OAAO,SAAS,IAAI,MAAM,GAAG,GAAG;GAC1C,SACG,IAAI,aAAa,IAAI,UAAU,IAChC;GACO;GACV"}
1
+ {"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/lakebase/client.ts"],"sourcesContent":["import { randomUUID } from \"node:crypto\";\nimport type { WorkspaceClient } from \"@databricks/sdk-experimental\";\nimport { ApiClient, Config } from \"@databricks/sdk-experimental\";\nimport pg from \"pg\";\nimport {\n type Counter,\n type Histogram,\n SpanStatusCode,\n TelemetryManager,\n type TelemetryProvider,\n} from \"@/telemetry\";\nimport { deepMerge } from \"../../utils\";\nimport { lakebaseDefaults } from \"./defaults\";\nimport type {\n LakebaseConfig,\n LakebaseConnectionConfig,\n LakebaseCredentials,\n} from \"./types\";\n\n/**\n * Enterprise-grade connector for Databricks Lakebase\n * @example Simplest - everything from env/context\n * ```typescript\n * const connector = new LakebaseConnector();\n * await connector.query('SELECT * FROM users');\n * ```\n *\n * @example With explicit connection string\n * ```typescript\n * const connector = new LakebaseConnector({\n * connectionString: 'postgresql://...'\n * });\n * ```\n */\nexport class LakebaseConnector {\n private readonly name: string = \"lakebase\";\n private readonly CACHE_BUFFER_MS = 2 * 60 * 1000;\n private readonly config: LakebaseConfig;\n private readonly connectionConfig: LakebaseConnectionConfig;\n private pool: pg.Pool | null = null;\n private credentials: LakebaseCredentials | null = null;\n\n // telemetry\n private readonly telemetry: TelemetryProvider;\n private readonly telemetryMetrics: {\n queryCount: Counter;\n queryDuration: Histogram;\n };\n\n constructor(userConfig?: Partial<LakebaseConfig>) {\n this.config = deepMerge(lakebaseDefaults, userConfig);\n this.connectionConfig = this.parseConnectionConfig();\n\n this.telemetry = TelemetryManager.getProvider(\n this.name,\n this.config.telemetry,\n );\n this.telemetryMetrics = {\n queryCount: this.telemetry\n .getMeter()\n .createCounter(\"lakebase.query.count\", {\n description: \"Total number of queries executed\",\n unit: \"1\",\n }),\n queryDuration: this.telemetry\n .getMeter()\n .createHistogram(\"lakebase.query.duration\", {\n description: \"Duration of queries executed\",\n unit: \"ms\",\n }),\n };\n\n // validate configuration\n if (this.config.maxPoolSize < 1) {\n throw new Error(\"maxPoolSize must be at least 1\");\n }\n }\n\n /**\n * Execute a SQL query\n *\n * @example\n * ```typescript\n * const users = await connector.query('SELECT * FROM users');\n * const user = await connector.query('SELECT * FROM users WHERE id = $1', [123]);\n * ```\n */\n async query<T extends pg.QueryResultRow>(\n sql: string,\n params?: any[],\n retryCount: number = 0,\n ): Promise<pg.QueryResult<T>> {\n const startTime = Date.now();\n\n return this.telemetry.startActiveSpan(\n \"lakebase.query\",\n {\n attributes: {\n \"db.system\": \"lakebase\",\n \"db.statement\": sql.substring(0, 500),\n \"db.retry_count\": retryCount,\n },\n },\n async (span) => {\n try {\n const pool = await this.getPool();\n const result = await pool.query<T>(sql, params);\n span.setAttribute(\"db.rows_affected\", result.rowCount ?? 0);\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n } catch (error) {\n // retry on auth failure\n if (this.isAuthError(error)) {\n span.addEvent(\"auth_error_retry\");\n await this.rotateCredentials();\n const newPool = await this.getPool();\n const result = await newPool.query<T>(sql, params);\n span.setAttribute(\"db.rows_affected\", result.rowCount ?? 0);\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n }\n\n // retry on transient errors, but only once\n if (this.isTransientError(error) && retryCount < 1) {\n span.addEvent(\"transient_error_retry\");\n await new Promise((resolve) => setTimeout(resolve, 100));\n return await this.query<T>(sql, params, retryCount + 1);\n }\n\n span.recordException(error as Error);\n span.setStatus({ code: SpanStatusCode.ERROR });\n\n throw error;\n } finally {\n const duration = Date.now() - startTime;\n this.telemetryMetrics.queryCount.add(1);\n this.telemetryMetrics.queryDuration.record(duration);\n span.end();\n }\n },\n );\n }\n\n /**\n * Execute a transaction\n *\n * COMMIT and ROLLBACK are automatically managed by the transaction function.\n *\n * @param callback - Callback function to execute within the transaction context\n * @example\n * ```typescript\n * await connector.transaction(async (client) => {\n * await client.query('INSERT INTO accounts (name) VALUES ($1)', ['Alice']);\n * await client.query('INSERT INTO logs (action) VALUES ($1)', ['Created Alice']);\n * });\n * ```\n */\n async transaction<T>(\n callback: (client: pg.PoolClient) => Promise<T>,\n retryCount: number = 0,\n ): Promise<T> {\n const startTime = Date.now();\n return this.telemetry.startActiveSpan(\n \"lakebase.transaction\",\n {\n attributes: {\n \"db.system\": \"lakebase\",\n \"db.retry_count\": retryCount,\n },\n },\n async (span) => {\n const pool = await this.getPool();\n const client = await pool.connect();\n try {\n await client.query(\"BEGIN\");\n const result = await callback(client);\n await client.query(\"COMMIT\");\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n } catch (error) {\n try {\n await client.query(\"ROLLBACK\");\n } catch {}\n // retry on auth failure\n if (this.isAuthError(error)) {\n span.addEvent(\"auth_error_retry\");\n client.release();\n await this.rotateCredentials();\n const newPool = await this.getPool();\n const retryClient = await newPool.connect();\n try {\n await client.query(\"BEGIN\");\n const result = await callback(retryClient);\n await client.query(\"COMMIT\");\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n } catch (retryError) {\n try {\n await retryClient.query(\"ROLLBACK\");\n } catch {}\n throw retryError;\n } finally {\n retryClient.release();\n }\n }\n\n // retry on transient errors, but only once\n if (this.isTransientError(error) && retryCount < 1) {\n span.addEvent(\"transaction_error_retry\");\n client.release();\n await new Promise((resolve) => setTimeout(resolve, 100));\n return await this.transaction<T>(callback, retryCount + 1);\n }\n span.recordException(error as Error);\n span.setStatus({ code: SpanStatusCode.ERROR });\n throw error;\n } finally {\n client.release();\n const duration = Date.now() - startTime;\n this.telemetryMetrics.queryCount.add(1);\n this.telemetryMetrics.queryDuration.record(duration);\n span.end();\n }\n },\n );\n }\n\n /** Check if database connection is healthy */\n async healthCheck(): Promise<boolean> {\n return this.telemetry.startActiveSpan(\n \"lakebase.healthCheck\",\n {},\n async (span) => {\n try {\n const result = await this.query<{ result: number }>(\n \"SELECT 1 as result\",\n );\n const healthy = result.rows[0]?.result === 1;\n span.setAttribute(\"db.healthy\", healthy);\n span.setStatus({ code: SpanStatusCode.OK });\n return healthy;\n } catch {\n span.setAttribute(\"db.healthy\", false);\n span.setStatus({ code: SpanStatusCode.ERROR });\n return false;\n } finally {\n span.end();\n }\n },\n );\n }\n\n /** Close connection pool (call on shutdown) */\n async close(): Promise<void> {\n if (this.pool) {\n await this.pool.end().catch((error: unknown) => {\n console.error(\"Error closing connection pool:\", error);\n });\n this.pool = null;\n }\n this.credentials = null;\n }\n\n /** Setup graceful shutdown to close connection pools */\n shutdown(): void {\n process.on(\"SIGTERM\", () => this.close());\n process.on(\"SIGINT\", () => this.close());\n this.close();\n }\n\n /** Get Databricks workspace client - from config or execution context */\n private getWorkspaceClient(): WorkspaceClient {\n if (this.config.workspaceClient) {\n return this.config.workspaceClient;\n }\n\n try {\n const { getWorkspaceClient: getClient } = require(\"../../context\");\n const client = getClient();\n\n // cache it for subsequent calls\n this.config.workspaceClient = client;\n return client;\n } catch (_error) {\n throw new Error(\n \"Databricks workspace client not available. Either pass it in config or ensure ServiceContext is initialized.\",\n );\n }\n }\n\n /** Get or create connection pool */\n private async getPool(): Promise<pg.Pool> {\n if (!this.connectionConfig) {\n throw new Error(\n \"Lakebase connection not configured. \" +\n \"Set PGHOST, PGDATABASE, PGAPPNAME env vars, provide a connectionString, or pass explicit config.\",\n );\n }\n\n if (!this.pool) {\n const creds = await this.getCredentials();\n this.pool = this.createPool(creds);\n }\n return this.pool;\n }\n\n /** Create PostgreSQL pool */\n private createPool(credentials: {\n username: string;\n password: string;\n }): pg.Pool {\n const { host, database, port, sslMode } = this.connectionConfig;\n\n const pool = new pg.Pool({\n host,\n port,\n database,\n user: credentials.username,\n password: credentials.password,\n max: this.config.maxPoolSize,\n idleTimeoutMillis: this.config.idleTimeoutMs,\n connectionTimeoutMillis: this.config.connectionTimeoutMs,\n ssl: sslMode === \"require\" ? { rejectUnauthorized: true } : false,\n });\n\n pool.on(\"error\", (error: Error & { code?: string }) => {\n console.error(\"Connection pool error:\", error.message, {\n code: error.code,\n });\n });\n\n return pool;\n }\n\n /** Get or fetch credentials with caching */\n private async getCredentials(): Promise<{\n username: string;\n password: string;\n }> {\n const now = Date.now();\n\n // return cached if still valid\n if (\n this.credentials &&\n now < this.credentials.expiresAt - this.CACHE_BUFFER_MS\n ) {\n return this.credentials;\n }\n\n // fetch new credentials\n const username = await this.fetchUsername();\n const { token, expiresAt } = await this.fetchPassword();\n\n this.credentials = {\n username,\n password: token,\n expiresAt,\n };\n\n return { username, password: token };\n }\n\n /** Rotate credentials and recreate pool */\n private async rotateCredentials(): Promise<void> {\n // clear cached credentials\n this.credentials = null;\n\n if (this.pool) {\n const oldPool = this.pool;\n this.pool = null;\n oldPool.end().catch((error: unknown) => {\n console.error(\n \"Error closing old connection pool during rotation:\",\n error,\n );\n });\n }\n }\n\n /** Fetch username from Databricks */\n private async fetchUsername(): Promise<string> {\n const workspaceClient = this.getWorkspaceClient();\n const user = await workspaceClient.currentUser.me();\n if (!user.userName) {\n throw new Error(\"Failed to get current user from Databricks workspace\");\n }\n return user.userName;\n }\n\n /** Fetch password (OAuth token) from Databricks */\n private async fetchPassword(): Promise<{ token: string; expiresAt: number }> {\n const workspaceClient = this.getWorkspaceClient();\n const config = new Config({ host: workspaceClient.config.host });\n const apiClient = new ApiClient(config);\n\n if (!this.connectionConfig.appName) {\n throw new Error(`Database app name not found in connection config`);\n }\n\n const credentials = await apiClient.request({\n path: `/api/2.0/database/credentials`,\n method: \"POST\",\n headers: new Headers(),\n raw: false,\n payload: {\n instance_names: [this.connectionConfig.appName],\n request_id: randomUUID(),\n },\n });\n\n if (!this.validateCredentials(credentials)) {\n throw new Error(\n `Failed to generate credentials for instance: ${this.connectionConfig.appName}`,\n );\n }\n\n const expiresAt = new Date(credentials.expiration_time).getTime();\n\n return { token: credentials.token, expiresAt };\n }\n\n /** Check if error is auth failure */\n private isAuthError(error: unknown): boolean {\n return (\n typeof error === \"object\" &&\n error !== null &&\n \"code\" in error &&\n (error as any).code === \"28P01\"\n );\n }\n\n /** Check if error is transient */\n private isTransientError(error: unknown): boolean {\n if (typeof error !== \"object\" || error === null || !(\"code\" in error)) {\n return false;\n }\n\n const code = (error as any).code;\n return (\n code === \"ECONNRESET\" ||\n code === \"ECONNREFUSED\" ||\n code === \"ETIMEDOUT\" ||\n code === \"57P01\" || // admin_shutdown\n code === \"57P03\" || // cannot_connect_now\n code === \"08006\" || // connection_failure\n code === \"08003\" || // connection_does_not_exist\n code === \"08000\" // connection_exception\n );\n }\n\n /** Type guard for credentials */\n private validateCredentials(\n value: unknown,\n ): value is { token: string; expiration_time: string } {\n if (typeof value !== \"object\" || value === null) {\n return false;\n }\n\n const credentials = value as { token: string; expiration_time: string };\n return (\n \"token\" in credentials &&\n typeof credentials.token === \"string\" &&\n \"expiration_time\" in credentials &&\n typeof credentials.expiration_time === \"string\" &&\n new Date(credentials.expiration_time).getTime() > Date.now()\n );\n }\n\n /** Parse connection configuration from config or environment */\n private parseConnectionConfig(): LakebaseConnectionConfig {\n if (this.config.connectionString) {\n return this.parseConnectionString(this.config.connectionString);\n }\n\n // get connection from config\n if (this.config.host && this.config.database && this.config.appName) {\n return {\n host: this.config.host,\n database: this.config.database,\n port: this.config.port ?? 5432,\n sslMode: this.config.sslMode ?? \"require\",\n appName: this.config.appName,\n };\n }\n\n // get connection from environment variables\n const pgHost = process.env.PGHOST;\n const pgDatabase = process.env.PGDATABASE;\n const pgAppName = process.env.PGAPPNAME;\n if (!pgHost || !pgDatabase || !pgAppName) {\n throw new Error(\n \"Lakebase connection not configured. Required env vars: PGHOST, PGDATABASE, PGAPPNAME. \" +\n \"Optional: PGPORT (default: 5432), PGSSLMODE (default: require).\",\n );\n }\n const pgPort = process.env.PGPORT;\n const port = pgPort ? parseInt(pgPort, 10) : 5432;\n\n if (Number.isNaN(port)) {\n throw new Error(`Invalid port: ${pgPort}. Must be a number.`);\n }\n\n const pgSSLMode = process.env.PGSSLMODE;\n const sslMode =\n (pgSSLMode as \"require\" | \"disable\" | \"prefer\") || \"require\";\n\n return {\n host: pgHost,\n database: pgDatabase,\n port,\n sslMode,\n appName: pgAppName,\n };\n }\n\n private parseConnectionString(\n connectionString: string,\n ): LakebaseConnectionConfig {\n const url = new URL(connectionString);\n const appName = url.searchParams.get(\"appName\");\n if (!appName) {\n throw new Error(\"Connection string must include appName parameter\");\n }\n\n return {\n host: url.hostname,\n database: url.pathname.slice(1), // remove leading slash\n port: url.port ? parseInt(url.port, 10) : 5432,\n sslMode:\n (url.searchParams.get(\"sslmode\") as \"require\" | \"disable\" | \"prefer\") ??\n \"require\",\n appName: appName,\n };\n }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;AAkCA,IAAa,oBAAb,MAA+B;CAe7B,YAAY,YAAsC;cAdlB;yBACG,MAAS;cAGb;qBACmB;AAUhD,OAAK,SAAS,UAAU,kBAAkB,WAAW;AACrD,OAAK,mBAAmB,KAAK,uBAAuB;AAEpD,OAAK,YAAY,iBAAiB,YAChC,KAAK,MACL,KAAK,OAAO,UACb;AACD,OAAK,mBAAmB;GACtB,YAAY,KAAK,UACd,UAAU,CACV,cAAc,wBAAwB;IACrC,aAAa;IACb,MAAM;IACP,CAAC;GACJ,eAAe,KAAK,UACjB,UAAU,CACV,gBAAgB,2BAA2B;IAC1C,aAAa;IACb,MAAM;IACP,CAAC;GACL;AAGD,MAAI,KAAK,OAAO,cAAc,EAC5B,OAAM,IAAI,MAAM,iCAAiC;;;;;;;;;;;CAarD,MAAM,MACJ,KACA,QACA,aAAqB,GACO;EAC5B,MAAM,YAAY,KAAK,KAAK;AAE5B,SAAO,KAAK,UAAU,gBACpB,kBACA,EACE,YAAY;GACV,aAAa;GACb,gBAAgB,IAAI,UAAU,GAAG,IAAI;GACrC,kBAAkB;GACnB,EACF,EACD,OAAO,SAAS;AACd,OAAI;IAEF,MAAM,SAAS,OADF,MAAM,KAAK,SAAS,EACP,MAAS,KAAK,OAAO;AAC/C,SAAK,aAAa,oBAAoB,OAAO,YAAY,EAAE;AAC3D,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AAEd,QAAI,KAAK,YAAY,MAAM,EAAE;AAC3B,UAAK,SAAS,mBAAmB;AACjC,WAAM,KAAK,mBAAmB;KAE9B,MAAM,SAAS,OADC,MAAM,KAAK,SAAS,EACP,MAAS,KAAK,OAAO;AAClD,UAAK,aAAa,oBAAoB,OAAO,YAAY,EAAE;AAC3D,UAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,YAAO;;AAIT,QAAI,KAAK,iBAAiB,MAAM,IAAI,aAAa,GAAG;AAClD,UAAK,SAAS,wBAAwB;AACtC,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,IAAI,CAAC;AACxD,YAAO,MAAM,KAAK,MAAS,KAAK,QAAQ,aAAa,EAAE;;AAGzD,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAE9C,UAAM;aACE;IACR,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,WAAW,IAAI,EAAE;AACvC,SAAK,iBAAiB,cAAc,OAAO,SAAS;AACpD,SAAK,KAAK;;IAGf;;;;;;;;;;;;;;;;CAiBH,MAAM,YACJ,UACA,aAAqB,GACT;EACZ,MAAM,YAAY,KAAK,KAAK;AAC5B,SAAO,KAAK,UAAU,gBACpB,wBACA,EACE,YAAY;GACV,aAAa;GACb,kBAAkB;GACnB,EACF,EACD,OAAO,SAAS;GAEd,MAAM,SAAS,OADF,MAAM,KAAK,SAAS,EACP,SAAS;AACnC,OAAI;AACF,UAAM,OAAO,MAAM,QAAQ;IAC3B,MAAM,SAAS,MAAM,SAAS,OAAO;AACrC,UAAM,OAAO,MAAM,SAAS;AAC5B,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AACd,QAAI;AACF,WAAM,OAAO,MAAM,WAAW;YACxB;AAER,QAAI,KAAK,YAAY,MAAM,EAAE;AAC3B,UAAK,SAAS,mBAAmB;AACjC,YAAO,SAAS;AAChB,WAAM,KAAK,mBAAmB;KAE9B,MAAM,cAAc,OADJ,MAAM,KAAK,SAAS,EACF,SAAS;AAC3C,SAAI;AACF,YAAM,OAAO,MAAM,QAAQ;MAC3B,MAAM,SAAS,MAAM,SAAS,YAAY;AAC1C,YAAM,OAAO,MAAM,SAAS;AAC5B,WAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,aAAO;cACA,YAAY;AACnB,UAAI;AACF,aAAM,YAAY,MAAM,WAAW;cAC7B;AACR,YAAM;eACE;AACR,kBAAY,SAAS;;;AAKzB,QAAI,KAAK,iBAAiB,MAAM,IAAI,aAAa,GAAG;AAClD,UAAK,SAAS,0BAA0B;AACxC,YAAO,SAAS;AAChB,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,IAAI,CAAC;AACxD,YAAO,MAAM,KAAK,YAAe,UAAU,aAAa,EAAE;;AAE5D,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,UAAM;aACE;AACR,WAAO,SAAS;IAChB,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,WAAW,IAAI,EAAE;AACvC,SAAK,iBAAiB,cAAc,OAAO,SAAS;AACpD,SAAK,KAAK;;IAGf;;;CAIH,MAAM,cAAgC;AACpC,SAAO,KAAK,UAAU,gBACpB,wBACA,EAAE,EACF,OAAO,SAAS;AACd,OAAI;IAIF,MAAM,WAHS,MAAM,KAAK,MACxB,qBACD,EACsB,KAAK,IAAI,WAAW;AAC3C,SAAK,aAAa,cAAc,QAAQ;AACxC,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;WACD;AACN,SAAK,aAAa,cAAc,MAAM;AACtC,SAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,WAAO;aACC;AACR,SAAK,KAAK;;IAGf;;;CAIH,MAAM,QAAuB;AAC3B,MAAI,KAAK,MAAM;AACb,SAAM,KAAK,KAAK,KAAK,CAAC,OAAO,UAAmB;AAC9C,YAAQ,MAAM,kCAAkC,MAAM;KACtD;AACF,QAAK,OAAO;;AAEd,OAAK,cAAc;;;CAIrB,WAAiB;AACf,UAAQ,GAAG,iBAAiB,KAAK,OAAO,CAAC;AACzC,UAAQ,GAAG,gBAAgB,KAAK,OAAO,CAAC;AACxC,OAAK,OAAO;;;CAId,AAAQ,qBAAsC;AAC5C,MAAI,KAAK,OAAO,gBACd,QAAO,KAAK,OAAO;AAGrB,MAAI;GACF,MAAM,EAAE,oBAAoB;GAC5B,MAAM,SAAS,WAAW;AAG1B,QAAK,OAAO,kBAAkB;AAC9B,UAAO;WACA,QAAQ;AACf,SAAM,IAAI,MACR,+GACD;;;;CAKL,MAAc,UAA4B;AACxC,MAAI,CAAC,KAAK,iBACR,OAAM,IAAI,MACR,uIAED;AAGH,MAAI,CAAC,KAAK,MAAM;GACd,MAAM,QAAQ,MAAM,KAAK,gBAAgB;AACzC,QAAK,OAAO,KAAK,WAAW,MAAM;;AAEpC,SAAO,KAAK;;;CAId,AAAQ,WAAW,aAGP;EACV,MAAM,EAAE,MAAM,UAAU,MAAM,YAAY,KAAK;EAE/C,MAAM,OAAO,IAAI,GAAG,KAAK;GACvB;GACA;GACA;GACA,MAAM,YAAY;GAClB,UAAU,YAAY;GACtB,KAAK,KAAK,OAAO;GACjB,mBAAmB,KAAK,OAAO;GAC/B,yBAAyB,KAAK,OAAO;GACrC,KAAK,YAAY,YAAY,EAAE,oBAAoB,MAAM,GAAG;GAC7D,CAAC;AAEF,OAAK,GAAG,UAAU,UAAqC;AACrD,WAAQ,MAAM,0BAA0B,MAAM,SAAS,EACrD,MAAM,MAAM,MACb,CAAC;IACF;AAEF,SAAO;;;CAIT,MAAc,iBAGX;EACD,MAAM,MAAM,KAAK,KAAK;AAGtB,MACE,KAAK,eACL,MAAM,KAAK,YAAY,YAAY,KAAK,gBAExC,QAAO,KAAK;EAId,MAAM,WAAW,MAAM,KAAK,eAAe;EAC3C,MAAM,EAAE,OAAO,cAAc,MAAM,KAAK,eAAe;AAEvD,OAAK,cAAc;GACjB;GACA,UAAU;GACV;GACD;AAED,SAAO;GAAE;GAAU,UAAU;GAAO;;;CAItC,MAAc,oBAAmC;AAE/C,OAAK,cAAc;AAEnB,MAAI,KAAK,MAAM;GACb,MAAM,UAAU,KAAK;AACrB,QAAK,OAAO;AACZ,WAAQ,KAAK,CAAC,OAAO,UAAmB;AACtC,YAAQ,MACN,sDACA,MACD;KACD;;;;CAKN,MAAc,gBAAiC;EAE7C,MAAM,OAAO,MADW,KAAK,oBAAoB,CACd,YAAY,IAAI;AACnD,MAAI,CAAC,KAAK,SACR,OAAM,IAAI,MAAM,uDAAuD;AAEzE,SAAO,KAAK;;;CAId,MAAc,gBAA+D;EAG3E,MAAM,YAAY,IAAI,UADP,IAAI,OAAO,EAAE,MADJ,KAAK,oBAAoB,CACC,OAAO,MAAM,CAAC,CACzB;AAEvC,MAAI,CAAC,KAAK,iBAAiB,QACzB,OAAM,IAAI,MAAM,mDAAmD;EAGrE,MAAM,cAAc,MAAM,UAAU,QAAQ;GAC1C,MAAM;GACN,QAAQ;GACR,SAAS,IAAI,SAAS;GACtB,KAAK;GACL,SAAS;IACP,gBAAgB,CAAC,KAAK,iBAAiB,QAAQ;IAC/C,YAAY,YAAY;IACzB;GACF,CAAC;AAEF,MAAI,CAAC,KAAK,oBAAoB,YAAY,CACxC,OAAM,IAAI,MACR,gDAAgD,KAAK,iBAAiB,UACvE;EAGH,MAAM,YAAY,IAAI,KAAK,YAAY,gBAAgB,CAAC,SAAS;AAEjE,SAAO;GAAE,OAAO,YAAY;GAAO;GAAW;;;CAIhD,AAAQ,YAAY,OAAyB;AAC3C,SACE,OAAO,UAAU,YACjB,UAAU,QACV,UAAU,SACT,MAAc,SAAS;;;CAK5B,AAAQ,iBAAiB,OAAyB;AAChD,MAAI,OAAO,UAAU,YAAY,UAAU,QAAQ,EAAE,UAAU,OAC7D,QAAO;EAGT,MAAM,OAAQ,MAAc;AAC5B,SACE,SAAS,gBACT,SAAS,kBACT,SAAS,eACT,SAAS,WACT,SAAS,WACT,SAAS,WACT,SAAS,WACT,SAAS;;;CAKb,AAAQ,oBACN,OACqD;AACrD,MAAI,OAAO,UAAU,YAAY,UAAU,KACzC,QAAO;EAGT,MAAM,cAAc;AACpB,SACE,WAAW,eACX,OAAO,YAAY,UAAU,YAC7B,qBAAqB,eACrB,OAAO,YAAY,oBAAoB,YACvC,IAAI,KAAK,YAAY,gBAAgB,CAAC,SAAS,GAAG,KAAK,KAAK;;;CAKhE,AAAQ,wBAAkD;AACxD,MAAI,KAAK,OAAO,iBACd,QAAO,KAAK,sBAAsB,KAAK,OAAO,iBAAiB;AAIjE,MAAI,KAAK,OAAO,QAAQ,KAAK,OAAO,YAAY,KAAK,OAAO,QAC1D,QAAO;GACL,MAAM,KAAK,OAAO;GAClB,UAAU,KAAK,OAAO;GACtB,MAAM,KAAK,OAAO,QAAQ;GAC1B,SAAS,KAAK,OAAO,WAAW;GAChC,SAAS,KAAK,OAAO;GACtB;EAIH,MAAM,SAAS,QAAQ,IAAI;EAC3B,MAAM,aAAa,QAAQ,IAAI;EAC/B,MAAM,YAAY,QAAQ,IAAI;AAC9B,MAAI,CAAC,UAAU,CAAC,cAAc,CAAC,UAC7B,OAAM,IAAI,MACR,wJAED;EAEH,MAAM,SAAS,QAAQ,IAAI;EAC3B,MAAM,OAAO,SAAS,SAAS,QAAQ,GAAG,GAAG;AAE7C,MAAI,OAAO,MAAM,KAAK,CACpB,OAAM,IAAI,MAAM,iBAAiB,OAAO,qBAAqB;AAO/D,SAAO;GACL,MAAM;GACN,UAAU;GACV;GACA,SARgB,QAAQ,IAAI,aAEuB;GAOnD,SAAS;GACV;;CAGH,AAAQ,sBACN,kBAC0B;EAC1B,MAAM,MAAM,IAAI,IAAI,iBAAiB;EACrC,MAAM,UAAU,IAAI,aAAa,IAAI,UAAU;AAC/C,MAAI,CAAC,QACH,OAAM,IAAI,MAAM,mDAAmD;AAGrE,SAAO;GACL,MAAM,IAAI;GACV,UAAU,IAAI,SAAS,MAAM,EAAE;GAC/B,MAAM,IAAI,OAAO,SAAS,IAAI,MAAM,GAAG,GAAG;GAC1C,SACG,IAAI,aAAa,IAAI,UAAU,IAChC;GACO;GACV"}
@@ -1 +1 @@
1
- {"version":3,"file":"defaults.js","names":["lakebaseDefaults: LakebaseConfig"],"sources":["../../../src/connectors/lakebase/defaults.ts"],"sourcesContent":["import type { LakebaseConfig } from \"./types\";\n\n/** Default configuration for Lakebase connector */\nexport const lakebaseDefaults: LakebaseConfig = {\n port: 5432,\n sslMode: \"require\",\n maxPoolSize: 10,\n idleTimeoutMs: 30_000,\n connectionTimeoutMs: 10_000,\n};\n"],"mappings":";;AAGA,MAAaA,mBAAmC;CAC9C,MAAM;CACN,SAAS;CACT,aAAa;CACb,eAAe;CACf,qBAAqB;CACtB"}
1
+ {"version":3,"file":"defaults.js","names":[],"sources":["../../../src/connectors/lakebase/defaults.ts"],"sourcesContent":["import type { LakebaseConfig } from \"./types\";\n\n/** Default configuration for Lakebase connector */\nexport const lakebaseDefaults: LakebaseConfig = {\n port: 5432,\n sslMode: \"require\",\n maxPoolSize: 10,\n idleTimeoutMs: 30_000,\n connectionTimeoutMs: 10_000,\n};\n"],"mappings":";;AAGA,MAAa,mBAAmC;CAC9C,MAAM;CACN,SAAS;CACT,aAAa;CACb,eAAe;CACf,qBAAqB;CACtB"}
@@ -1 +1 @@
1
- {"version":3,"file":"client.js","names":["body: sql.ExecuteStatementRequest","result:\n | sql.StatementResponse\n | { result: { statement_id: string; status: sql.StatementStatus } }","obj: Record<string, unknown>"],"sources":["../../../src/connectors/sql-warehouse/client.ts"],"sourcesContent":["import {\n Context,\n type sql,\n type WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport { ArrowStreamProcessor } from \"../../stream/arrow-stream-processor\";\nimport type { TelemetryOptions } from \"shared\";\nimport type { TelemetryProvider } from \"../../telemetry\";\nimport {\n type Counter,\n type Histogram,\n type Span,\n SpanKind,\n SpanStatusCode,\n TelemetryManager,\n} from \"../../telemetry\";\nimport { executeStatementDefaults } from \"./defaults\";\n\nexport interface SQLWarehouseConfig {\n timeout?: number;\n telemetry?: TelemetryOptions;\n}\n\nexport class SQLWarehouseConnector {\n private readonly name = \"sql-warehouse\";\n\n private config: SQLWarehouseConfig;\n\n // Lazy-initialized: only created when Arrow format is used\n private _arrowProcessor: ArrowStreamProcessor | null = null;\n // telemetry\n private readonly telemetry: TelemetryProvider;\n private readonly telemetryMetrics: {\n queryCount: Counter;\n queryDuration: Histogram;\n };\n\n constructor(config: SQLWarehouseConfig) {\n this.config = config;\n\n this.telemetry = TelemetryManager.getProvider(\n this.name,\n this.config.telemetry,\n );\n this.telemetryMetrics = {\n queryCount: this.telemetry.getMeter().createCounter(\"query.count\", {\n description: \"Total number of queries executed\",\n unit: \"1\",\n }),\n queryDuration: this.telemetry\n .getMeter()\n .createHistogram(\"query.duration\", {\n description: \"Duration of queries executed\",\n unit: \"ms\",\n }),\n };\n }\n\n /**\n * Lazily initializes and returns the ArrowStreamProcessor.\n * Only created on first Arrow format query to avoid unnecessary allocation.\n */\n private get arrowProcessor(): ArrowStreamProcessor {\n if (!this._arrowProcessor) {\n this._arrowProcessor = new ArrowStreamProcessor({\n timeout: this.config.timeout || executeStatementDefaults.timeout,\n maxConcurrentDownloads:\n ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS,\n retries: ArrowStreamProcessor.DEFAULT_RETRIES,\n });\n }\n return this._arrowProcessor;\n }\n\n async executeStatement(\n workspaceClient: WorkspaceClient,\n input: sql.ExecuteStatementRequest,\n signal?: AbortSignal,\n ) {\n const startTime = Date.now();\n let success = false;\n\n return this.telemetry.startActiveSpan(\n \"sql.query\",\n {\n kind: SpanKind.CLIENT,\n attributes: {\n \"db.system\": \"databricks\",\n \"db.warehouse_id\": input.warehouse_id || \"\",\n \"db.catalog\": input.catalog ?? \"\",\n \"db.schema\": input.schema ?? \"\",\n \"db.statement\": input.statement?.substring(0, 500) || \"\",\n \"db.has_parameters\": !!input.parameters,\n },\n },\n async (span: Span) => {\n try {\n // validate required fields\n if (!input.statement) {\n throw new Error(\n \"Statement is required: Please provide a SQL statement to execute\",\n );\n }\n\n if (!input.warehouse_id) {\n throw new Error(\n \"Warehouse ID is required: Please provide a warehouse_id to execute the statement\",\n );\n }\n\n const body: sql.ExecuteStatementRequest = {\n statement: input.statement,\n parameters: input.parameters,\n warehouse_id: input.warehouse_id,\n catalog: input.catalog,\n schema: input.schema,\n wait_timeout:\n input.wait_timeout || executeStatementDefaults.wait_timeout,\n disposition:\n input.disposition || executeStatementDefaults.disposition,\n format: input.format || executeStatementDefaults.format,\n byte_limit: input.byte_limit,\n row_limit: input.row_limit,\n on_wait_timeout:\n input.on_wait_timeout || executeStatementDefaults.on_wait_timeout,\n };\n\n span.addEvent(\"statement.submitting\", {\n \"db.warehouse_id\": input.warehouse_id,\n });\n\n const response =\n await workspaceClient.statementExecution.executeStatement(\n body,\n this._createContext(signal),\n );\n\n if (!response) {\n throw new Error(\"No response received from SQL Warehouse API\");\n }\n const status = response.status;\n const statementId = response.statement_id as string;\n\n span.setAttribute(\"db.statement_id\", statementId);\n span.addEvent(\"statement.submitted\", {\n \"db.statement_id\": response.statement_id,\n \"db.status\": status?.state,\n });\n\n let result:\n | sql.StatementResponse\n | { result: { statement_id: string; status: sql.StatementStatus } };\n\n switch (status?.state) {\n case \"RUNNING\":\n case \"PENDING\":\n span.addEvent(\"statement.polling_started\", {\n \"db.status\": response.status?.state,\n });\n result = await this._pollForStatementResult(\n workspaceClient,\n statementId,\n this.config.timeout,\n signal,\n );\n break;\n case \"SUCCEEDED\":\n result = this._transformDataArray(response);\n break;\n case \"FAILED\":\n throw new Error(\n `Statement failed: ${status.error?.message || \"Unknown error\"}`,\n );\n case \"CANCELED\":\n throw new Error(\"Statement was canceled\");\n case \"CLOSED\":\n throw new Error(\n \"Statement execution completed but results are no longer available (CLOSED state)\",\n );\n default:\n throw new Error(`Unknown statement state: ${status?.state}`);\n }\n\n const resultData = result.result as any;\n if (resultData?.data) {\n span.setAttribute(\"db.result.row_count\", resultData.data.length);\n } else if (resultData?.data_array) {\n span.setAttribute(\n \"db.result.row_count\",\n resultData.data_array.length,\n );\n }\n\n success = true;\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n } catch (error) {\n span.recordException(error as Error);\n span.setStatus({\n code: SpanStatusCode.ERROR,\n message: error instanceof Error ? error.message : String(error),\n });\n throw error;\n } finally {\n const duration = Date.now() - startTime;\n span.end();\n\n const attributes = {\n \"db.warehouse_id\": input.warehouse_id,\n \"db.catalog\": input.catalog ?? \"\",\n \"db.schema\": input.schema ?? \"\",\n \"db.statement\": input.statement?.substring(0, 500) || \"\",\n success: success.toString(),\n };\n\n this.telemetryMetrics.queryCount.add(1, attributes);\n this.telemetryMetrics.queryDuration.record(duration, attributes);\n }\n },\n { name: this.name, includePrefix: true },\n );\n }\n\n private async _pollForStatementResult(\n workspaceClient: WorkspaceClient,\n statementId: string,\n timeout = executeStatementDefaults.timeout,\n signal?: AbortSignal,\n ) {\n return this.telemetry.startActiveSpan(\n \"sql.poll\",\n {\n attributes: {\n \"db.statement_id\": statementId,\n \"db.polling.timeout\": timeout,\n },\n },\n async (span: Span) => {\n try {\n const startTime = Date.now();\n let delay = 1000;\n const maxDelayBetweenPolls = 5000; // max 5 seconds between polls\n let pollCount = 0;\n\n while (true) {\n pollCount++;\n span.setAttribute(\"db.polling.current_attempt\", pollCount);\n\n // check if timeout exceeded\n const elapsedTime = Date.now() - startTime;\n if (elapsedTime > timeout) {\n const error = new Error(\n `Statement polling timeout exceeded after ${timeout}ms (elapsed: ${elapsedTime}ms)`,\n );\n span.recordException(error);\n span.setStatus({ code: SpanStatusCode.ERROR });\n throw error;\n }\n\n if (signal?.aborted) {\n const error = new Error(\"Request aborted\");\n span.recordException(error);\n span.setStatus({ code: SpanStatusCode.ERROR });\n throw error;\n }\n\n span.addEvent(\"polling.attempt\", {\n \"poll.attempt\": pollCount,\n \"poll.delay_ms\": delay,\n \"poll.elapsed_ms\": elapsedTime,\n });\n\n const response =\n await workspaceClient.statementExecution.getStatement(\n {\n statement_id: statementId,\n },\n this._createContext(signal),\n );\n if (!response) {\n throw new Error(\"No response received from SQL Warehouse API\");\n }\n\n const status = response.status;\n\n span.addEvent(\"polling.status_check\", {\n \"db.status\": status?.state,\n \"poll.attempt\": pollCount,\n });\n\n switch (status?.state) {\n case \"PENDING\":\n case \"RUNNING\":\n // continue polling\n break;\n case \"SUCCEEDED\":\n span.setAttribute(\"db.polling.attempts\", pollCount);\n span.setAttribute(\"db.polling.total_duration_ms\", elapsedTime);\n span.addEvent(\"polling.completed\", {\n \"poll.attempts\": pollCount,\n \"poll.duration_ms\": elapsedTime,\n });\n span.setStatus({ code: SpanStatusCode.OK });\n return this._transformDataArray(response);\n case \"FAILED\":\n throw new Error(\n `Statement failed: ${\n status.error?.message || \"Unknown error\"\n }`,\n );\n case \"CANCELED\":\n throw new Error(\"Statement was canceled\");\n case \"CLOSED\":\n throw new Error(\n \"Statement execution completed but results are no longer available (CLOSED state)\",\n );\n default:\n throw new Error(`Unknown statement state: ${status?.state}`);\n }\n\n // continue polling after delay\n await new Promise((resolve) => setTimeout(resolve, delay));\n delay = Math.min(delay * 2, maxDelayBetweenPolls);\n }\n } catch (error) {\n span.recordException(error as Error);\n span.setStatus({\n code: SpanStatusCode.ERROR,\n message: error instanceof Error ? error.message : String(error),\n });\n throw error;\n } finally {\n span.end();\n }\n },\n { name: this.name, includePrefix: true },\n );\n }\n\n private _transformDataArray(response: sql.StatementResponse) {\n if (response.manifest?.format === \"ARROW_STREAM\") {\n return this.updateWithArrowStatus(response);\n }\n\n if (!response.result?.data_array || !response.manifest?.schema?.columns) {\n return response;\n }\n\n const columns = response.manifest.schema.columns;\n\n const transformedData = response.result.data_array.map((row) => {\n const obj: Record<string, unknown> = {};\n row.forEach((value, index) => {\n const column = columns[index];\n const columnName = column?.name || `column_${index}`;\n\n // attempt to parse JSON strings for string columns\n if (\n column?.type_name === \"STRING\" &&\n typeof value === \"string\" &&\n value &&\n (value[0] === \"{\" || value[0] === \"[\")\n ) {\n try {\n obj[columnName] = JSON.parse(value);\n } catch {\n // if parsing fails, keep as string\n obj[columnName] = value;\n }\n } else {\n obj[columnName] = value;\n }\n });\n return obj;\n });\n\n // remove data_array\n const { data_array: _data_array, ...restResult } = response.result;\n return {\n ...response,\n result: {\n ...restResult,\n data: transformedData,\n },\n };\n }\n\n private updateWithArrowStatus(response: sql.StatementResponse): {\n result: { statement_id: string; status: sql.StatementStatus };\n } {\n return {\n result: {\n statement_id: response.statement_id as string,\n status: {\n state: response.status?.state,\n error: response.status?.error,\n } as sql.StatementStatus,\n },\n };\n }\n\n async getArrowData(\n workspaceClient: WorkspaceClient,\n jobId: string,\n signal?: AbortSignal,\n ): Promise<ReturnType<typeof this.arrowProcessor.processChunks>> {\n const startTime = Date.now();\n\n return this.telemetry.startActiveSpan(\n \"arrow.getData\",\n {\n kind: SpanKind.CLIENT,\n attributes: {\n \"db.system\": \"databricks\",\n \"arrow.job_id\": jobId,\n },\n },\n async (span: Span) => {\n try {\n const response =\n await workspaceClient.statementExecution.getStatement(\n { statement_id: jobId },\n this._createContext(signal),\n );\n\n const chunks = response.result?.external_links;\n const schema = response.manifest?.schema;\n\n if (!chunks || !schema) {\n throw new Error(\"No chunks or schema found in response\");\n }\n\n span.setAttribute(\"arrow.chunk_count\", chunks.length);\n\n const result = await this.arrowProcessor.processChunks(\n chunks,\n schema,\n signal,\n );\n\n span.setAttribute(\"arrow.data_size_bytes\", result.data.length);\n span.setStatus({ code: SpanStatusCode.OK });\n\n const duration = Date.now() - startTime;\n this.telemetryMetrics.queryDuration.record(duration, {\n operation: \"arrow.getData\",\n status: \"success\",\n });\n\n return result;\n } catch (error) {\n span.setStatus({\n code: SpanStatusCode.ERROR,\n message: error instanceof Error ? error.message : \"Unknown error\",\n });\n span.recordException(error as Error);\n\n const duration = Date.now() - startTime;\n this.telemetryMetrics.queryDuration.record(duration, {\n operation: \"arrow.getData\",\n status: \"error\",\n });\n\n console.error(`Failed Arrow job: ${jobId}`, error);\n throw error;\n }\n },\n );\n }\n\n // create context for cancellation token\n private _createContext(signal?: AbortSignal) {\n return new Context({\n cancellationToken: {\n isCancellationRequested: signal?.aborted ?? false,\n onCancellationRequested: (cb: () => void) => {\n signal?.addEventListener(\"abort\", cb, { once: true });\n },\n },\n });\n }\n}\n"],"mappings":";;;;;;;AAuBA,IAAa,wBAAb,MAAmC;CAcjC,YAAY,QAA4B;cAbhB;yBAK+B;AASrD,OAAK,SAAS;AAEd,OAAK,YAAY,iBAAiB,YAChC,KAAK,MACL,KAAK,OAAO,UACb;AACD,OAAK,mBAAmB;GACtB,YAAY,KAAK,UAAU,UAAU,CAAC,cAAc,eAAe;IACjE,aAAa;IACb,MAAM;IACP,CAAC;GACF,eAAe,KAAK,UACjB,UAAU,CACV,gBAAgB,kBAAkB;IACjC,aAAa;IACb,MAAM;IACP,CAAC;GACL;;;;;;CAOH,IAAY,iBAAuC;AACjD,MAAI,CAAC,KAAK,gBACR,MAAK,kBAAkB,IAAI,qBAAqB;GAC9C,SAAS,KAAK,OAAO,WAAW,yBAAyB;GACzD,wBACE,qBAAqB;GACvB,SAAS,qBAAqB;GAC/B,CAAC;AAEJ,SAAO,KAAK;;CAGd,MAAM,iBACJ,iBACA,OACA,QACA;EACA,MAAM,YAAY,KAAK,KAAK;EAC5B,IAAI,UAAU;AAEd,SAAO,KAAK,UAAU,gBACpB,aACA;GACE,MAAM,SAAS;GACf,YAAY;IACV,aAAa;IACb,mBAAmB,MAAM,gBAAgB;IACzC,cAAc,MAAM,WAAW;IAC/B,aAAa,MAAM,UAAU;IAC7B,gBAAgB,MAAM,WAAW,UAAU,GAAG,IAAI,IAAI;IACtD,qBAAqB,CAAC,CAAC,MAAM;IAC9B;GACF,EACD,OAAO,SAAe;AACpB,OAAI;AAEF,QAAI,CAAC,MAAM,UACT,OAAM,IAAI,MACR,mEACD;AAGH,QAAI,CAAC,MAAM,aACT,OAAM,IAAI,MACR,mFACD;IAGH,MAAMA,OAAoC;KACxC,WAAW,MAAM;KACjB,YAAY,MAAM;KAClB,cAAc,MAAM;KACpB,SAAS,MAAM;KACf,QAAQ,MAAM;KACd,cACE,MAAM,gBAAgB,yBAAyB;KACjD,aACE,MAAM,eAAe,yBAAyB;KAChD,QAAQ,MAAM,UAAU,yBAAyB;KACjD,YAAY,MAAM;KAClB,WAAW,MAAM;KACjB,iBACE,MAAM,mBAAmB,yBAAyB;KACrD;AAED,SAAK,SAAS,wBAAwB,EACpC,mBAAmB,MAAM,cAC1B,CAAC;IAEF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,iBACvC,MACA,KAAK,eAAe,OAAO,CAC5B;AAEH,QAAI,CAAC,SACH,OAAM,IAAI,MAAM,8CAA8C;IAEhE,MAAM,SAAS,SAAS;IACxB,MAAM,cAAc,SAAS;AAE7B,SAAK,aAAa,mBAAmB,YAAY;AACjD,SAAK,SAAS,uBAAuB;KACnC,mBAAmB,SAAS;KAC5B,aAAa,QAAQ;KACtB,CAAC;IAEF,IAAIC;AAIJ,YAAQ,QAAQ,OAAhB;KACE,KAAK;KACL,KAAK;AACH,WAAK,SAAS,6BAA6B,EACzC,aAAa,SAAS,QAAQ,OAC/B,CAAC;AACF,eAAS,MAAM,KAAK,wBAClB,iBACA,aACA,KAAK,OAAO,SACZ,OACD;AACD;KACF,KAAK;AACH,eAAS,KAAK,oBAAoB,SAAS;AAC3C;KACF,KAAK,SACH,OAAM,IAAI,MACR,qBAAqB,OAAO,OAAO,WAAW,kBAC/C;KACH,KAAK,WACH,OAAM,IAAI,MAAM,yBAAyB;KAC3C,KAAK,SACH,OAAM,IAAI,MACR,mFACD;KACH,QACE,OAAM,IAAI,MAAM,4BAA4B,QAAQ,QAAQ;;IAGhE,MAAM,aAAa,OAAO;AAC1B,QAAI,YAAY,KACd,MAAK,aAAa,uBAAuB,WAAW,KAAK,OAAO;aACvD,YAAY,WACrB,MAAK,aACH,uBACA,WAAW,WAAW,OACvB;AAGH,cAAU;AACV,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AACd,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;KAChE,CAAC;AACF,UAAM;aACE;IACR,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,KAAK;IAEV,MAAM,aAAa;KACjB,mBAAmB,MAAM;KACzB,cAAc,MAAM,WAAW;KAC/B,aAAa,MAAM,UAAU;KAC7B,gBAAgB,MAAM,WAAW,UAAU,GAAG,IAAI,IAAI;KACtD,SAAS,QAAQ,UAAU;KAC5B;AAED,SAAK,iBAAiB,WAAW,IAAI,GAAG,WAAW;AACnD,SAAK,iBAAiB,cAAc,OAAO,UAAU,WAAW;;KAGpE;GAAE,MAAM,KAAK;GAAM,eAAe;GAAM,CACzC;;CAGH,MAAc,wBACZ,iBACA,aACA,UAAU,yBAAyB,SACnC,QACA;AACA,SAAO,KAAK,UAAU,gBACpB,YACA,EACE,YAAY;GACV,mBAAmB;GACnB,sBAAsB;GACvB,EACF,EACD,OAAO,SAAe;AACpB,OAAI;IACF,MAAM,YAAY,KAAK,KAAK;IAC5B,IAAI,QAAQ;IACZ,MAAM,uBAAuB;IAC7B,IAAI,YAAY;AAEhB,WAAO,MAAM;AACX;AACA,UAAK,aAAa,8BAA8B,UAAU;KAG1D,MAAM,cAAc,KAAK,KAAK,GAAG;AACjC,SAAI,cAAc,SAAS;MACzB,MAAM,wBAAQ,IAAI,MAChB,4CAA4C,QAAQ,eAAe,YAAY,KAChF;AACD,WAAK,gBAAgB,MAAM;AAC3B,WAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,YAAM;;AAGR,SAAI,QAAQ,SAAS;MACnB,MAAM,wBAAQ,IAAI,MAAM,kBAAkB;AAC1C,WAAK,gBAAgB,MAAM;AAC3B,WAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,YAAM;;AAGR,UAAK,SAAS,mBAAmB;MAC/B,gBAAgB;MAChB,iBAAiB;MACjB,mBAAmB;MACpB,CAAC;KAEF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,aACvC,EACE,cAAc,aACf,EACD,KAAK,eAAe,OAAO,CAC5B;AACH,SAAI,CAAC,SACH,OAAM,IAAI,MAAM,8CAA8C;KAGhE,MAAM,SAAS,SAAS;AAExB,UAAK,SAAS,wBAAwB;MACpC,aAAa,QAAQ;MACrB,gBAAgB;MACjB,CAAC;AAEF,aAAQ,QAAQ,OAAhB;MACE,KAAK;MACL,KAAK,UAEH;MACF,KAAK;AACH,YAAK,aAAa,uBAAuB,UAAU;AACnD,YAAK,aAAa,gCAAgC,YAAY;AAC9D,YAAK,SAAS,qBAAqB;QACjC,iBAAiB;QACjB,oBAAoB;QACrB,CAAC;AACF,YAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,cAAO,KAAK,oBAAoB,SAAS;MAC3C,KAAK,SACH,OAAM,IAAI,MACR,qBACE,OAAO,OAAO,WAAW,kBAE5B;MACH,KAAK,WACH,OAAM,IAAI,MAAM,yBAAyB;MAC3C,KAAK,SACH,OAAM,IAAI,MACR,mFACD;MACH,QACE,OAAM,IAAI,MAAM,4BAA4B,QAAQ,QAAQ;;AAIhE,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,MAAM,CAAC;AAC1D,aAAQ,KAAK,IAAI,QAAQ,GAAG,qBAAqB;;YAE5C,OAAO;AACd,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;KAChE,CAAC;AACF,UAAM;aACE;AACR,SAAK,KAAK;;KAGd;GAAE,MAAM,KAAK;GAAM,eAAe;GAAM,CACzC;;CAGH,AAAQ,oBAAoB,UAAiC;AAC3D,MAAI,SAAS,UAAU,WAAW,eAChC,QAAO,KAAK,sBAAsB,SAAS;AAG7C,MAAI,CAAC,SAAS,QAAQ,cAAc,CAAC,SAAS,UAAU,QAAQ,QAC9D,QAAO;EAGT,MAAM,UAAU,SAAS,SAAS,OAAO;EAEzC,MAAM,kBAAkB,SAAS,OAAO,WAAW,KAAK,QAAQ;GAC9D,MAAMC,MAA+B,EAAE;AACvC,OAAI,SAAS,OAAO,UAAU;IAC5B,MAAM,SAAS,QAAQ;IACvB,MAAM,aAAa,QAAQ,QAAQ,UAAU;AAG7C,QACE,QAAQ,cAAc,YACtB,OAAO,UAAU,YACjB,UACC,MAAM,OAAO,OAAO,MAAM,OAAO,KAElC,KAAI;AACF,SAAI,cAAc,KAAK,MAAM,MAAM;YAC7B;AAEN,SAAI,cAAc;;QAGpB,KAAI,cAAc;KAEpB;AACF,UAAO;IACP;EAGF,MAAM,EAAE,YAAY,aAAa,GAAG,eAAe,SAAS;AAC5D,SAAO;GACL,GAAG;GACH,QAAQ;IACN,GAAG;IACH,MAAM;IACP;GACF;;CAGH,AAAQ,sBAAsB,UAE5B;AACA,SAAO,EACL,QAAQ;GACN,cAAc,SAAS;GACvB,QAAQ;IACN,OAAO,SAAS,QAAQ;IACxB,OAAO,SAAS,QAAQ;IACzB;GACF,EACF;;CAGH,MAAM,aACJ,iBACA,OACA,QAC+D;EAC/D,MAAM,YAAY,KAAK,KAAK;AAE5B,SAAO,KAAK,UAAU,gBACpB,iBACA;GACE,MAAM,SAAS;GACf,YAAY;IACV,aAAa;IACb,gBAAgB;IACjB;GACF,EACD,OAAO,SAAe;AACpB,OAAI;IACF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,aACvC,EAAE,cAAc,OAAO,EACvB,KAAK,eAAe,OAAO,CAC5B;IAEH,MAAM,SAAS,SAAS,QAAQ;IAChC,MAAM,SAAS,SAAS,UAAU;AAElC,QAAI,CAAC,UAAU,CAAC,OACd,OAAM,IAAI,MAAM,wCAAwC;AAG1D,SAAK,aAAa,qBAAqB,OAAO,OAAO;IAErD,MAAM,SAAS,MAAM,KAAK,eAAe,cACvC,QACA,QACA,OACD;AAED,SAAK,aAAa,yBAAyB,OAAO,KAAK,OAAO;AAC9D,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;IAE3C,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,cAAc,OAAO,UAAU;KACnD,WAAW;KACX,QAAQ;KACT,CAAC;AAEF,WAAO;YACA,OAAO;AACd,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU;KACnD,CAAC;AACF,SAAK,gBAAgB,MAAe;IAEpC,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,cAAc,OAAO,UAAU;KACnD,WAAW;KACX,QAAQ;KACT,CAAC;AAEF,YAAQ,MAAM,qBAAqB,SAAS,MAAM;AAClD,UAAM;;IAGX;;CAIH,AAAQ,eAAe,QAAsB;AAC3C,SAAO,IAAI,QAAQ,EACjB,mBAAmB;GACjB,yBAAyB,QAAQ,WAAW;GAC5C,0BAA0B,OAAmB;AAC3C,YAAQ,iBAAiB,SAAS,IAAI,EAAE,MAAM,MAAM,CAAC;;GAExD,EACF,CAAC"}
1
+ {"version":3,"file":"client.js","names":[],"sources":["../../../src/connectors/sql-warehouse/client.ts"],"sourcesContent":["import {\n Context,\n type sql,\n type WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport { ArrowStreamProcessor } from \"../../stream/arrow-stream-processor\";\nimport type { TelemetryOptions } from \"shared\";\nimport type { TelemetryProvider } from \"../../telemetry\";\nimport {\n type Counter,\n type Histogram,\n type Span,\n SpanKind,\n SpanStatusCode,\n TelemetryManager,\n} from \"../../telemetry\";\nimport { executeStatementDefaults } from \"./defaults\";\n\nexport interface SQLWarehouseConfig {\n timeout?: number;\n telemetry?: TelemetryOptions;\n}\n\nexport class SQLWarehouseConnector {\n private readonly name = \"sql-warehouse\";\n\n private config: SQLWarehouseConfig;\n\n // Lazy-initialized: only created when Arrow format is used\n private _arrowProcessor: ArrowStreamProcessor | null = null;\n // telemetry\n private readonly telemetry: TelemetryProvider;\n private readonly telemetryMetrics: {\n queryCount: Counter;\n queryDuration: Histogram;\n };\n\n constructor(config: SQLWarehouseConfig) {\n this.config = config;\n\n this.telemetry = TelemetryManager.getProvider(\n this.name,\n this.config.telemetry,\n );\n this.telemetryMetrics = {\n queryCount: this.telemetry.getMeter().createCounter(\"query.count\", {\n description: \"Total number of queries executed\",\n unit: \"1\",\n }),\n queryDuration: this.telemetry\n .getMeter()\n .createHistogram(\"query.duration\", {\n description: \"Duration of queries executed\",\n unit: \"ms\",\n }),\n };\n }\n\n /**\n * Lazily initializes and returns the ArrowStreamProcessor.\n * Only created on first Arrow format query to avoid unnecessary allocation.\n */\n private get arrowProcessor(): ArrowStreamProcessor {\n if (!this._arrowProcessor) {\n this._arrowProcessor = new ArrowStreamProcessor({\n timeout: this.config.timeout || executeStatementDefaults.timeout,\n maxConcurrentDownloads:\n ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS,\n retries: ArrowStreamProcessor.DEFAULT_RETRIES,\n });\n }\n return this._arrowProcessor;\n }\n\n async executeStatement(\n workspaceClient: WorkspaceClient,\n input: sql.ExecuteStatementRequest,\n signal?: AbortSignal,\n ) {\n const startTime = Date.now();\n let success = false;\n\n return this.telemetry.startActiveSpan(\n \"sql.query\",\n {\n kind: SpanKind.CLIENT,\n attributes: {\n \"db.system\": \"databricks\",\n \"db.warehouse_id\": input.warehouse_id || \"\",\n \"db.catalog\": input.catalog ?? \"\",\n \"db.schema\": input.schema ?? \"\",\n \"db.statement\": input.statement?.substring(0, 500) || \"\",\n \"db.has_parameters\": !!input.parameters,\n },\n },\n async (span: Span) => {\n try {\n // validate required fields\n if (!input.statement) {\n throw new Error(\n \"Statement is required: Please provide a SQL statement to execute\",\n );\n }\n\n if (!input.warehouse_id) {\n throw new Error(\n \"Warehouse ID is required: Please provide a warehouse_id to execute the statement\",\n );\n }\n\n const body: sql.ExecuteStatementRequest = {\n statement: input.statement,\n parameters: input.parameters,\n warehouse_id: input.warehouse_id,\n catalog: input.catalog,\n schema: input.schema,\n wait_timeout:\n input.wait_timeout || executeStatementDefaults.wait_timeout,\n disposition:\n input.disposition || executeStatementDefaults.disposition,\n format: input.format || executeStatementDefaults.format,\n byte_limit: input.byte_limit,\n row_limit: input.row_limit,\n on_wait_timeout:\n input.on_wait_timeout || executeStatementDefaults.on_wait_timeout,\n };\n\n span.addEvent(\"statement.submitting\", {\n \"db.warehouse_id\": input.warehouse_id,\n });\n\n const response =\n await workspaceClient.statementExecution.executeStatement(\n body,\n this._createContext(signal),\n );\n\n if (!response) {\n throw new Error(\"No response received from SQL Warehouse API\");\n }\n const status = response.status;\n const statementId = response.statement_id as string;\n\n span.setAttribute(\"db.statement_id\", statementId);\n span.addEvent(\"statement.submitted\", {\n \"db.statement_id\": response.statement_id,\n \"db.status\": status?.state,\n });\n\n let result:\n | sql.StatementResponse\n | { result: { statement_id: string; status: sql.StatementStatus } };\n\n switch (status?.state) {\n case \"RUNNING\":\n case \"PENDING\":\n span.addEvent(\"statement.polling_started\", {\n \"db.status\": response.status?.state,\n });\n result = await this._pollForStatementResult(\n workspaceClient,\n statementId,\n this.config.timeout,\n signal,\n );\n break;\n case \"SUCCEEDED\":\n result = this._transformDataArray(response);\n break;\n case \"FAILED\":\n throw new Error(\n `Statement failed: ${status.error?.message || \"Unknown error\"}`,\n );\n case \"CANCELED\":\n throw new Error(\"Statement was canceled\");\n case \"CLOSED\":\n throw new Error(\n \"Statement execution completed but results are no longer available (CLOSED state)\",\n );\n default:\n throw new Error(`Unknown statement state: ${status?.state}`);\n }\n\n const resultData = result.result as any;\n if (resultData?.data) {\n span.setAttribute(\"db.result.row_count\", resultData.data.length);\n } else if (resultData?.data_array) {\n span.setAttribute(\n \"db.result.row_count\",\n resultData.data_array.length,\n );\n }\n\n success = true;\n span.setStatus({ code: SpanStatusCode.OK });\n return result;\n } catch (error) {\n span.recordException(error as Error);\n span.setStatus({\n code: SpanStatusCode.ERROR,\n message: error instanceof Error ? error.message : String(error),\n });\n throw error;\n } finally {\n const duration = Date.now() - startTime;\n span.end();\n\n const attributes = {\n \"db.warehouse_id\": input.warehouse_id,\n \"db.catalog\": input.catalog ?? \"\",\n \"db.schema\": input.schema ?? \"\",\n \"db.statement\": input.statement?.substring(0, 500) || \"\",\n success: success.toString(),\n };\n\n this.telemetryMetrics.queryCount.add(1, attributes);\n this.telemetryMetrics.queryDuration.record(duration, attributes);\n }\n },\n { name: this.name, includePrefix: true },\n );\n }\n\n private async _pollForStatementResult(\n workspaceClient: WorkspaceClient,\n statementId: string,\n timeout = executeStatementDefaults.timeout,\n signal?: AbortSignal,\n ) {\n return this.telemetry.startActiveSpan(\n \"sql.poll\",\n {\n attributes: {\n \"db.statement_id\": statementId,\n \"db.polling.timeout\": timeout,\n },\n },\n async (span: Span) => {\n try {\n const startTime = Date.now();\n let delay = 1000;\n const maxDelayBetweenPolls = 5000; // max 5 seconds between polls\n let pollCount = 0;\n\n while (true) {\n pollCount++;\n span.setAttribute(\"db.polling.current_attempt\", pollCount);\n\n // check if timeout exceeded\n const elapsedTime = Date.now() - startTime;\n if (elapsedTime > timeout) {\n const error = new Error(\n `Statement polling timeout exceeded after ${timeout}ms (elapsed: ${elapsedTime}ms)`,\n );\n span.recordException(error);\n span.setStatus({ code: SpanStatusCode.ERROR });\n throw error;\n }\n\n if (signal?.aborted) {\n const error = new Error(\"Request aborted\");\n span.recordException(error);\n span.setStatus({ code: SpanStatusCode.ERROR });\n throw error;\n }\n\n span.addEvent(\"polling.attempt\", {\n \"poll.attempt\": pollCount,\n \"poll.delay_ms\": delay,\n \"poll.elapsed_ms\": elapsedTime,\n });\n\n const response =\n await workspaceClient.statementExecution.getStatement(\n {\n statement_id: statementId,\n },\n this._createContext(signal),\n );\n if (!response) {\n throw new Error(\"No response received from SQL Warehouse API\");\n }\n\n const status = response.status;\n\n span.addEvent(\"polling.status_check\", {\n \"db.status\": status?.state,\n \"poll.attempt\": pollCount,\n });\n\n switch (status?.state) {\n case \"PENDING\":\n case \"RUNNING\":\n // continue polling\n break;\n case \"SUCCEEDED\":\n span.setAttribute(\"db.polling.attempts\", pollCount);\n span.setAttribute(\"db.polling.total_duration_ms\", elapsedTime);\n span.addEvent(\"polling.completed\", {\n \"poll.attempts\": pollCount,\n \"poll.duration_ms\": elapsedTime,\n });\n span.setStatus({ code: SpanStatusCode.OK });\n return this._transformDataArray(response);\n case \"FAILED\":\n throw new Error(\n `Statement failed: ${\n status.error?.message || \"Unknown error\"\n }`,\n );\n case \"CANCELED\":\n throw new Error(\"Statement was canceled\");\n case \"CLOSED\":\n throw new Error(\n \"Statement execution completed but results are no longer available (CLOSED state)\",\n );\n default:\n throw new Error(`Unknown statement state: ${status?.state}`);\n }\n\n // continue polling after delay\n await new Promise((resolve) => setTimeout(resolve, delay));\n delay = Math.min(delay * 2, maxDelayBetweenPolls);\n }\n } catch (error) {\n span.recordException(error as Error);\n span.setStatus({\n code: SpanStatusCode.ERROR,\n message: error instanceof Error ? error.message : String(error),\n });\n throw error;\n } finally {\n span.end();\n }\n },\n { name: this.name, includePrefix: true },\n );\n }\n\n private _transformDataArray(response: sql.StatementResponse) {\n if (response.manifest?.format === \"ARROW_STREAM\") {\n return this.updateWithArrowStatus(response);\n }\n\n if (!response.result?.data_array || !response.manifest?.schema?.columns) {\n return response;\n }\n\n const columns = response.manifest.schema.columns;\n\n const transformedData = response.result.data_array.map((row) => {\n const obj: Record<string, unknown> = {};\n row.forEach((value, index) => {\n const column = columns[index];\n const columnName = column?.name || `column_${index}`;\n\n // attempt to parse JSON strings for string columns\n if (\n column?.type_name === \"STRING\" &&\n typeof value === \"string\" &&\n value &&\n (value[0] === \"{\" || value[0] === \"[\")\n ) {\n try {\n obj[columnName] = JSON.parse(value);\n } catch {\n // if parsing fails, keep as string\n obj[columnName] = value;\n }\n } else {\n obj[columnName] = value;\n }\n });\n return obj;\n });\n\n // remove data_array\n const { data_array: _data_array, ...restResult } = response.result;\n return {\n ...response,\n result: {\n ...restResult,\n data: transformedData,\n },\n };\n }\n\n private updateWithArrowStatus(response: sql.StatementResponse): {\n result: { statement_id: string; status: sql.StatementStatus };\n } {\n return {\n result: {\n statement_id: response.statement_id as string,\n status: {\n state: response.status?.state,\n error: response.status?.error,\n } as sql.StatementStatus,\n },\n };\n }\n\n async getArrowData(\n workspaceClient: WorkspaceClient,\n jobId: string,\n signal?: AbortSignal,\n ): Promise<ReturnType<typeof this.arrowProcessor.processChunks>> {\n const startTime = Date.now();\n\n return this.telemetry.startActiveSpan(\n \"arrow.getData\",\n {\n kind: SpanKind.CLIENT,\n attributes: {\n \"db.system\": \"databricks\",\n \"arrow.job_id\": jobId,\n },\n },\n async (span: Span) => {\n try {\n const response =\n await workspaceClient.statementExecution.getStatement(\n { statement_id: jobId },\n this._createContext(signal),\n );\n\n const chunks = response.result?.external_links;\n const schema = response.manifest?.schema;\n\n if (!chunks || !schema) {\n throw new Error(\"No chunks or schema found in response\");\n }\n\n span.setAttribute(\"arrow.chunk_count\", chunks.length);\n\n const result = await this.arrowProcessor.processChunks(\n chunks,\n schema,\n signal,\n );\n\n span.setAttribute(\"arrow.data_size_bytes\", result.data.length);\n span.setStatus({ code: SpanStatusCode.OK });\n\n const duration = Date.now() - startTime;\n this.telemetryMetrics.queryDuration.record(duration, {\n operation: \"arrow.getData\",\n status: \"success\",\n });\n\n return result;\n } catch (error) {\n span.setStatus({\n code: SpanStatusCode.ERROR,\n message: error instanceof Error ? error.message : \"Unknown error\",\n });\n span.recordException(error as Error);\n\n const duration = Date.now() - startTime;\n this.telemetryMetrics.queryDuration.record(duration, {\n operation: \"arrow.getData\",\n status: \"error\",\n });\n\n console.error(`Failed Arrow job: ${jobId}`, error);\n throw error;\n }\n },\n );\n }\n\n // create context for cancellation token\n private _createContext(signal?: AbortSignal) {\n return new Context({\n cancellationToken: {\n isCancellationRequested: signal?.aborted ?? false,\n onCancellationRequested: (cb: () => void) => {\n signal?.addEventListener(\"abort\", cb, { once: true });\n },\n },\n });\n }\n}\n"],"mappings":";;;;;;;AAuBA,IAAa,wBAAb,MAAmC;CAcjC,YAAY,QAA4B;cAbhB;yBAK+B;AASrD,OAAK,SAAS;AAEd,OAAK,YAAY,iBAAiB,YAChC,KAAK,MACL,KAAK,OAAO,UACb;AACD,OAAK,mBAAmB;GACtB,YAAY,KAAK,UAAU,UAAU,CAAC,cAAc,eAAe;IACjE,aAAa;IACb,MAAM;IACP,CAAC;GACF,eAAe,KAAK,UACjB,UAAU,CACV,gBAAgB,kBAAkB;IACjC,aAAa;IACb,MAAM;IACP,CAAC;GACL;;;;;;CAOH,IAAY,iBAAuC;AACjD,MAAI,CAAC,KAAK,gBACR,MAAK,kBAAkB,IAAI,qBAAqB;GAC9C,SAAS,KAAK,OAAO,WAAW,yBAAyB;GACzD,wBACE,qBAAqB;GACvB,SAAS,qBAAqB;GAC/B,CAAC;AAEJ,SAAO,KAAK;;CAGd,MAAM,iBACJ,iBACA,OACA,QACA;EACA,MAAM,YAAY,KAAK,KAAK;EAC5B,IAAI,UAAU;AAEd,SAAO,KAAK,UAAU,gBACpB,aACA;GACE,MAAM,SAAS;GACf,YAAY;IACV,aAAa;IACb,mBAAmB,MAAM,gBAAgB;IACzC,cAAc,MAAM,WAAW;IAC/B,aAAa,MAAM,UAAU;IAC7B,gBAAgB,MAAM,WAAW,UAAU,GAAG,IAAI,IAAI;IACtD,qBAAqB,CAAC,CAAC,MAAM;IAC9B;GACF,EACD,OAAO,SAAe;AACpB,OAAI;AAEF,QAAI,CAAC,MAAM,UACT,OAAM,IAAI,MACR,mEACD;AAGH,QAAI,CAAC,MAAM,aACT,OAAM,IAAI,MACR,mFACD;IAGH,MAAM,OAAoC;KACxC,WAAW,MAAM;KACjB,YAAY,MAAM;KAClB,cAAc,MAAM;KACpB,SAAS,MAAM;KACf,QAAQ,MAAM;KACd,cACE,MAAM,gBAAgB,yBAAyB;KACjD,aACE,MAAM,eAAe,yBAAyB;KAChD,QAAQ,MAAM,UAAU,yBAAyB;KACjD,YAAY,MAAM;KAClB,WAAW,MAAM;KACjB,iBACE,MAAM,mBAAmB,yBAAyB;KACrD;AAED,SAAK,SAAS,wBAAwB,EACpC,mBAAmB,MAAM,cAC1B,CAAC;IAEF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,iBACvC,MACA,KAAK,eAAe,OAAO,CAC5B;AAEH,QAAI,CAAC,SACH,OAAM,IAAI,MAAM,8CAA8C;IAEhE,MAAM,SAAS,SAAS;IACxB,MAAM,cAAc,SAAS;AAE7B,SAAK,aAAa,mBAAmB,YAAY;AACjD,SAAK,SAAS,uBAAuB;KACnC,mBAAmB,SAAS;KAC5B,aAAa,QAAQ;KACtB,CAAC;IAEF,IAAI;AAIJ,YAAQ,QAAQ,OAAhB;KACE,KAAK;KACL,KAAK;AACH,WAAK,SAAS,6BAA6B,EACzC,aAAa,SAAS,QAAQ,OAC/B,CAAC;AACF,eAAS,MAAM,KAAK,wBAClB,iBACA,aACA,KAAK,OAAO,SACZ,OACD;AACD;KACF,KAAK;AACH,eAAS,KAAK,oBAAoB,SAAS;AAC3C;KACF,KAAK,SACH,OAAM,IAAI,MACR,qBAAqB,OAAO,OAAO,WAAW,kBAC/C;KACH,KAAK,WACH,OAAM,IAAI,MAAM,yBAAyB;KAC3C,KAAK,SACH,OAAM,IAAI,MACR,mFACD;KACH,QACE,OAAM,IAAI,MAAM,4BAA4B,QAAQ,QAAQ;;IAGhE,MAAM,aAAa,OAAO;AAC1B,QAAI,YAAY,KACd,MAAK,aAAa,uBAAuB,WAAW,KAAK,OAAO;aACvD,YAAY,WACrB,MAAK,aACH,uBACA,WAAW,WAAW,OACvB;AAGH,cAAU;AACV,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,WAAO;YACA,OAAO;AACd,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;KAChE,CAAC;AACF,UAAM;aACE;IACR,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,KAAK;IAEV,MAAM,aAAa;KACjB,mBAAmB,MAAM;KACzB,cAAc,MAAM,WAAW;KAC/B,aAAa,MAAM,UAAU;KAC7B,gBAAgB,MAAM,WAAW,UAAU,GAAG,IAAI,IAAI;KACtD,SAAS,QAAQ,UAAU;KAC5B;AAED,SAAK,iBAAiB,WAAW,IAAI,GAAG,WAAW;AACnD,SAAK,iBAAiB,cAAc,OAAO,UAAU,WAAW;;KAGpE;GAAE,MAAM,KAAK;GAAM,eAAe;GAAM,CACzC;;CAGH,MAAc,wBACZ,iBACA,aACA,UAAU,yBAAyB,SACnC,QACA;AACA,SAAO,KAAK,UAAU,gBACpB,YACA,EACE,YAAY;GACV,mBAAmB;GACnB,sBAAsB;GACvB,EACF,EACD,OAAO,SAAe;AACpB,OAAI;IACF,MAAM,YAAY,KAAK,KAAK;IAC5B,IAAI,QAAQ;IACZ,MAAM,uBAAuB;IAC7B,IAAI,YAAY;AAEhB,WAAO,MAAM;AACX;AACA,UAAK,aAAa,8BAA8B,UAAU;KAG1D,MAAM,cAAc,KAAK,KAAK,GAAG;AACjC,SAAI,cAAc,SAAS;MACzB,MAAM,wBAAQ,IAAI,MAChB,4CAA4C,QAAQ,eAAe,YAAY,KAChF;AACD,WAAK,gBAAgB,MAAM;AAC3B,WAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,YAAM;;AAGR,SAAI,QAAQ,SAAS;MACnB,MAAM,wBAAQ,IAAI,MAAM,kBAAkB;AAC1C,WAAK,gBAAgB,MAAM;AAC3B,WAAK,UAAU,EAAE,MAAM,eAAe,OAAO,CAAC;AAC9C,YAAM;;AAGR,UAAK,SAAS,mBAAmB;MAC/B,gBAAgB;MAChB,iBAAiB;MACjB,mBAAmB;MACpB,CAAC;KAEF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,aACvC,EACE,cAAc,aACf,EACD,KAAK,eAAe,OAAO,CAC5B;AACH,SAAI,CAAC,SACH,OAAM,IAAI,MAAM,8CAA8C;KAGhE,MAAM,SAAS,SAAS;AAExB,UAAK,SAAS,wBAAwB;MACpC,aAAa,QAAQ;MACrB,gBAAgB;MACjB,CAAC;AAEF,aAAQ,QAAQ,OAAhB;MACE,KAAK;MACL,KAAK,UAEH;MACF,KAAK;AACH,YAAK,aAAa,uBAAuB,UAAU;AACnD,YAAK,aAAa,gCAAgC,YAAY;AAC9D,YAAK,SAAS,qBAAqB;QACjC,iBAAiB;QACjB,oBAAoB;QACrB,CAAC;AACF,YAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;AAC3C,cAAO,KAAK,oBAAoB,SAAS;MAC3C,KAAK,SACH,OAAM,IAAI,MACR,qBACE,OAAO,OAAO,WAAW,kBAE5B;MACH,KAAK,WACH,OAAM,IAAI,MAAM,yBAAyB;MAC3C,KAAK,SACH,OAAM,IAAI,MACR,mFACD;MACH,QACE,OAAM,IAAI,MAAM,4BAA4B,QAAQ,QAAQ;;AAIhE,WAAM,IAAI,SAAS,YAAY,WAAW,SAAS,MAAM,CAAC;AAC1D,aAAQ,KAAK,IAAI,QAAQ,GAAG,qBAAqB;;YAE5C,OAAO;AACd,SAAK,gBAAgB,MAAe;AACpC,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU,OAAO,MAAM;KAChE,CAAC;AACF,UAAM;aACE;AACR,SAAK,KAAK;;KAGd;GAAE,MAAM,KAAK;GAAM,eAAe;GAAM,CACzC;;CAGH,AAAQ,oBAAoB,UAAiC;AAC3D,MAAI,SAAS,UAAU,WAAW,eAChC,QAAO,KAAK,sBAAsB,SAAS;AAG7C,MAAI,CAAC,SAAS,QAAQ,cAAc,CAAC,SAAS,UAAU,QAAQ,QAC9D,QAAO;EAGT,MAAM,UAAU,SAAS,SAAS,OAAO;EAEzC,MAAM,kBAAkB,SAAS,OAAO,WAAW,KAAK,QAAQ;GAC9D,MAAM,MAA+B,EAAE;AACvC,OAAI,SAAS,OAAO,UAAU;IAC5B,MAAM,SAAS,QAAQ;IACvB,MAAM,aAAa,QAAQ,QAAQ,UAAU;AAG7C,QACE,QAAQ,cAAc,YACtB,OAAO,UAAU,YACjB,UACC,MAAM,OAAO,OAAO,MAAM,OAAO,KAElC,KAAI;AACF,SAAI,cAAc,KAAK,MAAM,MAAM;YAC7B;AAEN,SAAI,cAAc;;QAGpB,KAAI,cAAc;KAEpB;AACF,UAAO;IACP;EAGF,MAAM,EAAE,YAAY,aAAa,GAAG,eAAe,SAAS;AAC5D,SAAO;GACL,GAAG;GACH,QAAQ;IACN,GAAG;IACH,MAAM;IACP;GACF;;CAGH,AAAQ,sBAAsB,UAE5B;AACA,SAAO,EACL,QAAQ;GACN,cAAc,SAAS;GACvB,QAAQ;IACN,OAAO,SAAS,QAAQ;IACxB,OAAO,SAAS,QAAQ;IACzB;GACF,EACF;;CAGH,MAAM,aACJ,iBACA,OACA,QAC+D;EAC/D,MAAM,YAAY,KAAK,KAAK;AAE5B,SAAO,KAAK,UAAU,gBACpB,iBACA;GACE,MAAM,SAAS;GACf,YAAY;IACV,aAAa;IACb,gBAAgB;IACjB;GACF,EACD,OAAO,SAAe;AACpB,OAAI;IACF,MAAM,WACJ,MAAM,gBAAgB,mBAAmB,aACvC,EAAE,cAAc,OAAO,EACvB,KAAK,eAAe,OAAO,CAC5B;IAEH,MAAM,SAAS,SAAS,QAAQ;IAChC,MAAM,SAAS,SAAS,UAAU;AAElC,QAAI,CAAC,UAAU,CAAC,OACd,OAAM,IAAI,MAAM,wCAAwC;AAG1D,SAAK,aAAa,qBAAqB,OAAO,OAAO;IAErD,MAAM,SAAS,MAAM,KAAK,eAAe,cACvC,QACA,QACA,OACD;AAED,SAAK,aAAa,yBAAyB,OAAO,KAAK,OAAO;AAC9D,SAAK,UAAU,EAAE,MAAM,eAAe,IAAI,CAAC;IAE3C,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,cAAc,OAAO,UAAU;KACnD,WAAW;KACX,QAAQ;KACT,CAAC;AAEF,WAAO;YACA,OAAO;AACd,SAAK,UAAU;KACb,MAAM,eAAe;KACrB,SAAS,iBAAiB,QAAQ,MAAM,UAAU;KACnD,CAAC;AACF,SAAK,gBAAgB,MAAe;IAEpC,MAAM,WAAW,KAAK,KAAK,GAAG;AAC9B,SAAK,iBAAiB,cAAc,OAAO,UAAU;KACnD,WAAW;KACX,QAAQ;KACT,CAAC;AAEF,YAAQ,MAAM,qBAAqB,SAAS,MAAM;AAClD,UAAM;;IAGX;;CAIH,AAAQ,eAAe,QAAsB;AAC3C,SAAO,IAAI,QAAQ,EACjB,mBAAmB;GACjB,yBAAyB,QAAQ,WAAW;GAC5C,0BAA0B,OAAmB;AAC3C,YAAQ,iBAAiB,SAAS,IAAI,EAAE,MAAM,MAAM,CAAC;;GAExD,EACF,CAAC"}
@@ -1 +1 @@
1
- {"version":3,"file":"defaults.js","names":["executeStatementDefaults: ExecuteStatementDefaults"],"sources":["../../../src/connectors/sql-warehouse/defaults.ts"],"sourcesContent":["import type { sql } from \"@databricks/sdk-experimental\";\n\ninterface ExecuteStatementDefaults {\n wait_timeout: string;\n disposition: sql.ExecuteStatementRequest[\"disposition\"];\n format: sql.ExecuteStatementRequest[\"format\"];\n on_wait_timeout: sql.ExecuteStatementRequest[\"on_wait_timeout\"];\n timeout: number;\n}\n\n// @TODO: Make these configurable globally and validate right values\nexport const executeStatementDefaults: ExecuteStatementDefaults = {\n wait_timeout: \"30s\",\n disposition: \"INLINE\",\n format: \"JSON_ARRAY\",\n on_wait_timeout: \"CONTINUE\",\n timeout: 60000,\n};\n"],"mappings":";AAWA,MAAaA,2BAAqD;CAChE,cAAc;CACd,aAAa;CACb,QAAQ;CACR,iBAAiB;CACjB,SAAS;CACV"}
1
+ {"version":3,"file":"defaults.js","names":[],"sources":["../../../src/connectors/sql-warehouse/defaults.ts"],"sourcesContent":["import type { sql } from \"@databricks/sdk-experimental\";\n\ninterface ExecuteStatementDefaults {\n wait_timeout: string;\n disposition: sql.ExecuteStatementRequest[\"disposition\"];\n format: sql.ExecuteStatementRequest[\"format\"];\n on_wait_timeout: sql.ExecuteStatementRequest[\"on_wait_timeout\"];\n timeout: number;\n}\n\n// @TODO: Make these configurable globally and validate right values\nexport const executeStatementDefaults: ExecuteStatementDefaults = {\n wait_timeout: \"30s\",\n disposition: \"INLINE\",\n format: \"JSON_ARRAY\",\n on_wait_timeout: \"CONTINUE\",\n timeout: 60000,\n};\n"],"mappings":";AAWA,MAAa,2BAAqD;CAChE,cAAc;CACd,aAAa;CACb,QAAQ;CACR,iBAAiB;CACjB,SAAS;CACV"}
@@ -0,0 +1,75 @@
1
+ import { __esmMin } from "../_virtual/rolldown_runtime.js";
2
+ import { ServiceContext, init_service_context } from "./service-context.js";
3
+ import { init_user_context, isUserContext } from "./user-context.js";
4
+ import { AsyncLocalStorage } from "node:async_hooks";
5
+
6
+ //#region src/context/execution-context.ts
7
+ /**
8
+ * Run a function in the context of a user.
9
+ * All calls within the function will have access to the user context.
10
+ *
11
+ * @param userContext - The user context to use
12
+ * @param fn - The function to run
13
+ * @returns The result of the function
14
+ */
15
+ function runInUserContext(userContext, fn) {
16
+ return executionContextStorage.run(userContext, fn);
17
+ }
18
+ /**
19
+ * Get the current execution context.
20
+ *
21
+ * - If running inside a user context (via asUser), returns the user context
22
+ * - Otherwise, returns the service context
23
+ *
24
+ * @throws Error if ServiceContext is not initialized
25
+ */
26
+ function getExecutionContext() {
27
+ const userContext = executionContextStorage.getStore();
28
+ if (userContext) return userContext;
29
+ return ServiceContext.get();
30
+ }
31
+ /**
32
+ * Get the current user ID for cache keying and telemetry.
33
+ *
34
+ * Returns the user ID if in user context, otherwise the service user ID.
35
+ */
36
+ function getCurrentUserId() {
37
+ const ctx = getExecutionContext();
38
+ if (isUserContext(ctx)) return ctx.userId;
39
+ return ctx.serviceUserId;
40
+ }
41
+ /**
42
+ * Get the WorkspaceClient for the current execution context.
43
+ */
44
+ function getWorkspaceClient() {
45
+ return getExecutionContext().client;
46
+ }
47
+ /**
48
+ * Get the warehouse ID promise.
49
+ */
50
+ function getWarehouseId() {
51
+ return getExecutionContext().warehouseId;
52
+ }
53
+ /**
54
+ * Get the workspace ID promise.
55
+ */
56
+ function getWorkspaceId() {
57
+ return getExecutionContext().workspaceId;
58
+ }
59
+ /**
60
+ * Check if currently running in a user context.
61
+ */
62
+ function isInUserContext() {
63
+ return executionContextStorage.getStore() !== void 0;
64
+ }
65
+ var executionContextStorage;
66
+ var init_execution_context = __esmMin((() => {
67
+ init_service_context();
68
+ init_user_context();
69
+ executionContextStorage = new AsyncLocalStorage();
70
+ }));
71
+
72
+ //#endregion
73
+ init_execution_context();
74
+ export { getCurrentUserId, getExecutionContext, getWarehouseId, getWorkspaceClient, getWorkspaceId, init_execution_context, isInUserContext, runInUserContext };
75
+ //# sourceMappingURL=execution-context.js.map
@@ -0,0 +1 @@
1
+ {"version":3,"file":"execution-context.js","names":[],"sources":["../../src/context/execution-context.ts"],"sourcesContent":["import { AsyncLocalStorage } from \"node:async_hooks\";\nimport { ServiceContext } from \"./service-context\";\nimport {\n isUserContext,\n type ExecutionContext,\n type UserContext,\n} from \"./user-context\";\n\n/**\n * AsyncLocalStorage for execution context.\n * Used to pass user context through the call stack without explicit parameters.\n */\nconst executionContextStorage = new AsyncLocalStorage<UserContext>();\n\n/**\n * Run a function in the context of a user.\n * All calls within the function will have access to the user context.\n *\n * @param userContext - The user context to use\n * @param fn - The function to run\n * @returns The result of the function\n */\nexport function runInUserContext<T>(userContext: UserContext, fn: () => T): T {\n return executionContextStorage.run(userContext, fn);\n}\n\n/**\n * Get the current execution context.\n *\n * - If running inside a user context (via asUser), returns the user context\n * - Otherwise, returns the service context\n *\n * @throws Error if ServiceContext is not initialized\n */\nexport function getExecutionContext(): ExecutionContext {\n const userContext = executionContextStorage.getStore();\n if (userContext) {\n return userContext;\n }\n return ServiceContext.get();\n}\n\n/**\n * Get the current user ID for cache keying and telemetry.\n *\n * Returns the user ID if in user context, otherwise the service user ID.\n */\nexport function getCurrentUserId(): string {\n const ctx = getExecutionContext();\n if (isUserContext(ctx)) {\n return ctx.userId;\n }\n return ctx.serviceUserId;\n}\n\n/**\n * Get the WorkspaceClient for the current execution context.\n */\nexport function getWorkspaceClient() {\n return getExecutionContext().client;\n}\n\n/**\n * Get the warehouse ID promise.\n */\nexport function getWarehouseId(): Promise<string> {\n return getExecutionContext().warehouseId;\n}\n\n/**\n * Get the workspace ID promise.\n */\nexport function getWorkspaceId(): Promise<string> {\n return getExecutionContext().workspaceId;\n}\n\n/**\n * Check if currently running in a user context.\n */\nexport function isInUserContext(): boolean {\n const ctx = executionContextStorage.getStore();\n return ctx !== undefined;\n}\n"],"mappings":";;;;;;;;;;;;;;AAsBA,SAAgB,iBAAoB,aAA0B,IAAgB;AAC5E,QAAO,wBAAwB,IAAI,aAAa,GAAG;;;;;;;;;;AAWrD,SAAgB,sBAAwC;CACtD,MAAM,cAAc,wBAAwB,UAAU;AACtD,KAAI,YACF,QAAO;AAET,QAAO,eAAe,KAAK;;;;;;;AAQ7B,SAAgB,mBAA2B;CACzC,MAAM,MAAM,qBAAqB;AACjC,KAAI,cAAc,IAAI,CACpB,QAAO,IAAI;AAEb,QAAO,IAAI;;;;;AAMb,SAAgB,qBAAqB;AACnC,QAAO,qBAAqB,CAAC;;;;;AAM/B,SAAgB,iBAAkC;AAChD,QAAO,qBAAqB,CAAC;;;;;AAM/B,SAAgB,iBAAkC;AAChD,QAAO,qBAAqB,CAAC;;;;;AAM/B,SAAgB,kBAA2B;AAEzC,QADY,wBAAwB,UAAU,KAC/B;;;;uBAhFkC;oBAK3B;CAMlB,0BAA0B,IAAI,mBAAgC"}
@@ -0,0 +1,27 @@
1
+ import { __esmMin, __exportAll } from "../_virtual/rolldown_runtime.js";
2
+ import { ServiceContext, init_service_context } from "./service-context.js";
3
+ import { init_user_context, isUserContext } from "./user-context.js";
4
+ import { getCurrentUserId, getExecutionContext, getWarehouseId, getWorkspaceClient, getWorkspaceId, init_execution_context, isInUserContext, runInUserContext } from "./execution-context.js";
5
+
6
+ //#region src/context/index.ts
7
+ var context_exports = /* @__PURE__ */ __exportAll({
8
+ ServiceContext: () => ServiceContext,
9
+ getCurrentUserId: () => getCurrentUserId,
10
+ getExecutionContext: () => getExecutionContext,
11
+ getWarehouseId: () => getWarehouseId,
12
+ getWorkspaceClient: () => getWorkspaceClient,
13
+ getWorkspaceId: () => getWorkspaceId,
14
+ isInUserContext: () => isInUserContext,
15
+ isUserContext: () => isUserContext,
16
+ runInUserContext: () => runInUserContext
17
+ });
18
+ var init_context = __esmMin((() => {
19
+ init_service_context();
20
+ init_user_context();
21
+ init_execution_context();
22
+ }));
23
+
24
+ //#endregion
25
+ init_context();
26
+ export { context_exports, init_context };
27
+ //# sourceMappingURL=index.js.map
@@ -0,0 +1 @@
1
+ {"version":3,"file":"index.js","names":[],"sources":["../../src/context/index.ts"],"sourcesContent":["export { ServiceContext, type ServiceContextState } from \"./service-context\";\n\nexport {\n isUserContext,\n type ExecutionContext,\n type UserContext,\n} from \"./user-context\";\n\nexport {\n getExecutionContext,\n getCurrentUserId,\n getWorkspaceClient,\n getWarehouseId,\n getWorkspaceId,\n isInUserContext,\n runInUserContext,\n} from \"./execution-context\";\n"],"mappings":";;;;;;;;;;;;;;;;;;uBAA6E;oBAMrD;yBAUK"}
@@ -0,0 +1,149 @@
1
+ import { __esmMin } from "../_virtual/rolldown_runtime.js";
2
+ import { name, version } from "../appkit/package.js";
3
+ import { WorkspaceClient } from "@databricks/sdk-experimental";
4
+ import { coerce } from "semver";
5
+
6
+ //#region src/context/service-context.ts
7
+ function getClientOptions() {
8
+ const isDev = process.env.NODE_ENV === "development";
9
+ return {
10
+ product: name,
11
+ productVersion: coerce(version)?.version ?? version,
12
+ ...isDev && { userAgentExtra: { mode: "dev" } }
13
+ };
14
+ }
15
+ var ServiceContext;
16
+ var init_service_context = __esmMin((() => {
17
+ ServiceContext = class ServiceContext {
18
+ static {
19
+ this.instance = null;
20
+ }
21
+ static {
22
+ this.initPromise = null;
23
+ }
24
+ /**
25
+ * Initialize the service context. Should be called once at app startup.
26
+ * Safe to call multiple times - will return the same instance.
27
+ */
28
+ static async initialize() {
29
+ if (ServiceContext.instance) return ServiceContext.instance;
30
+ if (ServiceContext.initPromise) return ServiceContext.initPromise;
31
+ ServiceContext.initPromise = ServiceContext.createContext();
32
+ ServiceContext.instance = await ServiceContext.initPromise;
33
+ return ServiceContext.instance;
34
+ }
35
+ /**
36
+ * Get the initialized service context.
37
+ * @throws Error if not initialized
38
+ */
39
+ static get() {
40
+ if (!ServiceContext.instance) throw new Error("ServiceContext not initialized. Call ServiceContext.initialize() first.");
41
+ return ServiceContext.instance;
42
+ }
43
+ /**
44
+ * Check if the service context has been initialized.
45
+ */
46
+ static isInitialized() {
47
+ return ServiceContext.instance !== null;
48
+ }
49
+ /**
50
+ * Create a user context from request headers.
51
+ *
52
+ * @param token - The user's access token from x-forwarded-access-token header
53
+ * @param userId - The user's ID from x-forwarded-user header
54
+ * @param userName - Optional user name
55
+ * @throws Error if token is not provided
56
+ */
57
+ static createUserContext(token, userId, userName) {
58
+ if (!token) throw new Error("User token is required to create user context");
59
+ const host = process.env.DATABRICKS_HOST;
60
+ if (!host) throw new Error("DATABRICKS_HOST environment variable is required for user context");
61
+ const serviceCtx = ServiceContext.get();
62
+ return {
63
+ client: new WorkspaceClient({
64
+ token,
65
+ host,
66
+ authType: "pat"
67
+ }, getClientOptions()),
68
+ userId,
69
+ userName,
70
+ warehouseId: serviceCtx.warehouseId,
71
+ workspaceId: serviceCtx.workspaceId,
72
+ isUserContext: true
73
+ };
74
+ }
75
+ /**
76
+ * Get the client options for WorkspaceClient.
77
+ * Exposed for testing purposes.
78
+ */
79
+ static getClientOptions() {
80
+ return getClientOptions();
81
+ }
82
+ static async createContext() {
83
+ const client = new WorkspaceClient({}, getClientOptions());
84
+ const warehouseId = ServiceContext.getWarehouseId(client);
85
+ const workspaceId = ServiceContext.getWorkspaceId(client);
86
+ const currentUser = await client.currentUser.me();
87
+ if (!currentUser.id) throw new Error("Service user ID not found");
88
+ return {
89
+ client,
90
+ serviceUserId: currentUser.id,
91
+ warehouseId,
92
+ workspaceId
93
+ };
94
+ }
95
+ static async getWorkspaceId(client) {
96
+ if (process.env.DATABRICKS_WORKSPACE_ID) return process.env.DATABRICKS_WORKSPACE_ID;
97
+ const response = await client.apiClient.request({
98
+ path: "/api/2.0/preview/scim/v2/Me",
99
+ method: "GET",
100
+ headers: new Headers(),
101
+ raw: false,
102
+ query: {},
103
+ responseHeaders: ["x-databricks-org-id"]
104
+ });
105
+ if (!response["x-databricks-org-id"]) throw new Error("Workspace ID not found");
106
+ return response["x-databricks-org-id"];
107
+ }
108
+ static async getWarehouseId(client) {
109
+ if (process.env.DATABRICKS_WAREHOUSE_ID) return process.env.DATABRICKS_WAREHOUSE_ID;
110
+ if (process.env.NODE_ENV === "development") {
111
+ const response = await client.apiClient.request({
112
+ path: "/api/2.0/sql/warehouses",
113
+ method: "GET",
114
+ headers: new Headers(),
115
+ raw: false,
116
+ query: { skip_cannot_use: "true" }
117
+ });
118
+ const priorities = {
119
+ RUNNING: 0,
120
+ STOPPED: 1,
121
+ STARTING: 2,
122
+ STOPPING: 3,
123
+ DELETED: 99,
124
+ DELETING: 99
125
+ };
126
+ const warehouses = (response.warehouses || []).sort((a, b) => {
127
+ return priorities[a.state] - priorities[b.state];
128
+ });
129
+ if (response.warehouses.length === 0) throw new Error("Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.");
130
+ const firstWarehouse = warehouses[0];
131
+ if (firstWarehouse.state === "DELETED" || firstWarehouse.state === "DELETING" || !firstWarehouse.id) throw new Error("Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.");
132
+ return firstWarehouse.id;
133
+ }
134
+ throw new Error("Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.");
135
+ }
136
+ /**
137
+ * Reset the service context. Only for testing purposes.
138
+ */
139
+ static reset() {
140
+ ServiceContext.instance = null;
141
+ ServiceContext.initPromise = null;
142
+ }
143
+ };
144
+ }));
145
+
146
+ //#endregion
147
+ init_service_context();
148
+ export { ServiceContext, init_service_context };
149
+ //# sourceMappingURL=service-context.js.map
@@ -0,0 +1 @@
1
+ {"version":3,"file":"service-context.js","names":["productName","productVersion"],"sources":["../../src/context/service-context.ts"],"sourcesContent":["import {\n type ClientOptions,\n type sql,\n WorkspaceClient,\n} from \"@databricks/sdk-experimental\";\nimport { coerce } from \"semver\";\nimport {\n name as productName,\n version as productVersion,\n} from \"../../package.json\";\nimport type { UserContext } from \"./user-context\";\n\n/**\n * Service context holds the service principal client and shared resources.\n * This is initialized once at app startup and shared across all requests.\n */\nexport interface ServiceContextState {\n /** WorkspaceClient authenticated as the service principal */\n client: WorkspaceClient;\n /** The service principal's user ID */\n serviceUserId: string;\n /** Promise that resolves to the warehouse ID */\n warehouseId: Promise<string>;\n /** Promise that resolves to the workspace ID */\n workspaceId: Promise<string>;\n}\n\nfunction getClientOptions(): ClientOptions {\n const isDev = process.env.NODE_ENV === \"development\";\n const semver = coerce(productVersion);\n const normalizedVersion = (semver?.version ??\n productVersion) as ClientOptions[\"productVersion\"];\n\n return {\n product: productName,\n productVersion: normalizedVersion,\n ...(isDev && { userAgentExtra: { mode: \"dev\" } }),\n };\n}\n\n/**\n * ServiceContext is a singleton that manages the service principal's\n * WorkspaceClient and shared resources like warehouse/workspace IDs.\n *\n * It's initialized once at app startup and provides the foundation\n * for both service principal and user context execution.\n */\nexport class ServiceContext {\n private static instance: ServiceContextState | null = null;\n private static initPromise: Promise<ServiceContextState> | null = null;\n\n /**\n * Initialize the service context. Should be called once at app startup.\n * Safe to call multiple times - will return the same instance.\n */\n static async initialize(): Promise<ServiceContextState> {\n if (ServiceContext.instance) {\n return ServiceContext.instance;\n }\n\n if (ServiceContext.initPromise) {\n return ServiceContext.initPromise;\n }\n\n ServiceContext.initPromise = ServiceContext.createContext();\n ServiceContext.instance = await ServiceContext.initPromise;\n return ServiceContext.instance;\n }\n\n /**\n * Get the initialized service context.\n * @throws Error if not initialized\n */\n static get(): ServiceContextState {\n if (!ServiceContext.instance) {\n throw new Error(\n \"ServiceContext not initialized. Call ServiceContext.initialize() first.\",\n );\n }\n return ServiceContext.instance;\n }\n\n /**\n * Check if the service context has been initialized.\n */\n static isInitialized(): boolean {\n return ServiceContext.instance !== null;\n }\n\n /**\n * Create a user context from request headers.\n *\n * @param token - The user's access token from x-forwarded-access-token header\n * @param userId - The user's ID from x-forwarded-user header\n * @param userName - Optional user name\n * @throws Error if token is not provided\n */\n static createUserContext(\n token: string,\n userId: string,\n userName?: string,\n ): UserContext {\n if (!token) {\n throw new Error(\"User token is required to create user context\");\n }\n\n const host = process.env.DATABRICKS_HOST;\n if (!host) {\n throw new Error(\n \"DATABRICKS_HOST environment variable is required for user context\",\n );\n }\n\n const serviceCtx = ServiceContext.get();\n\n // Create user client with the OAuth token from Databricks Apps\n // Note: We use authType: \"pat\" because the token is passed as a Bearer token\n // just like a PAT, even though it's technically an OAuth token\n const userClient = new WorkspaceClient(\n {\n token,\n host,\n authType: \"pat\",\n },\n getClientOptions(),\n );\n\n return {\n client: userClient,\n userId,\n userName,\n warehouseId: serviceCtx.warehouseId,\n workspaceId: serviceCtx.workspaceId,\n isUserContext: true,\n };\n }\n\n /**\n * Get the client options for WorkspaceClient.\n * Exposed for testing purposes.\n */\n static getClientOptions(): ClientOptions {\n return getClientOptions();\n }\n\n private static async createContext(): Promise<ServiceContextState> {\n const client = new WorkspaceClient({}, getClientOptions());\n\n const warehouseId = ServiceContext.getWarehouseId(client);\n const workspaceId = ServiceContext.getWorkspaceId(client);\n const currentUser = await client.currentUser.me();\n\n if (!currentUser.id) {\n throw new Error(\"Service user ID not found\");\n }\n\n return {\n client,\n serviceUserId: currentUser.id,\n warehouseId,\n workspaceId,\n };\n }\n\n private static async getWorkspaceId(\n client: WorkspaceClient,\n ): Promise<string> {\n if (process.env.DATABRICKS_WORKSPACE_ID) {\n return process.env.DATABRICKS_WORKSPACE_ID;\n }\n\n const response = (await client.apiClient.request({\n path: \"/api/2.0/preview/scim/v2/Me\",\n method: \"GET\",\n headers: new Headers(),\n raw: false,\n query: {},\n responseHeaders: [\"x-databricks-org-id\"],\n })) as { \"x-databricks-org-id\": string };\n\n if (!response[\"x-databricks-org-id\"]) {\n throw new Error(\"Workspace ID not found\");\n }\n\n return response[\"x-databricks-org-id\"];\n }\n\n private static async getWarehouseId(\n client: WorkspaceClient,\n ): Promise<string> {\n if (process.env.DATABRICKS_WAREHOUSE_ID) {\n return process.env.DATABRICKS_WAREHOUSE_ID;\n }\n\n if (process.env.NODE_ENV === \"development\") {\n const response = (await client.apiClient.request({\n path: \"/api/2.0/sql/warehouses\",\n method: \"GET\",\n headers: new Headers(),\n raw: false,\n query: {\n skip_cannot_use: \"true\",\n },\n })) as { warehouses: sql.EndpointInfo[] };\n\n const priorities: Record<sql.State, number> = {\n RUNNING: 0,\n STOPPED: 1,\n STARTING: 2,\n STOPPING: 3,\n DELETED: 99,\n DELETING: 99,\n };\n\n const warehouses = (response.warehouses || []).sort((a, b) => {\n return (\n priorities[a.state as sql.State] - priorities[b.state as sql.State]\n );\n });\n\n if (response.warehouses.length === 0) {\n throw new Error(\n \"Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.\",\n );\n }\n\n const firstWarehouse = warehouses[0];\n if (\n firstWarehouse.state === \"DELETED\" ||\n firstWarehouse.state === \"DELETING\" ||\n !firstWarehouse.id\n ) {\n throw new Error(\n \"Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.\",\n );\n }\n\n return firstWarehouse.id;\n }\n\n throw new Error(\n \"Warehouse ID not found. Please configure the DATABRICKS_WAREHOUSE_ID environment variable.\",\n );\n }\n\n /**\n * Reset the service context. Only for testing purposes.\n */\n static reset(): void {\n ServiceContext.instance = null;\n ServiceContext.initPromise = null;\n }\n}\n"],"mappings":";;;;;;AA2BA,SAAS,mBAAkC;CACzC,MAAM,QAAQ,QAAQ,IAAI,aAAa;AAKvC,QAAO;EACL,SAASA;EACT,gBANa,OAAOC,QAAe,EACF,WACjCA;EAKA,GAAI,SAAS,EAAE,gBAAgB,EAAE,MAAM,OAAO,EAAE;EACjD;;;;CAUU,iBAAb,MAAa,eAAe;;mBAC4B;;;sBACY;;;;;;EAMlE,aAAa,aAA2C;AACtD,OAAI,eAAe,SACjB,QAAO,eAAe;AAGxB,OAAI,eAAe,YACjB,QAAO,eAAe;AAGxB,kBAAe,cAAc,eAAe,eAAe;AAC3D,kBAAe,WAAW,MAAM,eAAe;AAC/C,UAAO,eAAe;;;;;;EAOxB,OAAO,MAA2B;AAChC,OAAI,CAAC,eAAe,SAClB,OAAM,IAAI,MACR,0EACD;AAEH,UAAO,eAAe;;;;;EAMxB,OAAO,gBAAyB;AAC9B,UAAO,eAAe,aAAa;;;;;;;;;;EAWrC,OAAO,kBACL,OACA,QACA,UACa;AACb,OAAI,CAAC,MACH,OAAM,IAAI,MAAM,gDAAgD;GAGlE,MAAM,OAAO,QAAQ,IAAI;AACzB,OAAI,CAAC,KACH,OAAM,IAAI,MACR,oEACD;GAGH,MAAM,aAAa,eAAe,KAAK;AAcvC,UAAO;IACL,QAViB,IAAI,gBACrB;KACE;KACA;KACA,UAAU;KACX,EACD,kBAAkB,CACnB;IAIC;IACA;IACA,aAAa,WAAW;IACxB,aAAa,WAAW;IACxB,eAAe;IAChB;;;;;;EAOH,OAAO,mBAAkC;AACvC,UAAO,kBAAkB;;EAG3B,aAAqB,gBAA8C;GACjE,MAAM,SAAS,IAAI,gBAAgB,EAAE,EAAE,kBAAkB,CAAC;GAE1D,MAAM,cAAc,eAAe,eAAe,OAAO;GACzD,MAAM,cAAc,eAAe,eAAe,OAAO;GACzD,MAAM,cAAc,MAAM,OAAO,YAAY,IAAI;AAEjD,OAAI,CAAC,YAAY,GACf,OAAM,IAAI,MAAM,4BAA4B;AAG9C,UAAO;IACL;IACA,eAAe,YAAY;IAC3B;IACA;IACD;;EAGH,aAAqB,eACnB,QACiB;AACjB,OAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;GAGrB,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;IAC/C,MAAM;IACN,QAAQ;IACR,SAAS,IAAI,SAAS;IACtB,KAAK;IACL,OAAO,EAAE;IACT,iBAAiB,CAAC,sBAAsB;IACzC,CAAC;AAEF,OAAI,CAAC,SAAS,uBACZ,OAAM,IAAI,MAAM,yBAAyB;AAG3C,UAAO,SAAS;;EAGlB,aAAqB,eACnB,QACiB;AACjB,OAAI,QAAQ,IAAI,wBACd,QAAO,QAAQ,IAAI;AAGrB,OAAI,QAAQ,IAAI,aAAa,eAAe;IAC1C,MAAM,WAAY,MAAM,OAAO,UAAU,QAAQ;KAC/C,MAAM;KACN,QAAQ;KACR,SAAS,IAAI,SAAS;KACtB,KAAK;KACL,OAAO,EACL,iBAAiB,QAClB;KACF,CAAC;IAEF,MAAM,aAAwC;KAC5C,SAAS;KACT,SAAS;KACT,UAAU;KACV,UAAU;KACV,SAAS;KACT,UAAU;KACX;IAED,MAAM,cAAc,SAAS,cAAc,EAAE,EAAE,MAAM,GAAG,MAAM;AAC5D,YACE,WAAW,EAAE,SAAsB,WAAW,EAAE;MAElD;AAEF,QAAI,SAAS,WAAW,WAAW,EACjC,OAAM,IAAI,MACR,6FACD;IAGH,MAAM,iBAAiB,WAAW;AAClC,QACE,eAAe,UAAU,aACzB,eAAe,UAAU,cACzB,CAAC,eAAe,GAEhB,OAAM,IAAI,MACR,6FACD;AAGH,WAAO,eAAe;;AAGxB,SAAM,IAAI,MACR,6FACD;;;;;EAMH,OAAO,QAAc;AACnB,kBAAe,WAAW;AAC1B,kBAAe,cAAc"}
@@ -0,0 +1,15 @@
1
+ import { __esmMin } from "../_virtual/rolldown_runtime.js";
2
+
3
+ //#region src/context/user-context.ts
4
+ /**
5
+ * Check if an execution context is a user context.
6
+ */
7
+ function isUserContext(ctx) {
8
+ return "isUserContext" in ctx && ctx.isUserContext === true;
9
+ }
10
+ var init_user_context = __esmMin((() => {}));
11
+
12
+ //#endregion
13
+ init_user_context();
14
+ export { init_user_context, isUserContext };
15
+ //# sourceMappingURL=user-context.js.map
@@ -0,0 +1 @@
1
+ {"version":3,"file":"user-context.js","names":[],"sources":["../../src/context/user-context.ts"],"sourcesContent":["import type { ServiceContextState } from \"./service-context\";\n\n/**\n * User execution context extends the service context with user-specific data.\n * Created on-demand when asUser(req) is called.\n */\nexport interface UserContext {\n /** WorkspaceClient authenticated as the user */\n client: ServiceContextState[\"client\"];\n /** The user's ID (from request headers) */\n userId: string;\n /** The user's name (from request headers) */\n userName?: string;\n /** Promise that resolves to the warehouse ID (inherited from service context) */\n warehouseId: Promise<string>;\n /** Promise that resolves to the workspace ID (inherited from service context) */\n workspaceId: Promise<string>;\n /** Flag indicating this is a user context */\n isUserContext: true;\n}\n\n/**\n * Execution context can be either service or user context.\n */\nexport type ExecutionContext = ServiceContextState | UserContext;\n\n/**\n * Check if an execution context is a user context.\n */\nexport function isUserContext(ctx: ExecutionContext): ctx is UserContext {\n return \"isUserContext\" in ctx && ctx.isUserContext === true;\n}\n"],"mappings":";;;;;;AA6BA,SAAgB,cAAc,KAA2C;AACvE,QAAO,mBAAmB,OAAO,IAAI,kBAAkB"}