From 37bc23c29d37ed9a1e1a480d0f03f7254b4655f5 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 1 Sep 2026 18:20:19 +0200 Subject: [PATCH 01/17] feat(appkit): migrate analytics to the modular @databricks/sdk-* MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Migrate the analytics stack (SQLWarehouseConnector + type-generator) off the legacy monolithic @databricks/sdk-experimental onto the new modular per-service @databricks/sdk-* SDK (v0.46.0, ESM-only), behind the existing workspace-client facade seam. The two services analytics depends on — warehouses and statementExecution — move together; every other service still routes through the legacy client (mixed state by design). - New packages/shared/src/workspace-client/modular.ts is the sole importer of @databricks/sdk-* (oxlint no-restricted-imports boundary), mirroring legacy.ts. Builds per-service WarehousesClient / StatementExecutionClient; maps wrapper options -> ClientOptions (host scheme-normalization, PAT empty-token guard, profile); stamps process-global client-info (sanitized, best-effort). - Connector + type-generator rewritten to the modular API: method renames (getStatement -> getStatementResult, getStatementResultChunkN -> getResultData), camelCase response model, CallOptions { signal }. - statementExecution relies on a pinned pnpm patch that restores the undocumented Reyden `attachment` field the SDK's unmarshal transform would otherwise strip. - Coerce the SDK's bigint row/byte counts back to number at the connector boundary so INLINE + ARROW_STREAM results stay JSON-serializable (cache / SSE frames). - Read the modular ApiError's `.code` (not only the legacy `.errorCode`) so the arrow disposition/format capability-rejection fallback still fires. Verified against live warehouses (standard + Reyden serverless): JSON and arrow (INLINE attachment + EXTERNAL_LINKS), OBO, warehouse auto-start, and metric views. Full appkit + shared suite green (3877 tests). Signed-off-by: MarioCadenas --- .oxlintrc.json | 7 +- .../api/appkit/Interface.WorkspaceClient.md | 8 +- .../connectors/sql-warehouse/arrow-schema.ts | 6 +- .../src/connectors/sql-warehouse/client.ts | 261 ++++++++++-------- .../src/connectors/sql-warehouse/defaults.ts | 14 +- .../sql-warehouse/tests/arrow-schema.test.ts | 26 +- .../sql-warehouse/tests/client.test.ts | 221 +++++++-------- .../sql-warehouse/warehouse-status-emitter.ts | 6 +- .../connectors/tests/sql-warehouse.test.ts | 227 +++++++++++++-- packages/appkit/src/evals/dataset.ts | 2 +- .../appkit/src/evals/tests/dataset.test.ts | 2 +- .../appkit/src/plugins/analytics/analytics.ts | 2 +- .../appkit/src/plugins/analytics/query.ts | 8 +- .../src/plugins/analytics/result-delivery.ts | 6 +- .../tests/analytics.integration.test.ts | 15 +- .../plugins/analytics/tests/analytics.test.ts | 14 +- .../tests/arrow-delivery.integration.test.ts | 6 +- .../plugins/analytics/tests/metric.test.ts | 4 +- .../src/stream/arrow-stream-processor.ts | 24 +- .../tests/arrow-stream-processor.test.ts | 15 +- packages/appkit/src/testing/fixtures.ts | 10 +- .../src/type-generator/statement-result.ts | 58 +++- .../tests/generate-queries.test.ts | 31 ++- .../src/type-generator/tests/index.test.ts | 74 +++-- .../type-generator/tests/mv-registry.test.ts | 52 +++- .../tests/statement-result.test.ts | 32 ++- .../tests/unreachable-warehouse-gate.test.ts | 2 +- .../tests/warehouse-status.test.ts | 10 +- .../src/type-generator/warehouse-status.ts | 6 +- packages/appkit/src/workspace-client/index.ts | 17 +- packages/shared/package.json | 5 + .../shared/src/workspace-client/client.ts | 24 +- packages/shared/src/workspace-client/index.ts | 2 + .../shared/src/workspace-client/modular.ts | 159 +++++++++++ .../workspace-client/tests/modular.test.ts | 113 ++++++++ packages/shared/src/workspace-client/types.ts | 19 +- ...ricks__sdk-statementexecution@0.46.0.patch | 34 +++ pnpm-lock.yaml | 84 ++++++ pnpm-workspace.yaml | 2 + 39 files changed, 1183 insertions(+), 425 deletions(-) create mode 100644 packages/shared/src/workspace-client/modular.ts create mode 100644 packages/shared/src/workspace-client/tests/modular.test.ts create mode 100644 patches/@databricks__sdk-statementexecution@0.46.0.patch diff --git a/.oxlintrc.json b/.oxlintrc.json index 96f5c64da..4d78b20bb 100644 --- a/.oxlintrc.json +++ b/.oxlintrc.json @@ -26,11 +26,8 @@ { "patterns": [ { - "group": [ - "@databricks/sdk-experimental", - "@databricks/sdk-experimental/**" - ], - "message": "Import the Databricks SDK only through the wrapper in packages/shared/src/workspace-client. Add a re-export there if you need a new symbol." + "group": ["@databricks/sdk-*", "@databricks/sdk-*/**"], + "message": "Import the Databricks SDK only through the wrapper in packages/shared/src/workspace-client (legacy.ts for @databricks/sdk-experimental, modular.ts for the modular @databricks/sdk-* packages). Add a re-export there if you need a new symbol." } ] } diff --git a/docs/docs/api/appkit/Interface.WorkspaceClient.md b/docs/docs/api/appkit/Interface.WorkspaceClient.md index bf508bbe4..26a680581 100644 --- a/docs/docs/api/appkit/Interface.WorkspaceClient.md +++ b/docs/docs/api/appkit/Interface.WorkspaceClient.md @@ -86,20 +86,20 @@ Serving Endpoints. ### statementExecution ```ts -readonly statementExecution: StatementExecutionService; +readonly statementExecution: StatementExecutionClient; ``` -Statement Execution. +Statement Execution (modular SDK). *** ### warehouses ```ts -readonly warehouses: WarehousesService; +readonly warehouses: WarehousesClient; ``` -SQL Warehouses. +SQL Warehouses (modular SDK). ## Methods diff --git a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts index 17d099e37..af5bfbb18 100644 --- a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts +++ b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts @@ -54,12 +54,12 @@ export function parseDatabricksType(typeText: string): DataType { export function buildEmptyArrowIPCBase64( columns: Array<{ name?: string; - type_text?: string; - type_name?: string; + typeText?: string; + typeName?: string; }>, ): string { const fields = columns.map((col, index) => { - const typeText = col.type_text ?? col.type_name ?? "STRING"; + const typeText = col.typeText ?? col.typeName ?? "STRING"; let dataType: DataType; try { dataType = parseDatabricksType(typeText); diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 1786f4f22..6afca34f8 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -21,10 +21,14 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; -import { - Context, - type sql, - type WorkspaceClient, +import type { + EndpointState, + ExecuteStatementRequest, + ExternalLink, + ResultData, + StatementResponse, + StatementStatus, + WorkspaceClient, } from "../../workspace-client"; import { buildEmptyArrowIPCBase64 } from "./arrow-schema"; import { executeStatementDefaults } from "./defaults"; @@ -40,9 +44,7 @@ const logger = createLogger("connectors:sql-warehouse"); * Arrow result to match the JSON path. Returns `undefined` when the manifest * carries no columns. */ -function arrowColumnNames( - response: sql.StatementResponse, -): string[] | undefined { +function arrowColumnNames(response: StatementResponse): string[] | undefined { const cols = response.manifest?.schema?.columns; if (!cols || cols.length === 0) return undefined; return cols.map((c, i) => @@ -50,6 +52,46 @@ function arrowColumnNames( ); } +/** + * Coerce the modular SDK's `bigint` row/byte counts back to `number` (the type + * the legacy SDK used). AppKit never does arithmetic on these — they are purely + * informational — but a stray `bigint` makes `JSON.stringify` throw ("Do not + * know how to serialize a BigInt") the instant the result is cached or written + * to an SSE frame. Reyden's INLINE + ARROW_STREAM result — which the analytics + * arrow path caches — carries them on `result`/`manifest`, so normalize every + * statement response at the SDK boundary. Mutates in place (the response is a + * fresh unmarshalled object, owned by the caller). + */ +const BIGINT_COUNT_FIELDS = ["rowOffset", "rowCount", "byteCount"] as const; + +function normalizeResultCounts(result: unknown): void { + if (!result || typeof result !== "object") return; + const r = result as Record; + for (const key of BIGINT_COUNT_FIELDS) { + if (typeof r[key] === "bigint") r[key] = Number(r[key]); + } + // EXTERNAL_LINKS entries carry the same count fields. + if (Array.isArray(r.externalLinks)) { + for (const link of r.externalLinks) normalizeResultCounts(link); + } +} + +function normalizeStatementCounts(response: T): T { + const manifest = response?.manifest as Record | undefined; + if (manifest) { + for (const key of ["totalRowCount", "totalByteCount"] as const) { + if (typeof manifest[key] === "bigint") + manifest[key] = Number(manifest[key]); + } + // Per-chunk `BaseChunkInfo` entries carry the same bigint count fields. + if (Array.isArray(manifest.chunks)) { + for (const chunk of manifest.chunks) normalizeResultCounts(chunk); + } + } + normalizeResultCounts(response?.result); + return response; +} + /** * Maximum size for inline Arrow IPC attachments (25 MiB decoded — the * Databricks Statement Execution API hard cap on INLINE responses). @@ -64,8 +106,8 @@ const MAX_INLINE_ATTACHMENT_BYTES = 25 * 1024 * 1024; /** * Safety cap on how many additional EXTERNAL_LINKS chunks * {@link SQLWarehouseConnector._resolveAllExternalLinks} will follow when the - * manifest omits `total_chunk_count`. High enough to cover any real result; - * only bounds a misbehaving warehouse with a cyclic `next_chunk_index`. + * manifest omits `totalChunkCount`. High enough to cover any real result; + * only bounds a misbehaving warehouse with a cyclic `nextChunkIndex`. */ const MAX_EXTERNAL_CHUNK_FOLLOWS = 10_000; @@ -105,7 +147,7 @@ const WAREHOUSE_RUNNING_CACHE_TTL_MS = 30_000; */ export interface WarehouseStatusUpdate { /** Current state from the SDK (RUNNING | STARTING | STOPPED | STOPPING | DELETED | DELETING). */ - state: sql.State; + state: EndpointState; /** Milliseconds elapsed since `ensureWarehouseRunning` was called. */ elapsedMs: number; /** 1-based attempt counter — useful for tests and telemetry. */ @@ -203,7 +245,7 @@ export class SQLWarehouseConnector { async executeStatement( workspaceClient: WorkspaceClient, - input: sql.ExecuteStatementRequest, + input: ExecuteStatementRequest, signal?: AbortSignal, ) { const startTime = Date.now(); @@ -220,7 +262,7 @@ export class SQLWarehouseConnector { kind: SpanKind.CLIENT, attributes: { "db.system": "databricks", - "db.warehouse_id": input.warehouse_id || "", + "db.warehouse_id": input.warehouseId || "", "db.catalog": input.catalog ?? "", "db.schema": input.schema ?? "", "db.statement": input.statement?.substring(0, 500) || "", @@ -252,52 +294,52 @@ export class SQLWarehouseConnector { throw ValidationError.missingField("statement"); } - if (!input.warehouse_id) { + if (!input.warehouseId) { throw ValidationError.missingField("warehouse_id"); } - const body: sql.ExecuteStatementRequest = { + const body: ExecuteStatementRequest = { statement: input.statement, parameters: input.parameters, - warehouse_id: input.warehouse_id, + warehouseId: input.warehouseId, catalog: input.catalog, schema: input.schema, - wait_timeout: - input.wait_timeout || executeStatementDefaults.wait_timeout, + waitTimeout: + input.waitTimeout || executeStatementDefaults.waitTimeout, disposition: input.disposition || executeStatementDefaults.disposition, format: input.format || executeStatementDefaults.format, - byte_limit: input.byte_limit, - row_limit: input.row_limit, - on_wait_timeout: - input.on_wait_timeout || executeStatementDefaults.on_wait_timeout, + byteLimit: input.byteLimit, + rowLimit: input.rowLimit, + onWaitTimeout: + input.onWaitTimeout || executeStatementDefaults.onWaitTimeout, }; span.addEvent("statement.submitting", { - "db.warehouse_id": input.warehouse_id, + "db.warehouse_id": input.warehouseId, }); const response = - await workspaceClient.statementExecution.executeStatement( - body, - this._createContext(signal), - ); + await workspaceClient.statementExecution.executeStatement(body, { + signal, + }); if (!response) { throw ConnectionError.apiFailure("SQL Warehouse"); } + normalizeStatementCounts(response); const status = response.status; - const statementId = response.statement_id as string; + const statementId = response.statementId as string; span.setAttribute("db.statement_id", statementId); span.addEvent("statement.submitted", { - "db.statement_id": response.statement_id, + "db.statement_id": response.statementId, "db.status": status?.state, }); let result: - | sql.StatementResponse - | { result: { statement_id: string; status: sql.StatementStatus } }; + | StatementResponse + | { result: { statement_id: string; status: StatementStatus } }; switch (status?.state) { case "RUNNING": @@ -322,7 +364,7 @@ export class SQLWarehouseConnector { case "FAILED": throw ExecutionError.statementFailed( status.error?.message, - status.error?.error_code, + status.error?.errorCode, ); case "CANCELED": throw ExecutionError.canceled(); @@ -336,7 +378,7 @@ export class SQLWarehouseConnector { const resultData = result.result as any; const rowCount = - resultData?.data?.length ?? resultData?.data_array?.length ?? 0; + resultData?.data?.length ?? resultData?.dataArray?.length ?? 0; if (rowCount > 0) { span.setAttribute("db.result.row_count", rowCount); @@ -344,7 +386,7 @@ export class SQLWarehouseConnector { const duration = Date.now() - startTime; logger.event()?.setContext("sql-warehouse", { - warehouse_id: input.warehouse_id, + warehouse_id: input.warehouseId, rows_returned: rowCount, query_duration_ms: duration, }); @@ -385,7 +427,7 @@ export class SQLWarehouseConnector { } const attributes = { - "db.warehouse_id": input.warehouse_id, + "db.warehouse_id": input.warehouseId, "db.catalog": input.catalog ?? "", "db.schema": input.schema ?? "", "db.statement": input.statement?.substring(0, 500) || "", @@ -622,11 +664,13 @@ export class SQLWarehouseConnector { ); } - let info: Awaited>; + let info: Awaited< + ReturnType + >; try { - info = await workspaceClient.warehouses.get( + info = await workspaceClient.warehouses.getWarehouse( { id: warehouseId }, - this._createContext(signal), + { signal }, ); } catch (error) { // A real cancellation must still surface as canceled, not be swallowed @@ -681,9 +725,9 @@ export class SQLWarehouseConnector { if (!didStart) { emitter.emit("STARTING", summary); onWarehouseStartIssued?.(); - await workspaceClient.warehouses.start( + await workspaceClient.warehouses.startWarehouse( { id: warehouseId }, - this._createContext(signal), + { signal }, ); didStart = true; } else { @@ -830,15 +874,14 @@ export class SQLWarehouseConnector { }); const response = - await workspaceClient.statementExecution.getStatement( - { - statement_id: statementId, - }, - this._createContext(signal), + await workspaceClient.statementExecution.getStatementResult( + { statementId }, + { signal }, ); if (!response) { throw ConnectionError.apiFailure("SQL Warehouse"); } + normalizeStatementCounts(response); const status = response.status; @@ -868,7 +911,7 @@ export class SQLWarehouseConnector { case "FAILED": throw ExecutionError.statementFailed( status.error?.message, - status.error?.error_code, + status.error?.errorCode, ); case "CANCELED": throw ExecutionError.canceled(); @@ -902,13 +945,13 @@ export class SQLWarehouseConnector { } private async _transformDataArray( - response: sql.StatementResponse, + response: StatementResponse, workspaceClient: WorkspaceClient, signal?: AbortSignal, ) { if (response.manifest?.format === "ARROW_STREAM") { const result = response.result as - | (sql.ResultData & { attachment?: string }) + | (ResultData & { attachment?: string }) | undefined; // Inline Arrow: pass the base64 IPC attachment through unmodified so @@ -924,20 +967,20 @@ export class SQLWarehouseConnector { // rather than omitting it) — it must NOT go down the streaming path // (`streamChunks([])` rejects), so fall through to synthesize an empty // Arrow table below. - if (result?.external_links && result.external_links.length > 0) { + if (result?.externalLinks && result.externalLinks.length > 0) { return this.updateWithArrowStatus(response, workspaceClient, signal); } // Empty result with a known schema: synthesize a zero-row Arrow IPC // attachment so the client always receives an Arrow Table for // ARROW_STREAM, regardless of whether the warehouse returned data. - // Note: an empty array (`data_array: []`) is truthy, so length-check + // Note: an empty array (`dataArray: []`) is truthy, so length-check // explicitly — otherwise zero-row responses fall through to the JSON // row transform below and return `[]` JSON rows instead of an Arrow // table. const hasNoRows = - !result?.data_array || - (Array.isArray(result.data_array) && result.data_array.length === 0); + !result?.dataArray || + (Array.isArray(result.dataArray) && result.dataArray.length === 0); if (hasNoRows && response.manifest?.schema?.columns) { const synthesized = buildEmptyArrowIPCBase64( response.manifest.schema.columns, @@ -948,19 +991,19 @@ export class SQLWarehouseConnector { }; } - // Inline data_array under ARROW_STREAM (rare): fall through to the + // Inline dataArray under ARROW_STREAM (rare): fall through to the // row transform below. The hook will receive `type: "result"` rows; // callers asking for ARROW_STREAM should not hit this path with // current Databricks warehouses. } - if (!response.result?.data_array || !response.manifest?.schema?.columns) { + if (!response.result?.dataArray || !response.manifest?.schema?.columns) { return response; } const columns = response.manifest.schema.columns; - const transformedData = response.result.data_array.map((row) => { + const transformedData = response.result.dataArray.map((row) => { const obj: Record = {}; row.forEach((value, index) => { const column = columns[index]; @@ -968,7 +1011,7 @@ export class SQLWarehouseConnector { // attempt to parse JSON strings for string columns if ( - column?.type_name === "STRING" && + column?.typeName === "STRING" && typeof value === "string" && value && (value[0] === "{" || value[0] === "[") @@ -986,8 +1029,8 @@ export class SQLWarehouseConnector { return obj; }); - // remove data_array - const { data_array: _data_array, ...restResult } = response.result; + // remove dataArray + const { dataArray: _dataArray, ...restResult } = response.result; return { ...response, result: { @@ -1009,7 +1052,7 @@ export class SQLWarehouseConnector { * mechanism used for both INLINE and EXTERNAL_LINKS. */ private _validateArrowAttachment( - response: sql.StatementResponse, + response: StatementResponse, attachment: string, ) { // Cap the size to protect against unbounded inline payloads from @@ -1037,14 +1080,16 @@ export class SQLWarehouseConnector { return { ...response, result: { - ...(response.result as sql.ResultData & { + ...(response.result as ResultData & { attachment?: string; columnNames?: string[]; }), - // `statement_id` is a top-level field, not on `ResultData` — carry it + // `statementId` is a top-level field, not on `ResultData` — carry it // onto the result (as the EXTERNAL_LINKS path does) so the route can // advertise it in `X-Appkit-Arrow-Columns-Ref` for wide inline schemas. - statement_id: response.statement_id, + // Kept as the synthetic `statement_id` key (the connector→route wire + // contract), sourced from the modular SDK's camelCase `statementId`. + statement_id: response.statementId, columnNames, }, }; @@ -1054,26 +1099,26 @@ export class SQLWarehouseConnector { } private async updateWithArrowStatus( - response: sql.StatementResponse, + response: StatementResponse, workspaceClient: WorkspaceClient, signal?: AbortSignal, ): Promise<{ result: { statement_id: string; - status: sql.StatementStatus; + status: StatementStatus; columnNames?: string[]; - external_links?: sql.ExternalLink[]; + external_links?: ExternalLink[]; refreshChunkLink?: RefreshChunkLink; }; }> { - const statementId = response.statement_id as string; + const statementId = response.statementId as string; return { result: { statement_id: statementId, status: { state: response.status?.state, error: response.status?.error, - } as sql.StatementStatus, + } as StatementStatus, columnNames: arrowColumnNames(response), // Resolve the pre-signed links for EVERY chunk in the caller's own // execution context. Streaming these directly (see @@ -1100,9 +1145,9 @@ export class SQLWarehouseConnector { /** * Resolve pre-signed links for EVERY chunk of an EXTERNAL_LINKS result. * - * The execute/getStatement response carries only the first chunk's links - * (each link, except the last, exposes `next_chunk_index`); the remaining - * chunks are fetched with `getStatementResultChunkN`. Runs in the caller's + * The execute/getStatementResult response carries only the first chunk's links + * (each link, except the last, exposes `nextChunkIndex`); the remaining + * chunks are fetched with `getResultData`. Runs in the caller's * identity context (user creds for `.obo.sql`), so there is no cross-identity * fetch. Only the tiny link metadata is resolved eagerly — the bytes still * stream one chunk at a time downstream. Without this a multi-chunk result @@ -1111,29 +1156,29 @@ export class SQLWarehouseConnector { private async _resolveAllExternalLinks( workspaceClient: WorkspaceClient, statementId: string, - response: sql.StatementResponse, + response: StatementResponse, signal?: AbortSignal, - ): Promise { - const first = response.result?.external_links; + ): Promise { + const first = response.result?.externalLinks; if (!first || first.length === 0) return first; - const links: sql.ExternalLink[] = [...first]; + const links: ExternalLink[] = [...first]; // Bound the follow loop so a warehouse returning a cyclic/never-ending // `next_chunk_index` can't spin forever. The manifest's chunk count is the // natural bound; fall back to a generous safety cap if it's absent (real // results still terminate earlier when `next_chunk_index` becomes null) so // a missing count doesn't silently truncate a genuine multi-chunk result. const maxFetches = - response.manifest?.total_chunk_count ?? MAX_EXTERNAL_CHUNK_FOLLOWS; + response.manifest?.totalChunkCount ?? MAX_EXTERNAL_CHUNK_FOLLOWS; let next = this._nextChunkIndex(first); for (let fetches = 0; next != null && fetches < maxFetches; fetches++) { if (signal?.aborted) throw ExecutionError.canceled(); - const chunk = - await workspaceClient.statementExecution.getStatementResultChunkN( - { statement_id: statementId, chunk_index: next }, - this._createContext(signal), - ); - const chunkLinks = chunk.external_links ?? []; + const chunk = await workspaceClient.statementExecution.getResultData( + { statementId, chunkIndex: next }, + { signal }, + ); + normalizeResultCounts(chunk); + const chunkLinks = chunk.externalLinks ?? []; if (chunkLinks.length === 0) break; links.push(...chunkLinks); next = this._nextChunkIndex(chunkLinks); @@ -1141,32 +1186,32 @@ export class SQLWarehouseConnector { return links; } - /** The `next_chunk_index` advertised by a chunk's links, if any. */ - private _nextChunkIndex(links: sql.ExternalLink[]): number | undefined { + /** The `nextChunkIndex` advertised by a chunk's links, if any. */ + private _nextChunkIndex(links: ExternalLink[]): number | undefined { for (const link of links) { - if (link.next_chunk_index != null) return link.next_chunk_index; + if (link.nextChunkIndex != null) return link.nextChunkIndex; } return undefined; } /** * A closure that re-mints a single chunk's pre-signed link via - * `getStatementResultChunkN`, bound to the caller's workspace client + + * `getResultData`, bound to the caller's workspace client + * statement id. Created here (in the caller's identity context) so the * streamer — which runs outside that context — can refresh an expired link - * for `.obo.sql` statements without a cross-identity `getStatement`. + * for `.obo.sql` statements without a cross-identity `getStatementResult`. */ private _makeChunkLinkRefresher( workspaceClient: WorkspaceClient, statementId: string, ): RefreshChunkLink { return async (chunkIndex, signal) => { - const chunk = - await workspaceClient.statementExecution.getStatementResultChunkN( - { statement_id: statementId, chunk_index: chunkIndex }, - this._createContext(signal), - ); - return chunk.external_links?.find((l) => l.chunk_index === chunkIndex); + const chunk = await workspaceClient.statementExecution.getResultData( + { statementId, chunkIndex }, + { signal }, + ); + normalizeResultCounts(chunk); + return chunk.externalLinks?.find((l) => l.chunkIndex === chunkIndex); }; } @@ -1178,7 +1223,7 @@ export class SQLWarehouseConnector { * the pre-signed URLs need no auth to download. */ streamExternalLinks( - chunks: sql.ExternalLink[], + chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, ): AsyncGenerator { @@ -1196,10 +1241,12 @@ export class SQLWarehouseConnector { jobId: string, signal?: AbortSignal, ): Promise { - const response = await workspaceClient.statementExecution.getStatement( - { statement_id: jobId }, - this._createContext(signal), - ); + const response = + await workspaceClient.statementExecution.getStatementResult( + { statementId: jobId }, + { signal }, + ); + normalizeStatementCounts(response); return arrowColumnNames(response); } @@ -1218,25 +1265,19 @@ export class SQLWarehouseConnector { if (error instanceof AppKitError) { throw error; } + // The legacy SDK exposed the Databricks error code as `errorCode`; the + // modular SDK's `ApiError` carries it as `code` (e.g. "INVALID_PARAMETER_VALUE"). + // Read either, so callers can still branch on the stable code — notably the + // analytics arrow disposition/format fallback, which keys on + // INVALID_PARAMETER_VALUE / NOT_IMPLEMENTED to switch INLINE↔EXTERNAL_LINKS. const sdkErrorCode = - error && typeof error === "object" && "errorCode" in error - ? (error as { errorCode?: unknown }).errorCode + error && typeof error === "object" + ? ((error as { errorCode?: unknown }).errorCode ?? + (error as { code?: unknown }).code) : undefined; throw ExecutionError.statementFailed( error instanceof Error ? error.message : String(error), typeof sdkErrorCode === "string" ? sdkErrorCode : undefined, ); } - - // create context for cancellation token - private _createContext(signal?: AbortSignal) { - return new Context({ - cancellationToken: { - isCancellationRequested: signal?.aborted ?? false, - onCancellationRequested: (cb: () => void) => { - signal?.addEventListener("abort", cb, { once: true }); - }, - }, - }); - } } diff --git a/packages/appkit/src/connectors/sql-warehouse/defaults.ts b/packages/appkit/src/connectors/sql-warehouse/defaults.ts index b046a5c4a..3a57c8058 100644 --- a/packages/appkit/src/connectors/sql-warehouse/defaults.ts +++ b/packages/appkit/src/connectors/sql-warehouse/defaults.ts @@ -1,18 +1,18 @@ -import type { sql } from "../../workspace-client"; +import type { ExecuteStatementRequest } from "../../workspace-client"; interface ExecuteStatementDefaults { - wait_timeout: string; - disposition: sql.ExecuteStatementRequest["disposition"]; - format: sql.ExecuteStatementRequest["format"]; - on_wait_timeout: sql.ExecuteStatementRequest["on_wait_timeout"]; + waitTimeout: string; + disposition: ExecuteStatementRequest["disposition"]; + format: ExecuteStatementRequest["format"]; + onWaitTimeout: ExecuteStatementRequest["onWaitTimeout"]; timeout: number; } // @TODO: Make these configurable globally and validate right values export const executeStatementDefaults: ExecuteStatementDefaults = { - wait_timeout: "30s", + waitTimeout: "30s", disposition: "INLINE", format: "JSON_ARRAY", - on_wait_timeout: "CONTINUE", + onWaitTimeout: "CONTINUE", timeout: 60000, }; diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts index d8f52f016..b7826e87e 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts @@ -428,11 +428,11 @@ describe("parseDatabricksType — error / robustness", () => { describe("buildEmptyArrowIPCBase64", () => { test("produces a decodable empty Arrow Table with the right schema", () => { const columns = [ - { name: "user_id", type_text: "BIGINT" }, - { name: "name", type_text: "STRING" }, - { name: "created_at", type_text: "TIMESTAMP" }, - { name: "balance", type_text: "DECIMAL(10,2)" }, - { name: "active", type_text: "BOOLEAN" }, + { name: "user_id", typeText: "BIGINT" }, + { name: "name", typeText: "STRING" }, + { name: "created_at", typeText: "TIMESTAMP" }, + { name: "balance", typeText: "DECIMAL(10,2)" }, + { name: "active", typeText: "BOOLEAN" }, ]; const b64 = buildEmptyArrowIPCBase64(columns); const buf = Buffer.from(b64, "base64"); @@ -463,9 +463,9 @@ describe("buildEmptyArrowIPCBase64", () => { test("round-trips nested types end-to-end", () => { const columns = [ - { name: "tags", type_text: "ARRAY" }, - { name: "meta", type_text: "STRUCT" }, - { name: "counts", type_text: "MAP" }, + { name: "tags", typeText: "ARRAY" }, + { name: "meta", typeText: "STRUCT" }, + { name: "counts", typeText: "MAP" }, ]; const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); const table = tableFromIPC(buf); @@ -476,8 +476,8 @@ describe("buildEmptyArrowIPCBase64", () => { expect(table.schema.fields[2]?.type).toBeInstanceOf(Map_); }); - test("falls back from type_text to type_name when type_text missing", () => { - const columns = [{ name: "id", type_name: "BIGINT" }]; + test("falls back from typeText to typeName when typeText missing", () => { + const columns = [{ name: "id", typeName: "BIGINT" }]; const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); const table = tableFromIPC(buf); expect( @@ -487,8 +487,8 @@ describe("buildEmptyArrowIPCBase64", () => { test("unknown type degrades to Utf8 without throwing", () => { const columns = [ - { name: "id", type_text: "BIGINT" }, - { name: "weird", type_text: "FUTURE_TYPE_NOT_YET_SUPPORTED" }, + { name: "id", typeText: "BIGINT" }, + { name: "weird", typeText: "FUTURE_TYPE_NOT_YET_SUPPORTED" }, ]; const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); const table = tableFromIPC(buf); @@ -499,7 +499,7 @@ describe("buildEmptyArrowIPCBase64", () => { }); test("missing column name gets a synthesized placeholder", () => { - const columns = [{ type_text: "STRING" }, { name: "", type_text: "INT" }]; + const columns = [{ typeText: "STRING" }, { name: "", typeText: "INT" }]; const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); const table = tableFromIPC(buf); expect(table.schema.fields[0]?.name).toBe("column_0"); diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts index 5a945b0cc..8344df2fc 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts @@ -1,7 +1,10 @@ import { tableFromIPC } from "apache-arrow"; import { describe, expect, test, vi } from "vitest"; -import type { sql } from "../../../workspace-client"; +import type { + ExternalLink, + StatementResponse, +} from "../../../workspace-client"; vi.mock("../../../telemetry", () => { const mockMeter = { @@ -40,11 +43,11 @@ function createConnector() { // `_transformDataArray` is async — it paginates multi-chunk EXTERNAL_LINKS // results. The workspace client is only touched when following -// `next_chunk_index`, so a bare stub suffices for the inline / JSON / +// `nextChunkIndex`, so a bare stub suffices for the inline / JSON / // single-chunk cases; the multi-chunk tests pass a real mock. function transform( connector: SQLWarehouseConnector, - response: sql.StatementResponse, + response: StatementResponse, workspaceClient: unknown = {}, ) { return (connector as any)._transformDataArray(response, workspaceClient); @@ -58,11 +61,11 @@ const REAL_ARROW_ATTACHMENT = describe("SQLWarehouseConnector._transformDataArray", () => { describe("classic warehouse (JSON_ARRAY + INLINE)", () => { - test("transforms data_array rows into named objects", async () => { + test("transforms dataArray rows into named objects", async () => { const connector = createConnector(); // Real response shape from classic warehouse: INLINE + JSON_ARRAY const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY", @@ -71,14 +74,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { columns: [ { name: "test_col", - type_text: "INT", - type_name: "INT", + typeText: "INT", + typeName: "INT", position: 0, }, { name: "test_col2", - type_text: "INT", - type_name: "INT", + typeText: "INT", + typeName: "INT", position: 1, }, ], @@ -87,33 +90,33 @@ describe("SQLWarehouseConnector._transformDataArray", () => { truncated: false, }, result: { - data_array: [["1", "2"]], + dataArray: [["1", "2"]], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.data).toEqual([{ test_col: "1", test_col2: "2" }]); - expect(result.result.data_array).toBeUndefined(); + expect(result.result.dataArray).toBeUndefined(); }); test("parses JSON strings in STRING columns", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY", schema: { columns: [ - { name: "id", type_name: "INT" }, - { name: "metadata", type_name: "STRING" }, + { name: "id", typeName: "INT" }, + { name: "metadata", typeName: "STRING" }, ], }, }, result: { - data_array: [["1", '{"key":"value"}']], + dataArray: [["1", '{"key":"value"}']], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.data[0].metadata).toEqual({ key: "value" }); @@ -125,26 +128,26 @@ describe("SQLWarehouseConnector._transformDataArray", () => { const connector = createConnector(); // Real response shape from classic warehouse: EXTERNAL_LINKS + ARROW_STREAM const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "test_col", type_name: "INT" }, - { name: "test_col2", type_name: "INT" }, + { name: "test_col", typeName: "INT" }, + { name: "test_col2", typeName: "INT" }, ], }, }, result: { - external_links: [ + externalLinks: [ { - external_link: "https://storage.example.com/chunk0", + externalLink: "https://storage.example.com/chunk0", expiration: "2026-04-15T00:00:00Z", }, ], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.statement_id).toBe("stmt-1"); @@ -156,9 +159,9 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("passes attachment through unchanged for client-side decoding", async () => { const connector = createConnector(); // Real response shape from serverless warehouse: INLINE + ARROW_STREAM - // Data arrives in result.attachment as base64-encoded Arrow IPC, not data_array. + // Data arrives in result.attachment as base64-encoded Arrow IPC, not dataArray. const response = { - statement_id: "00000001-test-stmt", + statementId: "00000001-test-stmt", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", @@ -167,30 +170,30 @@ describe("SQLWarehouseConnector._transformDataArray", () => { columns: [ { name: "test_col", - type_text: "INT", - type_name: "INT", + typeText: "INT", + typeName: "INT", position: 0, }, { name: "test_col2", - type_text: "INT", - type_name: "INT", + typeText: "INT", + typeName: "INT", position: 1, }, ], - total_chunk_count: 1, - chunks: [{ chunk_index: 0, row_offset: 0, row_count: 1 }], + totalChunkCount: 1, + chunks: [{ chunkIndex: 0, row_offset: 0, row_count: 1 }], total_row_count: 1, }, truncated: false, }, result: { - chunk_index: 0, + chunkIndex: 0, row_offset: 0, row_count: 1, attachment: REAL_ARROW_ATTACHMENT, }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); @@ -206,56 +209,56 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("preserves manifest and status alongside attachment", async () => { const connector = createConnector(); const response = { - statement_id: "00000001-test-stmt", + statementId: "00000001-test-stmt", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "test_col", type_name: "INT" }, - { name: "test_col2", type_name: "INT" }, + { name: "test_col", typeName: "INT" }, + { name: "test_col2", typeName: "INT" }, ], }, }, result: { - chunk_index: 0, + chunkIndex: 0, row_count: 1, attachment: REAL_ARROW_ATTACHMENT, }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); // Manifest, statement_id, and attachment are all preserved expect(result.manifest.format).toBe("ARROW_STREAM"); - expect(result.statement_id).toBe("00000001-test-stmt"); + expect(result.statementId).toBe("00000001-test-stmt"); expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); }); test("synthesizes an empty Arrow IPC attachment for empty results so the client always gets a Table", async () => { const connector = createConnector(); - // Empty result: no attachment, no data_array, no external_links — but + // Empty result: no attachment, no dataArray, no external_links — but // the manifest still describes the schema. The connector should fill in // `attachment` with a zero-row Arrow IPC matching the schema. const response = { - statement_id: "stmt-empty", + statementId: "stmt-empty", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "user_id", type_text: "BIGINT", type_name: "BIGINT" }, - { name: "name", type_text: "STRING", type_name: "STRING" }, + { name: "user_id", typeText: "BIGINT", typeName: "BIGINT" }, + { name: "name", typeText: "STRING", typeName: "STRING" }, { name: "balance", - type_text: "DECIMAL(10,2)", - type_name: "DECIMAL", + typeText: "DECIMAL(10,2)", + typeName: "DECIMAL", }, ], }, total_row_count: 0, }, result: {}, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const transformed = await transform(connector, response); const attachment: string = transformed.result.attachment; @@ -275,18 +278,18 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("does NOT synthesize an attachment when external_links are present", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-ext", + statementId: "stmt-ext", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", - schema: { columns: [{ name: "x", type_text: "INT" }] }, + schema: { columns: [{ name: "x", typeText: "INT" }] }, }, result: { - external_links: [ - { external_link: "https://example.com/x", expiration: "9999" }, + externalLinks: [ + { externalLink: "https://example.com/x", expiration: "9999" }, ], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const transformed = await transform(connector, response); // External-links path returns the statement_id projection — no attachment. @@ -296,21 +299,21 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("empty external_links array is a zero-row result → synthesizes an empty table (not the streaming path)", async () => { const connector = createConnector(); - // Some warehouses emit `external_links: []` for a zero-row result rather + // Some warehouses emit `externalLinks: []` for a zero-row result rather // than omitting it. An empty array must NOT go down the streaming path // (streamChunks([]) rejects) — synthesize an empty Arrow table instead. const response = { - statement_id: "stmt-empty-ext", + statementId: "stmt-empty-ext", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { - columns: [{ name: "x", type_text: "INT", type_name: "INT" }], + columns: [{ name: "x", typeText: "INT", typeName: "INT" }], }, total_row_count: 0, }, - result: { external_links: [] }, - } as unknown as sql.StatementResponse; + result: { externalLinks: [] }, + } as unknown as StatementResponse; const transformed = await transform(connector, response); const attachment: string = transformed.result.attachment; @@ -323,11 +326,11 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("does NOT synthesize an attachment when schema is missing", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-no-schema", + statementId: "stmt-no-schema", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: {}, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const transformed = await transform(connector, response); // Without a schema we cannot build a Table — pass through unchanged. @@ -340,11 +343,11 @@ describe("SQLWarehouseConnector._transformDataArray", () => { // base64 chars decodes to ~27 MiB, comfortably above the limit. const oversized = "A".repeat(36 * 1024 * 1024); const response = { - statement_id: "stmt-oversized", + statementId: "stmt-oversized", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: oversized }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; await expect(transform(connector, response)).rejects.toThrow( /exceeds maximum size/, @@ -352,28 +355,28 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }); }); - describe("ARROW_STREAM with data_array (hypothetical inline variant)", () => { - test("transforms data_array like JSON_ARRAY path", async () => { + describe("ARROW_STREAM with dataArray (hypothetical inline variant)", () => { + test("transforms dataArray like JSON_ARRAY path", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "id", type_name: "INT" }, - { name: "value", type_name: "STRING" }, + { name: "id", typeName: "INT" }, + { name: "value", typeName: "STRING" }, ], }, }, result: { - data_array: [ + dataArray: [ ["1", "hello"], ["2", "world"], ], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.data).toEqual([ @@ -384,88 +387,88 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }); describe("edge cases", () => { - test("returns response unchanged when no data_array, attachment, or schema", async () => { + test("returns response unchanged when no dataArray, attachment, or schema", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY" }, result: {}, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result).toBe(response); }); - test("attachment takes priority over data_array when both present", async () => { + test("attachment takes priority over dataArray when both present", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "test_col", type_name: "INT" }, - { name: "test_col2", type_name: "INT" }, + { name: "test_col", typeName: "INT" }, + { name: "test_col2", typeName: "INT" }, ], }, }, result: { attachment: REAL_ARROW_ATTACHMENT, - data_array: [["999", "999"]], + dataArray: [["999", "999"]], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); - // Should pass attachment through (client decodes), not transform data_array + // Should pass attachment through (client decodes), not transform dataArray expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); expect(result.result.data).toBeUndefined(); }); }); describe("multi-chunk EXTERNAL_LINKS pagination", () => { - function multiChunkResponse(totalChunks: number): sql.StatementResponse { + function multiChunkResponse(totalChunks: number): StatementResponse { return { - statement_id: "stmt-multi", + statementId: "stmt-multi", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", - total_chunk_count: totalChunks, - schema: { columns: [{ name: "x", type_name: "INT" }] }, + totalChunkCount: totalChunks, + schema: { columns: [{ name: "x", typeName: "INT" }] }, }, result: { - external_links: [ + externalLinks: [ { - chunk_index: 0, - external_link: "https://example.com/chunk0", - next_chunk_index: 1, + chunkIndex: 0, + externalLink: "https://example.com/chunk0", + nextChunkIndex: 1, }, ], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; } - test("follows next_chunk_index to resolve every chunk's links", async () => { + test("follows nextChunkIndex to resolve every chunk's links", async () => { const connector = createConnector(); - const getStatementResultChunkN = vi + const getResultData = vi .fn() .mockResolvedValueOnce({ - external_links: [ + externalLinks: [ { - chunk_index: 1, - external_link: "https://example.com/chunk1", - next_chunk_index: 2, + chunkIndex: 1, + externalLink: "https://example.com/chunk1", + nextChunkIndex: 2, }, ], }) .mockResolvedValueOnce({ - external_links: [ - { chunk_index: 2, external_link: "https://example.com/chunk2" }, + externalLinks: [ + { chunkIndex: 2, externalLink: "https://example.com/chunk2" }, ], }); const workspaceClient = { - statementExecution: { getStatementResultChunkN }, + statementExecution: { getResultData }, }; const result = await transform( @@ -474,16 +477,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { workspaceClient, ); - expect(getStatementResultChunkN).toHaveBeenCalledTimes(2); - expect(getStatementResultChunkN).toHaveBeenNthCalledWith( + expect(getResultData).toHaveBeenCalledTimes(2); + expect(getResultData).toHaveBeenNthCalledWith( 1, - expect.objectContaining({ statement_id: "stmt-multi", chunk_index: 1 }), + expect.objectContaining({ statementId: "stmt-multi", chunkIndex: 1 }), expect.anything(), ); expect( - result.result.external_links.map( - (l: sql.ExternalLink) => l.external_link, - ), + result.result.external_links.map((l: ExternalLink) => l.externalLink), ).toEqual([ "https://example.com/chunk0", "https://example.com/chunk1", @@ -491,20 +492,20 @@ describe("SQLWarehouseConnector._transformDataArray", () => { ]); }); - test("is bounded by total_chunk_count when next_chunk_index never terminates", async () => { + test("is bounded by totalChunkCount when nextChunkIndex never terminates", async () => { const connector = createConnector(); // Misbehaving warehouse: always advertises another chunk. - const getStatementResultChunkN = vi.fn().mockResolvedValue({ - external_links: [ + const getResultData = vi.fn().mockResolvedValue({ + externalLinks: [ { - chunk_index: 1, - external_link: "https://example.com/loop", - next_chunk_index: 99, + chunkIndex: 1, + externalLink: "https://example.com/loop", + nextChunkIndex: 99, }, ], }); const workspaceClient = { - statementExecution: { getStatementResultChunkN }, + statementExecution: { getResultData }, }; const result = await transform( @@ -513,8 +514,8 @@ describe("SQLWarehouseConnector._transformDataArray", () => { workspaceClient, ); - // Terminates (no hang) — capped at total_chunk_count fetches. - expect(getStatementResultChunkN).toHaveBeenCalledTimes(2); + // Terminates (no hang) — capped at totalChunkCount fetches. + expect(getResultData).toHaveBeenCalledTimes(2); expect(result.result.external_links.length).toBeGreaterThan(0); }); }); diff --git a/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts b/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts index aba8488c3..a9061f61d 100644 --- a/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts +++ b/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts @@ -1,5 +1,5 @@ import type { Span } from "../../telemetry"; -import type { sql } from "../../workspace-client"; +import type { EndpointState } from "../../workspace-client"; import type { WarehouseStatusUpdate } from "./client"; /** @@ -10,7 +10,7 @@ import type { WarehouseStatusUpdate } from "./client"; */ export class WarehouseStatusEmitter { attempt = 0; - private lastEmittedState: sql.State | null = null; + private lastEmittedState: EndpointState | null = null; constructor( private readonly span: Span, @@ -18,7 +18,7 @@ export class WarehouseStatusEmitter { private readonly onStatus: (update: WarehouseStatusUpdate) => void, ) {} - emit(state: sql.State, summary: string | undefined): void { + emit(state: EndpointState, summary: string | undefined): void { this.attempt += 1; this.span.addEvent("warehouse.status", { "db.warehouse.state": state, diff --git a/packages/appkit/src/connectors/tests/sql-warehouse.test.ts b/packages/appkit/src/connectors/tests/sql-warehouse.test.ts index ce4b1b717..00c626878 100644 --- a/packages/appkit/src/connectors/tests/sql-warehouse.test.ts +++ b/packages/appkit/src/connectors/tests/sql-warehouse.test.ts @@ -61,7 +61,7 @@ describe("SQLWarehouseConnector", () => { await expect( connector.executeStatement(mockWorkspaceClient as any, { statement: sensitiveStatement, - warehouse_id: "test-warehouse", + warehouseId: "test-warehouse", }), ).rejects.toThrow(); @@ -89,7 +89,9 @@ describe("SQLWarehouseConnector", () => { statement_id: "stmt-123", status: { state: "RUNNING" }, }), - getStatement: vi.fn().mockRejectedValue(new Error("polling timeout")), + getStatementResult: vi + .fn() + .mockRejectedValue(new Error("polling timeout")), }, config: { host: "https://test.databricks.com" }, }; @@ -97,7 +99,7 @@ describe("SQLWarehouseConnector", () => { await expect( connector.executeStatement(mockWorkspaceClient as any, { statement: "SELECT secret_data FROM vault", - warehouse_id: "test-warehouse", + warehouseId: "test-warehouse", }), ).rejects.toThrow(); @@ -118,6 +120,141 @@ describe("SQLWarehouseConnector", () => { }); }); + describe("statement error-code propagation", () => { + let connector: SQLWarehouseConnector; + + beforeEach(() => { + vi.clearAllMocks(); + connector = new SQLWarehouseConnector({ timeout: 5000 }); + }); + + // Regression: the modular `@databricks/sdk-core` `ApiError` carries the + // Databricks error code on `.code`, whereas the legacy SDK used + // `.errorCode`. The analytics arrow disposition/format fallback keys on + // this code ("INVALID_PARAMETER_VALUE" / "NOT_IMPLEMENTED") to switch + // INLINE→EXTERNAL_LINKS, so the connector MUST surface either field as + // `ExecutionError.errorCode` — reading only `.errorCode` broke every arrow + // query (the INLINE+ARROW_STREAM probe rejection went unrecognized). + test("surfaces the modular SDK ApiError.code as ExecutionError.errorCode", async () => { + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + class FakeApiError extends Error { + readonly code = "INVALID_PARAMETER_VALUE"; + } + const mockWorkspaceClient = { + statementExecution: { + executeStatement: vi + .fn() + .mockRejectedValue( + new FakeApiError( + "Incompatible parameters: The format field must be JSON_ARRAY when the disposition field is INLINE.", + ), + ), + }, + config: { host: "https://test.databricks.com" }, + }; + + await expect( + connector.executeStatement(mockWorkspaceClient as any, { + statement: "SELECT 1", + warehouseId: "test-warehouse", + disposition: "INLINE", + format: "ARROW_STREAM", + }), + ).rejects.toMatchObject({ errorCode: "INVALID_PARAMETER_VALUE" }); + + errorSpy.mockRestore(); + }); + + // A failed statement STATUS (not a thrown ApiError) still carries the code + // on `status.error.errorCode` — the SDK unmarshals `error_code` there, so + // that path was already correct and must stay so. + test("surfaces status.error.errorCode from a FAILED statement status", async () => { + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + const mockWorkspaceClient = { + statementExecution: { + executeStatement: vi.fn().mockResolvedValue({ + statementId: "stmt-123", + status: { + state: "FAILED", + error: { + errorCode: "INVALID_PARAMETER_VALUE", + message: "bad parameter", + }, + }, + }), + }, + config: { host: "https://test.databricks.com" }, + }; + + await expect( + connector.executeStatement(mockWorkspaceClient as any, { + statement: "SELECT 1", + warehouseId: "test-warehouse", + }), + ).rejects.toMatchObject({ errorCode: "INVALID_PARAMETER_VALUE" }); + + errorSpy.mockRestore(); + }); + }); + + describe("bigint count normalization", () => { + let connector: SQLWarehouseConnector; + + beforeEach(() => { + vi.clearAllMocks(); + connector = new SQLWarehouseConnector({ timeout: 5000 }); + }); + + // Regression: the modular SDK types rowCount/byteCount/rowOffset (and the + // per-chunk BaseChunkInfo counts) as `bigint`, whereas the legacy SDK used + // `number`. Reyden's INLINE+ARROW_STREAM result is cached by the analytics + // arrow path, and `JSON.stringify` throws ("Do not know how to serialize a + // BigInt") on any surviving bigint — which broke EVERY query on Reyden. The + // connector must coerce these to `number` at the SDK boundary so the result + // stays serializable for the cache / SSE frames. + test("coerces bigint manifest/result/chunk counts so the result is JSON-serializable", async () => { + const mockWorkspaceClient = { + statementExecution: { + executeStatement: vi.fn().mockResolvedValue({ + statementId: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { + format: "JSON_ARRAY", + totalRowCount: 2n, + totalByteCount: 100n, + chunks: [ + { chunkIndex: 0, rowOffset: 0n, rowCount: 2n, byteCount: 100n }, + ], + schema: { columns: [{ name: "id", typeName: "INT" }] }, + }, + result: { + dataArray: [["1"], ["2"]], + rowOffset: 0n, + rowCount: 2n, + byteCount: 100n, + }, + }), + }, + config: { host: "https://test.databricks.com" }, + }; + + const out: any = await connector.executeStatement( + mockWorkspaceClient as any, + { statement: "SELECT id FROM t", warehouseId: "reyden" }, + ); + + // The arrow cache serializes exactly this — it must not throw. + expect(() => JSON.stringify(out)).not.toThrow(); + // Counts are coerced to number (legacy parity), including per-chunk ones. + expect(typeof out.manifest.totalRowCount).toBe("number"); + expect(typeof out.manifest.totalByteCount).toBe("number"); + expect(typeof out.manifest.chunks[0].byteCount).toBe("number"); + expect(typeof out.result.rowCount).toBe("number"); + }); + }); + describe("ensureWarehouseRunning", () => { let connector: SQLWarehouseConnector; @@ -137,7 +274,9 @@ describe("SQLWarehouseConnector", () => { test("emits a single RUNNING update and returns when warehouse is already running", async () => { const get = vi.fn().mockResolvedValue({ state: "RUNNING" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; await connector.ensureWarehouseRunning(wsClient as any, "wh-1", { @@ -158,7 +297,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); const start = vi.fn().mockResolvedValue(undefined); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; const promise = connector.ensureWarehouseRunning( @@ -191,7 +332,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; const promise = connector.ensureWarehouseRunning( @@ -212,7 +355,9 @@ describe("SQLWarehouseConnector", () => { test("rejects when warehouse is DELETED", async () => { const get = vi.fn().mockResolvedValue({ state: "DELETED" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; await expect( @@ -228,7 +373,9 @@ describe("SQLWarehouseConnector", () => { test("rejects when warehouse is DELETING", async () => { const get = vi.fn().mockResolvedValue({ state: "DELETING" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; await expect( @@ -243,7 +390,9 @@ describe("SQLWarehouseConnector", () => { test("aborts immediately when signal is already aborted", async () => { const get = vi.fn(); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const controller = new AbortController(); controller.abort(); @@ -258,7 +407,9 @@ describe("SQLWarehouseConnector", () => { test("times out if warehouse never reaches RUNNING", async () => { const get = vi.fn().mockResolvedValue({ state: "STARTING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const promise = connector.ensureWarehouseRunning( wsClient as any, @@ -281,7 +432,7 @@ describe("SQLWarehouseConnector", () => { test("rejects when warehouse_id is empty", async () => { const wsClient = { - warehouses: { get: vi.fn(), start: vi.fn() }, + warehouses: { getWarehouse: vi.fn(), startWarehouse: vi.fn() }, }; await expect( @@ -293,7 +444,9 @@ describe("SQLWarehouseConnector", () => { test("skips the SDK round-trip on a subsequent call within the recently-running TTL", async () => { const get = vi.fn().mockResolvedValue({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const updates1: any[] = []; await connector.ensureWarehouseRunning(wsClient as any, "wh-cache", { @@ -314,7 +467,9 @@ describe("SQLWarehouseConnector", () => { test("rejects with ConfigurationError when STOPPED and autoStart is false", async () => { const get = vi.fn().mockResolvedValue({ state: "STOPPED" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; await expect( connector.ensureWarehouseRunning(wsClient as any, "wh-no-auto", { @@ -332,7 +487,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const updates: any[] = []; const promise = connector.ensureWarehouseRunning( @@ -361,7 +518,9 @@ describe("SQLWarehouseConnector", () => { .fn() .mockRejectedValue(new Error("permission denied reading warehouse")); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; await expect( @@ -385,7 +544,9 @@ describe("SQLWarehouseConnector", () => { err.name = "AbortError"; return Promise.reject(err); }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; await expect( connector.ensureWarehouseRunning(wsClient as any, "wh-probe-abort", { @@ -398,12 +559,14 @@ describe("SQLWarehouseConnector", () => { test("does not leak raw SDK error text in the rethrown error", async () => { const sensitive = "getaddrinfo ENOTFOUND adb-1234567890.10.azuredatabricks.net"; - // The status probe (get) no longer blocks the query, so exercise a path - // that still throws — an auto-start (`start`) failure — to verify raw SDK + // The status probe (getWarehouse) no longer blocks the query, so exercise a path + // that still throws — an auto-start (`startWarehouse`) failure — to verify raw SDK // text is sanitized out of the rethrown readiness error. const get = vi.fn().mockResolvedValue({ state: "STOPPED" }); const start = vi.fn().mockRejectedValue(new Error(sensitive)); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; await expect( connector.ensureWarehouseRunning(wsClient as any, "wh-leak", { @@ -428,7 +591,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); const start = vi.fn().mockResolvedValue(undefined); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const allUpdates = [0, 1, 2].map(() => [] as { state: string }[]); const waits = allUpdates.map((updates) => @@ -453,7 +618,9 @@ describe("SQLWarehouseConnector", () => { .fn() .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const controller = new AbortController(); const aborted = connector.ensureWarehouseRunning( @@ -485,7 +652,9 @@ describe("SQLWarehouseConnector", () => { .fn() .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const mount1 = new AbortController(); const first = connector.ensureWarehouseRunning( @@ -514,7 +683,9 @@ describe("SQLWarehouseConnector", () => { test("orphan before warehouses.start is aborted on the next microtask", async () => { const get = vi.fn().mockResolvedValue({ state: "STARTING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const controller = new AbortController(); const only = connector.ensureWarehouseRunning( @@ -532,7 +703,7 @@ describe("SQLWarehouseConnector", () => { await Promise.resolve(); expect(get).toHaveBeenCalledTimes(1); - expect(wsClient.warehouses.start).not.toHaveBeenCalled(); + expect(wsClient.warehouses.startWarehouse).not.toHaveBeenCalled(); }); test("orphan after warehouses.start runs poll to completion", async () => { @@ -542,7 +713,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); const start = vi.fn().mockResolvedValue(undefined); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const controller = new AbortController(); const only = connector.ensureWarehouseRunning( @@ -579,7 +752,9 @@ describe("SQLWarehouseConnector", () => { .fn() .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; let callCount = 0; const promise = connector.ensureWarehouseRunning( diff --git a/packages/appkit/src/evals/dataset.ts b/packages/appkit/src/evals/dataset.ts index b6cdc0dac..f49270fcc 100644 --- a/packages/appkit/src/evals/dataset.ts +++ b/packages/appkit/src/evals/dataset.ts @@ -83,7 +83,7 @@ export async function readEvalDataset( : ""; const connector = new SQLWarehouseConnector({}); const response = await connector.executeStatement(client, { - warehouse_id: options.warehouseId, + warehouseId: options.warehouseId, statement: `SELECT inputs, expectations FROM ${options.table}${limit}`, }); diff --git a/packages/appkit/src/evals/tests/dataset.test.ts b/packages/appkit/src/evals/tests/dataset.test.ts index 12214d1ac..76d113b4d 100644 --- a/packages/appkit/src/evals/tests/dataset.test.ts +++ b/packages/appkit/src/evals/tests/dataset.test.ts @@ -45,7 +45,7 @@ describe("readEvalDataset", () => { // SELECT targets the table; no LIMIT when unset. const [, input] = executeStatement.mock.calls[0]; - expect(input.warehouse_id).toBe("wh1"); + expect(input.warehouseId).toBe("wh1"); expect(input.statement).toBe( "SELECT inputs, expectations FROM main.default.eval_ds", ); diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 4e8bdf077..511fd643e 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1070,7 +1070,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { workspaceClient, { statement, - warehouse_id: warehouseId, + warehouseId, parameters: sqlParameters, ...formatParameters, }, diff --git a/packages/appkit/src/plugins/analytics/query.ts b/packages/appkit/src/plugins/analytics/query.ts index bcd77a817..4b57b051b 100644 --- a/packages/appkit/src/plugins/analytics/query.ts +++ b/packages/appkit/src/plugins/analytics/query.ts @@ -4,7 +4,7 @@ import { isSQLTypeMarker, type SQLTypeMarker, sql as sqlHelpers } from "shared"; import { getWorkspaceId } from "../../context"; import { ValidationError } from "../../errors"; -import type { sql } from "../../workspace-client"; +import type { StatementParameter } from "../../workspace-client"; type SQLParameterValue = SQLTypeMarker | null | undefined; @@ -37,8 +37,8 @@ export class QueryProcessor { convertToSQLParameters( query: string, parameters?: Record, - ): { statement: string; parameters: sql.StatementParameterListItem[] } { - const sqlParameters: sql.StatementParameterListItem[] = []; + ): { statement: string; parameters: StatementParameter[] } { + const sqlParameters: StatementParameter[] = []; if (parameters) { // extract all params from the query @@ -72,7 +72,7 @@ export class QueryProcessor { private _createParameter( key: string, value: SQLParameterValue, - ): sql.StatementParameterListItem | null { + ): StatementParameter | null { if (value === null || value === undefined) { return null; } diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts index a0435513c..3f40439d2 100644 --- a/packages/appkit/src/plugins/analytics/result-delivery.ts +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -4,7 +4,7 @@ import type { SQLTypeMarker } from "shared"; import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import type { RefreshChunkLink } from "../../stream/arrow-stream-processor"; -import type { sql } from "../../workspace-client"; +import type { ExternalLink } from "../../workspace-client"; /** * Centralized disposition/format fallback for analytics result delivery. @@ -39,7 +39,7 @@ export interface QueryExecutor { | { attachment?: string; data?: Record[]; - external_links?: sql.ExternalLink[]; + external_links?: ExternalLink[]; columnNames?: string[]; statement_id?: string; status?: unknown; @@ -52,7 +52,7 @@ export interface QueryExecutor { /** Streams already-resolved EXTERNAL_LINKS chunks; the connector provides it. */ export interface ArrowChunkStreamer { streamExternalLinks( - chunks: sql.ExternalLink[], + chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, ): AsyncGenerator; diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts index c69a82e83..0265feeea 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts @@ -26,7 +26,7 @@ describe("Analytics Plugin Integration", () => { let app: TestApp<[ReturnType]>; /** The SQL mock the analytics route drives, via the harness's client. */ let executeStatement: ReturnType; - let getStatement: ReturnType; + let getStatementResult: ReturnType; beforeAll(async () => { // The harness owns the env setup, the singleton resets, the mock client, the @@ -36,7 +36,10 @@ describe("Analytics Plugin Integration", () => { app.client, "statementExecution.executeStatement", ); - getStatement = getMock(app.client, "statementExecution.getStatement"); + getStatementResult = getMock( + app.client, + "statementExecution.getStatementResult", + ); }); afterAll(async () => { @@ -48,7 +51,7 @@ describe("Analytics Plugin Integration", () => { // Reset drops the built-in canned SUCCEEDED default too, matching the // "script it yourself" semantics this suite relied on before. executeStatement.mockReset(); - getStatement.mockReset(); + getStatementResult.mockReset(); getAppQuerySpy.mockReset(); }); @@ -60,8 +63,8 @@ describe("Analytics Plugin Integration", () => { ["Bob", "25"], ]; const mockColumns = [ - { name: "name", type_name: "STRING" }, - { name: "age", type_name: "STRING" }, + { name: "name", typeName: "STRING" }, + { name: "age", typeName: "STRING" }, ]; getAppQuerySpy.mockResolvedValueOnce({ @@ -93,7 +96,7 @@ describe("Analytics Plugin Integration", () => { expect(executeStatement).toHaveBeenCalledWith( expect.objectContaining({ statement: testQuery, - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.anything(), ); diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index a25e18daa..3104d710b 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -145,7 +145,7 @@ describe("Analytics Plugin", () => { expect.anything(), expect.objectContaining({ statement: "SELECT * FROM test", - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.any(AbortSignal), ); @@ -222,7 +222,7 @@ describe("Analytics Plugin", () => { expect.anything(), expect.objectContaining({ statement: "SELECT * FROM users WHERE id = :user_id", - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.any(AbortSignal), ); @@ -619,7 +619,7 @@ describe("Analytics Plugin", () => { expect.objectContaining({ statement: "SELECT * FROM test", parameters: [], - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.any(AbortSignal), ); @@ -654,7 +654,7 @@ describe("Analytics Plugin", () => { expect.anything(), expect.objectContaining({ statement: "SELECT * FROM test", - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", disposition: "INLINE", format: "ARROW_STREAM", }), @@ -1667,7 +1667,7 @@ describe("Analytics Plugin", () => { result: { data: [] }, }), }, - warehouses: { get: warehouseGet, start: vi.fn() }, + warehouses: { getWarehouse: warehouseGet, startWarehouse: vi.fn() }, }, }); const mockReq = createMockRequest({ @@ -1724,7 +1724,7 @@ describe("Analytics Plugin", () => { serviceContextMock = await mockServiceContext({ serviceDatabricksClient: { statementExecution: { executeStatement: vi.fn() }, - warehouses: { get: warehouseGet, start: vi.fn() }, + warehouses: { getWarehouse: warehouseGet, startWarehouse: vi.fn() }, }, }); const mockReq = createMockRequest({ @@ -1769,7 +1769,7 @@ describe("Analytics Plugin", () => { serviceContextMock = await mockServiceContext({ serviceDatabricksClient: { statementExecution: { executeStatement: vi.fn() }, - warehouses: { get: warehouseGet, start: vi.fn() }, + warehouses: { getWarehouse: warehouseGet, startWarehouse: vi.fn() }, }, }); const mockReq = createMockRequest({ diff --git a/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts b/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts index 934e83c28..404719440 100644 --- a/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts @@ -38,9 +38,9 @@ describe.runIf(!!warehouseId)("arrow delivery (live warehouse)", () => { client, { statement, - warehouse_id: warehouseId as string, - wait_timeout: "50s", - on_wait_timeout: "CONTINUE", + warehouseId: warehouseId as string, + waitTimeout: "50s", + onWaitTimeout: "CONTINUE", disposition: fp.disposition as never, format: fp.format as never, }, diff --git a/packages/appkit/src/plugins/analytics/tests/metric.test.ts b/packages/appkit/src/plugins/analytics/tests/metric.test.ts index abb8952cd..07b83e449 100644 --- a/packages/appkit/src/plugins/analytics/tests/metric.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/metric.test.ts @@ -700,7 +700,7 @@ describe("analytics metric route", () => { expect.objectContaining({ statement: "SELECT MEASURE(`arr`) AS `arr` FROM `cat`.`sch`.`revenue_metrics`", - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.any(AbortSignal), ); @@ -748,7 +748,7 @@ describe("analytics metric route", () => { result: { data: [] }, }), }, - warehouses: { get: warehouseGet, start: vi.fn() }, + warehouses: { getWarehouse: warehouseGet, startWarehouse: vi.fn() }, }, }); const mockReq = createMockRequest({ diff --git a/packages/appkit/src/stream/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index 62cab4df4..e063b6abb 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -1,11 +1,9 @@ import { ExecutionError, ValidationError } from "../errors"; import { createLogger } from "../logging/logger"; -import type { sql } from "../workspace-client"; +import type { ExternalLink } from "../workspace-client"; const logger = createLogger("stream:arrow"); -type ExternalLink = sql.ExternalLink; - /** * Re-mint a chunk's pre-signed URL. DBSQL external links expire in <= 15 min, * so a large result whose tail chunks are reached after the earlier chunks @@ -83,11 +81,11 @@ export class ArrowStreamProcessor { signal?: AbortSignal, refresh?: RefreshChunkLink, ): AsyncGenerator { - let externalLink = chunk.external_link; + let externalLink = chunk.externalLink; if (!externalLink) { // A missing link cannot be fixed by retrying — fail loudly. throw ExecutionError.statementFailed( - `External link missing for chunk ${chunk.chunk_index}`, + `External link missing for chunk ${chunk.chunkIndex}`, ); } @@ -114,7 +112,7 @@ export class ArrowStreamProcessor { clearTimeout(timer); if (!r.ok) { throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunk_index}: ${r.status} ${r.statusText}`, + `Failed to download chunk ${chunk.chunkIndex}: ${r.status} ${r.statusText}`, ); } // Keep this attempt's controller alive to drive the body read + idle @@ -134,16 +132,16 @@ export class ArrowStreamProcessor { // chunk's link — a stale URL would just 403 again on the same address. // Only meaningful before any bytes are yielded (below), which is why // this lives in the establish-response loop. - if (refresh && chunk.chunk_index != null) { + if (refresh && chunk.chunkIndex != null) { try { - const fresh = await refresh(chunk.chunk_index, signal); - if (fresh?.external_link) externalLink = fresh.external_link; + const fresh = await refresh(chunk.chunkIndex, signal); + if (fresh?.externalLink) externalLink = fresh.externalLink; } catch (refreshError) { // Keep retrying the current URL; surface the original error if // all attempts fail. logger.warn( "Failed to re-resolve link for chunk %s: %O", - chunk.chunk_index, + chunk.chunkIndex, refreshError, ); } @@ -154,7 +152,7 @@ export class ArrowStreamProcessor { if (!response || !controller) { throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunk_index} after ${this.options.retries} attempts: ${ + `Failed to download chunk ${chunk.chunkIndex} after ${this.options.retries} attempts: ${ lastError instanceof Error ? lastError.message : String(lastError) }`, ); @@ -194,13 +192,13 @@ export class ArrowStreamProcessor { if (signal?.aborted) throw ExecutionError.canceled(); logger.error( "Failed streaming chunk %s body: %O", - chunk.chunk_index, + chunk.chunkIndex, error, ); throw error instanceof ExecutionError ? error : ExecutionError.statementFailed( - `Failed streaming chunk ${chunk.chunk_index}: ${ + `Failed streaming chunk ${chunk.chunkIndex}: ${ error instanceof Error ? error.message : String(error) }`, ); diff --git a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts index 555f87339..84d1ee031 100644 --- a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts +++ b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts @@ -1,6 +1,5 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; -import type { sql } from "../../workspace-client"; import { ArrowStreamProcessor } from "../arrow-stream-processor"; /** A ReadableStream that emits the given pieces in order, then closes. */ @@ -15,10 +14,10 @@ function streamOf(...pieces: Uint8Array[]): ReadableStream { function mockChunks(count: number) { return Array.from({ length: count }, (_, i) => ({ - chunk_index: i, - external_link: `https://example.com/chunk-${i}`, - row_offset: i * 100, - row_count: 100, + chunkIndex: i, + externalLink: `https://example.com/chunk-${i}`, + rowOffset: BigInt(i * 100), + rowCount: 100n, })); } @@ -166,7 +165,7 @@ describe("ArrowStreamProcessor.streamChunks", () => { }); test("throws immediately when a chunk has no external_link", async () => { - const chunks = [{ chunk_index: 0 }] as any; + const chunks = [{ chunkIndex: 0 }] as any; await expect(drain(processor.streamChunks(chunks))).rejects.toThrow( /External link missing/, ); @@ -218,8 +217,8 @@ describe("ArrowStreamProcessor.streamChunks", () => { globalThis.fetch = fetchMock; const refresh = vi.fn(async (chunkIndex: number) => ({ - chunk_index: chunkIndex, - external_link: "https://example.com/fresh-link", + chunkIndex, + externalLink: "https://example.com/fresh-link", })); const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); diff --git a/packages/appkit/src/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index 5570c90f7..6f5c53075 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -640,19 +640,19 @@ export async function runWithRequestContext( */ export function createSuccessfulSQLResponse( data: Any[][], - columns: Array<{ name: string; type_name?: string }>, + columns: Array<{ name: string; typeName?: string }>, ) { return { status: { state: "SUCCEEDED" }, - statement_id: `stmt-${Date.now()}`, + statementId: `stmt-${Date.now()}`, result: { - data_array: data, + dataArray: data, }, manifest: { schema: { columns: columns.map((col) => ({ name: col.name, - type_name: col.type_name ?? "STRING", + typeName: col.typeName ?? "STRING", })), }, }, @@ -668,7 +668,7 @@ export function createFailedSQLResponse(errorMessage: string) { message: errorMessage, }, }, - statement_id: `stmt-${Date.now()}`, + statementId: `stmt-${Date.now()}`, }; } diff --git a/packages/appkit/src/type-generator/statement-result.ts b/packages/appkit/src/type-generator/statement-result.ts index 7ae091aaf..24620988b 100644 --- a/packages/appkit/src/type-generator/statement-result.ts +++ b/packages/appkit/src/type-generator/statement-result.ts @@ -1,5 +1,5 @@ import { createLogger } from "../logging/logger"; -import type { WorkspaceClient } from "../workspace-client"; +import type { StatementResponse, WorkspaceClient } from "../workspace-client"; import { getErrorMessage } from "./errors"; import type { DatabricksStatementExecutionResponse } from "./types"; @@ -147,6 +147,42 @@ function isFormatRejection( ); } +/** + * Adapt the modular SDK's camelCase {@link StatementResponse} onto the + * type-generator's own snake_case {@link DatabricksStatementExecutionResponse} + * — the shape every downstream DESCRIBE parser (and every mocked test) reads. + * Keeping the boundary here means only this mapper touches the SDK shape; + * {@link normalizeResultRows} and the parsers stay unchanged. `attachment` + * survives thanks to the pinned pnpm patch on `@databricks/sdk-statementexecution`. + */ +function toDescribeResponse( + r: StatementResponse, +): DatabricksStatementExecutionResponse { + return { + statement_id: r.statementId ?? "", + status: { + state: r.status?.state ?? "", + error: r.status?.error + ? { + error_code: r.status.error.errorCode, + message: r.status.error.message, + } + : undefined, + }, + manifest: r.manifest ? { format: r.manifest.format } : undefined, + result: r.result + ? { + // DESCRIBE rows are always string/null cells. Local key stays + // snake_case (`data_array`); value is the SDK's camelCase `dataArray`. + data_array: r.result.dataArray as (string | null)[][] | undefined, + attachment: r.result.attachment, + next_chunk_index: r.result.nextChunkIndex, + next_chunk_internal_link: r.result.nextChunkInternalLink, + } + : undefined, + }; +} + /** * Run a DESCRIBE and return a response whose rows are readable via * `result.data_array`, adapting to the warehouse's result-format capability. @@ -175,15 +211,17 @@ export async function describeAdaptive( let lastError: unknown; for (const format of formats) { try { - const response = (await client.statementExecution.executeStatement({ - statement, - warehouse_id: warehouseId, - // Synchronous wait: without it the call can return PENDING/RUNNING with - // no rows, which downstream misreads as a no-result degrade. - wait_timeout: "30s", - format, - disposition: "INLINE", - })) as DatabricksStatementExecutionResponse; + const response = toDescribeResponse( + await client.statementExecution.executeStatement({ + statement, + warehouseId, + // Synchronous wait: without it the call can return PENDING/RUNNING with + // no rows, which downstream misreads as a no-result degrade. + waitTimeout: "30s", + format, + disposition: "INLINE", + }), + ); const normalized = await normalizeResultRows(response); if ( normalized.status?.state === "FAILED" && diff --git a/packages/appkit/src/type-generator/tests/generate-queries.test.ts b/packages/appkit/src/type-generator/tests/generate-queries.test.ts index 4e9b69bf0..9ad1b1dce 100644 --- a/packages/appkit/src/type-generator/tests/generate-queries.test.ts +++ b/packages/appkit/src/type-generator/tests/generate-queries.test.ts @@ -34,7 +34,10 @@ vi.mock("../../workspace-client", async (importOriginal) => { mocks.createWorkspaceClient(...args); return { statementExecution: { executeStatement: mocks.executeStatement }, - warehouses: { get: mocks.getWarehouse, start: mocks.startWarehouse }, + warehouses: { + getWarehouse: mocks.getWarehouse, + startWarehouse: mocks.startWarehouse, + }, }; }, }; @@ -87,9 +90,9 @@ const lastSavedQueries = () => function succeededResult(columns: [string, string, string | null][]) { return { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, - result: { data_array: columns }, + result: { dataArray: columns }, }; } @@ -113,7 +116,7 @@ async function succeededArrowAttachmentResult( "base64", ); return { - statement_id: "stmt-arrow", + statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, // No data_array — rows live in the attachment, like a real INLINE Arrow @@ -227,8 +230,8 @@ describe("generateQueriesFromDescribe", () => { expect(mocks.executeStatement).toHaveBeenCalledTimes(1); expect(mocks.executeStatement.mock.calls[0][0]).toMatchObject({ - warehouse_id: "wh-123", - wait_timeout: "30s", + warehouseId: "wh-123", + waitTimeout: "30s", format: "JSON_ARRAY", disposition: "INLINE", }); @@ -238,7 +241,7 @@ describe("generateQueriesFromDescribe", () => { mocks.readdir.mockResolvedValue(["bad_table.sql"]); mocks.readFile.mockResolvedValue("SELECT * FROM bad_table"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-2", + statementId: "stmt-2", status: { state: "FAILED", error: { message: "Table or view not found: bad_table" }, @@ -258,7 +261,7 @@ describe("generateQueriesFromDescribe", () => { mocks.readdir.mockResolvedValue(["query.sql"]); mocks.readFile.mockResolvedValue("SELECT 1"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-3", + statementId: "stmt-3", status: { state: "FAILED" }, }); @@ -280,7 +283,7 @@ describe("generateQueriesFromDescribe", () => { mocks.executeStatement .mockResolvedValueOnce(succeededResult([["id", "INT", null]])) .mockResolvedValueOnce({ - statement_id: "stmt-fail", + statementId: "stmt-fail", status: { state: "FAILED", error: { message: "Table not found" }, @@ -312,7 +315,7 @@ describe("generateQueriesFromDescribe", () => { mocks.executeStatement .mockRejectedValueOnce(new Error("Connection refused")) .mockResolvedValueOnce({ - statement_id: "stmt-fail-2", + statementId: "stmt-fail-2", status: { state: "FAILED", error: { message: "Table not found" } }, }); @@ -492,7 +495,7 @@ describe("generateQueriesFromDescribe", () => { .mockResolvedValueOnce("SELECT * FROM whatever"); mocks.executeStatement .mockResolvedValueOnce({ - statement_id: "stmt-syntax", + statementId: "stmt-syntax", status: { state: "FAILED", error: { message: "Table not found" }, @@ -649,7 +652,7 @@ describe("generateQueriesFromDescribe", () => { // state with no result rows. Must degrade like a transient outage, not be // misreported as EMPTY (which would discard a good cached type). mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "PENDING" }, }); @@ -684,7 +687,7 @@ describe("generateQueriesFromDescribe", () => { mocks.executeStatement .mockResolvedValueOnce(succeededResult([["id", "INT", null]])) .mockResolvedValueOnce({ - statement_id: "stmt-pending", + statementId: "stmt-pending", status: { state: "RUNNING" }, }); @@ -711,7 +714,7 @@ describe("generateQueriesFromDescribe", () => { mocks.readdir.mockResolvedValue(["broken.sql"]); mocks.readFile.mockResolvedValue("SELECT * FROM missing"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { message: "Table or view not found: missing" }, diff --git a/packages/appkit/src/type-generator/tests/index.test.ts b/packages/appkit/src/type-generator/tests/index.test.ts index 21b5a7e0c..6fa33ed48 100644 --- a/packages/appkit/src/type-generator/tests/index.test.ts +++ b/packages/appkit/src/type-generator/tests/index.test.ts @@ -11,8 +11,37 @@ import { vi, } from "vitest"; +import type { StatementResponse } from "../../workspace-client"; import type { DatabricksStatementExecutionResponse } from "../types"; +/** + * Adapt a local snake_case describe fixture to the modular SDK's camelCase + * `StatementResponse` — the shape the mocked `executeStatement` now returns. + * `describeAdaptive` maps it back to the local shape via `toDescribeResponse`, + * so fixtures stay authored in the type-generator's own domain shape. + */ +function asSdkResponse( + r: DatabricksStatementExecutionResponse, +): StatementResponse { + return { + statementId: r.statement_id, + status: r.status && { + state: r.status.state, + error: r.status.error && { + errorCode: r.status.error.error_code, + message: r.status.error.message, + }, + }, + manifest: r.manifest && { format: r.manifest.format }, + result: r.result && { + dataArray: r.result.data_array, + attachment: r.result.attachment, + nextChunkIndex: r.result.next_chunk_index, + nextChunkInternalLink: r.result.next_chunk_internal_link, + }, + } as unknown as StatementResponse; +} + const mocks = vi.hoisted(() => ({ generateQueriesFromDescribe: vi.fn(), getWarehouseState: vi.fn(), @@ -553,7 +582,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("non-blocking + RUNNING warehouse: DESCRIBEs run and land full schemas", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ @@ -568,7 +597,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { expect(mocks.executeStatement).toHaveBeenCalledWith( expect.objectContaining({ statement: "DESCRIBE TABLE EXTENDED `demo`.`sales`.`revenue` AS JSON", - warehouse_id: "wh-1", + warehouseId: "wh-1", }), ); const declarations = fs.readFileSync(metricFile, "utf-8"); @@ -582,7 +611,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("blocking + RUNNING: one preflight probe, no start/wait, DESCRIBEs run", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ @@ -756,7 +785,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { mocks.getWarehouseState.mockResolvedValue("STOPPED"); mocks.startWarehouse.mockResolvedValue(undefined); mocks.waitUntilRunning.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ @@ -1015,7 +1044,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("non-blocking + RUNNING with the default fetcher: probe and DESCRIBEs share exactly one client", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ @@ -1154,24 +1183,23 @@ describe("generateFromEntryPoint — metric cache section", () => { const outFile = path.join(cacheTestDir, "generated", "analytics.d.ts"); const metricFile = path.join(cacheTestDir, "generated", "metric-views.d.ts"); - const describeResponseFor = ( - measure: string, - ): DatabricksStatementExecutionResponse => ({ - statement_id: "stmt-mock", - status: { state: "SUCCEEDED" }, - result: { - data_array: [ - [ - JSON.stringify({ - columns: [ - { name: measure, type: "DECIMAL(38,2)", is_measure: true }, - { name: "region", type: "STRING", is_measure: false }, - ], - }), + const describeResponseFor = (measure: string): StatementResponse => + asSdkResponse({ + statement_id: "stmt-mock", + status: { state: "SUCCEEDED" }, + result: { + data_array: [ + [ + JSON.stringify({ + columns: [ + { name: measure, type: "DECIMAL(38,2)", is_measure: true }, + { name: "region", type: "STRING", is_measure: false }, + ], + }), + ], ], - ], - }, - }); + }, + }); const writeConfig = ( metricViews: Record< @@ -2081,7 +2109,7 @@ describe("generateFromEntryPoint — anti-clobber for blocking mode", () => { }; mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ diff --git a/packages/appkit/src/type-generator/tests/mv-registry.test.ts b/packages/appkit/src/type-generator/tests/mv-registry.test.ts index 90b07185c..e6d90bb75 100644 --- a/packages/appkit/src/type-generator/tests/mv-registry.test.ts +++ b/packages/appkit/src/type-generator/tests/mv-registry.test.ts @@ -10,6 +10,7 @@ import { afterEach, beforeEach, describe, expect, test } from "vitest"; // imports it from there. import { quoteFqnForSql } from "../../../../shared/src/schemas/metric-fqn"; import { metricSourceSchema } from "../../../../shared/src/schemas/metric-source"; +import type { StatementResponse } from "../../workspace-client"; import { readMetricConfig, resolveMetricConfig } from "../mv-registry/config"; import { createWorkspaceDescribeFetcher, @@ -53,6 +54,35 @@ function mockDescribeResponse( }; } +/** + * Adapt a local snake_case describe fixture to the modular SDK's camelCase + * `StatementResponse` — the shape a mocked `executeStatement` (consumed by the + * real `createWorkspaceDescribeFetcher` → `describeAdaptive`) now returns. + * Direct `syncMetrics(resolution, fetcher)` fixtures stay in the local snake + * shape (they bypass the SDK), so only the executeStatement mocks wrap with this. + */ +function asSdkResponse( + r: DatabricksStatementExecutionResponse, +): StatementResponse { + return { + statementId: r.statement_id, + status: r.status && { + state: r.status.state, + error: r.status.error && { + errorCode: r.status.error.error_code, + message: r.status.error.message, + }, + }, + manifest: r.manifest && { format: r.manifest.format }, + result: r.result && { + dataArray: r.result.data_array, + attachment: r.result.attachment, + nextChunkIndex: r.result.next_chunk_index, + nextChunkInternalLink: r.result.next_chunk_internal_link, + }, + } as unknown as StatementResponse; +} + /** * Real Arrow IPC attachment captured live from dogfood: * DESCRIBE TABLE EXTENDED `appkit_demo`.`public`.`revenue_metrics` AS JSON @@ -326,9 +356,11 @@ describe("resolveMetricConfig — FQN naming (UC-accurate)", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return mockDescribeResponse({ - columns: [{ name: "arr", type: "DECIMAL", is_measure: true }], - }); + return asSdkResponse( + mockDescribeResponse({ + columns: [{ name: "arr", type: "DECIMAL", is_measure: true }], + }), + ); }, }, } as unknown as Parameters[0]; @@ -664,7 +696,7 @@ describe("createWorkspaceDescribeFetcher", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return mockDescribeResponse(payload); + return asSdkResponse(mockDescribeResponse(payload)); }, }, } as unknown as Parameters[0]; @@ -680,8 +712,8 @@ describe("createWorkspaceDescribeFetcher", () => { expect(statements).toHaveLength(1); expect(statements[0]).toMatchObject({ statement: "DESCRIBE TABLE EXTENDED `demo`.`sales`.`revenue` AS JSON", - warehouse_id: "wh-1", - wait_timeout: "30s", + warehouseId: "wh-1", + waitTimeout: "30s", // describeAdaptive tries JSON_ARRAY first (standard DBSQL); it falls back // to ARROW_STREAM only if the warehouse rejects that format. format: "JSON_ARRAY", @@ -700,13 +732,13 @@ describe("createWorkspaceDescribeFetcher", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return { + return asSdkResponse({ statement_id: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, // Only an attachment — no data_array (the bug's trigger condition). result: { attachment: ARROW_ATTACHMENT_B64 }, - } as DatabricksStatementExecutionResponse; + }); }, }, } as unknown as Parameters[0]; @@ -1016,7 +1048,7 @@ describe("syncMetrics", () => { const client = { statementExecution: { executeStatement: async () => - ({ + asSdkResponse({ statement_id: "stmt-chunked", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, @@ -1024,7 +1056,7 @@ describe("syncMetrics", () => { attachment: ARROW_ATTACHMENT_B64, next_chunk_index: 1, }, - }) as DatabricksStatementExecutionResponse, + }), }, } as unknown as Parameters[0]; const fetcher = createWorkspaceDescribeFetcher(client, "wh-1"); diff --git a/packages/appkit/src/type-generator/tests/statement-result.test.ts b/packages/appkit/src/type-generator/tests/statement-result.test.ts index 4221cd705..d545e49ce 100644 --- a/packages/appkit/src/type-generator/tests/statement-result.test.ts +++ b/packages/appkit/src/type-generator/tests/statement-result.test.ts @@ -3,7 +3,10 @@ import path from "node:path"; import { describe, expect, test } from "vitest"; -import type { WorkspaceClient } from "../../workspace-client"; +import type { + StatementResponse, + WorkspaceClient, +} from "../../workspace-client"; import { type DescribeFormatMemo, describeAdaptive, @@ -270,6 +273,31 @@ describe("describeAdaptive", () => { | DatabricksStatementExecutionResponse | Promise; + // Adapt a local snake_case fixture to the modular SDK's camelCase + // StatementResponse — the shape executeStatement now returns; describeAdaptive + // maps it back to the local shape via toDescribeResponse. + function asSdkResponse( + r: DatabricksStatementExecutionResponse, + ): StatementResponse { + return { + statementId: r.statement_id, + status: r.status && { + state: r.status.state, + error: r.status.error && { + errorCode: r.status.error.error_code, + message: r.status.error.message, + }, + }, + manifest: r.manifest && { format: r.manifest.format }, + result: r.result && { + dataArray: r.result.data_array, + attachment: r.result.attachment, + nextChunkIndex: r.result.next_chunk_index, + nextChunkInternalLink: r.result.next_chunk_internal_link, + }, + } as unknown as StatementResponse; + } + // Minimal WorkspaceClient stub: records the formats requested and delegates // each executeStatement to behavior(format), which may resolve or throw. function stubClient(behavior: StubBehavior) { @@ -278,7 +306,7 @@ describe("describeAdaptive", () => { statementExecution: { executeStatement: async (req: { format: string }) => { formats.push(req.format); - return behavior(req.format); + return asSdkResponse(await behavior(req.format)); }, }, } as unknown as WorkspaceClient; diff --git a/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts b/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts index 6134f8348..ef0b81519 100644 --- a/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts +++ b/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts @@ -21,7 +21,7 @@ vi.mock("../../workspace-client", async (importOriginal) => { ...actual, createWorkspaceClient: () => ({ statementExecution: { executeStatement: mocks.executeStatement }, - warehouses: { get: mocks.getWarehouse, start: vi.fn() }, + warehouses: { getWarehouse: mocks.getWarehouse, startWarehouse: vi.fn() }, }), }; }); diff --git a/packages/appkit/src/type-generator/tests/warehouse-status.test.ts b/packages/appkit/src/type-generator/tests/warehouse-status.test.ts index 882188e34..4a1b0f7b3 100644 --- a/packages/appkit/src/type-generator/tests/warehouse-status.test.ts +++ b/packages/appkit/src/type-generator/tests/warehouse-status.test.ts @@ -8,15 +8,15 @@ import { } from "../warehouse-status"; /** - * Build a minimal WorkspaceClient stub exposing only `warehouses.get`, the one - * method these helpers touch. Cast through `unknown` to the SDK type so callers - * type-check without us constructing a real client. + * Build a minimal WorkspaceClient stub exposing only `warehouses.getWarehouse`, + * the one method these helpers touch. Cast through `unknown` to the SDK type so + * callers type-check without us constructing a real client. */ function makeClient(get: ReturnType): WorkspaceClient { - return { warehouses: { get } } as unknown as WorkspaceClient; + return { warehouses: { getWarehouse: get } } as unknown as WorkspaceClient; } -/** A warehouses.get resolution carrying a given lifecycle state. */ +/** A warehouses.getWarehouse resolution carrying a given lifecycle state. */ const stateResponse = (state: WarehouseState) => ({ state }); describe("getWarehouseState", () => { diff --git a/packages/appkit/src/type-generator/warehouse-status.ts b/packages/appkit/src/type-generator/warehouse-status.ts index 27a0afeb5..8aae70e5e 100644 --- a/packages/appkit/src/type-generator/warehouse-status.ts +++ b/packages/appkit/src/type-generator/warehouse-status.ts @@ -71,14 +71,14 @@ export async function getWarehouseState( client: WorkspaceClient, warehouseId: string, ): Promise { - const response = await client.warehouses.get({ id: warehouseId }); + const response = await client.warehouses.getWarehouse({ id: warehouseId }); return response.state as WarehouseState; } /** * Initiate a start of a stopped/stopping SQL warehouse. * - * Only KICKS OFF the start: the SDK's `start()` returns a Waiter, but we + * Only KICKS OFF the start: the SDK's `startWarehouse()` returns a Waiter, but we * deliberately do not `.wait()` on it. Blocking on the full cold-start isn't our * job here — {@link waitUntilRunning} is the poller that watches the warehouse * the rest of the way to RUNNING. We just nudge it out of the stopped state. @@ -90,7 +90,7 @@ export async function startWarehouse( client: WorkspaceClient, warehouseId: string, ): Promise { - await client.warehouses.start({ id: warehouseId }); + await client.warehouses.startWarehouse({ id: warehouseId }); } /** diff --git a/packages/appkit/src/workspace-client/index.ts b/packages/appkit/src/workspace-client/index.ts index 581cb79a8..a7d7e1c34 100644 --- a/packages/appkit/src/workspace-client/index.ts +++ b/packages/appkit/src/workspace-client/index.ts @@ -13,15 +13,8 @@ export { Time, TimeUnits, } from "shared"; -export type { - CancellationToken, - ClientOptions, - files, - GenieMessage, - jobs, - serving, - sql, - Waiter, - WorkspaceClient, - WorkspaceClientOptions, -} from "shared/workspace-client"; +// Forwards every wrapper type — legacy service namespaces (files/jobs/serving), +// the client option/waiter types, and the modular SDK client + model types +// (warehouses, statementExecution). `sql` is gone: its statement + warehouse +// types now come from the modular SDK. +export type * from "shared/workspace-client"; diff --git a/packages/shared/package.json b/packages/shared/package.json index 811031c4c..7dd98e48b 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -48,7 +48,12 @@ "dependencies": { "@ast-grep/napi": "0.37.0", "@clack/prompts": "1.0.1", + "@databricks/sdk-auth": "0.46.0", + "@databricks/sdk-core": "0.46.0", "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-options": "0.46.0", + "@databricks/sdk-statementexecution": "0.46.0", + "@databricks/sdk-warehouses": "0.46.0", "@standard-schema/spec": "1.1.0", "commander": "12.1.0", "dotenv": "16.6.1", diff --git a/packages/shared/src/workspace-client/client.ts b/packages/shared/src/workspace-client/client.ts index 18ef76a17..adf30c8db 100644 --- a/packages/shared/src/workspace-client/client.ts +++ b/packages/shared/src/workspace-client/client.ts @@ -12,11 +12,19 @@ import { type LegacyWorkspaceClient, type WorkspaceClientOptions, } from "./legacy"; +import { + buildStatementExecutionClient, + buildWarehousesClient, + type StatementExecutionClient, + type WarehousesClient, +} from "./modular"; import type { WorkspaceClient } from "./types"; export class AppKitWorkspaceClient implements WorkspaceClient { readonly #opts: WorkspaceClientOptions; #legacy?: LegacyWorkspaceClient; + #warehouses?: WarehousesClient; + #statementExecution?: StatementExecutionClient; constructor(opts: WorkspaceClientOptions) { this.#opts = opts; @@ -26,8 +34,12 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#getLegacy().files; } - get warehouses() { - return this.#getLegacy().warehouses; + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get warehouses(): WarehousesClient { + if (!this.#warehouses) { + this.#warehouses = buildWarehousesClient(this.#opts); + } + return this.#warehouses; } get genie() { @@ -38,8 +50,12 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#getLegacy().jobs; } - get statementExecution() { - return this.#getLegacy().statementExecution; + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get statementExecution(): StatementExecutionClient { + if (!this.#statementExecution) { + this.#statementExecution = buildStatementExecutionClient(this.#opts); + } + return this.#statementExecution; } get servingEndpoints() { diff --git a/packages/shared/src/workspace-client/index.ts b/packages/shared/src/workspace-client/index.ts index 91921efeb..2b981ebec 100644 --- a/packages/shared/src/workspace-client/index.ts +++ b/packages/shared/src/workspace-client/index.ts @@ -23,3 +23,5 @@ export { TimeUnits, } from "./legacy"; export type { files, jobs, serving, sql, WorkspaceClient } from "./types"; +// Modular SDK client + model types (warehouses). +export type * from "./modular"; diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts new file mode 100644 index 000000000..a9cdef32b --- /dev/null +++ b/packages/shared/src/workspace-client/modular.ts @@ -0,0 +1,159 @@ +/** + * The single module allowed to import the modular `@databricks/sdk-*` SDK + * directly — the new-SDK sibling of {@link ./legacy.ts}. Every other AppKit + * module reaches these clients through the {@link WorkspaceClient} facade and + * the type re-exports below, so the modular SDK stays isolated exactly like the + * legacy one (the oxlint `no-restricted-imports` boundary walls `@databricks/sdk-*` + * off everywhere outside `packages/shared/src/workspace-client/`). + * + * Migrated services are built here as per-service clients; the facade delegates + * their accessors to these instead of the legacy monolithic client. Currently + * `warehouses` and `statementExecution` are migrated; every other service still + * routes through `legacy.ts`. + * + * NOTE: statementExecution relies on a pinned pnpm patch + * (`patches/@databricks__sdk-statementexecution@0.46.0.patch`) that restores the + * undocumented Reyden `attachment` response field, which the SDK's generated + * unmarshal transform would otherwise strip. + */ +import { newPatCredentials } from "@databricks/sdk-auth/credentials"; +import { addToDefault, setProduct } from "@databricks/sdk-core/clientinfo"; +import type { ClientOptions } from "@databricks/sdk-options/client"; +import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; +import { WarehousesClient } from "@databricks/sdk-warehouses/v1"; + +import type { WorkspaceClientOptions } from "./legacy"; + +/** + * Prepend `https://` to a scheme-less host. The legacy SDK normalized the host + * this way; the modular SDK does NOT — it passes the host straight into `fetch`, + * so a bare `DATABRICKS_HOST=my-workspace.cloud.databricks.com` (the common form, + * and what the Databricks Apps runtime sets) yields `TypeError: Invalid URL`. + */ +function normalizeHost(host: string | undefined): string | undefined { + const trimmed = host?.trim(); + if (!trimmed) return undefined; + return /^https?:\/\//i.test(trimmed) ? trimmed : `https://${trimmed}`; +} + +/** + * Map wrapper options onto the modular SDK's `ClientOptions`. Mirrors + * `buildLegacyWorkspaceClient`'s auth resolution verbatim, including the + * privilege-escalation guard: check `token !== undefined` (NOT truthiness) so an + * explicitly-passed token — even an empty string — pins the PAT path and fails + * loudly at request time rather than silently authenticating as the service + * principal via the default chain (which would be an OBO privilege escalation). + */ +function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { + const clientOptions: ClientOptions = {}; + // Resolve + scheme-normalize the host the way the legacy SDK did. Explicit + // `opts.host` wins; otherwise fall back to `DATABRICKS_HOST` (env is where the + // Apps runtime and dev set it). When a profile is selected without an explicit + // host, defer to the SDK's profile-file resolution instead of the env. + const host = normalizeHost( + opts.host ?? (opts.profile ? undefined : process.env.DATABRICKS_HOST), + ); + if (host) { + clientOptions.host = host; + } + if (opts.token !== undefined) { + clientOptions.credentials = newPatCredentials(opts.token); + } else if (opts.profile) { + clientOptions.profileOptions = { profile: opts.profile }; + } + // Neither token nor profile → leave credentials unset so the SDK walks its + // default auth chain (env vars + ~/.databrickscfg), matching the legacy `{}` case. + return clientOptions; +} + +// The modular SDK has no per-client User-Agent option; product/client-info is a +// process-global set once via `setProduct`/`addToDefault` before any client is +// built. The AppKit product/version/userAgentExtra arrive on `opts.clientOptions` +// (from `getClientOptions()`); build-time callers omit them and are left unstamped, +// preserving the legacy behavior where build-time clients carry no AppKit UA. The +// flag latches only once we actually stamp, so a first (unstamped) build-time +// client never blocks a later runtime client from stamping. +let clientInfoStamped = false; + +/** + * Coerce an arbitrary string into a valid client-info segment. The modular SDK + * validates keys as simple tokens and throws `ClientInfoError` on anything else, + * so the legacy product name `@databricks/appkit` (with `@` and `/`) is rejected + * — collapse invalid runs to `-` and trim the ends (`@databricks/appkit` → + * `databricks-appkit`). + */ +function toClientInfoKey(value: string): string { + return value.replace(/[^A-Za-z0-9._-]+/g, "-").replace(/^-+|-+$/g, ""); +} + +function ensureClientInfo(opts: WorkspaceClientOptions): void { + if (clientInfoStamped) { + return; + } + const co = opts.clientOptions; + if (!co?.product || !co?.productVersion) { + return; + } + // User-Agent stamping is best-effort: a value the SDK's client-info validator + // rejects must NEVER break client construction (the legacy SDK stamped the UA + // without validating). On failure the outbound request just carries the SDK's + // default User-Agent. + try { + setProduct(toClientInfoKey(co.product), co.productVersion); + if (co.userAgentExtra) { + for (const [key, value] of Object.entries(co.userAgentExtra)) { + addToDefault(toClientInfoKey(key), String(value)); + } + } + clientInfoStamped = true; + } catch { + clientInfoStamped = true; + } +} + +/** Build a modular Warehouses client from wrapper options. */ +export function buildWarehousesClient( + opts: WorkspaceClientOptions, +): WarehousesClient { + ensureClientInfo(opts); + return new WarehousesClient(mapToClientOptions(opts)); +} + +/** Build a modular Statement Execution client from wrapper options. */ +export function buildStatementExecutionClient( + opts: WorkspaceClientOptions, +): StatementExecutionClient { + ensureClientInfo(opts); + return new StatementExecutionClient(mapToClientOptions(opts)); +} + +// ── Client type re-exports (for the facade accessor types) ─────────────── +export type { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; +export type { WarehousesClient } from "@databricks/sdk-warehouses/v1"; + +// ── Model type re-exports ──────────────────────────────────────────────── +// AppKit modules import request/response/enum types from the wrapper rather +// than the SDK, so the import boundary holds. Type-only: the connector compares +// state against string literals, which satisfy the SDK's `Enum | (string & {})` +// field unions — no runtime enum values needed. +export type { + ColumnInfo, + Disposition, + ExecuteStatementRequest, + ExternalLink, + Format, + ResultData, + ResultManifest, + Schema, + ServiceError, + StatementParameter, + StatementResponse, + StatementStatus, + StatementStatus_State, +} from "@databricks/sdk-statementexecution/v1"; +export type { + EndpointHealth, + EndpointInfo, + EndpointState, + GetWarehouseResponse, +} from "@databricks/sdk-warehouses/v1"; diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts new file mode 100644 index 000000000..d6d092bcd --- /dev/null +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -0,0 +1,113 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +// The wrapper's own tests are the one place allowed to mock the SDK directly. +// Capture the `ClientOptions` the modular `WarehousesClient` constructor receives +// so we can assert how wrapper options map onto the modular SDK's config. +const { ctorOpts, patTokens, productCalls } = vi.hoisted(() => ({ + ctorOpts: [] as Array>, + patTokens: [] as string[], + productCalls: [] as Array<[string, string]>, +})); + +vi.mock("@databricks/sdk-warehouses/v1", () => ({ + WarehousesClient: vi.fn().mockImplementation((opts) => { + ctorOpts.push(opts); + return { opts }; + }), +})); +vi.mock("@databricks/sdk-statementexecution/v1", () => ({ + StatementExecutionClient: vi.fn().mockImplementation((opts) => ({ opts })), +})); +vi.mock("@databricks/sdk-auth/credentials", () => ({ + newPatCredentials: vi.fn((token: string) => { + patTokens.push(token); + return { kind: "pat", token }; + }), +})); +vi.mock("@databricks/sdk-core/clientinfo", () => ({ + setProduct: vi.fn((name: string, version: string) => { + // Mirror the real SDK: reject client-info keys that aren't simple tokens. + if (/[^A-Za-z0-9._-]/.test(name)) { + throw new Error(`Invalid key: ${name}.`); + } + productCalls.push([name, version]); + }), + addToDefault: vi.fn(), +})); + +import { buildWarehousesClient } from "../modular"; + +describe("modular mapToClientOptions (via buildWarehousesClient)", () => { + const originalHost = process.env.DATABRICKS_HOST; + + beforeEach(() => { + ctorOpts.length = 0; + patTokens.length = 0; + productCalls.length = 0; + delete process.env.DATABRICKS_HOST; + }); + + afterEach(() => { + if (originalHost === undefined) delete process.env.DATABRICKS_HOST; + else process.env.DATABRICKS_HOST = originalHost; + }); + + test("prepends https:// to a scheme-less explicit host", () => { + buildWarehousesClient({ host: "ws.cloud.databricks.com" }); + expect(ctorOpts[0].host).toBe("https://ws.cloud.databricks.com"); + }); + + test("leaves an explicit host that already has a scheme unchanged", () => { + buildWarehousesClient({ host: "https://ws.cloud.databricks.com" }); + expect(ctorOpts[0].host).toBe("https://ws.cloud.databricks.com"); + }); + + test("falls back to DATABRICKS_HOST (scheme-normalized) when no host is passed", () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + buildWarehousesClient({}); + expect(ctorOpts[0].host).toBe("https://envhost.cloud.databricks.com"); + }); + + test("a token takes the PAT path and pins the resolved host", () => { + buildWarehousesClient({ token: "abc", host: "https://x" }); + expect(patTokens).toEqual(["abc"]); + expect(ctorOpts[0].host).toBe("https://x"); + expect(ctorOpts[0].credentials).toEqual({ kind: "pat", token: "abc" }); + }); + + test("an empty-string token still uses PAT (no silent fall-through to default auth)", () => { + buildWarehousesClient({ token: "", host: "https://x" }); + expect(patTokens).toEqual([""]); + expect(ctorOpts[0].credentials).toEqual({ kind: "pat", token: "" }); + }); + + test("a profile sets profileOptions and defers host to the SDK (ignores env)", () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + buildWarehousesClient({ profile: "myprofile" }); + expect(ctorOpts[0].profileOptions).toEqual({ profile: "myprofile" }); + expect(ctorOpts[0].host).toBeUndefined(); + expect(patTokens).toEqual([]); + }); + + test("no host, no token, no profile, no env → empty options (SDK default chain)", () => { + buildWarehousesClient({}); + expect(ctorOpts[0].host).toBeUndefined(); + expect(ctorOpts[0].credentials).toBeUndefined(); + expect(ctorOpts[0].profileOptions).toBeUndefined(); + }); + + test("client-info: sanitizes an invalid product name (e.g. @databricks/appkit) rather than crashing the client build", () => { + // Regression: the modular SDK's `setProduct` rejects `@databricks/appkit` + // (INVALID_KEY), which the legacy SDK accepted. UA stamping must be + // best-effort — a bad product string must never break client construction. + const client = buildWarehousesClient({ + clientOptions: { + product: "@databricks/appkit", + productVersion: "0.64.0", + userAgentExtra: { mode: "dev" }, + }, + } as never); + expect(client).toBeDefined(); + expect(productCalls[0]).toEqual(["databricks-appkit", "0.64.0"]); + }); +}); diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index 398a4afe0..6d9865ace 100644 --- a/packages/shared/src/workspace-client/types.ts +++ b/packages/shared/src/workspace-client/types.ts @@ -14,10 +14,17 @@ * as each service migrates. */ import type { LegacyWorkspaceClient } from "./legacy"; +import type { StatementExecutionClient, WarehousesClient } from "./modular"; -// SDK type namespaces, re-exported so AppKit modules import them from the -// wrapper rather than the SDK directly. +// Legacy SDK type namespaces for un-migrated services, re-exported so AppKit +// modules import them from the wrapper rather than the SDK directly. `sql` +// stays only for the dev-mode warehouse listing in service-context, which reads +// the raw (snake_case) `/api/2.0/sql/warehouses` body via the still-legacy +// `apiClient` and types it as `sql.EndpointInfo[]`. Statement + warehouse +// service types now come from `./modular`. export type { files, jobs, serving, sql } from "@databricks/sdk-experimental"; +// Modular SDK client + model types (warehouses, statementExecution). +export type * from "./modular"; /** * AppKit's workspace client facade. Mirrors the multi-client shape of the @@ -31,8 +38,8 @@ export interface WorkspaceClient { /** UC Volumes / Files API. */ readonly files: LegacyWorkspaceClient["files"]; - /** SQL Warehouses. */ - readonly warehouses: LegacyWorkspaceClient["warehouses"]; + /** SQL Warehouses (modular SDK). */ + readonly warehouses: WarehousesClient; /** Genie / dashboards. */ readonly genie: LegacyWorkspaceClient["genie"]; @@ -40,8 +47,8 @@ export interface WorkspaceClient { /** Jobs. */ readonly jobs: LegacyWorkspaceClient["jobs"]; - /** Statement Execution. */ - readonly statementExecution: LegacyWorkspaceClient["statementExecution"]; + /** Statement Execution (modular SDK). */ + readonly statementExecution: StatementExecutionClient; /** Serving Endpoints. */ readonly servingEndpoints: LegacyWorkspaceClient["servingEndpoints"]; diff --git a/patches/@databricks__sdk-statementexecution@0.46.0.patch b/patches/@databricks__sdk-statementexecution@0.46.0.patch new file mode 100644 index 000000000..c206b65c4 --- /dev/null +++ b/patches/@databricks__sdk-statementexecution@0.46.0.patch @@ -0,0 +1,34 @@ +diff --git a/dist/v1/model.d.ts b/dist/v1/model.d.ts +index e8d95659ea348b384a3d32b6a3d4f754287b38b6..705b9bed203981a2f3cde5417ed8019ff3a7065c 100644 +--- a/dist/v1/model.d.ts ++++ b/dist/v1/model.d.ts +@@ -385,6 +385,8 @@ interface QueryTag { + * link is returned.) + */ + interface ResultData { ++ /** PATCH(appkit): Reyden's non-standard INLINE ARROW_STREAM payload (base64 Arrow IPC). */ ++ attachment?: string | undefined; + externalLinks?: ExternalLink[] | undefined; + /** + * The `JSON_ARRAY` format is an array of arrays of values, where each non-null value is +diff --git a/dist/v1/model.js b/dist/v1/model.js +index fc35e28bbad5e7e873c14f6492696f9086ec9280..3bd78f3ad2dea39cdc883fa0daddb940e99e83b9 100644 +--- a/dist/v1/model.js ++++ b/dist/v1/model.js +@@ -177,10 +177,15 @@ const unmarshalResultDataSchema = z.object({ + z.string() + ]).transform((v) => BigInt(v)).optional(), + next_chunk_index: z.number().optional(), +- next_chunk_internal_link: z.string().optional() ++ next_chunk_internal_link: z.string().optional(), ++ // PATCH(appkit): preserve Reyden's non-standard INLINE ARROW_STREAM `attachment` ++ // (base64 Arrow IPC). The generated schema + rebuild-transform would otherwise ++ // strip it, breaking the inline-arrow delivery path. See patches/ for rationale. ++ attachment: z.string().optional() + }).transform((d) => ({ + externalLinks: d.external_links, + dataArray: d.data_array, ++ attachment: d.attachment, + chunkIndex: d.chunk_index, + rowOffset: d.row_offset, + rowCount: d.row_count, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4ced623e4..63cfb5da7 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -11,6 +11,9 @@ overrides: qs@<6.15.2: 6.15.2 size-sensor: 1.0.3 +patchedDependencies: + '@databricks/sdk-statementexecution@0.46.0': a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d + importers: .: @@ -591,9 +594,24 @@ importers: '@clack/prompts': specifier: 1.0.1 version: 1.0.1 + '@databricks/sdk-auth': + specifier: 0.46.0 + version: 0.46.0 + '@databricks/sdk-core': + specifier: 0.46.0 + version: 0.46.0 '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-options': + specifier: 0.46.0 + version: 0.46.0 + '@databricks/sdk-statementexecution': + specifier: 0.46.0 + version: 0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d) + '@databricks/sdk-warehouses': + specifier: 0.46.0 + version: 0.46.0 '@standard-schema/spec': specifier: 1.1.0 version: 1.1.0 @@ -1959,6 +1977,14 @@ packages: engines: {node: ^20 || ^22 || ^24 || ^25, pnpm: '>=10'} hasBin: true + '@databricks/sdk-auth@0.46.0': + resolution: {integrity: sha512-cMrwxsFtpiEKFxta5dKHchKdrgmHkQ6upJ2C4OacmlHrOHJ+ChzQBVTpAiVtjX62xj+YjNgp/29IpSdKKYUVDA==} + engines: {node: '>=22.0.0'} + + '@databricks/sdk-core@0.46.0': + resolution: {integrity: sha512-Q2LAGWYIi+jyeKR9OIqvkgyde2GdzqfSG8lewxA9Xu/C9RJBBFbSfg5Nh8ZC66TKElGIosVOecoEJdbxnNMuxw==} + engines: {node: '>=22.0.0'} + '@databricks/sdk-experimental@0.15.0': resolution: {integrity: sha512-HkoMiF7dNDt6WRW0xhi7oPlBJQfxJ9suJhEZRFt08VwLMaWcw2PiF8monfHlkD4lkufEYV6CTxi5njQkciqiHA==} engines: {node: '>=22.0', npm: '>=10.0.0'} @@ -1967,6 +1993,18 @@ packages: resolution: {integrity: sha512-dOJIt4F2nBk6HKObnv7Xbmy/qLYTy2835qhXSuW0Qw1QAXui9plmCet1KqG3yeQcMTyncWGbnhjGdQi8GEGQSA==} engines: {node: '>=22.0', npm: '>=10.0.0'} + '@databricks/sdk-options@0.46.0': + resolution: {integrity: sha512-UtADlR+41rYEoOCycZvJh1g96uDN6GVWgqQk+72cHBzcxi+koxKSJXHcYRI997Sc4Fcc1d2oyAC2I2ddVhurjA==} + engines: {node: '>=22.0.0'} + + '@databricks/sdk-statementexecution@0.46.0': + resolution: {integrity: sha512-VJA3e7UHmxRxN42/mV5VtKeINME0vCz3Na3hrwmta3tZqJWZbBU7XTfUdD1yOQ5Z1JU5UIM65OlWq8gc4IzHFg==} + engines: {node: '>=22.0.0'} + + '@databricks/sdk-warehouses@0.46.0': + resolution: {integrity: sha512-9r/gbdTb6ASiWCiibCwAOF8QizqNacidIw78uwzJYKK9dbdqmWIfNK0pF/jt3BG1sUiXz2b1I6URPdX7Qi0oLg==} + engines: {node: '>=22.0.0'} + '@date-fns/tz@1.4.1': resolution: {integrity: sha512-P5LUNhtbj6YfI3iJjw5EL9eUAG6OitD0W3fWQcpQjDRc/QIsL0tRNuO1PcDvPccWL1fSTXXdE1ds+l95DV/OFA==} @@ -2633,6 +2671,10 @@ packages: '@js-sdsl/ordered-map@4.4.2': resolution: {integrity: sha512-iUKgm52T8HOE/makSxjqoWhe95ZJA1/G1sYsGev2JDKUSS14KAgg1LHb+Ba+IPow0xflbnSkOsZcO08C7w1gYw==} + '@js-temporal/polyfill@0.5.1': + resolution: {integrity: sha512-hloP58zRVCRSpgDxmqCWJNlizAlUgJFqG2ypq79DCvyv9tHjRYMDOcPFjzfl/A1/YxDvRCZz8wvZvmapQnKwFQ==} + engines: {node: '>=12'} + '@jsep-plugin/assignment@1.3.0': resolution: {integrity: sha512-VVgV+CXrhbMI3aSusQyclHkenWSAm95WaiKrMxRFam3JSUiIaQjoMIw2sEs/OX4XifnqeQUN4DYbJjlA8EfktQ==} engines: {node: '>= 10.16.0'} @@ -8685,6 +8727,9 @@ packages: resolution: {integrity: sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==} hasBin: true + jsbi@4.3.2: + resolution: {integrity: sha512-9fqMSQbhJykSeii05nxKl4m6Eqn2P6rOlYiS+C5Dr/HPIU/7yZxu5qzbs40tgaFORiw2Amd0mirjxatXYMkIew==} + jsdom@27.0.0: resolution: {integrity: sha512-lIHeR1qlIRrIN5VMccd8tI2Sgw6ieYXSVktcSHaNe3Z5nE/tcPQYQWOq00wxMvYOsz+73eAkNenVvmPC6bba9A==} engines: {node: '>=20'} @@ -14228,6 +14273,16 @@ snapshots: transitivePeerDependencies: - supports-color + '@databricks/sdk-auth@0.46.0': + dependencies: + '@databricks/sdk-core': 0.46.0 + zod: 4.3.6 + + '@databricks/sdk-core@0.46.0': + dependencies: + json-bigint: 1.0.0 + zod: 4.3.6 + '@databricks/sdk-experimental@0.15.0': dependencies: google-auth-library: 10.5.0 @@ -14246,6 +14301,29 @@ snapshots: transitivePeerDependencies: - supports-color + '@databricks/sdk-options@0.46.0': + dependencies: + '@databricks/sdk-auth': 0.46.0 + '@databricks/sdk-core': 0.46.0 + + '@databricks/sdk-statementexecution@0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d)': + dependencies: + '@databricks/sdk-auth': 0.46.0 + '@databricks/sdk-core': 0.46.0 + '@databricks/sdk-options': 0.46.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + + '@databricks/sdk-warehouses@0.46.0': + dependencies: + '@databricks/sdk-auth': 0.46.0 + '@databricks/sdk-core': 0.46.0 + '@databricks/sdk-options': 0.46.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + '@date-fns/tz@1.4.1': {} '@discoveryjs/json-ext@0.5.7': {} @@ -15444,6 +15522,10 @@ snapshots: '@js-sdsl/ordered-map@4.4.2': {} + '@js-temporal/polyfill@0.5.1': + dependencies: + jsbi: 4.3.2 + '@jsep-plugin/assignment@1.3.0(jsep@1.4.0)': dependencies: jsep: 1.4.0 @@ -22048,6 +22130,8 @@ snapshots: dependencies: argparse: 2.0.1 + jsbi@4.3.2: {} + jsdom@27.0.0(postcss@8.5.6): dependencies: '@asamuzakjp/dom-selector': 6.6.2 diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 3b88ac7e2..9a8221175 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -18,3 +18,5 @@ allowBuilds: core-js-pure: false esbuild: true protobufjs: false +patchedDependencies: + '@databricks/sdk-statementexecution@0.46.0': patches/@databricks__sdk-statementexecution@0.46.0.patch From 94ebe11a7f7cd735b1d46f02f5c0b3138ec57fa6 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 11:40:05 +0200 Subject: [PATCH 02/17] chore: fixup --- packages/appkit/src/plugins/analytics/analytics.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 511fd643e..21d0985d9 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -29,7 +29,6 @@ import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; import { getWarehouseId } from "../../resources"; -import type { WorkspaceClient } from "../../workspace-client"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { From 7e80a2a0697f6f51c3d77d9f8eb8647c6083dd61 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 12:32:51 +0200 Subject: [PATCH 03/17] refactor(appkit): camelCase type-generator domain type; drop SDK response mappers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the analytics SDK migration: DatabricksStatementExecutionResponse is now camelCase (aligned with the modular SDK's StatementResponse), so the translation shims the migration introduced are no longer needed. - Delete `toDescribeResponse` — `describeAdaptive` narrows the SDK response onto the domain type directly (one cast: the SDK types DESCRIBE cells as JsonValue[][], but for a DESCRIBE they are always string/null). - Delete the three duplicated `asSdkResponse` test helpers; fixtures are now authored in the camelCase domain shape and feed both the SDK-mock and the snake-free parsers directly. - `error_code` stays snake ONLY where it reads the raw Databricks error wire shape (errors.ts auth classification, `{"error_code":...}` JSON bodies) — a network payload, not our type. Internal type only (not exported, not in API docs) — no breaking change. Net -117 LOC. Verified: typecheck, 4355 tests, build (attw + publint), lint, format. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../type-generator/mv-registry/describe.ts | 2 +- .../src/type-generator/query-registry.ts | 10 +- .../src/type-generator/statement-result.ts | 86 +++------- .../tests/generate-queries.test.ts | 6 +- .../src/type-generator/tests/index.test.ts | 110 +++++-------- .../type-generator/tests/mv-registry.test.ts | 113 +++++-------- .../tests/query-registry.test.ts | 28 ++-- .../tests/statement-result.test.ts | 153 ++++++++---------- .../tests/sync-metric-views-types.test.ts | 6 +- .../tests/unreachable-warehouse-gate.test.ts | 4 +- packages/appkit/src/type-generator/types.ts | 27 ++-- 11 files changed, 214 insertions(+), 331 deletions(-) diff --git a/packages/appkit/src/type-generator/mv-registry/describe.ts b/packages/appkit/src/type-generator/mv-registry/describe.ts index db2302761..a01f26190 100644 --- a/packages/appkit/src/type-generator/mv-registry/describe.ts +++ b/packages/appkit/src/type-generator/mv-registry/describe.ts @@ -25,7 +25,7 @@ export function parseDescribeTableExtendedJson( throw new Error(`DESCRIBE TABLE EXTENDED failed: ${msg}`); } - const rows = response.result?.data_array ?? []; + const rows = response.result?.dataArray ?? []; if (rows.length === 0) { throw new Error( "DESCRIBE TABLE EXTENDED returned no rows. Verify the FQN points to a metric view.", diff --git a/packages/appkit/src/type-generator/query-registry.ts b/packages/appkit/src/type-generator/query-registry.ts index dc6fc9d5f..d91b8cabe 100644 --- a/packages/appkit/src/type-generator/query-registry.ts +++ b/packages/appkit/src/type-generator/query-registry.ts @@ -158,7 +158,7 @@ function formatParametersType(sql: string): string { /** * Decode a base64 Arrow IPC attachment from a DESCRIBE QUERY response and * extract column metadata. Returns the same shape as rows parsed from the - * legacy data_array path. + * legacy dataArray path. * * IMPORTANT: a DESCRIBE QUERY response is itself a result *table* with rows * shaped like `(col_name, data_type, comment)` describing the user query's @@ -196,7 +196,7 @@ export function convertToQueryType( sql: string, queryName: string, ): { type: string; hasResults: boolean } { - const dataRows = result.result?.data_array || []; + const dataRows = result.result?.dataArray || []; let columns = dataRows.map((row) => ({ name: row[0] || "", type_name: row[1]?.toUpperCase() || "STRING", @@ -204,10 +204,10 @@ export function convertToQueryType( })); // Fallback: serverless warehouses return ARROW_STREAM format with an inline - // base64 attachment instead of data_array. Decode the Arrow IPC rows (the + // base64 attachment instead of dataArray. Decode the Arrow IPC rows (the // DESCRIBE QUERY result table) to extract column names and types. if (columns.length === 0 && result.result?.attachment) { - logger.debug("data_array empty, decoding Arrow IPC attachment for schema"); + logger.debug("dataArray empty, decoding Arrow IPC attachment for schema"); try { columns = columnsFromArrowAttachment(result.result.attachment); } catch (err) { @@ -849,7 +849,7 @@ export async function generateQueriesFromDescribe( "DESCRIBE result for %s: state=%s, rows=%d, hasAttachment=%s", queryName, result.status.state, - result.result?.data_array?.length ?? 0, + result.result?.dataArray?.length ?? 0, !!result.result?.attachment, ); diff --git a/packages/appkit/src/type-generator/statement-result.ts b/packages/appkit/src/type-generator/statement-result.ts index 24620988b..9f5f791fe 100644 --- a/packages/appkit/src/type-generator/statement-result.ts +++ b/packages/appkit/src/type-generator/statement-result.ts @@ -1,5 +1,5 @@ import { createLogger } from "../logging/logger"; -import type { StatementResponse, WorkspaceClient } from "../workspace-client"; +import type { WorkspaceClient } from "../workspace-client"; import { getErrorMessage } from "./errors"; import type { DatabricksStatementExecutionResponse } from "./types"; @@ -7,18 +7,18 @@ const logger = createLogger("type-generator:statement-result"); /** * Normalize a Statement Execution response so downstream parsers can always - * read rows from `result.data_array`, regardless of the wire format the + * read rows from `result.dataArray`, regardless of the wire format the * warehouse chose. * * `@databricks/sdk-experimental`'s `executeStatement` defaults to an * `ARROW_STREAM` disposition. With an `INLINE` disposition the single * DESCRIBE row is returned as a base64-encoded Arrow IPC stream in - * `result.attachment` and `result.data_array` is left undefined. The metric - * and query type generators only ever read `result.data_array`, so without + * `result.attachment` and `result.dataArray` is left undefined. The metric + * and query type generators only ever read `result.dataArray`, so without * this normalization an Arrow response reads as "returned no rows" — the * registry ships empty and the runtime fail-closed gate 503s every affected * metric/query. (A warehouse configured to return `JSON_ARRAY` populates - * `data_array` directly and needs no decoding — that path, and every mocked + * `dataArray` directly and needs no decoding — that path, and every mocked * test, flows through here unchanged.) */ export async function normalizeResultRows( @@ -30,18 +30,18 @@ export async function normalizeResultRows( // types. A deliberate throw — unlike the best-effort decode below — that both // callers catch per-entry as a loud per-key/per-query failure. if ( - response.result?.next_chunk_index != null || - response.result?.next_chunk_internal_link != null + response.result?.nextChunkIndex != null || + response.result?.nextChunkInternalLink != null ) { throw new Error( - "DESCRIBE result is multi-chunk (truncated); refusing to emit partial types — see next_chunk_index", + "DESCRIBE result is multi-chunk (truncated); refusing to emit partial types — see nextChunkIndex", ); } // Passthrough: rows already materialized (JSON_ARRAY warehouses + every - // mocked test). `data_array` being an empty array still counts as present — + // mocked test). `dataArray` being an empty array still counts as present — // that is a genuine "no rows" answer we must not overwrite with a decode. - if (response.result?.data_array !== undefined) { + if (response.result?.dataArray !== undefined) { return response; } @@ -78,7 +78,7 @@ export async function normalizeResultRows( ...response, result: { ...response.result, - data_array: dataArray, + dataArray: dataArray, }, }; } catch (err) { @@ -147,45 +147,9 @@ function isFormatRejection( ); } -/** - * Adapt the modular SDK's camelCase {@link StatementResponse} onto the - * type-generator's own snake_case {@link DatabricksStatementExecutionResponse} - * — the shape every downstream DESCRIBE parser (and every mocked test) reads. - * Keeping the boundary here means only this mapper touches the SDK shape; - * {@link normalizeResultRows} and the parsers stay unchanged. `attachment` - * survives thanks to the pinned pnpm patch on `@databricks/sdk-statementexecution`. - */ -function toDescribeResponse( - r: StatementResponse, -): DatabricksStatementExecutionResponse { - return { - statement_id: r.statementId ?? "", - status: { - state: r.status?.state ?? "", - error: r.status?.error - ? { - error_code: r.status.error.errorCode, - message: r.status.error.message, - } - : undefined, - }, - manifest: r.manifest ? { format: r.manifest.format } : undefined, - result: r.result - ? { - // DESCRIBE rows are always string/null cells. Local key stays - // snake_case (`data_array`); value is the SDK's camelCase `dataArray`. - data_array: r.result.dataArray as (string | null)[][] | undefined, - attachment: r.result.attachment, - next_chunk_index: r.result.nextChunkIndex, - next_chunk_internal_link: r.result.nextChunkInternalLink, - } - : undefined, - }; -} - /** * Run a DESCRIBE and return a response whose rows are readable via - * `result.data_array`, adapting to the warehouse's result-format capability. + * `result.dataArray`, adapting to the warehouse's result-format capability. * * No single format is portable: standard DBSQL (PRO/CLASSIC) serves * `INLINE`+`JSON_ARRAY` and rejects `INLINE`+`ARROW_STREAM`; the Reyden engine @@ -211,23 +175,25 @@ export async function describeAdaptive( let lastError: unknown; for (const format of formats) { try { - const response = toDescribeResponse( - await client.statementExecution.executeStatement({ - statement, - warehouseId, - // Synchronous wait: without it the call can return PENDING/RUNNING with - // no rows, which downstream misreads as a no-result degrade. - waitTimeout: "30s", - format, - disposition: "INLINE", - }), - ); + // Narrow the modular SDK's camelCase StatementResponse straight onto our + // subset. The only gap is `dataArray` cells (the SDK types them as + // `JsonValue[][]`); for a DESCRIBE they are always string/null, so the + // assertion is safe. `attachment` survives via the pinned pnpm patch. + const response = (await client.statementExecution.executeStatement({ + statement, + warehouseId, + // Synchronous wait: without it the call can return PENDING/RUNNING with + // no rows, which downstream misreads as a no-result degrade. + waitTimeout: "30s", + format, + disposition: "INLINE", + })) as DatabricksStatementExecutionResponse; const normalized = await normalizeResultRows(response); if ( normalized.status?.state === "FAILED" && isFormatRejection( normalized.status.error?.message, - normalized.status.error?.error_code, + normalized.status.error?.errorCode, ) ) { lastResponse = normalized; diff --git a/packages/appkit/src/type-generator/tests/generate-queries.test.ts b/packages/appkit/src/type-generator/tests/generate-queries.test.ts index 9ad1b1dce..2fe29a072 100644 --- a/packages/appkit/src/type-generator/tests/generate-queries.test.ts +++ b/packages/appkit/src/type-generator/tests/generate-queries.test.ts @@ -98,7 +98,7 @@ function succeededResult(columns: [string, string, string | null][]) { /** * Build a SUCCEEDED DESCRIBE QUERY response whose rows arrive only as a base64 - * Arrow IPC `attachment` (no `data_array`) — the ARROW_STREAM/INLINE wire shape + * Arrow IPC `attachment` (no `dataArray`) — the ARROW_STREAM/INLINE wire shape * the fetcher now requests. The describeOne path pipes this through * normalizeResultRows, which decodes the attachment so convertToQueryType can * read the columns. Each [name, type, comment] triple becomes one DESCRIBE row. @@ -119,7 +119,7 @@ async function succeededArrowAttachmentResult( statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, - // No data_array — rows live in the attachment, like a real INLINE Arrow + // No dataArray — rows live in the attachment, like a real INLINE Arrow // response. This is the condition the silent-degrade bug left unread. result: { attachment }, }; @@ -184,7 +184,7 @@ describe("generateQueriesFromDescribe", () => { test("ARROW attachment path — decodes Arrow rows into a real query schema", async () => { // The warehouse answers ARROW_STREAM/INLINE: columns arrive only as a - // base64 Arrow IPC attachment with data_array undefined. describeOne pipes + // base64 Arrow IPC attachment with dataArray undefined. describeOne pipes // this through normalizeResultRows before convertToQueryType, so the schema // resolves to real columns instead of the degraded `result: unknown`. mocks.readdir.mockResolvedValue(["users.sql"]); diff --git a/packages/appkit/src/type-generator/tests/index.test.ts b/packages/appkit/src/type-generator/tests/index.test.ts index 6fa33ed48..088b2bcf5 100644 --- a/packages/appkit/src/type-generator/tests/index.test.ts +++ b/packages/appkit/src/type-generator/tests/index.test.ts @@ -11,37 +11,8 @@ import { vi, } from "vitest"; -import type { StatementResponse } from "../../workspace-client"; import type { DatabricksStatementExecutionResponse } from "../types"; -/** - * Adapt a local snake_case describe fixture to the modular SDK's camelCase - * `StatementResponse` — the shape the mocked `executeStatement` now returns. - * `describeAdaptive` maps it back to the local shape via `toDescribeResponse`, - * so fixtures stay authored in the type-generator's own domain shape. - */ -function asSdkResponse( - r: DatabricksStatementExecutionResponse, -): StatementResponse { - return { - statementId: r.statement_id, - status: r.status && { - state: r.status.state, - error: r.status.error && { - errorCode: r.status.error.error_code, - message: r.status.error.message, - }, - }, - manifest: r.manifest && { format: r.manifest.format }, - result: r.result && { - dataArray: r.result.data_array, - attachment: r.result.attachment, - nextChunkIndex: r.result.next_chunk_index, - nextChunkInternalLink: r.result.next_chunk_internal_link, - }, - } as unknown as StatementResponse; -} - const mocks = vi.hoisted(() => ({ generateQueriesFromDescribe: vi.fn(), getWarehouseState: vi.fn(), @@ -327,10 +298,10 @@ describe("generateFromEntryPoint — metric-view emission", () => { const metricFile = path.join(metricsDir, "generated", "metric-views.d.ts"); const describeResponse: DatabricksStatementExecutionResponse = { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ [ JSON.stringify({ columns: [ @@ -489,7 +460,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { // the statement still PENDING — no rows yet. Previously this fell // into the "returned no rows" failure with per-key warns. metricFetcher: async () => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }), }), @@ -582,7 +553,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("non-blocking + RUNNING warehouse: DESCRIBEs run and land full schemas", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ @@ -611,7 +582,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("blocking + RUNNING: one preflight probe, no start/wait, DESCRIBEs run", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ @@ -737,7 +708,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { warehouseId: "wh-1", mode: "blocking", metricFetcher: async () => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }), }), @@ -785,7 +756,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { mocks.getWarehouseState.mockResolvedValue("STOPPED"); mocks.startWarehouse.mockResolvedValue(undefined); mocks.waitUntilRunning.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ @@ -922,7 +893,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { // The fall-through DESCRIBE hits a still-cold warehouse: non-terminal // response, which classifies as degraded (never an error). mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }); const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -1044,7 +1015,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("non-blocking + RUNNING with the default fetcher: probe and DESCRIBEs share exactly one client", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ @@ -1183,23 +1154,24 @@ describe("generateFromEntryPoint — metric cache section", () => { const outFile = path.join(cacheTestDir, "generated", "analytics.d.ts"); const metricFile = path.join(cacheTestDir, "generated", "metric-views.d.ts"); - const describeResponseFor = (measure: string): StatementResponse => - asSdkResponse({ - statement_id: "stmt-mock", - status: { state: "SUCCEEDED" }, - result: { - data_array: [ - [ - JSON.stringify({ - columns: [ - { name: measure, type: "DECIMAL(38,2)", is_measure: true }, - { name: "region", type: "STRING", is_measure: false }, - ], - }), - ], + const describeResponseFor = ( + measure: string, + ): DatabricksStatementExecutionResponse => ({ + statementId: "stmt-mock", + status: { state: "SUCCEEDED" }, + result: { + dataArray: [ + [ + JSON.stringify({ + columns: [ + { name: measure, type: "DECIMAL(38,2)", is_measure: true }, + { name: "region", type: "STRING", is_measure: false }, + ], + }), ], - }, - }); + ], + }, + }); const writeConfig = ( metricViews: Record< @@ -1522,26 +1494,26 @@ describe("generateFromEntryPoint — metric cache section", () => { } if (statement.includes("failed_stmt")) { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }; } if (statement.includes("no_rows")) { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }; } if (statement.includes("no_columns")) { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [[JSON.stringify({ unrelated: true })]] }, + result: { dataArray: [[JSON.stringify({ unrelated: true })]] }, }; } if (statement.includes("pending")) { - return { statement_id: "stmt-mock", status: { state: "PENDING" } }; + return { statementId: "stmt-mock", status: { state: "PENDING" } }; } return describeResponseFor("total_revenue"); }, @@ -1601,7 +1573,7 @@ describe("generateFromEntryPoint — metric cache section", () => { writeConfig({ revenue: { source: "demo.sales.revenue" } }); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const firstWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -1638,7 +1610,7 @@ describe("generateFromEntryPoint — metric cache section", () => { writeConfig({ revenue: { source: "demo.sales.revenue" } }); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -1655,7 +1627,7 @@ describe("generateFromEntryPoint — metric cache section", () => { vi.clearAllMocks(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const error = await run({ mode: "blocking" }).then( @@ -1693,7 +1665,7 @@ describe("generateFromEntryPoint — metric cache section", () => { writeConfig({ revenue: { source: "demo.sales.revenue" } }); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -1735,7 +1707,7 @@ describe("generateFromEntryPoint — metric cache section", () => { writeConfig({ revenue: { source: "demo.sales.revenue" } }); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -2041,7 +2013,7 @@ describe("generateFromEntryPoint — anti-clobber for blocking mode", () => { warehouseId: "wh-1", mode: "blocking", metricFetcher: async () => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }), }), @@ -2088,10 +2060,10 @@ describe("generateFromEntryPoint — anti-clobber for blocking mode", () => { ); const describeResponse: DatabricksStatementExecutionResponse = { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ [ JSON.stringify({ columns: [ @@ -2109,7 +2081,7 @@ describe("generateFromEntryPoint — anti-clobber for blocking mode", () => { }; mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ diff --git a/packages/appkit/src/type-generator/tests/mv-registry.test.ts b/packages/appkit/src/type-generator/tests/mv-registry.test.ts index e6d90bb75..50c69fe1c 100644 --- a/packages/appkit/src/type-generator/tests/mv-registry.test.ts +++ b/packages/appkit/src/type-generator/tests/mv-registry.test.ts @@ -10,7 +10,6 @@ import { afterEach, beforeEach, describe, expect, test } from "vitest"; // imports it from there. import { quoteFqnForSql } from "../../../../shared/src/schemas/metric-fqn"; import { metricSourceSchema } from "../../../../shared/src/schemas/metric-source"; -import type { StatementResponse } from "../../workspace-client"; import { readMetricConfig, resolveMetricConfig } from "../mv-registry/config"; import { createWorkspaceDescribeFetcher, @@ -46,43 +45,14 @@ function mockDescribeResponse( payload: unknown, ): DatabricksStatementExecutionResponse { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, result: { - data_array: [[JSON.stringify(payload)]], + dataArray: [[JSON.stringify(payload)]], }, }; } -/** - * Adapt a local snake_case describe fixture to the modular SDK's camelCase - * `StatementResponse` — the shape a mocked `executeStatement` (consumed by the - * real `createWorkspaceDescribeFetcher` → `describeAdaptive`) now returns. - * Direct `syncMetrics(resolution, fetcher)` fixtures stay in the local snake - * shape (they bypass the SDK), so only the executeStatement mocks wrap with this. - */ -function asSdkResponse( - r: DatabricksStatementExecutionResponse, -): StatementResponse { - return { - statementId: r.statement_id, - status: r.status && { - state: r.status.state, - error: r.status.error && { - errorCode: r.status.error.error_code, - message: r.status.error.message, - }, - }, - manifest: r.manifest && { format: r.manifest.format }, - result: r.result && { - dataArray: r.result.data_array, - attachment: r.result.attachment, - nextChunkIndex: r.result.next_chunk_index, - nextChunkInternalLink: r.result.next_chunk_internal_link, - }, - } as unknown as StatementResponse; -} - /** * Real Arrow IPC attachment captured live from dogfood: * DESCRIBE TABLE EXTENDED `appkit_demo`.`public`.`revenue_metrics` AS JSON @@ -356,11 +326,9 @@ describe("resolveMetricConfig — FQN naming (UC-accurate)", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return asSdkResponse( - mockDescribeResponse({ - columns: [{ name: "arr", type: "DECIMAL", is_measure: true }], - }), - ); + return mockDescribeResponse({ + columns: [{ name: "arr", type: "DECIMAL", is_measure: true }], + }); }, }, } as unknown as Parameters[0]; @@ -391,7 +359,7 @@ describe("resolveMetricConfig — FQN naming (UC-accurate)", () => { // crashing the pass — exactly the pre-existing degrade behavior. const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const { schemas, failures } = await syncMetrics(resolution, fetcher); @@ -588,7 +556,7 @@ describe("parseDescribeTableExtendedJson", () => { test("throws on a FAILED status", () => { expect(() => parseDescribeTableExtendedJson({ - statement_id: "x", + statementId: "x", status: { state: "FAILED", error: { message: "no such table" } }, }), ).toThrowError(/no such table/); @@ -597,9 +565,9 @@ describe("parseDescribeTableExtendedJson", () => { test("throws when the response is empty", () => { expect(() => parseDescribeTableExtendedJson({ - statement_id: "x", + statementId: "x", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }), ).toThrowError(/no rows/); }); @@ -607,9 +575,9 @@ describe("parseDescribeTableExtendedJson", () => { test("throws when the cell is not a JSON string", () => { expect(() => parseDescribeTableExtendedJson({ - statement_id: "x", + statementId: "x", status: { state: "SUCCEEDED" }, - result: { data_array: [[null]] }, + result: { dataArray: [[null]] }, }), ).toThrowError(/JSON string/); }); @@ -696,7 +664,7 @@ describe("createWorkspaceDescribeFetcher", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return asSdkResponse(mockDescribeResponse(payload)); + return mockDescribeResponse(payload); }, }, } as unknown as Parameters[0]; @@ -723,7 +691,7 @@ describe("createWorkspaceDescribeFetcher", () => { test("decodes an Arrow attachment-only response into parseable columns (fetcher → normalizer → parser)", async () => { // The warehouse answers ARROW_STREAM/INLINE: rows arrive as a base64 Arrow - // IPC attachment with `data_array` undefined. Before the normalizer was + // IPC attachment with `dataArray` undefined. Before the normalizer was // wired in, parseDescribeTableExtendedJson read this as "no rows" and the // metric shipped degraded. Now the fetcher pipes the response through // normalizeResultRows, so the real describe doc is recovered end-to-end. @@ -732,13 +700,13 @@ describe("createWorkspaceDescribeFetcher", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return asSdkResponse({ - statement_id: "stmt-arrow", + return { + statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, - // Only an attachment — no data_array (the bug's trigger condition). + // Only an attachment — no dataArray (the bug's trigger condition). result: { attachment: ARROW_ATTACHMENT_B64 }, - }); + }; }, }, } as unknown as Parameters[0]; @@ -747,7 +715,7 @@ describe("createWorkspaceDescribeFetcher", () => { const response = await fetcher("appkit_demo.public.revenue_metrics"); // The fetcher decoded the attachment: rows are now readable. - expect(response.result?.data_array).toBeDefined(); + expect(response.result?.dataArray).toBeDefined(); const parsed = parseDescribeTableExtendedJson(response); const cols = extractMetricColumns(parsed); // The real revenue_metrics describe doc carries measures and dimensions. @@ -1037,7 +1005,7 @@ describe("syncMetrics", () => { test("a multi-chunk (truncated) DESCRIBE surfaces as a loud failure, not a crash (fetcher → normalizer → syncMetrics)", async () => { // End-to-end loudness check for the truncation guard. The warehouse paginates - // the DESCRIBE result (sets next_chunk_index on the first chunk); the fetcher + // the DESCRIBE result (sets nextChunkIndex on the first chunk); the fetcher // pipes the response through normalizeResultRows, which THROWS rather than // emit partial types. That throw must be caught inside describeOne and // recorded as a MetricSyncFailure — never an uncaught crash that aborts the @@ -1047,16 +1015,15 @@ describe("syncMetrics", () => { }); const client = { statementExecution: { - executeStatement: async () => - asSdkResponse({ - statement_id: "stmt-chunked", - status: { state: "SUCCEEDED" }, - manifest: { format: "ARROW_STREAM" }, - result: { - attachment: ARROW_ATTACHMENT_B64, - next_chunk_index: 1, - }, - }), + executeStatement: async () => ({ + statementId: "stmt-chunked", + status: { state: "SUCCEEDED" }, + manifest: { format: "ARROW_STREAM" }, + result: { + attachment: ARROW_ATTACHMENT_B64, + nextChunkIndex: 1, + }, + }), }, } as unknown as Parameters[0]; const fetcher = createWorkspaceDescribeFetcher(client, "wh-1"); @@ -1166,24 +1133,24 @@ describe("syncMetrics — failure transience (D′)", () => { [ "a FAILED statement", { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }, ], [ "a SUCCEEDED statement with zero rows", { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }, ], [ "an unparseable payload", { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [["{not json"]] }, + result: { dataArray: [["{not json"]] }, }, ], ["zero extracted columns", mockDescribeResponse({ unrelated: true })], @@ -1230,7 +1197,7 @@ describe("syncMetrics — DESCRIBE state classification", () => { test(`a non-terminal ${state} response degrades the schema without recording a failure`, async () => { const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state }, }); @@ -1254,7 +1221,7 @@ describe("syncMetrics — DESCRIBE state classification", () => { test("a FAILED response stays a genuine failure (and its schema is degraded)", async () => { const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); @@ -1273,9 +1240,9 @@ describe("syncMetrics — DESCRIBE state classification", () => { // wrong FQN, not warehouse readiness. const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }); const { schemas, failures } = await syncMetrics( @@ -1462,7 +1429,7 @@ describe("syncMetrics — bounded-concurrency scheduling", () => { throw new Error(`boom ${key}`); } if (key === nonTerminal) { - return { statement_id: "stmt-mock", status: { state: "PENDING" } }; + return { statementId: "stmt-mock", status: { state: "PENDING" } }; } return mockDescribeResponse({ columns: [ @@ -1634,7 +1601,7 @@ describe("generateMetricTypeDeclarations — snapshot", () => { ): Promise => fqn.endsWith("cold_metric") ? // Stopped/cold warehouse: wait_timeout elapsed → non-terminal, no rows. - { statement_id: "stmt-mock", status: { state: "PENDING" } } + { statementId: "stmt-mock", status: { state: "PENDING" } } : // Genuinely measure-less view: SUCCEEDED with dimension columns only. mockDescribeResponse({ columns: [{ name: "region", type: "STRING", is_measure: false }], @@ -1799,7 +1766,7 @@ describe("metric metadata bundle", () => { // Non-terminal DESCRIBE → degraded schema (empty column arrays). const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }); const { schemas } = await syncMetrics(resolution, fetcher); diff --git a/packages/appkit/src/type-generator/tests/query-registry.test.ts b/packages/appkit/src/type-generator/tests/query-registry.test.ts index 1f8156a2f..00d78e5fb 100644 --- a/packages/appkit/src/type-generator/tests/query-registry.test.ts +++ b/packages/appkit/src/type-generator/tests/query-registry.test.ts @@ -307,10 +307,10 @@ describe("defaultForType", () => { describe("convertToQueryType", () => { // DESCRIBE QUERY returns rows as [col_name, data_type, comment] const mockResponse: DatabricksStatementExecutionResponse = { - statement_id: "test-123", + statementId: "test-123", status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ ["id", "STRING", null], ["name", "STRING", null], ["count", "INT", null], @@ -373,10 +373,10 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` test("uses column comment when available", () => { const responseWithComment: DatabricksStatementExecutionResponse = { - statement_id: "test-123", + statementId: "test-123", status: { state: "SUCCEEDED" }, result: { - data_array: [["total", "DECIMAL", "Total amount in USD"]], + dataArray: [["total", "DECIMAL", "Total amount in USD"]], }, }; @@ -391,10 +391,10 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` test("quotes invalid column identifiers", () => { const responseWithInvalidName: DatabricksStatementExecutionResponse = { - statement_id: "test-123", + statementId: "test-123", status: { state: "SUCCEEDED" }, result: { - data_array: [["(1 = 1)", "BOOLEAN", null]], + dataArray: [["(1 = 1)", "BOOLEAN", null]], }, }; @@ -414,9 +414,9 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` test("returns hasResults: false when no columns exist", () => { const emptyResponse: DatabricksStatementExecutionResponse = { - statement_id: "test-123", + statementId: "test-123", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }; const { hasResults } = convertToQueryType( emptyResponse, @@ -439,7 +439,7 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` { col_name: "active", data_type: "BOOLEAN", comment: null }, ]); const response: DatabricksStatementExecutionResponse = { - statement_id: "test-arrow", + statementId: "test-arrow", status: { state: "SUCCEEDED" }, result: { attachment }, }; @@ -467,7 +467,7 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` { col_name: "id", data_type: "int", comment: null }, ]); const response: DatabricksStatementExecutionResponse = { - statement_id: "test-arrow", + statementId: "test-arrow", status: { state: "SUCCEEDED" }, result: { attachment }, }; @@ -477,15 +477,15 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` expect(type).toContain("id: number"); }); - test("prefers data_array over attachment when both are present", () => { + test("prefers dataArray over attachment when both are present", () => { const attachment = describeQueryAttachment([ { col_name: "from_arrow", data_type: "STRING", comment: null }, ]); const response: DatabricksStatementExecutionResponse = { - statement_id: "test-both", + statementId: "test-both", status: { state: "SUCCEEDED" }, result: { - data_array: [["from_data_array", "INT", null]], + dataArray: [["from_data_array", "INT", null]], attachment, }, }; @@ -498,7 +498,7 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` test("logs a warning and yields the unknown-result fallback on malformed attachment", () => { mockLoggerWarn.mockClear(); const response: DatabricksStatementExecutionResponse = { - statement_id: "test-bad", + statementId: "test-bad", status: { state: "SUCCEEDED" }, result: { attachment: "not-valid-arrow-ipc" }, }; diff --git a/packages/appkit/src/type-generator/tests/statement-result.test.ts b/packages/appkit/src/type-generator/tests/statement-result.test.ts index d545e49ce..1d1f43a6e 100644 --- a/packages/appkit/src/type-generator/tests/statement-result.test.ts +++ b/packages/appkit/src/type-generator/tests/statement-result.test.ts @@ -3,10 +3,7 @@ import path from "node:path"; import { describe, expect, test } from "vitest"; -import type { - StatementResponse, - WorkspaceClient, -} from "../../workspace-client"; +import type { WorkspaceClient } from "../../workspace-client"; import { type DescribeFormatMemo, describeAdaptive, @@ -46,9 +43,9 @@ const ARROW_REORDERED_FIELDS_B64 = fs.readFileSync( ); describe("normalizeResultRows", () => { - test("decodes an Arrow attachment into data_array (real fixture)", async () => { + test("decodes an Arrow attachment into dataArray (real fixture)", async () => { const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-arrow", + statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: ARROW_ATTACHMENT_B64 }, @@ -57,10 +54,10 @@ describe("normalizeResultRows", () => { const normalized = await normalizeResultRows(response); // One row, one cell — the JSON-string DESCRIBE payload. - expect(normalized.result?.data_array).toHaveLength(1); - expect(normalized.result?.data_array?.[0]).toHaveLength(1); + expect(normalized.result?.dataArray).toHaveLength(1); + expect(normalized.result?.dataArray?.[0]).toHaveLength(1); - const cell = normalized.result?.data_array?.[0]?.[0]; + const cell = normalized.result?.dataArray?.[0]?.[0]; expect(typeof cell).toBe("string"); // The real describe doc parses to an object with a non-empty `columns` array. @@ -69,9 +66,9 @@ describe("normalizeResultRows", () => { expect(parsed.columns.length).toBeGreaterThan(0); }); - test("preserves status, statement_id, and manifest when decoding", async () => { + test("preserves status, statementId, and manifest when decoding", async () => { const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-arrow", + statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: ARROW_ATTACHMENT_B64 }, @@ -79,46 +76,46 @@ describe("normalizeResultRows", () => { const normalized = await normalizeResultRows(response); - expect(normalized.statement_id).toBe("stmt-arrow"); + expect(normalized.statementId).toBe("stmt-arrow"); expect(normalized.status.state).toBe("SUCCEEDED"); expect(normalized.manifest?.format).toBe("ARROW_STREAM"); - // The attachment is left in place; only data_array is added. + // The attachment is left in place; only dataArray is added. expect(normalized.result?.attachment).toBe(ARROW_ATTACHMENT_B64); }); - test("passes through unchanged when data_array is already present", async () => { + test("passes through unchanged when dataArray is already present", async () => { // JSON_ARRAY warehouses (and every mocked test) take this path: no decode. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-json", + statementId: "stmt-json", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY" }, - result: { data_array: [['{"columns":[]}']] }, + result: { dataArray: [['{"columns":[]}']] }, }; const normalized = await normalizeResultRows(response); expect(normalized).toBe(response); - expect(normalized.result?.data_array).toEqual([['{"columns":[]}']]); + expect(normalized.result?.dataArray).toEqual([['{"columns":[]}']]); }); - test("treats an empty data_array as present (genuine no-rows, no decode)", async () => { + test("treats an empty dataArray as present (genuine no-rows, no decode)", async () => { // An empty array is a real "no rows" answer — it must not be overwritten by // an attachment decode even if an attachment is somehow also present. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-empty", + statementId: "stmt-empty", status: { state: "SUCCEEDED" }, - result: { data_array: [], attachment: ARROW_ATTACHMENT_B64 }, + result: { dataArray: [], attachment: ARROW_ATTACHMENT_B64 }, }; const normalized = await normalizeResultRows(response); expect(normalized).toBe(response); - expect(normalized.result?.data_array).toEqual([]); + expect(normalized.result?.dataArray).toEqual([]); }); - test("returns response unchanged when neither data_array nor attachment is present", async () => { + test("returns response unchanged when neither dataArray nor attachment is present", async () => { const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-bare", + statementId: "stmt-bare", status: { state: "SUCCEEDED" }, result: {}, }; @@ -126,12 +123,12 @@ describe("normalizeResultRows", () => { const normalized = await normalizeResultRows(response); expect(normalized).toBe(response); - expect(normalized.result?.data_array).toBeUndefined(); + expect(normalized.result?.dataArray).toBeUndefined(); }); test("returns response unchanged when result is entirely absent", async () => { const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-noresult", + statementId: "stmt-noresult", status: { state: "RUNNING" }, }; @@ -144,13 +141,13 @@ describe("normalizeResultRows", () => { test("does not throw when Arrow decoding rejects; degrades to no usable rows", async () => { // Bytes that look like an Arrow IPC header but aren't make `tableFromIPC` // throw. The decoder must swallow that so the generation pass does not - // crash — it leaves data_array absent and the downstream "returned no + // crash — it leaves dataArray absent and the downstream "returned no // rows" degrade fires instead. const notArrow = Buffer.from( "hello world this is plainly not an arrow ipc stream", ).toString("base64"); const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-corrupt", + statementId: "stmt-corrupt", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: notArrow }, @@ -163,19 +160,19 @@ describe("normalizeResultRows", () => { })(), ).resolves.toBeUndefined(); - // No fabricated rows: decode rejected, so data_array stays absent. - expect(normalized.result?.data_array).toBeUndefined(); + // No fabricated rows: decode rejected, so dataArray stays absent. + expect(normalized.result?.dataArray).toBeUndefined(); // The (bad) attachment is preserved; nothing was invented. expect(normalized.result?.attachment).toBe(notArrow); }); - test("decodes garbage that yields an empty Arrow table to an empty data_array", async () => { + test("decodes garbage that yields an empty Arrow table to an empty dataArray", async () => { // Some malformed payloads decode without throwing into a zero-row table - // (e.g. truncated/garbage bytes). That surfaces as an empty data_array — + // (e.g. truncated/garbage bytes). That surfaces as an empty dataArray — // which is itself a valid "no rows" answer and degrades correctly // downstream, never a fabricated row. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-garbage", + statementId: "stmt-garbage", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: "not-valid-base64-arrow-ipc!!!" }, @@ -185,58 +182,57 @@ describe("normalizeResultRows", () => { // Either absent or empty — both mean "no usable rows". Crucially: no // non-empty fabricated row. - expect(normalized.result?.data_array ?? []).toHaveLength(0); + expect(normalized.result?.dataArray ?? []).toHaveLength(0); }); - test("throws on a multi-chunk result flagged by next_chunk_index", async () => { + test("throws on a multi-chunk result flagged by nextChunkIndex", async () => { // A DESCRIBE result that exceeds INLINE's size limit is paginated. The - // first chunk carries `next_chunk_index`; decoding it alone would silently + // first chunk carries `nextChunkIndex`; decoding it alone would silently // cache partial types. The normalizer must throw (loud) rather than degrade // — distinct from the malformed-attachment path which degrades silently. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-chunked-json", + statementId: "stmt-chunked-json", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY" }, - // data_array present (first chunk) — but the guard runs ABOVE the + // dataArray present (first chunk) — but the guard runs ABOVE the // passthrough, so truncation still throws instead of returning rows. result: { - data_array: [["col_a", "STRING", null]], - next_chunk_index: 1, + dataArray: [["col_a", "STRING", null]], + nextChunkIndex: 1, }, }; await expect(normalizeResultRows(response)).rejects.toThrow(/multi-chunk/i); await expect(normalizeResultRows(response)).rejects.toThrow( - /next_chunk_index/, + /nextChunkIndex/, ); }); - test("throws on a multi-chunk result flagged by next_chunk_internal_link", async () => { + test("throws on a multi-chunk result flagged by nextChunkInternalLink", async () => { // The attachment transport can paginate too: first chunk arrives as an - // Arrow attachment with `next_chunk_internal_link` set. The guard runs + // Arrow attachment with `nextChunkInternalLink` set. The guard runs // before the decode, so this throws rather than emitting first-chunk types. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-chunked-arrow", + statementId: "stmt-chunked-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: ARROW_ATTACHMENT_B64, - next_chunk_internal_link: - "/api/2.0/sql/statements/stmt/result/chunks/1", + nextChunkInternalLink: "/api/2.0/sql/statements/stmt/result/chunks/1", }, }; await expect(normalizeResultRows(response)).rejects.toThrow(/multi-chunk/i); }); - test("throws on a multi-chunk result with neither data_array nor attachment", async () => { + test("throws on a multi-chunk result with neither dataArray nor attachment", async () => { // Even when the first chunk somehow carries no inline rows, the chunk // markers alone mean the answer is truncated — refuse, do not fall through // to the "no rows" degrade. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-chunked-bare", + statementId: "stmt-chunked-bare", status: { state: "SUCCEEDED" }, - result: { next_chunk_index: 2 }, + result: { nextChunkIndex: 2 }, }; await expect(normalizeResultRows(response)).rejects.toThrow( @@ -252,7 +248,7 @@ describe("normalizeResultRows", () => { // scrambling the [col_name, data_type, comment] triple. The `[...row]` // iterator preserves field order. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-reordered", + statementId: "stmt-reordered", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: ARROW_REORDERED_FIELDS_B64 }, @@ -260,7 +256,7 @@ describe("normalizeResultRows", () => { const normalized = await normalizeResultRows(response); - expect(normalized.result?.data_array).toEqual([ + expect(normalized.result?.dataArray).toEqual([ ["revenue", "DOUBLE", "total revenue"], ]); }); @@ -273,40 +269,17 @@ describe("describeAdaptive", () => { | DatabricksStatementExecutionResponse | Promise; - // Adapt a local snake_case fixture to the modular SDK's camelCase - // StatementResponse — the shape executeStatement now returns; describeAdaptive - // maps it back to the local shape via toDescribeResponse. - function asSdkResponse( - r: DatabricksStatementExecutionResponse, - ): StatementResponse { - return { - statementId: r.statement_id, - status: r.status && { - state: r.status.state, - error: r.status.error && { - errorCode: r.status.error.error_code, - message: r.status.error.message, - }, - }, - manifest: r.manifest && { format: r.manifest.format }, - result: r.result && { - dataArray: r.result.data_array, - attachment: r.result.attachment, - nextChunkIndex: r.result.next_chunk_index, - nextChunkInternalLink: r.result.next_chunk_internal_link, - }, - } as unknown as StatementResponse; - } - // Minimal WorkspaceClient stub: records the formats requested and delegates - // each executeStatement to behavior(format), which may resolve or throw. + // each executeStatement to behavior(format), which may resolve or throw. The + // fixtures are the camelCase domain type — the same shape executeStatement + // returns — so describeAdaptive consumes them directly (no adapter needed). function stubClient(behavior: StubBehavior) { const formats: string[] = []; const client = { statementExecution: { executeStatement: async (req: { format: string }) => { formats.push(req.format); - return asSdkResponse(await behavior(req.format)); + return await behavior(req.format); }, }, } as unknown as WorkspaceClient; @@ -316,9 +289,9 @@ describe("describeAdaptive", () => { const rows = ( data: (string | null)[][], ): DatabricksStatementExecutionResponse => ({ - statement_id: "stmt", + statementId: "stmt", status: { state: "SUCCEEDED" }, - result: { data_array: data }, + result: { dataArray: data }, }); test("standard DBSQL: JSON_ARRAY succeeds, memoized, no fallback", async () => { @@ -335,7 +308,7 @@ describe("describeAdaptive", () => { memo, ); - expect(result.result?.data_array).toEqual([["schema"]]); + expect(result.result?.dataArray).toEqual([["schema"]]); expect(memo.format).toBe("JSON_ARRAY"); expect(formats).toEqual(["JSON_ARRAY"]); }); @@ -356,7 +329,7 @@ describe("describeAdaptive", () => { memo, ); - expect(result.result?.data_array).toEqual([["arrow-decoded"]]); + expect(result.result?.dataArray).toEqual([["arrow-decoded"]]); expect(memo.format).toBe("ARROW_STREAM"); expect(formats).toEqual(["JSON_ARRAY", "ARROW_STREAM"]); }); @@ -366,7 +339,7 @@ describe("describeAdaptive", () => { const { client, formats } = stubClient((format) => { if (format === "JSON_ARRAY") { return { - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { message: "merge_json_arrays" } }, result: {}, } as DatabricksStatementExecutionResponse; @@ -381,7 +354,7 @@ describe("describeAdaptive", () => { memo, ); - expect(result.result?.data_array).toEqual([["arrow-decoded"]]); + expect(result.result?.dataArray).toEqual([["arrow-decoded"]]); expect(memo.format).toBe("ARROW_STREAM"); expect(formats).toEqual(["JSON_ARRAY", "ARROW_STREAM"]); }); @@ -403,7 +376,7 @@ describe("describeAdaptive", () => { const { client, formats } = stubClient((format) => { if (format === "JSON_ARRAY") { return { - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { message: "[TABLE_OR_VIEW_NOT_FOUND]" }, @@ -438,11 +411,11 @@ describe("describeAdaptive", () => { const { client, formats } = stubClient((format) => { if (format === "JSON_ARRAY") { return { - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { - error_code: "TABLE_OR_VIEW_NOT_FOUND", + errorCode: "TABLE_OR_VIEW_NOT_FOUND", message: "table x has no disposition column; format unknown", }, }, @@ -461,7 +434,7 @@ describe("describeAdaptive", () => { // The real diagnostic survives unmasked, and no second format was probed. expect(result.status.state).toBe("FAILED"); - expect(result.status.error?.error_code).toBe("TABLE_OR_VIEW_NOT_FOUND"); + expect(result.status.error?.errorCode).toBe("TABLE_OR_VIEW_NOT_FOUND"); expect(memo.format).toBeUndefined(); expect(formats).toEqual(["JSON_ARRAY"]); }); @@ -474,11 +447,11 @@ describe("describeAdaptive", () => { const { client, formats } = stubClient((format) => { if (format === "JSON_ARRAY") { return { - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { - error_code: "INVALID_PARAMETER_VALUE", + errorCode: "INVALID_PARAMETER_VALUE", message: "disposition must be one of INLINE, EXTERNAL_LINKS; format must be JSON_ARRAY, ARROW_STREAM", }, @@ -496,7 +469,7 @@ describe("describeAdaptive", () => { memo, ); - expect(result.result?.data_array).toEqual([["arrow-decoded"]]); + expect(result.result?.dataArray).toEqual([["arrow-decoded"]]); expect(memo.format).toBe("ARROW_STREAM"); expect(formats).toEqual(["JSON_ARRAY", "ARROW_STREAM"]); }); diff --git a/packages/appkit/src/type-generator/tests/sync-metric-views-types.test.ts b/packages/appkit/src/type-generator/tests/sync-metric-views-types.test.ts index a0a4e9995..27a2580b9 100644 --- a/packages/appkit/src/type-generator/tests/sync-metric-views-types.test.ts +++ b/packages/appkit/src/type-generator/tests/sync-metric-views-types.test.ts @@ -63,9 +63,9 @@ function mockDescribeResponse( payload: unknown, ): DatabricksStatementExecutionResponse { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [[JSON.stringify(payload)]] }, + result: { dataArray: [[JSON.stringify(payload)]] }, }; } @@ -232,7 +232,7 @@ describe("syncMetricViewsTypes", () => { mode: "blocking", suppressDegradedWrite: true, metricFetcher: async () => ({ - statement_id: "stmt-pending", + statementId: "stmt-pending", status: { state: "PENDING" }, }), }); diff --git a/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts b/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts index ef0b81519..0ffc0ed71 100644 --- a/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts +++ b/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts @@ -140,7 +140,7 @@ describe("--wait gate: environmental query failures (real query path)", () => { test("non-terminal DESCRIBE + no committed types → crashes instead of silently exiting 0", async () => { mocks.getWarehouse.mockResolvedValue({ state: "RUNNING" }); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-pending", + statementId: "stmt-pending", status: { state: "PENDING" }, }); @@ -162,7 +162,7 @@ describe("--wait gate: environmental query failures (real query path)", () => { test("non-terminal DESCRIBE + committed types → warns unavailable and keeps them", async () => { mocks.getWarehouse.mockResolvedValue({ state: "RUNNING" }); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-pending", + statementId: "stmt-pending", status: { state: "RUNNING" }, }); fs.mkdirSync(path.dirname(outFile), { recursive: true }); diff --git a/packages/appkit/src/type-generator/types.ts b/packages/appkit/src/type-generator/types.ts index 954bde706..124a990c5 100644 --- a/packages/appkit/src/type-generator/types.ts +++ b/packages/appkit/src/type-generator/types.ts @@ -2,34 +2,39 @@ * Databricks statement execution response interface for DESCRIBE QUERY / * DESCRIBE TABLE EXTENDED. * + * A hand-written camelCase subset of the modular SDK's `StatementResponse` — + * only the fields the type generators read. `describeAdaptive` narrows the SDK + * response into this type directly (the sole difference is `dataArray`, which + * the SDK types as `JsonValue[][]`; DESCRIBE cells are always string/null). + * * Two result shapes matter here: - * - `result.data_array` — rows already materialized as JSON arrays. Present + * - `result.dataArray` — rows already materialized as JSON arrays. Present * when the warehouse returns `JSON_ARRAY` (and what every mocked test * builds). * - `result.attachment` — a base64-encoded Arrow IPC stream. Present when the * statement runs with `format: "ARROW_STREAM"` + `disposition: "INLINE"`, * which is the SDK's default disposition. The single row lands here and - * `data_array` is left undefined. {@link normalizeResultRows} decodes this - * back into `data_array` so downstream parsers stay shape-agnostic. + * `dataArray` is left undefined. {@link normalizeResultRows} decodes this + * back into `dataArray` so downstream parsers stay shape-agnostic. * - * @property statement_id - the id of the statement + * @property statementId - the id of the statement * @property status - the status of the statement * @property manifest - result metadata; `manifest.format` echoes the wire * format (`ARROW_STREAM`, `JSON_ARRAY`, ...) the warehouse chose. - * @property result - the result; either `data_array` (rows as + * @property result - the result; either `dataArray` (rows as * `[col_name, data_type, comment]` arrays) or `attachment` (base64 Arrow IPC) */ export interface DatabricksStatementExecutionResponse { - statement_id: string; + statementId: string; status: { state: string; - error?: { error_code?: string; message?: string }; + error?: { errorCode?: string; message?: string }; }; manifest?: { format?: string; }; result?: { - data_array?: (string | null)[][]; + dataArray?: (string | null)[][]; /** Base64-encoded Arrow IPC stream (ARROW_STREAM + INLINE disposition). */ attachment?: string; /** @@ -37,9 +42,9 @@ export interface DatabricksStatementExecutionResponse { * limit). Its presence means this response holds only the FIRST chunk; * {@link normalizeResultRows} throws rather than emit truncated types. */ - next_chunk_index?: number; - /** Companion to {@link next_chunk_index}: link to fetch the next chunk. */ - next_chunk_internal_link?: string; + nextChunkIndex?: number; + /** Companion to {@link nextChunkIndex}: link to fetch the next chunk. */ + nextChunkInternalLink?: string; }; } From 00647042e1568cf412c04c8e25f0831653b3b39a Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 15:22:46 +0200 Subject: [PATCH 04/17] fix(appkit): point #540 never-crash mock at the modular warehouse API Rebasing the analytics SDK migration onto #540 (createTestApp + never-crash mock client) surfaced that #540's mock hardcodes the LEGACY warehouse method names. Point them at the modular API so warehouse-readiness calls resolve: - mock-workspace-client.ts DEFAULT_RESPONSES: warehouses.get/start -> getWarehouse/startWarehouse. - mock-workspace-client.test.ts: never-crash table + convergence asserts. (fixtures.ts + analytics.integration.test.ts adaptations landed inside the migration commit during rebase conflict resolution.) Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/src/testing/mock-workspace-client.ts | 9 +++++---- .../src/testing/tests/mock-workspace-client.test.ts | 10 ++++++---- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/packages/appkit/src/testing/mock-workspace-client.ts b/packages/appkit/src/testing/mock-workspace-client.ts index 7b9298379..59ec97635 100644 --- a/packages/appkit/src/testing/mock-workspace-client.ts +++ b/packages/appkit/src/testing/mock-workspace-client.ts @@ -43,8 +43,9 @@ export type MockWorkspaceClient = WorkspaceClient; /** * Applied beneath caller-supplied `responses`. * - * `statementExecution.executeStatement`, `warehouses.get` and `warehouses.start` - * must stay byte-identical to the old `fixtures.ts` values — suites reach them + * `statementExecution.executeStatement`, `warehouses.getWarehouse` and + * `warehouses.startWarehouse` must stay byte-identical to the old `fixtures.ts` + * values — suites reach them * implicitly through `mockServiceContext`. `currentUser.me` is * required: `ServiceContext.createContext` reads `.id`, so `createApp({ client })` * cannot boot without it. @@ -54,8 +55,8 @@ const DEFAULT_RESPONSES: Record = { status: { state: "SUCCEEDED" }, result: { data: [] }, }, - "warehouses.get": { state: "RUNNING" }, - "warehouses.start": undefined, + "warehouses.getWarehouse": { state: "RUNNING" }, + "warehouses.startWarehouse": undefined, "currentUser.me": { id: "test-service-user", userName: "test-service-user", diff --git a/packages/appkit/src/testing/tests/mock-workspace-client.test.ts b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts index a2cbeda1a..5ec90a0d9 100644 --- a/packages/appkit/src/testing/tests/mock-workspace-client.test.ts +++ b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts @@ -24,8 +24,8 @@ describe("createMockWorkspaceClient", () => { ["genie", "getMessage", undefined], ["jobs", "getRun", undefined], ["servingEndpoints", "get", undefined], - ["warehouses", "get", { state: "RUNNING" }], - ["warehouses", "start", undefined], + ["warehouses", "getWarehouse", { state: "RUNNING" }], + ["warehouses", "startWarehouse", undefined], ["statementExecution", "executeStatement", SUCCEEDED], ["currentUser", "me", TEST_USER], ])("%s.%s resolves its default", async (service, method, expected) => { @@ -248,11 +248,13 @@ describe("createMockWorkspaceClient", () => { await expect( client.statementExecution.executeStatement({} as never), ).resolves.toEqual(SUCCEEDED); - await expect(client.warehouses.get({} as never)).resolves.toEqual({ + await expect( + client.warehouses.getWarehouse({} as never), + ).resolves.toEqual({ state: "RUNNING", }); await expect( - client.warehouses.start({} as never), + client.warehouses.startWarehouse({} as never), ).resolves.toBeUndefined(); }); From 8604199e0c616b5fb8aa908a836b22a38c33a549 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 15:51:53 +0200 Subject: [PATCH 05/17] fix(appkit): declare modular @databricks/sdk-* packages as dependencies appkit bundles `shared` inline, whose code imports the modular @databricks/sdk-* packages, but only @databricks/sdk-experimental was declared. A published `@databricks/appkit` install (npm, incl. the Databricks Apps runtime) would therefore fail at runtime with "Cannot find module @databricks/sdk-statementexecution" on the analytics path. Declare sdk-auth/core/options/warehouses/statementexecution at 0.46.0 (mirroring shared) and exempt them from knip's unused check like sdk-experimental (they are imported only by the inlined `shared` workspace, which knip does not analyze). Known follow-up: the pnpm patch restoring Reyden's stripped `attachment` field is a workspace-only mechanism, does not reach consumers, and cannot be bundled under tsdown `unbundle` mode. Standard warehouses (dataArray + external_links) are unaffected; Reyden INLINE+ARROW_STREAM needs the upstream SDK `attachment` fix. See SDK_MIGRATION_GAPS.md. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- knip.json | 5 +++++ packages/appkit/package.json | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/knip.json b/knip.json index 02320c52b..21f7b9ca0 100644 --- a/knip.json +++ b/knip.json @@ -10,7 +10,12 @@ "packages/appkit": { "ignoreDependencies": [ "vitest", + "@databricks/sdk-auth", + "@databricks/sdk-core", "@databricks/sdk-experimental", + "@databricks/sdk-options", + "@databricks/sdk-statementexecution", + "@databricks/sdk-warehouses", "@mlflow/core", "autoevals" ] diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 488c98da2..d6fbb9daa 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -71,7 +71,12 @@ "dependencies": { "@ast-grep/napi": "0.37.0", "@databricks/lakebase": "workspace:*", + "@databricks/sdk-auth": "0.46.0", + "@databricks/sdk-core": "0.46.0", "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-options": "0.46.0", + "@databricks/sdk-statementexecution": "0.46.0", + "@databricks/sdk-warehouses": "0.46.0", "@opentelemetry/api": "1.9.0", "@opentelemetry/api-logs": "0.219.0", "@opentelemetry/auto-instrumentations-node": "0.77.0", From 7b312527dfc99d73b39d9c1d524ffc345f4f64ef Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 18:41:24 +0200 Subject: [PATCH 06/17] fix(shared): resolve service-principal credentials from env in modular SDK client The modular @databricks/sdk-* default credential chain resolves auth only from a ~/.databrickscfg profile and reads no DATABRICKS_* env vars, unlike the legacy sdk-experimental. The Databricks Apps runtime injects the app's service-principal credentials via env vars only (no config file), so a deployed app built a modular client with no credentials and every request failed ("Warehouse readiness check failed"). Resolve the service principal from the environment in mapToClientOptions when no explicit token/profile is given: M2M from DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET (what Apps injects), else PAT from DATABRICKS_TOKEN, else fall through to the profile default chain for local dev. The explicit token path (asUser OBO) is unchanged and still guarded by token !== undefined so an empty/invalid OBO token fails loudly instead of silently falling through to the service principal. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../shared/src/workspace-client/modular.ts | 44 ++++++++-- .../workspace-client/tests/modular.test.ts | 84 +++++++++++++++++-- 2 files changed, 114 insertions(+), 14 deletions(-) diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index a9cdef32b..2714b8dbd 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -16,7 +16,10 @@ * undocumented Reyden `attachment` response field, which the SDK's generated * unmarshal transform would otherwise strip. */ -import { newPatCredentials } from "@databricks/sdk-auth/credentials"; +import { + newM2mCredentials, + newPatCredentials, +} from "@databricks/sdk-auth/credentials"; import { addToDefault, setProduct } from "@databricks/sdk-core/clientinfo"; import type { ClientOptions } from "@databricks/sdk-options/client"; import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; @@ -37,12 +40,13 @@ function normalizeHost(host: string | undefined): string | undefined { } /** - * Map wrapper options onto the modular SDK's `ClientOptions`. Mirrors - * `buildLegacyWorkspaceClient`'s auth resolution verbatim, including the - * privilege-escalation guard: check `token !== undefined` (NOT truthiness) so an - * explicitly-passed token — even an empty string — pins the PAT path and fails - * loudly at request time rather than silently authenticating as the service - * principal via the default chain (which would be an OBO privilege escalation). + * Map wrapper options onto the modular SDK's `ClientOptions`, reproducing the + * legacy SDK's auth resolution: explicit token → PAT (the OBO path); profile → + * profile file; otherwise the service principal from the environment. It carries + * the privilege-escalation guard — check `token !== undefined` (NOT truthiness) + * so an explicitly-passed token, even an empty string, pins the PAT path and + * fails loudly at request time rather than silently falling through to the + * service-principal env credentials (which would be an OBO privilege escalation). */ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { const clientOptions: ClientOptions = {}; @@ -57,12 +61,34 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { clientOptions.host = host; } if (opts.token !== undefined) { + // Explicit token (this is the OBO path: `asUser` passes the user's token). clientOptions.credentials = newPatCredentials(opts.token); } else if (opts.profile) { clientOptions.profileOptions = { profile: opts.profile }; + } else { + // No token, no profile: authenticate as the service principal from the + // environment, the way the legacy SDK did. The modular SDK's default auth + // chain resolves ONLY from a `~/.databrickscfg` profile — it reads no + // `DATABRICKS_*` env vars — so on the Databricks Apps runtime (which injects + // the app's SP credentials via env, with no config file) it would find no + // credentials and every request would fail. Resolve them here instead: + // M2M (client id + secret, what Apps injects) first, then a PAT, else fall + // through to the default chain for local dev with a config file. + const clientId = process.env.DATABRICKS_CLIENT_ID; + const clientSecret = process.env.DATABRICKS_CLIENT_SECRET; + const envToken = process.env.DATABRICKS_TOKEN; + if (host && clientId && clientSecret) { + clientOptions.credentials = newM2mCredentials({ + host, + clientId, + clientSecret, + }); + } else if (envToken) { + clientOptions.credentials = newPatCredentials(envToken); + } + // Otherwise leave credentials unset and let the SDK walk its profile-based + // default chain (local dev with `~/.databrickscfg`). } - // Neither token nor profile → leave credentials unset so the SDK walks its - // default auth chain (env vars + ~/.databrickscfg), matching the legacy `{}` case. return clientOptions; } diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index d6d092bcd..07af9ef07 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -3,9 +3,10 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; // The wrapper's own tests are the one place allowed to mock the SDK directly. // Capture the `ClientOptions` the modular `WarehousesClient` constructor receives // so we can assert how wrapper options map onto the modular SDK's config. -const { ctorOpts, patTokens, productCalls } = vi.hoisted(() => ({ +const { ctorOpts, patTokens, m2mOpts, productCalls } = vi.hoisted(() => ({ ctorOpts: [] as Array>, patTokens: [] as string[], + m2mOpts: [] as Array>, productCalls: [] as Array<[string, string]>, })); @@ -23,6 +24,10 @@ vi.mock("@databricks/sdk-auth/credentials", () => ({ patTokens.push(token); return { kind: "pat", token }; }), + newM2mCredentials: vi.fn((opts: Record) => { + m2mOpts.push(opts); + return { kind: "m2m", ...opts }; + }), })); vi.mock("@databricks/sdk-core/clientinfo", () => ({ setProduct: vi.fn((name: string, version: string) => { @@ -38,18 +43,32 @@ vi.mock("@databricks/sdk-core/clientinfo", () => ({ import { buildWarehousesClient } from "../modular"; describe("modular mapToClientOptions (via buildWarehousesClient)", () => { - const originalHost = process.env.DATABRICKS_HOST; + // Auth resolution reads these env vars; snapshot + clear them so the dev + // machine's own DATABRICKS_* values never leak into a case. + const AUTH_ENV = [ + "DATABRICKS_HOST", + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CLIENT_SECRET", + "DATABRICKS_TOKEN", + ] as const; + const originalEnv: Record = {}; beforeEach(() => { ctorOpts.length = 0; patTokens.length = 0; + m2mOpts.length = 0; productCalls.length = 0; - delete process.env.DATABRICKS_HOST; + for (const key of AUTH_ENV) { + originalEnv[key] = process.env[key]; + delete process.env[key]; + } }); afterEach(() => { - if (originalHost === undefined) delete process.env.DATABRICKS_HOST; - else process.env.DATABRICKS_HOST = originalHost; + for (const key of AUTH_ENV) { + if (originalEnv[key] === undefined) delete process.env[key]; + else process.env[key] = originalEnv[key]; + } }); test("prepends https:// to a scheme-less explicit host", () => { @@ -96,6 +115,61 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(ctorOpts[0].profileOptions).toBeUndefined(); }); + test("service-principal by default: DATABRICKS_CLIENT_ID/SECRET + host env → M2M creds", () => { + // The Databricks Apps runtime injects the app's SP credentials this way + // (env only, no config file). The modular SDK's default chain reads no env, + // so we must map them to M2M credentials ourselves. + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + buildWarehousesClient({}); + expect(m2mOpts).toEqual([ + { + host: "https://envhost.cloud.databricks.com", + clientId: "sp-client-id", + clientSecret: "sp-secret", + }, + ]); + expect(ctorOpts[0].credentials).toEqual({ + kind: "m2m", + host: "https://envhost.cloud.databricks.com", + clientId: "sp-client-id", + clientSecret: "sp-secret", + }); + expect(patTokens).toEqual([]); + }); + + test("falls back to DATABRICKS_TOKEN (PAT) when no client id/secret is set", () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + process.env.DATABRICKS_TOKEN = "env-pat"; + buildWarehousesClient({}); + expect(patTokens).toEqual(["env-pat"]); + expect(ctorOpts[0].credentials).toEqual({ kind: "pat", token: "env-pat" }); + expect(m2mOpts).toEqual([]); + }); + + test("an explicit (OBO) token wins over env SP credentials — no escalation", () => { + // asUser passes the user's token; it must NOT be shadowed by the SP env + // creds the deployed runtime also sets. + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + buildWarehousesClient({ token: "user-token", host: "https://x" }); + expect(patTokens).toEqual(["user-token"]); + expect(ctorOpts[0].credentials).toEqual({ + kind: "pat", + token: "user-token", + }); + expect(m2mOpts).toEqual([]); + }); + + test("M2M needs a host: client id/secret with no resolvable host falls through to the default chain", () => { + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + buildWarehousesClient({}); + expect(m2mOpts).toEqual([]); + expect(ctorOpts[0].credentials).toBeUndefined(); + }); + test("client-info: sanitizes an invalid product name (e.g. @databricks/appkit) rather than crashing the client build", () => { // Regression: the modular SDK's `setProduct` rejects `@databricks/appkit` // (INVALID_KEY), which the legacy SDK accepted. UA stamping must be From 1209120a1a27894285435adcd7129ac12ce6034e Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 18:48:15 +0200 Subject: [PATCH 07/17] chore(deps): sync pnpm-lock with appkit modular SDK dependencies The 5 modular @databricks/sdk-* deps were added to packages/appkit/package.json but the appkit importer in pnpm-lock.yaml was never regenerated, so CI's `pnpm install --frozen-lockfile` failed with ERR_PNPM_OUTDATED_LOCKFILE. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- pnpm-lock.yaml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 63cfb5da7..b58231fa1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -279,9 +279,24 @@ importers: '@databricks/lakebase': specifier: workspace:* version: link:../lakebase + '@databricks/sdk-auth': + specifier: 0.46.0 + version: 0.46.0 + '@databricks/sdk-core': + specifier: 0.46.0 + version: 0.46.0 '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-options': + specifier: 0.46.0 + version: 0.46.0 + '@databricks/sdk-statementexecution': + specifier: 0.46.0 + version: 0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d) + '@databricks/sdk-warehouses': + specifier: 0.46.0 + version: 0.46.0 '@opentelemetry/api': specifier: 1.9.0 version: 1.9.0 From 3599b2f91706159ff1508c8019d16d7a5ba0aa01 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 15 Sep 2026 14:49:53 +0200 Subject: [PATCH 08/17] fix(shared): preserve the @databricks/appkit User-Agent on modular clients MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The modular SDK's client-info `setProduct` validates the product as a token and rejects `@databricks/appkit` (the `@`/`/`), so the migrated clients sent a sanitized `databricks-appkit` User-Agent. Databricks-side dashboards filter AppKit's analytics/warehouse traffic on the literal `@databricks/appkit` UA, so that silently dropped AppKit out of them. Set the User-Agent on a per-client `ClientOptions.httpClient` transport wrapper (`buildHttpClient`) that prepends `@databricks/appkit/` (+ userAgentExtra) to each request, replacing the process-global `setProduct` sanitization/latch. This keeps the exact legacy string and ships inside appkit's bundled dist, so it reaches deployed apps — a pnpm patch of the validator would not (pnpm patches apply only at workspace install, not to the npm-installed tarball). Also corrects the service-principal auth comments: the SDK's default chain does read DATABRICKS_* env; the deployed failure was its M2M strategy using the raw scheme-less DATABRICKS_HOST for OAuth discovery, which the normalized-host resolution already works around. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../shared/src/workspace-client/modular.ts | 95 ++++++++++--------- .../workspace-client/tests/modular.test.ts | 81 ++++++++++++---- 2 files changed, 109 insertions(+), 67 deletions(-) diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index 2714b8dbd..44a31fbec 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -20,7 +20,7 @@ import { newM2mCredentials, newPatCredentials, } from "@databricks/sdk-auth/credentials"; -import { addToDefault, setProduct } from "@databricks/sdk-core/clientinfo"; +import { type HttpClient, newFetchHttpClient } from "@databricks/sdk-core/http"; import type { ClientOptions } from "@databricks/sdk-options/client"; import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; import { WarehousesClient } from "@databricks/sdk-warehouses/v1"; @@ -67,13 +67,16 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { clientOptions.profileOptions = { profile: opts.profile }; } else { // No token, no profile: authenticate as the service principal from the - // environment, the way the legacy SDK did. The modular SDK's default auth - // chain resolves ONLY from a `~/.databrickscfg` profile — it reads no - // `DATABRICKS_*` env vars — so on the Databricks Apps runtime (which injects - // the app's SP credentials via env, with no config file) it would find no - // credentials and every request would fail. Resolve them here instead: - // M2M (client id + secret, what Apps injects) first, then a PAT, else fall - // through to the default chain for local dev with a config file. + // environment. The SDK's own default chain DOES read the DATABRICKS_* env + // vars (host, client id/secret, token) — but its M2M strategy feeds the RAW + // `DATABRICKS_HOST` straight into OAuth token-endpoint discovery, and the + // Databricks Apps runtime sets that host scheme-less (e.g. + // `x.cloud.databricks.com`), so discovery fails with `Invalid URL` and every + // request dies as "Warehouse readiness check failed". Resolve the SP here + // with the scheme-normalized `host` instead: M2M from client id + secret + // (what Apps injects), else PAT from `DATABRICKS_TOKEN`, else fall through to + // the SDK default chain (local dev, where a `~/.databrickscfg` host already + // carries a scheme). const clientId = process.env.DATABRICKS_CLIENT_ID; const clientSecret = process.env.DATABRICKS_CLIENT_SECRET; const envToken = process.env.DATABRICKS_TOKEN; @@ -89,59 +92,60 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { // Otherwise leave credentials unset and let the SDK walk its profile-based // default chain (local dev with `~/.databrickscfg`). } + const httpClient = buildHttpClient(opts); + if (httpClient) { + clientOptions.httpClient = httpClient; + } return clientOptions; } -// The modular SDK has no per-client User-Agent option; product/client-info is a -// process-global set once via `setProduct`/`addToDefault` before any client is -// built. The AppKit product/version/userAgentExtra arrive on `opts.clientOptions` -// (from `getClientOptions()`); build-time callers omit them and are left unstamped, -// preserving the legacy behavior where build-time clients carry no AppKit UA. The -// flag latches only once we actually stamp, so a first (unstamped) build-time -// client never blocks a later runtime client from stamping. -let clientInfoStamped = false; - /** - * Coerce an arbitrary string into a valid client-info segment. The modular SDK - * validates keys as simple tokens and throws `ClientInfoError` on anything else, - * so the legacy product name `@databricks/appkit` (with `@` and `/`) is rejected - * — collapse invalid runs to `-` and trim the ends (`@databricks/appkit` → - * `databricks-appkit`). + * Wrap the SDK's default fetch transport to prepend AppKit's product segment to + * the outgoing `User-Agent`, preserving the exact legacy string (e.g. + * `@databricks/appkit/0.75.1`) that Databricks-side dashboards match on. + * + * Why the transport and not `setProduct`: the modular SDK's client-info API + * validates the product as a simple token and rejects `@databricks/appkit` (the + * `@`/`/`), and it is process-global. Setting the header on the `httpClient` + * instead keeps the literal product name, is per-client, and — unlike a pnpm + * patch — ships inside appkit's bundled `dist`, so it also reaches deployed apps + * (npm-installed from the tarball, where pnpm patches do not apply). The SDK's + * own client-info (`sdk-js-core/…`, runtime) is already on `request.headers`, so + * prepending keeps it intact after AppKit's segment. + * + * Returns `undefined` when no product is configured (build-time callers), leaving + * the SDK's default User-Agent untouched — matching the legacy behavior where + * build-time clients carried no AppKit UA. */ -function toClientInfoKey(value: string): string { - return value.replace(/[^A-Za-z0-9._-]+/g, "-").replace(/^-+|-+$/g, ""); -} - -function ensureClientInfo(opts: WorkspaceClientOptions): void { - if (clientInfoStamped) { - return; - } +function buildHttpClient(opts: WorkspaceClientOptions): HttpClient | undefined { const co = opts.clientOptions; if (!co?.product || !co?.productVersion) { - return; + return undefined; } - // User-Agent stamping is best-effort: a value the SDK's client-info validator - // rejects must NEVER break client construction (the legacy SDK stamped the UA - // without validating). On failure the outbound request just carries the SDK's - // default User-Agent. - try { - setProduct(toClientInfoKey(co.product), co.productVersion); - if (co.userAgentExtra) { - for (const [key, value] of Object.entries(co.userAgentExtra)) { - addToDefault(toClientInfoKey(key), String(value)); - } + const segments = [`${co.product}/${co.productVersion}`]; + if (co.userAgentExtra) { + for (const [key, value] of Object.entries(co.userAgentExtra)) { + segments.push(`${key}/${String(value)}`); } - clientInfoStamped = true; - } catch { - clientInfoStamped = true; } + const appkitUserAgent = segments.join(" "); + const base = newFetchHttpClient(); + return { + send(request) { + const existing = request.headers.get("User-Agent"); + request.headers.set( + "User-Agent", + existing ? `${appkitUserAgent} ${existing}` : appkitUserAgent, + ); + return base.send(request); + }, + }; } /** Build a modular Warehouses client from wrapper options. */ export function buildWarehousesClient( opts: WorkspaceClientOptions, ): WarehousesClient { - ensureClientInfo(opts); return new WarehousesClient(mapToClientOptions(opts)); } @@ -149,7 +153,6 @@ export function buildWarehousesClient( export function buildStatementExecutionClient( opts: WorkspaceClientOptions, ): StatementExecutionClient { - ensureClientInfo(opts); return new StatementExecutionClient(mapToClientOptions(opts)); } diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index 07af9ef07..51814a293 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -3,11 +3,10 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; // The wrapper's own tests are the one place allowed to mock the SDK directly. // Capture the `ClientOptions` the modular `WarehousesClient` constructor receives // so we can assert how wrapper options map onto the modular SDK's config. -const { ctorOpts, patTokens, m2mOpts, productCalls } = vi.hoisted(() => ({ +const { ctorOpts, patTokens, m2mOpts } = vi.hoisted(() => ({ ctorOpts: [] as Array>, patTokens: [] as string[], m2mOpts: [] as Array>, - productCalls: [] as Array<[string, string]>, })); vi.mock("@databricks/sdk-warehouses/v1", () => ({ @@ -29,19 +28,42 @@ vi.mock("@databricks/sdk-auth/credentials", () => ({ return { kind: "m2m", ...opts }; }), })); -vi.mock("@databricks/sdk-core/clientinfo", () => ({ - setProduct: vi.fn((name: string, version: string) => { - // Mirror the real SDK: reject client-info keys that aren't simple tokens. - if (/[^A-Za-z0-9._-]/.test(name)) { - throw new Error(`Invalid key: ${name}.`); - } - productCalls.push([name, version]); - }), - addToDefault: vi.fn(), +// The default transport: its `send` echoes the final request headers so tests +// can assert the User-Agent the wrapper set before delegating. +vi.mock("@databricks/sdk-core/http", () => ({ + newFetchHttpClient: vi.fn(() => ({ + send: vi.fn((request: { headers: Headers }) => + Promise.resolve({ + statusCode: 200, + headers: request.headers, + body: null, + }), + ), + })), })); import { buildWarehousesClient } from "../modular"; +/** Drive the wrapped httpClient with one request and return the UA it set. */ +async function sentUserAgent( + httpClient: unknown, + seedUserAgent?: string, +): Promise { + const headers = new Headers( + seedUserAgent ? { "User-Agent": seedUserAgent } : undefined, + ); + await ( + httpClient as { + send: (r: { + url: string; + method: string; + headers: Headers; + }) => Promise; + } + ).send({ url: "https://x", method: "GET", headers }); + return headers.get("User-Agent"); +} + describe("modular mapToClientOptions (via buildWarehousesClient)", () => { // Auth resolution reads these env vars; snapshot + clear them so the dev // machine's own DATABRICKS_* values never leak into a case. @@ -57,7 +79,6 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { ctorOpts.length = 0; patTokens.length = 0; m2mOpts.length = 0; - productCalls.length = 0; for (const key of AUTH_ENV) { originalEnv[key] = process.env[key]; delete process.env[key]; @@ -116,9 +137,10 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { }); test("service-principal by default: DATABRICKS_CLIENT_ID/SECRET + host env → M2M creds", () => { - // The Databricks Apps runtime injects the app's SP credentials this way - // (env only, no config file). The modular SDK's default chain reads no env, - // so we must map them to M2M credentials ourselves. + // The Databricks Apps runtime injects the app's SP credentials via env. The + // SDK's default chain would read them too, but its M2M strategy uses the raw + // scheme-less DATABRICKS_HOST for OAuth discovery (→ Invalid URL); we resolve + // M2M here with the scheme-normalized host so discovery succeeds. process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; @@ -170,18 +192,35 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(ctorOpts[0].credentials).toBeUndefined(); }); - test("client-info: sanitizes an invalid product name (e.g. @databricks/appkit) rather than crashing the client build", () => { + test("User-Agent: prepends the exact @databricks/appkit product segment (dashboards match on it)", async () => { // Regression: the modular SDK's `setProduct` rejects `@databricks/appkit` - // (INVALID_KEY), which the legacy SDK accepted. UA stamping must be - // best-effort — a bad product string must never break client construction. - const client = buildWarehousesClient({ + // (the `@`/`/`). We set the UA on the httpClient transport instead, keeping + // the literal legacy product string that Databricks-side dashboards match. + buildWarehousesClient({ clientOptions: { product: "@databricks/appkit", productVersion: "0.64.0", userAgentExtra: { mode: "dev" }, }, } as never); - expect(client).toBeDefined(); - expect(productCalls[0]).toEqual(["databricks-appkit", "0.64.0"]); + // The SDK's own client-info UA is already on the request; ours prepends. + const ua = await sentUserAgent(ctorOpts[0].httpClient, "sdk-js-core/1.0.0"); + expect(ua).toBe("@databricks/appkit/0.64.0 mode/dev sdk-js-core/1.0.0"); + }); + + test("User-Agent: sets the product segment even when the request has no prior UA", async () => { + buildWarehousesClient({ + clientOptions: { + product: "@databricks/appkit", + productVersion: "0.64.0", + }, + } as never); + const ua = await sentUserAgent(ctorOpts[0].httpClient); + expect(ua).toBe("@databricks/appkit/0.64.0"); + }); + + test("no product configured (build-time) → no httpClient override (SDK default UA)", () => { + buildWarehousesClient({ host: "https://x" }); + expect(ctorOpts[0].httpClient).toBeUndefined(); }); }); From 7466b81d456529afe9bc9d326b6902f9c0ddcf16 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 7 Oct 2026 12:53:31 +0200 Subject: [PATCH 09/17] docs(shared): point the legacy sql type note at resources/warehouse.ts The dev-mode warehouse discovery moved from service-context to resources/warehouse.ts (#592). Note why it stays on the legacy apiClient: the modular listWarehouses request has no skip_cannot_use filter. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/shared/src/workspace-client/types.ts | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index 6d9865ace..87071b004 100644 --- a/packages/shared/src/workspace-client/types.ts +++ b/packages/shared/src/workspace-client/types.ts @@ -18,10 +18,12 @@ import type { StatementExecutionClient, WarehousesClient } from "./modular"; // Legacy SDK type namespaces for un-migrated services, re-exported so AppKit // modules import them from the wrapper rather than the SDK directly. `sql` -// stays only for the dev-mode warehouse listing in service-context, which reads -// the raw (snake_case) `/api/2.0/sql/warehouses` body via the still-legacy -// `apiClient` and types it as `sql.EndpointInfo[]`. Statement + warehouse -// service types now come from `./modular`. +// stays only for the dev-mode warehouse discovery in appkit's +// `resources/warehouse.ts`, which reads the raw (snake_case) +// `/api/2.0/sql/warehouses` body via the still-legacy `apiClient` and types it +// as `sql.EndpointInfo[]`. It is not on the modular `listWarehouses` because +// that request has no `skip_cannot_use` filter, so it could pick a warehouse +// the caller can't use. Statement + warehouse service types come from `./modular`. export type { files, jobs, serving, sql } from "@databricks/sdk-experimental"; // Modular SDK client + model types (warehouses, statementExecution). export type * from "./modular"; From 2111d766d63161eef4ca8630013671e246c64623 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 7 Oct 2026 14:48:37 +0200 Subject: [PATCH 10/17] chore(deps): bump modular @databricks/sdk-* and re-create the attachment patch sdk-core/sdk-auth/sdk-options 0.46.0 -> 0.51.0, sdk-warehouses -> 0.53.0, sdk-statementexecution -> 0.52.0. The Reyden `attachment` field is still stripped by statementexecution 0.52.0, so the patch is re-created against that version (same three edits: schema, transform, ResultData type). These versions normalize scheme-less hosts in the request path and in OAuth discovery; AppKit's own normalizeHost/env-M2M resolution is kept for now. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/package.json | 10 +-- packages/shared/package.json | 10 +-- ...icks__sdk-statementexecution@0.52.0.patch} | 4 +- pnpm-lock.yaml | 90 +++++++++---------- pnpm-workspace.yaml | 2 +- 5 files changed, 58 insertions(+), 58 deletions(-) rename patches/{@databricks__sdk-statementexecution@0.46.0.patch => @databricks__sdk-statementexecution@0.52.0.patch} (90%) diff --git a/packages/appkit/package.json b/packages/appkit/package.json index d6fbb9daa..555f9da55 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -71,12 +71,12 @@ "dependencies": { "@ast-grep/napi": "0.37.0", "@databricks/lakebase": "workspace:*", - "@databricks/sdk-auth": "0.46.0", - "@databricks/sdk-core": "0.46.0", + "@databricks/sdk-auth": "0.51.0", + "@databricks/sdk-core": "0.51.0", "@databricks/sdk-experimental": "0.17.0", - "@databricks/sdk-options": "0.46.0", - "@databricks/sdk-statementexecution": "0.46.0", - "@databricks/sdk-warehouses": "0.46.0", + "@databricks/sdk-options": "0.51.0", + "@databricks/sdk-statementexecution": "0.52.0", + "@databricks/sdk-warehouses": "0.53.0", "@opentelemetry/api": "1.9.0", "@opentelemetry/api-logs": "0.219.0", "@opentelemetry/auto-instrumentations-node": "0.77.0", diff --git a/packages/shared/package.json b/packages/shared/package.json index 7dd98e48b..85d2a0672 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -48,12 +48,12 @@ "dependencies": { "@ast-grep/napi": "0.37.0", "@clack/prompts": "1.0.1", - "@databricks/sdk-auth": "0.46.0", - "@databricks/sdk-core": "0.46.0", + "@databricks/sdk-auth": "0.51.0", + "@databricks/sdk-core": "0.51.0", "@databricks/sdk-experimental": "0.17.0", - "@databricks/sdk-options": "0.46.0", - "@databricks/sdk-statementexecution": "0.46.0", - "@databricks/sdk-warehouses": "0.46.0", + "@databricks/sdk-options": "0.51.0", + "@databricks/sdk-statementexecution": "0.52.0", + "@databricks/sdk-warehouses": "0.53.0", "@standard-schema/spec": "1.1.0", "commander": "12.1.0", "dotenv": "16.6.1", diff --git a/patches/@databricks__sdk-statementexecution@0.46.0.patch b/patches/@databricks__sdk-statementexecution@0.52.0.patch similarity index 90% rename from patches/@databricks__sdk-statementexecution@0.46.0.patch rename to patches/@databricks__sdk-statementexecution@0.52.0.patch index c206b65c4..168212482 100644 --- a/patches/@databricks__sdk-statementexecution@0.46.0.patch +++ b/patches/@databricks__sdk-statementexecution@0.52.0.patch @@ -1,8 +1,8 @@ diff --git a/dist/v1/model.d.ts b/dist/v1/model.d.ts -index e8d95659ea348b384a3d32b6a3d4f754287b38b6..705b9bed203981a2f3cde5417ed8019ff3a7065c 100644 +index a43af9275af5fed230b9d428ba30b979e4d6aa61..3f6b302871a886158ad18d7d957439a229a8c549 100644 --- a/dist/v1/model.d.ts +++ b/dist/v1/model.d.ts -@@ -385,6 +385,8 @@ interface QueryTag { +@@ -384,6 +384,8 @@ interface QueryTag { * link is returned.) */ interface ResultData { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b58231fa1..929d9c01e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -12,7 +12,7 @@ overrides: size-sensor: 1.0.3 patchedDependencies: - '@databricks/sdk-statementexecution@0.46.0': a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d + '@databricks/sdk-statementexecution@0.52.0': 9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c importers: @@ -280,23 +280,23 @@ importers: specifier: workspace:* version: link:../lakebase '@databricks/sdk-auth': - specifier: 0.46.0 - version: 0.46.0 + specifier: 0.51.0 + version: 0.51.0 '@databricks/sdk-core': - specifier: 0.46.0 - version: 0.46.0 + specifier: 0.51.0 + version: 0.51.0 '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 '@databricks/sdk-options': - specifier: 0.46.0 - version: 0.46.0 + specifier: 0.51.0 + version: 0.51.0 '@databricks/sdk-statementexecution': - specifier: 0.46.0 - version: 0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d) + specifier: 0.52.0 + version: 0.52.0(patch_hash=9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c) '@databricks/sdk-warehouses': - specifier: 0.46.0 - version: 0.46.0 + specifier: 0.53.0 + version: 0.53.0 '@opentelemetry/api': specifier: 1.9.0 version: 1.9.0 @@ -610,23 +610,23 @@ importers: specifier: 1.0.1 version: 1.0.1 '@databricks/sdk-auth': - specifier: 0.46.0 - version: 0.46.0 + specifier: 0.51.0 + version: 0.51.0 '@databricks/sdk-core': - specifier: 0.46.0 - version: 0.46.0 + specifier: 0.51.0 + version: 0.51.0 '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 '@databricks/sdk-options': - specifier: 0.46.0 - version: 0.46.0 + specifier: 0.51.0 + version: 0.51.0 '@databricks/sdk-statementexecution': - specifier: 0.46.0 - version: 0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d) + specifier: 0.52.0 + version: 0.52.0(patch_hash=9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c) '@databricks/sdk-warehouses': - specifier: 0.46.0 - version: 0.46.0 + specifier: 0.53.0 + version: 0.53.0 '@standard-schema/spec': specifier: 1.1.0 version: 1.1.0 @@ -1992,12 +1992,12 @@ packages: engines: {node: ^20 || ^22 || ^24 || ^25, pnpm: '>=10'} hasBin: true - '@databricks/sdk-auth@0.46.0': - resolution: {integrity: sha512-cMrwxsFtpiEKFxta5dKHchKdrgmHkQ6upJ2C4OacmlHrOHJ+ChzQBVTpAiVtjX62xj+YjNgp/29IpSdKKYUVDA==} + '@databricks/sdk-auth@0.51.0': + resolution: {integrity: sha512-CswN+io9XB7rrcHo9BG/2qIdRCUMb8gFjiZmqVDG2qRmyhzY7Wv1tvwzA+U8uBplmdzAcdhuxMvHQ7vKGFWYqQ==} engines: {node: '>=22.0.0'} - '@databricks/sdk-core@0.46.0': - resolution: {integrity: sha512-Q2LAGWYIi+jyeKR9OIqvkgyde2GdzqfSG8lewxA9Xu/C9RJBBFbSfg5Nh8ZC66TKElGIosVOecoEJdbxnNMuxw==} + '@databricks/sdk-core@0.51.0': + resolution: {integrity: sha512-XzsdLcC5vXB+z5P68IlZ4emnXyp2frvY+aJnIwlpc2i3F/ZCR53a8HDrq7ruZPQG/FgGUR75MSFhku0nnYOysw==} engines: {node: '>=22.0.0'} '@databricks/sdk-experimental@0.15.0': @@ -2008,16 +2008,16 @@ packages: resolution: {integrity: sha512-dOJIt4F2nBk6HKObnv7Xbmy/qLYTy2835qhXSuW0Qw1QAXui9plmCet1KqG3yeQcMTyncWGbnhjGdQi8GEGQSA==} engines: {node: '>=22.0', npm: '>=10.0.0'} - '@databricks/sdk-options@0.46.0': - resolution: {integrity: sha512-UtADlR+41rYEoOCycZvJh1g96uDN6GVWgqQk+72cHBzcxi+koxKSJXHcYRI997Sc4Fcc1d2oyAC2I2ddVhurjA==} + '@databricks/sdk-options@0.51.0': + resolution: {integrity: sha512-p5uBh64Y1onnvwlEwfRsmKmGRI23Y4Ly8fZiCgNUwVDe4Z9d0ctI0mo1v4gnDqlmBaBml3TZfqn5BJBROCXsEw==} engines: {node: '>=22.0.0'} - '@databricks/sdk-statementexecution@0.46.0': - resolution: {integrity: sha512-VJA3e7UHmxRxN42/mV5VtKeINME0vCz3Na3hrwmta3tZqJWZbBU7XTfUdD1yOQ5Z1JU5UIM65OlWq8gc4IzHFg==} + '@databricks/sdk-statementexecution@0.52.0': + resolution: {integrity: sha512-ymDGDW67PJ/0VujM89mUrE6mbbrabsKJP5LvPv8ZbgzPR0bm0DwXMcQzvS0XUt0UGITjV2gqbx5jQwP3ijtOmg==} engines: {node: '>=22.0.0'} - '@databricks/sdk-warehouses@0.46.0': - resolution: {integrity: sha512-9r/gbdTb6ASiWCiibCwAOF8QizqNacidIw78uwzJYKK9dbdqmWIfNK0pF/jt3BG1sUiXz2b1I6URPdX7Qi0oLg==} + '@databricks/sdk-warehouses@0.53.0': + resolution: {integrity: sha512-+NlLn/HP3+5YUxTKlp0+zwsbAGaU2GHl8C+8JlWk04VLmBZhJg/PJnr4WlNrQVTAzqH01dy4DmK7MfqyOetHRg==} engines: {node: '>=22.0.0'} '@date-fns/tz@1.4.1': @@ -14288,12 +14288,12 @@ snapshots: transitivePeerDependencies: - supports-color - '@databricks/sdk-auth@0.46.0': + '@databricks/sdk-auth@0.51.0': dependencies: - '@databricks/sdk-core': 0.46.0 + '@databricks/sdk-core': 0.51.0 zod: 4.3.6 - '@databricks/sdk-core@0.46.0': + '@databricks/sdk-core@0.51.0': dependencies: json-bigint: 1.0.0 zod: 4.3.6 @@ -14316,25 +14316,25 @@ snapshots: transitivePeerDependencies: - supports-color - '@databricks/sdk-options@0.46.0': + '@databricks/sdk-options@0.51.0': dependencies: - '@databricks/sdk-auth': 0.46.0 - '@databricks/sdk-core': 0.46.0 + '@databricks/sdk-auth': 0.51.0 + '@databricks/sdk-core': 0.51.0 - '@databricks/sdk-statementexecution@0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d)': + '@databricks/sdk-statementexecution@0.52.0(patch_hash=9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c)': dependencies: - '@databricks/sdk-auth': 0.46.0 - '@databricks/sdk-core': 0.46.0 - '@databricks/sdk-options': 0.46.0 + '@databricks/sdk-auth': 0.51.0 + '@databricks/sdk-core': 0.51.0 + '@databricks/sdk-options': 0.51.0 '@js-temporal/polyfill': 0.5.1 json-bigint: 1.0.0 zod: 4.3.6 - '@databricks/sdk-warehouses@0.46.0': + '@databricks/sdk-warehouses@0.53.0': dependencies: - '@databricks/sdk-auth': 0.46.0 - '@databricks/sdk-core': 0.46.0 - '@databricks/sdk-options': 0.46.0 + '@databricks/sdk-auth': 0.51.0 + '@databricks/sdk-core': 0.51.0 + '@databricks/sdk-options': 0.51.0 '@js-temporal/polyfill': 0.5.1 json-bigint: 1.0.0 zod: 4.3.6 diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 9a8221175..bc863eab6 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -19,4 +19,4 @@ allowBuilds: esbuild: true protobufjs: false patchedDependencies: - '@databricks/sdk-statementexecution@0.46.0': patches/@databricks__sdk-statementexecution@0.46.0.patch + '@databricks/sdk-statementexecution@0.52.0': patches/@databricks__sdk-statementexecution@0.52.0.patch From ee9a37fcd33025ae775c4b2e98dc7ed979293d50 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 8 Oct 2026 15:57:28 +0200 Subject: [PATCH 11/17] fix(shared): cache env-M2M OAuth tokens until shortly before expiry sdk-auth 0.51.0's newM2mCredentials caches only the token endpoint and mints a new OAuth token on every request, unlike the legacy SDK which reused it until expiry. Wrap the env service-principal credentials in a cache that refreshes 40s before expiry (the legacy margin) and shares one in-flight fetch across concurrent callers. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../shared/src/workspace-client/modular.ts | 45 ++++++++++++++-- .../workspace-client/tests/modular.test.ts | 52 ++++++++++++++++--- 2 files changed, 84 insertions(+), 13 deletions(-) diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index 44a31fbec..c22590de0 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -16,6 +16,12 @@ * undocumented Reyden `attachment` response field, which the SDK's generated * unmarshal transform would otherwise strip. */ +import { + newTokenCredentials, + type Token, + type TokenCredentials, + tokenProviderFn, +} from "@databricks/sdk-auth"; import { newM2mCredentials, newPatCredentials, @@ -39,6 +45,37 @@ function normalizeHost(host: string | undefined): string | undefined { return /^https?:\/\//i.test(trimmed) ? trimmed : `https://${trimmed}`; } +// Same margin as the legacy SDK: refresh 40s early, since Azure Databricks +// rejects tokens that expire in 30s or less. +const TOKEN_REFRESH_MARGIN_MS = 40_000; + +/** + * Cache a token until shortly before it expires. The modular `newM2mCredentials` + * caches only the token endpoint and mints a fresh OAuth token on EVERY request; + * the legacy SDK reused it until expiry. Concurrent callers share one in-flight + * fetch. Like the legacy SDK, a token without an expiry is reused indefinitely. + */ +function withTokenCache(credentials: TokenCredentials): TokenCredentials { + let current: Token | undefined; + let inflight: Promise | undefined; + const isFresh = (t: Token) => + t.expiry === undefined || + t.expiry.getTime() - TOKEN_REFRESH_MARGIN_MS > Date.now(); + return newTokenCredentials( + credentials.name(), + tokenProviderFn(async () => { + if (current && isFresh(current)) return current; + inflight ??= credentials + .token() + .then((t) => (current = t)) + .finally(() => { + inflight = undefined; + }); + return inflight; + }), + ); +} + /** * Map wrapper options onto the modular SDK's `ClientOptions`, reproducing the * legacy SDK's auth resolution: explicit token → PAT (the OBO path); profile → @@ -81,11 +118,9 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { const clientSecret = process.env.DATABRICKS_CLIENT_SECRET; const envToken = process.env.DATABRICKS_TOKEN; if (host && clientId && clientSecret) { - clientOptions.credentials = newM2mCredentials({ - host, - clientId, - clientSecret, - }); + clientOptions.credentials = withTokenCache( + newM2mCredentials({ host, clientId, clientSecret }), + ); } else if (envToken) { clientOptions.credentials = newPatCredentials(envToken); } diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index 51814a293..ae0bda9d2 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -3,10 +3,12 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; // The wrapper's own tests are the one place allowed to mock the SDK directly. // Capture the `ClientOptions` the modular `WarehousesClient` constructor receives // so we can assert how wrapper options map onto the modular SDK's config. -const { ctorOpts, patTokens, m2mOpts } = vi.hoisted(() => ({ +const { ctorOpts, patTokens, m2mOpts, m2mMint } = vi.hoisted(() => ({ ctorOpts: [] as Array>, patTokens: [] as string[], m2mOpts: [] as Array>, + // Each M2M `token()` call mints a new token, like the real (uncached) SDK. + m2mMint: { count: 0, ttlMs: 60 * 60 * 1000 }, })); vi.mock("@databricks/sdk-warehouses/v1", () => ({ @@ -25,7 +27,13 @@ vi.mock("@databricks/sdk-auth/credentials", () => ({ }), newM2mCredentials: vi.fn((opts: Record) => { m2mOpts.push(opts); - return { kind: "m2m", ...opts }; + return { + name: () => "oauth-m2m", + token: async () => ({ + value: `m2m-token-${++m2mMint.count}`, + expiry: new Date(Date.now() + m2mMint.ttlMs), + }), + }; }), })); // The default transport: its `send` echoes the final request headers so tests @@ -79,6 +87,8 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { ctorOpts.length = 0; patTokens.length = 0; m2mOpts.length = 0; + m2mMint.count = 0; + m2mMint.ttlMs = 60 * 60 * 1000; for (const key of AUTH_ENV) { originalEnv[key] = process.env[key]; delete process.env[key]; @@ -152,15 +162,41 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { clientSecret: "sp-secret", }, ]); - expect(ctorOpts[0].credentials).toEqual({ - kind: "m2m", - host: "https://envhost.cloud.databricks.com", - clientId: "sp-client-id", - clientSecret: "sp-secret", - }); + // Wrapped in the token cache, so assert the strategy, not identity. + expect((ctorOpts[0].credentials as { name(): string }).name()).toBe( + "oauth-m2m", + ); expect(patTokens).toEqual([]); }); + test("env M2M: caches the OAuth token until 40s before expiry, then refreshes", async () => { + // sdk-auth's newM2mCredentials mints a new OAuth token on every request; + // the legacy SDK reused it until expiry. + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + buildWarehousesClient({}); + const creds = ctorOpts[0].credentials as { + authHeaders(): Promise>; + }; + const bearer = async () => + (await creds.authHeaders()).find((h) => h.key === "Authorization")?.value; + expect(await bearer()).toBe("Bearer m2m-token-1"); + expect(await bearer()).toBe("Bearer m2m-token-1"); + expect(m2mMint.count).toBe(1); + + // A token inside the 40s refresh margin is replaced on the next request. + m2mMint.ttlMs = 30_000; + vi.useFakeTimers({ toFake: ["Date"] }); + try { + vi.setSystemTime(Date.now() + 60 * 60 * 1000); + expect(await bearer()).toBe("Bearer m2m-token-2"); + expect(await bearer()).toBe("Bearer m2m-token-3"); + } finally { + vi.useRealTimers(); + } + }); + test("falls back to DATABRICKS_TOKEN (PAT) when no client id/secret is set", () => { process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; process.env.DATABRICKS_TOKEN = "env-pat"; From 12db322dd20e629e9e6b1b872895aa1775dd723f Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 14 Oct 2026 16:31:20 +0200 Subject: [PATCH 12/17] fix(appkit): read the modular SDK's httpStatusCode for OBO expiry and HTTP status Typed modular clients throw sdk-core's ApiError, which carries the HTTP status in httpStatusCode rather than statusCode. isUnauthorized missed those 401s, so an expired user token no longer became IdentityExpiredError, and Plugin.execute and the server error handler returned 500 for typed 4xx failures. Both now read either field through a shared httpStatusOf. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/src/errors/http-status.ts | 19 ++++++++ .../appkit/src/errors/identity-expired.ts | 3 ++ .../src/errors/tests/http-status.test.ts | 46 +++++++++++++++++++ packages/appkit/src/plugin/plugin.ts | 22 ++------- packages/appkit/src/plugins/server/index.ts | 29 ++++-------- .../server/tests/error-middleware.test.ts | 16 +++++++ packages/shared/src/workspace-client/index.ts | 1 + .../shared/src/workspace-client/modular.ts | 4 ++ 8 files changed, 102 insertions(+), 38 deletions(-) create mode 100644 packages/appkit/src/errors/http-status.ts create mode 100644 packages/appkit/src/errors/tests/http-status.test.ts diff --git a/packages/appkit/src/errors/http-status.ts b/packages/appkit/src/errors/http-status.ts new file mode 100644 index 000000000..dfa0b9c9b --- /dev/null +++ b/packages/appkit/src/errors/http-status.ts @@ -0,0 +1,19 @@ +/** + * The HTTP status carried by a thrown Error: `statusCode` (the workspace-client + * and legacy SDK `ApiError`) or `httpStatusCode` (the modular SDK's sdk-core + * `ApiError`, which reports -1 when there was no HTTP response). + */ +export function httpStatusOf(error: unknown): number | undefined { + if (!(error instanceof Error)) return undefined; + const { statusCode, httpStatusCode } = error as { + statusCode?: unknown; + httpStatusCode?: unknown; + }; + const status = typeof statusCode === "number" ? statusCode : httpStatusCode; + return typeof status === "number" && + Number.isInteger(status) && + status >= 100 && + status <= 599 + ? status + : undefined; +} diff --git a/packages/appkit/src/errors/identity-expired.ts b/packages/appkit/src/errors/identity-expired.ts index 9cfa7f1ff..72c535198 100644 --- a/packages/appkit/src/errors/identity-expired.ts +++ b/packages/appkit/src/errors/identity-expired.ts @@ -26,6 +26,8 @@ export function isUnauthorized(error: unknown): boolean { const value = current as { status?: unknown; statusCode?: unknown; + // Modular SDK (sdk-core) ApiError. + httpStatusCode?: unknown; response?: { status?: unknown }; cause?: unknown; }; @@ -33,6 +35,7 @@ export function isUnauthorized(error: unknown): boolean { !(current instanceof AppKitError) && (value.status === 401 || value.statusCode === 401 || + value.httpStatusCode === 401 || value.response?.status === 401) ) return true; diff --git a/packages/appkit/src/errors/tests/http-status.test.ts b/packages/appkit/src/errors/tests/http-status.test.ts new file mode 100644 index 000000000..99898dd8e --- /dev/null +++ b/packages/appkit/src/errors/tests/http-status.test.ts @@ -0,0 +1,46 @@ +import { ModularApiError } from "shared/workspace-client"; +import { describe, expect, test } from "vitest"; + +import { httpStatusOf } from "../http-status"; +import { isUnauthorized } from "../identity-expired"; + +/** What a typed modular SDK call throws: the status lives in `httpStatusCode`. */ +function modularError(status: number) { + const body = new TextEncoder().encode( + JSON.stringify({ error_code: "X", message: `status ${status}` }), + ); + const error = ModularApiError.fromHttpError(status, undefined, body); + if (!error) throw new Error(`no ApiError for ${status}`); + return error; +} + +describe("modular SDK errors", () => { + test("carry the status in httpStatusCode only", () => { + const error = modularError(401); + expect(error.httpStatusCode).toBe(401); + expect("statusCode" in error).toBe(false); + }); + + test("isUnauthorized recognizes a modular 401, directly and as a cause", () => { + expect(isUnauthorized(modularError(401))).toBe(true); + expect( + isUnauthorized(new Error("wrapped", { cause: modularError(401) })), + ).toBe(true); + expect(isUnauthorized(modularError(403))).toBe(false); + }); + + test("httpStatusOf reads statusCode or httpStatusCode", () => { + expect(httpStatusOf(modularError(404))).toBe(404); + expect( + httpStatusOf(Object.assign(new Error("x"), { statusCode: 403 })), + ).toBe(403); + }); + + test("httpStatusOf ignores non-HTTP values", () => { + expect( + httpStatusOf(Object.assign(new Error("x"), { httpStatusCode: -1 })), + ).toBeUndefined(); + expect(httpStatusOf({ statusCode: 404 })).toBeUndefined(); + expect(httpStatusOf(new Error("plain"))).toBeUndefined(); + }); +}); diff --git a/packages/appkit/src/plugin/plugin.ts b/packages/appkit/src/plugin/plugin.ts index 7eed23826..b00baed66 100644 --- a/packages/appkit/src/plugin/plugin.ts +++ b/packages/appkit/src/plugin/plugin.ts @@ -27,6 +27,7 @@ import { import { scopePlugin } from "../context/scoped-api"; import type { PluginContext } from "../core/plugin-context"; import { AppKitError, AuthenticationError } from "../errors"; +import { httpStatusOf } from "../errors/http-status"; import { IdentityExpiredError } from "../errors/identity-expired"; import { createLogger } from "../logging/logger"; import { StreamManager } from "../stream"; @@ -53,20 +54,6 @@ const logger = createLogger("plugin"); export { isPlainObject } from "../utils/is-plain-object"; export { isDevOboFallback } from "../context/request-scope"; -/** - * Narrow an unknown thrown value to an Error that carries a numeric - * `statusCode` property (e.g. `ApiError` from `@databricks/sdk-experimental`). - */ -function hasHttpStatusCode( - error: unknown, -): error is Error & { statusCode: number } { - return ( - error instanceof Error && - "statusCode" in error && - typeof (error as Record).statusCode === "number" - ); -} - /** * Base abstract class for creating AppKit plugins. * @@ -511,12 +498,13 @@ export abstract class Plugin< }; } - if (hasHttpStatusCode(error)) { + const status = httpStatusOf(error); + if (status !== undefined && error instanceof Error) { const isDev = process.env.NODE_ENV !== "production"; - const isClientError = error.statusCode >= 400 && error.statusCode < 500; + const isClientError = status >= 400 && status < 500; return { ok: false, - status: error.statusCode, + status, message: isDev || isClientError ? error.message : "Server error", }; } diff --git a/packages/appkit/src/plugins/server/index.ts b/packages/appkit/src/plugins/server/index.ts index 19951a040..ff8efcd0d 100644 --- a/packages/appkit/src/plugins/server/index.ts +++ b/packages/appkit/src/plugins/server/index.ts @@ -9,6 +9,7 @@ import type { PluginClientConfigs, PluginPhase } from "shared"; import { camelToKebab } from "shared"; import { AppKitError, ServerError } from "../../errors"; +import { httpStatusOf } from "../../errors/http-status"; import { TelemetryReporter } from "../../internal-telemetry"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; @@ -562,24 +563,6 @@ export function requestMetricsMiddleware( next(); } -/** - * Narrow an unknown thrown value to an Error that carries a numeric - * `statusCode` property in the HTTP error range (e.g. `ApiError` from - * `@databricks/sdk-experimental`). - */ -function hasHttpStatusCode( - error: unknown, -): error is Error & { statusCode: number } { - if (!(error instanceof Error) || !("statusCode" in error)) return false; - const statusCode = error.statusCode; - return ( - typeof statusCode === "number" && - Number.isInteger(statusCode) && - statusCode >= 400 && - statusCode <= 599 - ); -} - /** * Terminal Express error-handling middleware. * @@ -602,9 +585,13 @@ export function errorHandlerMiddleware( res: express.Response, next: express.NextFunction, ) { - const httpError = - err instanceof AppKitError || hasHttpStatusCode(err) ? err : null; - const statusCode = httpError ? httpError.statusCode : 500; + const httpStatus = httpStatusOf(err); + const statusCode = + err instanceof AppKitError + ? err.statusCode + : httpStatus !== undefined && httpStatus >= 400 + ? httpStatus + : 500; const isClientError = statusCode < 500; if (isClientError) { diff --git a/packages/appkit/src/plugins/server/tests/error-middleware.test.ts b/packages/appkit/src/plugins/server/tests/error-middleware.test.ts index 67335058a..000eee446 100644 --- a/packages/appkit/src/plugins/server/tests/error-middleware.test.ts +++ b/packages/appkit/src/plugins/server/tests/error-middleware.test.ts @@ -1,3 +1,4 @@ +import { ModularApiError } from "shared/workspace-client"; import { afterEach, describe, expect, test, vi } from "vitest"; import { @@ -40,6 +41,21 @@ describe("errorHandlerMiddleware", () => { vi.clearAllMocks(); }); + test("maps a modular SDK ApiError (httpStatusCode) to its status", () => { + const res = makeRes(); + const error = ModularApiError.fromHttpError( + 404, + undefined, + new TextEncoder().encode( + JSON.stringify({ error_code: "NOT_FOUND", message: "No such space" }), + ), + ); + + errorHandlerMiddleware(error, makeReq(), res, vi.fn()); + + expect(res.status).toHaveBeenCalledWith(404); + }); + test("maps AppKitError to its statusCode with its message", () => { const res = makeRes(); const next = vi.fn(); diff --git a/packages/shared/src/workspace-client/index.ts b/packages/shared/src/workspace-client/index.ts index 2b981ebec..585bc46c5 100644 --- a/packages/shared/src/workspace-client/index.ts +++ b/packages/shared/src/workspace-client/index.ts @@ -5,6 +5,7 @@ * behind a stable facade. */ export { ApiError } from "./errors"; +export { ModularApiError } from "./modular"; export { createWorkspaceClient } from "./factory"; export type { CancellationToken, diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index c22590de0..e6ad93724 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -191,6 +191,10 @@ export function buildStatementExecutionClient( return new StatementExecutionClient(mapToClientOptions(opts)); } +// The modular SDK's error class, for tests that must throw what typed calls +// really throw (status in `httpStatusCode`, not the wrapper's `statusCode`). +export { ApiError as ModularApiError } from "@databricks/sdk-core/apierror"; + // ── Client type re-exports (for the facade accessor types) ─────────────── export type { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; export type { WarehousesClient } from "@databricks/sdk-warehouses/v1"; From 85c7ec574245a7958799cc2eb5c9fcd28f01e15b Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 15 Oct 2026 17:05:54 +0200 Subject: [PATCH 13/17] fix(shared): cache tokens for profile and default-chain auth The modular SDK's default chain mints a token on every request: its CLI strategy shells out to `databricks auth token` per call, and its M2M strategy fetches a new OAuth token per call. The legacy SDK reused the token until expiry. Local dev with a CLI profile and CI typegen with an M2M profile paid this on every API call. Resolve the profile once and pick PAT, M2M or CLI in the SDK's order, wrapping M2M and CLI in the existing token cache. Profiles that pin an authType or a groupId keep the SDK's own uncached selection. Deployed apps are unaffected: env M2M was already cached and OBO uses the forwarded token. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../shared/src/workspace-client/modular.ts | 71 +++++++++- .../workspace-client/tests/modular.test.ts | 128 ++++++++++++++++-- 2 files changed, 187 insertions(+), 12 deletions(-) diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index e6ad93724..27f7400c0 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -17,16 +17,20 @@ * unmarshal transform would otherwise strip. */ import { + type Credentials, newTokenCredentials, type Token, type TokenCredentials, tokenProviderFn, } from "@databricks/sdk-auth"; import { + defaultCredentials, newM2mCredentials, newPatCredentials, + newU2mCredentials, } from "@databricks/sdk-auth/credentials"; import { type HttpClient, newFetchHttpClient } from "@databricks/sdk-core/http"; +import { type ProfileOptions, resolve } from "@databricks/sdk-core/profiles"; import type { ClientOptions } from "@databricks/sdk-options/client"; import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; import { WarehousesClient } from "@databricks/sdk-warehouses/v1"; @@ -76,6 +80,65 @@ function withTokenCache(credentials: TokenCredentials): TokenCredentials { ); } +/** + * The SDK default chain (profile / `~/.databrickscfg`), with the token cache the + * chain lacks: its CLI strategy shells out to `databricks auth token` and its + * M2M strategy mints a new OAuth token on EVERY request. `DefaultCredentials` + * hides the token behind `authHeaders()`, so mirror its strategy order (PAT → + * M2M → CLI) to get a cacheable `TokenCredentials`. Profiles pinning an + * `authType` or a `groupId` keep the SDK's own (uncached) selection logic. + * TODO: drop once sdk-auth caches tokens in its M2M and CLI credentials. + */ +function cachedDefaultCredentials( + profileOptions?: ProfileOptions, +): Credentials { + const select = async (): Promise => { + const profile = await resolve(profileOptions); + // Same as env M2M: OAuth discovery needs a scheme, which a profile may omit. + const host = normalizeHost(profile.host)?.replace(/\/+$/, ""); + const { name, token, clientId, clientSecret } = profile; + if (host && !profile.authType && !profile.groupId) { + if (token) return newPatCredentials(token.value); + if (clientId && clientSecret) { + return withTokenCache( + newM2mCredentials({ + host, + clientId, + clientSecret: clientSecret.value, + ...(profile.accountId !== undefined && { + accountId: profile.accountId, + }), + }), + ); + } + if (name) { + return withTokenCache( + newU2mCredentials({ + profile: name, + ...(profile.databricksCliPath !== undefined && { + cliPath: profile.databricksCliPath, + }), + }), + ); + } + } + return defaultCredentials({ + profile: { ...profile, ...(host && { host }) }, + }); + }; + let selected: Promise | undefined; + return { + name: () => "default", + async authHeaders() { + selected ??= select().catch((e) => { + selected = undefined; + throw e; + }); + return (await selected).authHeaders(); + }, + }; +} + /** * Map wrapper options onto the modular SDK's `ClientOptions`, reproducing the * legacy SDK's auth resolution: explicit token → PAT (the OBO path); profile → @@ -102,6 +165,9 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { clientOptions.credentials = newPatCredentials(opts.token); } else if (opts.profile) { clientOptions.profileOptions = { profile: opts.profile }; + clientOptions.credentials = cachedDefaultCredentials({ + profile: opts.profile, + }); } else { // No token, no profile: authenticate as the service principal from the // environment. The SDK's own default chain DOES read the DATABRICKS_* env @@ -123,9 +189,10 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { ); } else if (envToken) { clientOptions.credentials = newPatCredentials(envToken); + } else { + // Local dev: the profile-based default chain, with token caching. + clientOptions.credentials = cachedDefaultCredentials(); } - // Otherwise leave credentials unset and let the SDK walk its profile-based - // default chain (local dev with `~/.databrickscfg`). } const httpClient = buildHttpClient(opts); if (httpClient) { diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index ae0bda9d2..77945e7a8 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -3,13 +3,21 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; // The wrapper's own tests are the one place allowed to mock the SDK directly. // Capture the `ClientOptions` the modular `WarehousesClient` constructor receives // so we can assert how wrapper options map onto the modular SDK's config. -const { ctorOpts, patTokens, m2mOpts, m2mMint } = vi.hoisted(() => ({ - ctorOpts: [] as Array>, - patTokens: [] as string[], - m2mOpts: [] as Array>, - // Each M2M `token()` call mints a new token, like the real (uncached) SDK. - m2mMint: { count: 0, ttlMs: 60 * 60 * 1000 }, -})); +const { ctorOpts, patTokens, m2mOpts, m2mMint, u2mMint, profileState } = + vi.hoisted(() => ({ + ctorOpts: [] as Array>, + patTokens: [] as string[], + m2mOpts: [] as Array>, + // Each M2M `token()` call mints a new token, like the real (uncached) SDK. + m2mMint: { count: 0, ttlMs: 60 * 60 * 1000 }, + // Each CLI `token()` call is a `databricks auth token` subprocess. + u2mMint: { count: 0 }, + // What the mocked `~/.databrickscfg` profile resolution returns. + profileState: { + profile: {} as Record, + resolveCalls: [] as unknown[], + }, + })); vi.mock("@databricks/sdk-warehouses/v1", () => ({ WarehousesClient: vi.fn().mockImplementation((opts) => { @@ -35,6 +43,20 @@ vi.mock("@databricks/sdk-auth/credentials", () => ({ }), }; }), + newU2mCredentials: vi.fn(() => ({ + name: () => "databricks-cli", + token: async () => ({ + value: `cli-token-${++u2mMint.count}`, + expiry: new Date(Date.now() + 60 * 60 * 1000), + }), + })), + defaultCredentials: vi.fn(() => ({ kind: "sdk-default" })), +})); +vi.mock("@databricks/sdk-core/profiles", () => ({ + resolve: vi.fn(async (opts: unknown) => { + profileState.resolveCalls.push(opts); + return profileState.profile; + }), })); // The default transport: its `send` echoes the final request headers so tests // can assert the User-Agent the wrapper set before delegating. @@ -89,6 +111,9 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { m2mOpts.length = 0; m2mMint.count = 0; m2mMint.ttlMs = 60 * 60 * 1000; + u2mMint.count = 0; + profileState.profile = {}; + profileState.resolveCalls.length = 0; for (const key of AUTH_ENV) { originalEnv[key] = process.env[key]; delete process.env[key]; @@ -139,10 +164,12 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(patTokens).toEqual([]); }); - test("no host, no token, no profile, no env → empty options (SDK default chain)", () => { + test("no host, no token, no profile, no env → cached SDK default chain", () => { buildWarehousesClient({}); expect(ctorOpts[0].host).toBeUndefined(); - expect(ctorOpts[0].credentials).toBeUndefined(); + expect((ctorOpts[0].credentials as { name(): string }).name()).toBe( + "default", + ); expect(ctorOpts[0].profileOptions).toBeUndefined(); }); @@ -225,7 +252,88 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; buildWarehousesClient({}); expect(m2mOpts).toEqual([]); - expect(ctorOpts[0].credentials).toBeUndefined(); + expect((ctorOpts[0].credentials as { name(): string }).name()).toBe( + "default", + ); + }); + + describe("default chain (profile / ~/.databrickscfg) token cache", () => { + const bearerOf = (creds: unknown) => async () => + ( + await ( + creds as { + authHeaders(): Promise>; + } + ).authHeaders() + ).find((h) => h.key === "Authorization")?.value; + + test("CLI profile: 5 calls → one `databricks auth token`, profile resolved once", async () => { + profileState.profile = { name: "dev", host: "https://ws" }; + buildWarehousesClient({ profile: "dev" }); + const bearer = bearerOf(ctorOpts[0].credentials); + for (let i = 0; i < 5; i++) { + expect(await bearer()).toBe("Bearer cli-token-1"); + } + expect(u2mMint.count).toBe(1); + expect(profileState.resolveCalls).toEqual([{ profile: "dev" }]); + }); + + test("M2M profile: 5 calls → one OAuth token", async () => { + profileState.profile = { + name: "ci", + host: "ws.cloud.databricks.com/", + clientId: "id", + clientSecret: { value: "secret" }, + }; + buildWarehousesClient({ profile: "ci" }); + const bearer = bearerOf(ctorOpts[0].credentials); + for (let i = 0; i < 5; i++) { + expect(await bearer()).toBe("Bearer m2m-token-1"); + } + expect(m2mMint.count).toBe(1); + expect(m2mOpts).toEqual([ + // Scheme-normalized, trailing slash dropped, like env M2M. + { + host: "https://ws.cloud.databricks.com", + clientId: "id", + clientSecret: "secret", + }, + ]); + }); + + test("PAT profile: uses the profile token", async () => { + profileState.profile = { + name: "p", + host: "https://ws", + token: { value: "profile-pat" }, + }; + buildWarehousesClient({}); + await (ctorOpts[0].credentials as { authHeaders(): Promise }) + .authHeaders() + .catch(() => undefined); + expect(patTokens).toEqual(["profile-pat"]); + expect(u2mMint.count).toBe(0); + }); + + test("authType / groupId profiles keep the SDK's own selection", async () => { + const { defaultCredentials } = + await import("@databricks/sdk-auth/credentials"); + profileState.profile = { + name: "g", + host: "https://ws", + clientId: "id", + clientSecret: { value: "s" }, + groupId: "grp", + }; + buildWarehousesClient({ profile: "g" }); + await (ctorOpts[0].credentials as { authHeaders(): Promise }) + .authHeaders() + .catch(() => undefined); + expect(defaultCredentials).toHaveBeenCalledWith({ + profile: profileState.profile, + }); + expect(m2mOpts).toEqual([]); + }); }); test("User-Agent: prepends the exact @databricks/appkit product segment (dashboards match on it)", async () => { From 58d0459bc36f2ffac87c9e32b4e06ef533d67db0 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 16 Oct 2026 11:07:36 +0200 Subject: [PATCH 14/17] fix(shared): honor a profile's auth_type in the cached default chain `databricks auth login` writes `auth_type = databricks-cli`, and the cached default chain skipped every profile with an auth type, so the common local setup still ran `databricks auth token` on every request. Pin the named strategy (pat, oauth-m2m or databricks-cli) and cache it. Other auth types and groupId profiles still use the SDK's own selection. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../shared/src/workspace-client/modular.ts | 16 ++--- .../workspace-client/tests/modular.test.ts | 61 +++++++++++++++---- 2 files changed, 57 insertions(+), 20 deletions(-) diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index 27f7400c0..39198afa6 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -85,8 +85,9 @@ function withTokenCache(credentials: TokenCredentials): TokenCredentials { * chain lacks: its CLI strategy shells out to `databricks auth token` and its * M2M strategy mints a new OAuth token on EVERY request. `DefaultCredentials` * hides the token behind `authHeaders()`, so mirror its strategy order (PAT → - * M2M → CLI) to get a cacheable `TokenCredentials`. Profiles pinning an - * `authType` or a `groupId` keep the SDK's own (uncached) selection logic. + * M2M → CLI) to get a cacheable `TokenCredentials`. A profile `authType` naming + * one of these (`databricks auth login` writes `databricks-cli`) pins that + * strategy; other auth types and `groupId` keep the SDK's own (uncached) logic. * TODO: drop once sdk-auth caches tokens in its M2M and CLI credentials. */ function cachedDefaultCredentials( @@ -96,10 +97,11 @@ function cachedDefaultCredentials( const profile = await resolve(profileOptions); // Same as env M2M: OAuth discovery needs a scheme, which a profile may omit. const host = normalizeHost(profile.host)?.replace(/\/+$/, ""); - const { name, token, clientId, clientSecret } = profile; - if (host && !profile.authType && !profile.groupId) { - if (token) return newPatCredentials(token.value); - if (clientId && clientSecret) { + const { name, token, clientId, clientSecret, authType } = profile; + const allows = (strategy: string) => !authType || authType === strategy; + if (host && !profile.groupId) { + if (token && allows("pat")) return newPatCredentials(token.value); + if (clientId && clientSecret && allows("oauth-m2m")) { return withTokenCache( newM2mCredentials({ host, @@ -111,7 +113,7 @@ function cachedDefaultCredentials( }), ); } - if (name) { + if (name && allows("databricks-cli")) { return withTokenCache( newU2mCredentials({ profile: name, diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index 77945e7a8..7fc01238d 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -315,23 +315,58 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(u2mMint.count).toBe(0); }); - test("authType / groupId profiles keep the SDK's own selection", async () => { - const { defaultCredentials } = - await import("@databricks/sdk-auth/credentials"); + test("auth_type = databricks-cli (what `databricks auth login` writes) is cached", async () => { + // A PAT in the same profile must not win over the pinned CLI strategy. + profileState.profile = { + name: "dev", + host: "https://ws", + authType: "databricks-cli", + token: { value: "stale-pat" }, + }; + buildWarehousesClient({ profile: "dev" }); + const bearer = bearerOf(ctorOpts[0].credentials); + for (let i = 0; i < 5; i++) { + expect(await bearer()).toBe("Bearer cli-token-1"); + } + expect(u2mMint.count).toBe(1); + expect(patTokens).toEqual([]); + }); + + test("auth_type = oauth-m2m is cached", async () => { profileState.profile = { - name: "g", + name: "ci", host: "https://ws", + authType: "oauth-m2m", clientId: "id", - clientSecret: { value: "s" }, - groupId: "grp", + clientSecret: { value: "secret" }, }; - buildWarehousesClient({ profile: "g" }); - await (ctorOpts[0].credentials as { authHeaders(): Promise }) - .authHeaders() - .catch(() => undefined); - expect(defaultCredentials).toHaveBeenCalledWith({ - profile: profileState.profile, - }); + buildWarehousesClient({ profile: "ci" }); + const bearer = bearerOf(ctorOpts[0].credentials); + expect(await bearer()).toBe("Bearer m2m-token-1"); + expect(await bearer()).toBe("Bearer m2m-token-1"); + expect(m2mMint.count).toBe(1); + }); + + test("other auth types and groupId profiles keep the SDK's own selection", async () => { + const { defaultCredentials } = + await import("@databricks/sdk-auth/credentials"); + for (const extra of [{ authType: "azure-cli" }, { groupId: "grp" }]) { + profileState.profile = { + name: "g", + host: "https://ws", + clientId: "id", + clientSecret: { value: "s" }, + ...extra, + }; + ctorOpts.length = 0; + buildWarehousesClient({ profile: "g" }); + await (ctorOpts[0].credentials as { authHeaders(): Promise }) + .authHeaders() + .catch(() => undefined); + expect(defaultCredentials).toHaveBeenLastCalledWith({ + profile: profileState.profile, + }); + } expect(m2mOpts).toEqual([]); }); }); From db97339209bdaea57aca086f0b66fd83327272e1 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 16 Oct 2026 12:18:45 +0200 Subject: [PATCH 15/17] fix(appkit): keep modular SDK API errors' status through the cache The cache interceptor passed through only AppKitError and the legacy SDK's ApiError; anything else became ExecutionError.statementFailed (500). The modular SDK's ApiError is a different class, so a cached call failing with a 4xx (e.g. a jobs run that doesn't exist, 400) reached the client as a 500. Also pass through any error that carries an HTTP status (httpStatusOf). Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/src/cache/index.ts | 9 ++++++++- .../src/cache/tests/cache-manager.test.ts | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/packages/appkit/src/cache/index.ts b/packages/appkit/src/cache/index.ts index fcaffc637..bc18c13aa 100644 --- a/packages/appkit/src/cache/index.ts +++ b/packages/appkit/src/cache/index.ts @@ -6,6 +6,7 @@ import { createLakebasePool } from "../connectors/lakebase"; import { getClientOptions } from "../context/client-options"; import { getCurrentPrincipalKey } from "../context/execution-context"; import { AppKitError, ExecutionError, InitializationError } from "../errors"; +import { httpStatusOf } from "../errors/http-status"; import { createLogger } from "../logging/logger"; import type { Counter, TelemetryProvider } from "../telemetry"; import { SpanStatusCode, TelemetryManager } from "../telemetry"; @@ -339,7 +340,13 @@ export class CacheManager { if (sharedController.signal.aborted) { throw error; } - if (error instanceof AppKitError || error instanceof ApiError) { + // Keep any error that carries an HTTP status: the modular SDK's + // ApiError is a different class from the legacy one. + if ( + error instanceof AppKitError || + error instanceof ApiError || + httpStatusOf(error) !== undefined + ) { throw error; } throw ExecutionError.statementFailed( diff --git a/packages/appkit/src/cache/tests/cache-manager.test.ts b/packages/appkit/src/cache/tests/cache-manager.test.ts index 1565978ce..fe1028fba 100644 --- a/packages/appkit/src/cache/tests/cache-manager.test.ts +++ b/packages/appkit/src/cache/tests/cache-manager.test.ts @@ -997,6 +997,25 @@ describe("CacheManager", () => { ); }); + test("should re-throw a modular SDK ApiError without wrapping (keeps its status)", async () => { + const { ModularApiError } = await import("shared/workspace-client"); + const cache = await CacheManager.getInstance({ + storage: createMockStorage(), + }); + const body = new TextEncoder().encode( + JSON.stringify({ + error_code: "INVALID_PARAMETER_VALUE", + message: "Run 1 does not exist.", + }), + ); + const apiError = ModularApiError.fromHttpError(400, undefined, body); + const fn = vi.fn().mockRejectedValue(apiError); + + await expect(cache.getOrExecute(["key"], fn, "user1")).rejects.toBe( + apiError, + ); + }); + test("should allow retry after error", async () => { const cache = await CacheManager.getInstance({ storage: createMockStorage(), From 0990e36a90afaff9789f683a56e58ba9a12422b5 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 16 Oct 2026 15:28:08 +0200 Subject: [PATCH 16/17] fix(shared): restore legacy auth, retry and ApiError behavior on modular clients Review findings against the legacy SDK: - Host-only config read ~/.databrickscfg anyway, so a [DEFAULT] profile's credential could be sent to a different DATABRICKS_HOST, and a host-only .env failed to boot. Like the legacy SDK, a configured host now skips the config file and uses the Databricks CLI token for that host (cached). - DATABRICKS_AUTH_TYPE was ignored (pat still ran as the SP). It now pins the strategy; an unsupported value fails with an actionable message. - Modular clients do not retry. The transport now retries 429 and the legacy transient errors with the legacy backoff (up to 300s). Streamed bodies are never retried. - Modular errors were not instanceof the exported ApiError and had no statusCode, so the retry interceptor retried 4xx. The exported ApiError now matches both classes and modular errors expose statusCode/errorCode. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../src/errors/tests/http-status.test.ts | 4 +- .../appkit/src/plugin/tests/retry.test.ts | 14 ++ .../shared/src/workspace-client/errors.ts | 48 +++- .../shared/src/workspace-client/modular.ts | 155 ++++++++++-- .../workspace-client/tests/modular.test.ts | 230 ++++++++++++++++-- 5 files changed, 393 insertions(+), 58 deletions(-) diff --git a/packages/appkit/src/errors/tests/http-status.test.ts b/packages/appkit/src/errors/tests/http-status.test.ts index 99898dd8e..59f2c018f 100644 --- a/packages/appkit/src/errors/tests/http-status.test.ts +++ b/packages/appkit/src/errors/tests/http-status.test.ts @@ -15,10 +15,10 @@ function modularError(status: number) { } describe("modular SDK errors", () => { - test("carry the status in httpStatusCode only", () => { + test("carry the status in httpStatusCode, mirrored as the legacy statusCode", () => { const error = modularError(401); expect(error.httpStatusCode).toBe(401); - expect("statusCode" in error).toBe(false); + expect((error as unknown as { statusCode: number }).statusCode).toBe(401); }); test("isUnauthorized recognizes a modular 401, directly and as a cause", () => { diff --git a/packages/appkit/src/plugin/tests/retry.test.ts b/packages/appkit/src/plugin/tests/retry.test.ts index e1bb61d0f..1d74cfeee 100644 --- a/packages/appkit/src/plugin/tests/retry.test.ts +++ b/packages/appkit/src/plugin/tests/retry.test.ts @@ -1,4 +1,5 @@ import type { RetryConfig } from "shared"; +import { ModularApiError } from "shared/workspace-client"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { AuthenticationError } from "../../errors/authentication"; @@ -37,6 +38,19 @@ describe("RetryInterceptor", () => { expect(fn).toHaveBeenCalledTimes(1); }); + test("does not retry a modular SDK 4xx (status is read from httpStatusCode)", async () => { + const interceptor = new RetryInterceptor({ enabled: true, attempts: 3 }); + const body = new TextEncoder().encode( + JSON.stringify({ error_code: "INVALID_PARAMETER_VALUE", message: "no" }), + ); + const fn = vi + .fn() + .mockRejectedValue(ModularApiError.fromHttpError(400, undefined, body)); + + await expect(interceptor.intercept(fn, context)).rejects.toThrow("no"); + expect(fn).toHaveBeenCalledTimes(1); + }); + test("should retry on failure up to max attempts", async () => { const config: RetryConfig = { enabled: true, diff --git a/packages/shared/src/workspace-client/errors.ts b/packages/shared/src/workspace-client/errors.ts index 01427e191..408c6db4c 100644 --- a/packages/shared/src/workspace-client/errors.ts +++ b/packages/shared/src/workspace-client/errors.ts @@ -1,14 +1,42 @@ /** * API error type surfaced by the AppKit workspace client. * - * Re-exports the legacy `@databricks/sdk-experimental` `ApiError` so AppKit - * catch sites (`instanceof ApiError`, `.statusCode`) and the one construction - * site in `connectors/files/client.ts` keep working unchanged while importing - * from the wrapper instead of the SDK. - * - * When the first modular service lands (throwing `@databricks/sdk-core`'s - * distinct `ApiError`), add an `isApiError(e)` / `getApiErrorStatusCode(e)` - * predicate here to unify the two shapes — `instanceof` can't, since they're - * classes from different packages. + * Migrated services throw the modular `@databricks/sdk-core` `ApiError`, a + * different class from the legacy `@databricks/sdk-experimental` one, with the + * status in `httpStatusCode` and the code in `code`. To keep `instanceof + * ApiError`, `.statusCode` and `.errorCode` working for AppKit and its + * consumers, the exported `ApiError` matches both classes, and modular errors + * also expose the legacy field names. */ -export { ApiError } from "@databricks/sdk-experimental"; +import { ApiError as ModularApiError } from "@databricks/sdk-core/apierror"; +import { ApiError as LegacyApiError } from "@databricks/sdk-experimental"; + +// Legacy-name aliases on modular errors. Setters store an own value, so call +// sites that assign `statusCode` (e.g. via Object.assign) keep working. +for (const [name, source] of [ + ["statusCode", "httpStatusCode"], + ["errorCode", "code"], +] as const) { + if (!Object.hasOwn(ModularApiError.prototype, name)) { + Object.defineProperty(ModularApiError.prototype, name, { + configurable: true, + get(this: ModularApiError) { + return this[source]; + }, + set(this: ModularApiError, value: unknown) { + Object.defineProperty(this, name, { + value, + writable: true, + enumerable: true, + configurable: true, + }); + }, + }); + } +} + +export class ApiError extends LegacyApiError { + static [Symbol.hasInstance](value: unknown): boolean { + return value instanceof LegacyApiError || value instanceof ModularApiError; + } +} diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index 39198afa6..33ec26f53 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -16,6 +16,9 @@ * undocumented Reyden `attachment` response field, which the SDK's generated * unmarshal transform would otherwise strip. */ +import { execFile } from "node:child_process"; +import { promisify } from "node:util"; + import { type Credentials, newTokenCredentials, @@ -35,6 +38,7 @@ import type { ClientOptions } from "@databricks/sdk-options/client"; import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; import { WarehousesClient } from "@databricks/sdk-warehouses/v1"; +import "./errors"; import type { WorkspaceClientOptions } from "./legacy"; /** @@ -141,6 +145,55 @@ function cachedDefaultCredentials( }; } +const execFileAsync = promisify(execFile); + +/** + * Databricks CLI auth for a host, not a profile: what the legacy SDK used when + * only a host was configured. The modular CLI strategy needs a profile name. + */ +function cliHostCredentials(host: string): TokenCredentials { + return withTokenCache( + newTokenCredentials( + "databricks-cli", + tokenProviderFn(async () => { + let stdout: string; + try { + ({ stdout } = await execFileAsync("databricks", [ + "auth", + "token", + "--host", + host, + ])); + } catch (e) { + const err = e as { stderr?: unknown; message?: string }; + const detail = String(err.stderr ?? err.message ?? e).trim(); + throw new Error( + `cannot get an access token for ${host} from the Databricks CLI: ${detail}`, + ); + } + const parsed = JSON.parse(stdout) as { + access_token: string; + expiry?: string; + }; + return { + value: parsed.access_token, + ...(parsed.expiry && { expiry: new Date(parsed.expiry) }), + }; + }), + ), + ); +} + +/** Credentials that fail on use with `message`, so boot surfaces it clearly. */ +function unsupportedCredentials(message: string): Credentials { + return { + name: () => "unsupported", + authHeaders: async () => { + throw new Error(message); + }, + }; +} + /** * Map wrapper options onto the modular SDK's `ClientOptions`, reproducing the * legacy SDK's auth resolution: explicit token → PAT (the OBO path); profile → @@ -182,24 +235,32 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { // (what Apps injects), else PAT from `DATABRICKS_TOKEN`, else fall through to // the SDK default chain (local dev, where a `~/.databrickscfg` host already // carries a scheme). + // `DATABRICKS_AUTH_TYPE`, like in the legacy SDK, pins the strategy. const clientId = process.env.DATABRICKS_CLIENT_ID; const clientSecret = process.env.DATABRICKS_CLIENT_SECRET; const envToken = process.env.DATABRICKS_TOKEN; - if (host && clientId && clientSecret) { + const authType = process.env.DATABRICKS_AUTH_TYPE?.trim() || undefined; + const allows = (strategy: string) => !authType || authType === strategy; + if (host && clientId && clientSecret && allows("oauth-m2m")) { clientOptions.credentials = withTokenCache( newM2mCredentials({ host, clientId, clientSecret }), ); - } else if (envToken) { + } else if (envToken && allows("pat")) { clientOptions.credentials = newPatCredentials(envToken); - } else { - // Local dev: the profile-based default chain, with token caching. + } else if (host && allows("databricks-cli")) { + // A configured host means ~/.databrickscfg is NOT consulted (legacy + // rule): a profile's credentials must never be sent to another host. + clientOptions.credentials = cliHostCredentials(host); + } else if (!host && !authType) { + // Local dev with no host anywhere: the profile-based default chain. clientOptions.credentials = cachedDefaultCredentials(); + } else { + clientOptions.credentials = unsupportedCredentials( + `DATABRICKS_AUTH_TYPE=${authType ?? ""} cannot be satisfied: supported types are pat (DATABRICKS_TOKEN), oauth-m2m (DATABRICKS_CLIENT_ID/SECRET) and databricks-cli (\`databricks auth login\`).`, + ); } } - const httpClient = buildHttpClient(opts); - if (httpClient) { - clientOptions.httpClient = httpClient; - } + clientOptions.httpClient = buildHttpClient(opts); return clientOptions; } @@ -221,31 +282,79 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { * the SDK's default User-Agent untouched — matching the legacy behavior where * build-time clients carried no AppKit UA. */ -function buildHttpClient(opts: WorkspaceClientOptions): HttpClient | undefined { +function buildHttpClient(opts: WorkspaceClientOptions): HttpClient { const co = opts.clientOptions; - if (!co?.product || !co?.productVersion) { - return undefined; - } - const segments = [`${co.product}/${co.productVersion}`]; - if (co.userAgentExtra) { - for (const [key, value] of Object.entries(co.userAgentExtra)) { - segments.push(`${key}/${String(value)}`); + let appkitUserAgent: string | undefined; + if (co?.product && co?.productVersion) { + const segments = [`${co.product}/${co.productVersion}`]; + if (co.userAgentExtra) { + for (const [key, value] of Object.entries(co.userAgentExtra)) { + segments.push(`${key}/${String(value)}`); + } } + appkitUserAgent = segments.join(" "); } - const appkitUserAgent = segments.join(" "); const base = newFetchHttpClient(); return { send(request) { - const existing = request.headers.get("User-Agent"); - request.headers.set( - "User-Agent", - existing ? `${appkitUserAgent} ${existing}` : appkitUserAgent, - ); - return base.send(request); + if (appkitUserAgent) { + const existing = request.headers.get("User-Agent"); + request.headers.set( + "User-Agent", + existing ? `${appkitUserAgent} ${existing}` : appkitUserAgent, + ); + } + return sendWithRetry(base, request); }, }; } +// The legacy SDK's retry policy: 429, or an error body naming a transient +// failure; backoff of `attempt` seconds plus jitter, capped at 10s, for up to +// 300s in total. The modular SDK retries nothing outside long-running waits. +const RETRY_TIMEOUT_MS = 300_000; +const TRANSIENT_ERROR_MESSAGES = [ + "com.databricks.backend.manager.util.UnknownWorkerEnvironmentException", + "does not have any associated worker environments", + "There is no worker environment with id", + "Unknown worker environment", + "ClusterNotReadyException", + "connection reset by peer", + "TLS handshake timeout", + "connection refused", + "Unexpected error", + "i/o timeout", +]; + +async function sendWithRetry( + base: HttpClient, + request: Parameters[0], +): ReturnType { + // A streamed body (e.g. a file upload) can only be sent once. + if (request.body instanceof ReadableStream) return base.send(request); + const deadline = Date.now() + RETRY_TIMEOUT_MS; + for (let attempt = 1; ; attempt++) { + const response = await base.send(request); + if (response.statusCode < 400) return response; + const body = new Uint8Array( + await new Response(response.body).arrayBuffer(), + ); + const text = new TextDecoder().decode(body); + const retryable = + response.statusCode === 429 || + TRANSIENT_ERROR_MESSAGES.some((m) => text.includes(m)); + const waitMs = Math.min(attempt * 1000 + 50 + Math.random() * 700, 10_000); + if ( + !retryable || + request.signal?.aborted || + Date.now() + waitMs > deadline + ) { + return { ...response, body: new Response(body).body }; + } + await new Promise((resolve) => setTimeout(resolve, waitMs)); + } +} + /** Build a modular Warehouses client from wrapper options. */ export function buildWarehousesClient( opts: WorkspaceClientOptions, diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index 7fc01238d..f78b65cc4 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -3,21 +3,33 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; // The wrapper's own tests are the one place allowed to mock the SDK directly. // Capture the `ClientOptions` the modular `WarehousesClient` constructor receives // so we can assert how wrapper options map onto the modular SDK's config. -const { ctorOpts, patTokens, m2mOpts, m2mMint, u2mMint, profileState } = - vi.hoisted(() => ({ - ctorOpts: [] as Array>, - patTokens: [] as string[], - m2mOpts: [] as Array>, - // Each M2M `token()` call mints a new token, like the real (uncached) SDK. - m2mMint: { count: 0, ttlMs: 60 * 60 * 1000 }, - // Each CLI `token()` call is a `databricks auth token` subprocess. - u2mMint: { count: 0 }, - // What the mocked `~/.databrickscfg` profile resolution returns. - profileState: { - profile: {} as Record, - resolveCalls: [] as unknown[], - }, - })); +const { + ctorOpts, + patTokens, + m2mOpts, + m2mMint, + u2mMint, + profileState, + cliCalls, + responses, +} = vi.hoisted(() => ({ + ctorOpts: [] as Array>, + patTokens: [] as string[], + m2mOpts: [] as Array>, + // Each M2M `token()` call mints a new token, like the real (uncached) SDK. + m2mMint: { count: 0, ttlMs: 60 * 60 * 1000 }, + // Each CLI `token()` call is a `databricks auth token` subprocess. + u2mMint: { count: 0 }, + // What the mocked `~/.databrickscfg` profile resolution returns. + profileState: { + profile: {} as Record, + resolveCalls: [] as unknown[], + }, + // Each `databricks auth token --host` subprocess run by the host CLI path. + cliCalls: [] as string[][], + // Queued transport responses; when empty the transport answers 200. + responses: [] as Array<{ statusCode: number; body: string }>, +})); vi.mock("@databricks/sdk-warehouses/v1", () => ({ WarehousesClient: vi.fn().mockImplementation((opts) => { @@ -62,15 +74,33 @@ vi.mock("@databricks/sdk-core/profiles", () => ({ // can assert the User-Agent the wrapper set before delegating. vi.mock("@databricks/sdk-core/http", () => ({ newFetchHttpClient: vi.fn(() => ({ - send: vi.fn((request: { headers: Headers }) => - Promise.resolve({ - statusCode: 200, + send: vi.fn((request: { headers: Headers }) => { + const next = responses.shift(); + return Promise.resolve({ + statusCode: next?.statusCode ?? 200, headers: request.headers, - body: null, - }), - ), + body: next ? new Response(next.body).body : null, + }); + }), })), })); +vi.mock("node:child_process", () => ({ + execFile: vi.fn( + ( + _cmd: string, + args: string[], + cb: (err: null, out: { stdout: string }) => void, + ) => { + cliCalls.push(args); + cb(null, { + stdout: JSON.stringify({ + access_token: `host-cli-token-${cliCalls.length}`, + expiry: new Date(Date.now() + 3_600_000).toISOString(), + }), + }); + }, + ), +})); import { buildWarehousesClient } from "../modular"; @@ -102,6 +132,7 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { "DATABRICKS_CLIENT_ID", "DATABRICKS_CLIENT_SECRET", "DATABRICKS_TOKEN", + "DATABRICKS_AUTH_TYPE", ] as const; const originalEnv: Record = {}; @@ -112,6 +143,8 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { m2mMint.count = 0; m2mMint.ttlMs = 60 * 60 * 1000; u2mMint.count = 0; + cliCalls.length = 0; + responses.length = 0; profileState.profile = {}; profileState.resolveCalls.length = 0; for (const key of AUTH_ENV) { @@ -398,8 +431,159 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(ua).toBe("@databricks/appkit/0.64.0"); }); - test("no product configured (build-time) → no httpClient override (SDK default UA)", () => { + test("no product configured (build-time) → transport adds no AppKit User-Agent", async () => { buildWarehousesClient({ host: "https://x" }); - expect(ctorOpts[0].httpClient).toBeUndefined(); + expect(await sentUserAgent(ctorOpts[0].httpClient)).toBeNull(); + }); +}); + +describe("migration review fixes (host-only auth, DATABRICKS_AUTH_TYPE, retries, ApiError)", () => { + const ENV = [ + "DATABRICKS_HOST", + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CLIENT_SECRET", + "DATABRICKS_TOKEN", + "DATABRICKS_AUTH_TYPE", + ] as const; + const saved: Record = {}; + beforeEach(() => { + ctorOpts.length = 0; + patTokens.length = 0; + m2mOpts.length = 0; + cliCalls.length = 0; + responses.length = 0; + profileState.profile = { + name: "DEFAULT", + host: "https://other-workspace", + token: { value: "default-profile-pat" }, + }; + profileState.resolveCalls.length = 0; + for (const k of ENV) { + saved[k] = process.env[k]; + delete process.env[k]; + } + }); + afterEach(() => { + vi.useRealTimers(); + for (const k of ENV) { + if (saved[k] === undefined) delete process.env[k]; + else process.env[k] = saved[k]; + } + }); + + const bearer = async (creds: unknown) => + ( + await ( + creds as { + authHeaders(): Promise>; + } + ).authHeaders() + ).find((h) => h.key === "Authorization")?.value; + + test("host-only env never uses another host's [DEFAULT] profile; uses the CLI for that host, cached", async () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + buildWarehousesClient({}); + for (let i = 0; i < 3; i++) { + expect(await bearer(ctorOpts[0].credentials)).toBe( + "Bearer host-cli-token-1", + ); + } + expect(cliCalls).toEqual([ + ["auth", "token", "--host", "https://envhost.cloud.databricks.com"], + ]); + expect(profileState.resolveCalls).toEqual([]); + expect(patTokens).toEqual([]); + }); + + test("DATABRICKS_AUTH_TYPE=pat wins over SP client id/secret", async () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + process.env.DATABRICKS_TOKEN = "user-pat"; + process.env.DATABRICKS_AUTH_TYPE = "pat"; + buildWarehousesClient({}); + expect(patTokens).toEqual(["user-pat"]); + expect(m2mOpts).toEqual([]); + }); + + test("an unsupported DATABRICKS_AUTH_TYPE fails with an actionable message", async () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + process.env.DATABRICKS_AUTH_TYPE = "azure-cli"; + buildWarehousesClient({}); + await expect(bearer(ctorOpts[0].credentials)).rejects.toThrow( + /DATABRICKS_AUTH_TYPE=azure-cli cannot be satisfied/, + ); + expect(cliCalls).toEqual([]); + }); + + async function send(body?: unknown) { + buildWarehousesClient({ host: "https://x" }); + const client = ctorOpts[0].httpClient as { + send( + r: unknown, + ): Promise<{ statusCode: number; body: ReadableStream | null }>; + }; + return client.send({ + url: "https://x/api", + method: "POST", + headers: new Headers(), + ...(body !== undefined && { body }), + }); + } + + test("transport retries a 429 (legacy behavior) and returns the eventual success", async () => { + vi.useFakeTimers(); + responses.push({ + statusCode: 429, + body: '{"error_code":"TOO_MANY_REQUESTS"}', + }); + const pending = send("{}"); + await vi.runAllTimersAsync(); + const res = await pending; + expect(res.statusCode).toBe(200); + expect(responses).toEqual([]); + }); + + test("transport does not retry a plain 400 and keeps its body readable", async () => { + responses.push({ + statusCode: 400, + body: '{"message":"Run 1 does not exist."}', + }); + responses.push({ statusCode: 200, body: "{}" }); + const res = await send("{}"); + expect(res.statusCode).toBe(400); + expect(await new Response(res.body).text()).toContain( + "Run 1 does not exist", + ); + expect(responses).toHaveLength(1); + }); + + test("transport never retries a streamed body (it can only be sent once)", async () => { + responses.push({ statusCode: 429, body: "{}" }); + const res = await send(new Response("data").body); + expect(res.statusCode).toBe(429); + }); + + test("modular SDK errors are instanceof the exported ApiError and expose legacy fields", async () => { + const { ApiError, ModularApiError } = await import("../index"); + const body = new TextEncoder().encode( + JSON.stringify({ + error_code: "INVALID_PARAMETER_VALUE", + message: "nope", + }), + ); + const err = ModularApiError.fromHttpError(400, undefined, body); + expect(err).toBeInstanceOf(ApiError); + expect((err as unknown as { statusCode: number }).statusCode).toBe(400); + expect((err as unknown as { errorCode: string }).errorCode).toBe( + "INVALID_PARAMETER_VALUE", + ); + // Assigning the legacy field (as some connectors do) must not throw. + Object.assign(err as object, { statusCode: 418 }); + expect((err as unknown as { statusCode: number }).statusCode).toBe(418); + expect(new Error("x")).not.toBeInstanceOf(ApiError); + expect(new ApiError("m", "CODE", 404, undefined, [])).toBeInstanceOf( + ApiError, + ); }); }); From db1d298582f0223a690aa47684c8d61ed1250cbc Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 16 Oct 2026 16:39:03 +0200 Subject: [PATCH 17/17] docs(shared): describe the workspace client facade as partly migrated The header still said the seam migrates no service, which is no longer true once warehouses and statementExecution use the modular SDK. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/shared/src/workspace-client/types.ts | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index 87071b004..9ca7ea189 100644 --- a/packages/shared/src/workspace-client/types.ts +++ b/packages/shared/src/workspace-client/types.ts @@ -2,11 +2,11 @@ * The wrapper facade type and the SDK type-namespace re-exports used across * AppKit. * - * Every service accessor is currently typed against the legacy SDK and - * delegated through `toLegacyWorkspaceClient()` — this PR introduces the seam - * only, it does not migrate any service to the modular SDK. Migrating a service - * later means changing that accessor's type here + its getter in `client.ts` - * and updating the one connector that consumes it. + * Accessors for migrated services are typed against their modular + * `@databricks/sdk-*` clients (built in `modular.ts`); the rest are still the + * legacy SDK's, delegated through `toLegacyWorkspaceClient()`. Migrating a + * service means changing its accessor's type here, its getter in `client.ts`, + * and the connector that consumes it. * * Type-namespace re-exports (`files`, `jobs`, `serving`, `sql`) point at * `@databricks/sdk-experimental` so existing AppKit call-site shapes stay