diff --git a/apps/dev-playground/tests/utils/test-utils.ts b/apps/dev-playground/tests/utils/test-utils.ts index a4449023..f608a0e0 100644 --- a/apps/dev-playground/tests/utils/test-utils.ts +++ b/apps/dev-playground/tests/utils/test-utils.ts @@ -1,4 +1,5 @@ import type { Page, Request } from "@playwright/test"; +import { tableFromJSON, tableToIPC } from "apache-arrow"; import { mockAnalyticsData, mockReconnectMessages, @@ -27,65 +28,90 @@ function getSSEHeaders() { }; } +/** Maps a query URL to its mock rows (matched by query-key substring). */ +const ANALYTICS_MOCK: Array = [ + ["spend_summary", mockAnalyticsData.spendSummary], + ["apps_list", mockAnalyticsData.appsList], + ["untagged_apps", mockAnalyticsData.untaggedApps], + ["spend_data", mockAnalyticsData.spendData], + ["top_contributors", mockAnalyticsData.topContributors], + ["app_activity_heatmap", mockAnalyticsData.appActivityHeatmap], + ["sql_helpers_test", mockAnalyticsData.sqlHelpersTest], +]; + +function resolveAnalyticsMock(url: string): readonly unknown[] { + return ANALYTICS_MOCK.find(([key]) => url.includes(key))?.[1] ?? []; +} + +/** + * Build raw Arrow IPC stream bytes from mock rows — the format `fetchArrowDirect` + * expects for `format: "ARROW_STREAM"` (raw bytes on the response body, not SSE). + * + * Mirrors a real Databricks warehouse: the Arrow schema is encoded + * *positionally* (`col_0`, `col_1`, …) and the real names are returned + * separately for the `X-Appkit-Arrow-Columns` header, so the client's relabel + * path (positional schema → real names) is exercised end-to-end — that path + * was the live bug and is otherwise only unit-tested. + * + * Array/object cells are stringified so Arrow can infer a primitive column + * type, matching how the warehouse serializes complex types; chart-relevant + * columns (names, numbers) are unaffected. + */ +function toArrowIPC(rows: readonly unknown[]): { + body: Buffer; + columnNames: string[]; +} { + const columnNames = Object.keys(rows[0] as Record); + const positional = rows.map((row) => { + const values = Object.values(row as Record); + return Object.fromEntries( + values.map((v, i) => [ + `col_${i}`, + v !== null && typeof v === "object" ? JSON.stringify(v) : v, + ]), + ); + }); + return { + body: Buffer.from(tableToIPC(tableFromJSON(positional), "stream")), + columnNames, + }; +} + export async function setupMockAPI(page: Page) { await page.route("**/api/analytics/query/**", async (route) => { - const url = route.request().url(); + const data = resolveAnalyticsMock(route.request().url()); - if (url.includes("spend_summary")) { - return route.fulfill({ - status: 200, - headers: getSSEHeaders(), - body: createSSEResponse(mockAnalyticsData.spendSummary), - }); - } - if (url.includes("apps_list")) { - return route.fulfill({ - status: 200, - headers: getSSEHeaders(), - body: createSSEResponse(mockAnalyticsData.appsList), - }); - } - if (url.includes("untagged_apps")) { - return route.fulfill({ - status: 200, - headers: getSSEHeaders(), - body: createSSEResponse(mockAnalyticsData.untaggedApps), - }); - } - if (url.includes("spend_data")) { - return route.fulfill({ - status: 200, - headers: getSSEHeaders(), - body: createSSEResponse(mockAnalyticsData.spendData), - }); + // ARROW_STREAM requests read raw Arrow IPC bytes off the response body + // (see fetchArrowDirect); everything else consumes the SSE/JSON result. + let requestFormat: string | undefined; + try { + requestFormat = route.request().postDataJSON()?.format; + } catch { + // Non-JSON body (e.g. a GET) — treat as the SSE/JSON path. } - if (url.includes("top_contributors")) { - return route.fulfill({ - status: 200, - headers: getSSEHeaders(), - body: createSSEResponse(mockAnalyticsData.topContributors), - }); - } - if (url.includes("app_activity_heatmap")) { - return route.fulfill({ - status: 200, - headers: getSSEHeaders(), - body: createSSEResponse(mockAnalyticsData.appActivityHeatmap), - }); - } - if (url.includes("sql_helpers_test")) { + + if (requestFormat === "ARROW_STREAM" && data.length > 0) { + const { body, columnNames } = toArrowIPC(data); return route.fulfill({ status: 200, - headers: getSSEHeaders(), - body: createSSEResponse(mockAnalyticsData.sqlHelpersTest), + headers: { + "Content-Type": "application/vnd.apache.arrow.stream", + "Cache-Control": "no-store", + // Real warehouses encode the schema positionally (col_N) and send + // the aliased names in this header; the client relabels the decoded + // Table. Mirror that so the relabel path is covered. + "X-Appkit-Arrow-Columns": encodeURIComponent( + JSON.stringify(columnNames), + ), + }, + body, }); } - // Default empty response for unknown queries return route.fulfill({ status: 200, headers: getSSEHeaders(), - body: createSSEResponse([]), + body: createSSEResponse(data), }); }); diff --git a/bundle-size-baseline.json b/bundle-size-baseline.json index e07f325c..ea7156ce 100644 --- a/bundle-size-baseline.json +++ b/bundle-size-baseline.json @@ -3,25 +3,25 @@ { "name": "@databricks/appkit", "tarball": { - "packed": 674146, - "unpacked": 2386195 + "packed": 727720, + "unpacked": 2550712 }, "dist": { "total": { - "raw": 2381750, - "gzip": 797897 + "raw": 2546267, + "gzip": 855702 }, "js": { - "raw": 702696, - "gzip": 245275 + "raw": 750162, + "gzip": 262624 }, "types": { - "raw": 274230, - "gzip": 92925 + "raw": 287585, + "gzip": 98642 }, "maps": { - "raw": 1394029, - "gzip": 455879 + "raw": 1497725, + "gzip": 490618 }, "css": { "raw": 0, @@ -31,22 +31,22 @@ "raw": 10795, "gzip": 3818 }, - "fileCount": 518 + "fileCount": 533 }, "entries": [ { "id": ".", - "gzip": 78098, + "gzip": 83090, "composition": { - "initialGzip": 75524, + "initialGzip": 80516, "lazyGzip": 2574, - "totalGzip": 78098, - "own": 248907, + "totalGzip": 83090, + "own": 263971, "nodeModules": null, "chunks": [ { "label": "index.js", - "gzip": 71426, + "gzip": 76418, "kind": "initial" }, { @@ -64,17 +64,17 @@ }, { "id": "./beta", - "gzip": 39886, + "gzip": 40063, "composition": { - "initialGzip": 39655, + "initialGzip": 39832, "lazyGzip": 231, - "totalGzip": 39886, - "own": 119365, + "totalGzip": 40063, + "own": 119920, "nodeModules": null, "chunks": [ { "label": "beta.js", - "gzip": 30505, + "gzip": 30577, "kind": "initial" }, { @@ -84,7 +84,7 @@ }, { "label": "service-context.js", - "gzip": 3123, + "gzip": 3228, "kind": "initial" }, { @@ -107,17 +107,17 @@ }, { "id": "./type-generator", - "gzip": 18895, + "gzip": 19068, "composition": { - "initialGzip": 18895, + "initialGzip": 19068, "lazyGzip": 0, - "totalGzip": 18895, - "own": 54388, + "totalGzip": 19068, + "own": 55003, "nodeModules": null, "chunks": [ { "label": "index.js", - "gzip": 18895, + "gzip": 19068, "kind": "initial" } ] @@ -128,25 +128,25 @@ { "name": "@databricks/appkit-ui", "tarball": { - "packed": 296998, - "unpacked": 1254198 + "packed": 313597, + "unpacked": 1314080 }, "dist": { "total": { - "raw": 1250252, - "gzip": 413942 + "raw": 1310134, + "gzip": 434210 }, "js": { - "raw": 354361, - "gzip": 117408 + "raw": 377132, + "gzip": 125150 }, "types": { - "raw": 203935, - "gzip": 73265 + "raw": 209874, + "gzip": 75430 }, "maps": { - "raw": 675096, - "gzip": 219923 + "raw": 706268, + "gzip": 230284 }, "css": { "raw": 16860, @@ -156,22 +156,22 @@ "raw": 0, "gzip": 0 }, - "fileCount": 462 + "fileCount": 478 }, "entries": [ { "id": "./js", - "gzip": 4155, + "gzip": 4254, "composition": { - "initialGzip": 4309, + "initialGzip": 4410, "lazyGzip": 50587, - "totalGzip": 54896, - "own": 11585, + "totalGzip": 54997, + "own": 11865, "nodeModules": 213288, "chunks": [ { "label": "index.js", - "gzip": 4189, + "gzip": 4290, "kind": "initial" }, { @@ -207,17 +207,17 @@ }, { "id": "./react", - "gzip": 45534, + "gzip": 45960, "composition": { - "initialGzip": 437612, + "initialGzip": 438067, "lazyGzip": 49772, - "totalGzip": 487384, - "own": 165847, - "nodeModules": 1402932, + "totalGzip": 487839, + "own": 167329, + "nodeModules": 1402934, "chunks": [ { "label": "index.js", - "gzip": 435462, + "gzip": 435917, "kind": "initial" }, { diff --git a/docs/docs/api/appkit/Class.AppKitError.md b/docs/docs/api/appkit/Class.AppKitError.md index c2803d91..18b852d2 100644 --- a/docs/docs/api/appkit/Class.AppKitError.md +++ b/docs/docs/api/appkit/Class.AppKitError.md @@ -43,6 +43,7 @@ console.error(error.toJSON()); // Safe for logging, sensitive values redacted ```ts new AppKitError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; }): AppKitError; ``` @@ -52,8 +53,9 @@ new AppKitError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | #### Returns @@ -68,6 +70,24 @@ Error.constructor ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +*** + ### cause? ```ts @@ -122,6 +142,23 @@ abstract readonly statusCode: number; HTTP status code suggestion (can be overridden) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + ## Methods ### toJSON() diff --git a/docs/docs/api/appkit/Class.AuthenticationError.md b/docs/docs/api/appkit/Class.AuthenticationError.md index 6d1d57f7..bc70ceed 100644 --- a/docs/docs/api/appkit/Class.AuthenticationError.md +++ b/docs/docs/api/appkit/Class.AuthenticationError.md @@ -21,6 +21,7 @@ throw new AuthenticationError("Failed to generate credentials", { cause: origina ```ts new AuthenticationError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; }): AuthenticationError; ``` @@ -30,8 +31,9 @@ new AuthenticationError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | #### Returns @@ -44,6 +46,28 @@ new AuthenticationError(message: string, options?: { ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + ### cause? ```ts @@ -112,6 +136,27 @@ HTTP status code suggestion (can be overridden) [`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() diff --git a/docs/docs/api/appkit/Class.ConfigurationError.md b/docs/docs/api/appkit/Class.ConfigurationError.md index ef71bf3a..5734349d 100644 --- a/docs/docs/api/appkit/Class.ConfigurationError.md +++ b/docs/docs/api/appkit/Class.ConfigurationError.md @@ -21,6 +21,7 @@ throw new ConfigurationError("Warehouse ID not found", { context: { env: "produc ```ts new ConfigurationError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; }): ConfigurationError; ``` @@ -30,8 +31,9 @@ new ConfigurationError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | #### Returns @@ -44,6 +46,28 @@ new ConfigurationError(message: string, options?: { ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + ### cause? ```ts @@ -112,6 +136,27 @@ HTTP status code suggestion (can be overridden) [`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() diff --git a/docs/docs/api/appkit/Class.ConnectionError.md b/docs/docs/api/appkit/Class.ConnectionError.md index 12a60739..afb6e6e3 100644 --- a/docs/docs/api/appkit/Class.ConnectionError.md +++ b/docs/docs/api/appkit/Class.ConnectionError.md @@ -21,6 +21,7 @@ throw new ConnectionError("No response received from SQL Warehouse API"); ```ts new ConnectionError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; }): ConnectionError; ``` @@ -30,8 +31,9 @@ new ConnectionError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | #### Returns @@ -44,6 +46,28 @@ new ConnectionError(message: string, options?: { ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + ### cause? ```ts @@ -112,6 +136,27 @@ HTTP status code suggestion (can be overridden) [`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() diff --git a/docs/docs/api/appkit/Class.ExecutionError.md b/docs/docs/api/appkit/Class.ExecutionError.md index 75886c4d..e7c6f457 100644 --- a/docs/docs/api/appkit/Class.ExecutionError.md +++ b/docs/docs/api/appkit/Class.ExecutionError.md @@ -21,7 +21,9 @@ throw new ExecutionError("Statement was canceled"); ```ts new ExecutionError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; + errorCode?: string; }): ExecutionError; ``` @@ -30,20 +32,44 @@ new ExecutionError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; `errorCode?`: `string`; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | +| `options.errorCode?` | `string` | #### Returns `ExecutionError` -#### Inherited from +#### Overrides [`AppKitError`](Class.AppKitError.md).[`constructor`](Class.AppKitError.md#constructor) ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + ### cause? ```ts @@ -86,6 +112,19 @@ Additional context for the error *** +### errorCode? + +```ts +readonly optional errorCode: string; +``` + +Structured error code from the upstream source (typically the warehouse's +`error_code` for statement-level failures, or the SDK's `ApiError.errorCode` +for HTTP failures). Preserved through wrapping so callers can branch on a +stable identifier without substring-matching the message. + +*** + ### isRetryable ```ts @@ -112,6 +151,30 @@ HTTP status code suggestion (can be overridden) [`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Execution errors default to a generic message — the raw warehouse / +SDK text in `.message` often includes statement fragments, internal +paths, and correlation IDs. UI code should branch on `errorCode` +(`RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`, etc.) and not on +the human string. + +##### Returns + +`string` + +#### Overrides + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() @@ -202,16 +265,21 @@ Create an execution error for closed/expired results ### statementFailed() ```ts -static statementFailed(errorMessage?: string): ExecutionError; +static statementFailed( + errorMessage?: string, + errorCode?: string, + clientMessage?: string): ExecutionError; ``` -Create an execution error for statement failure +Create an execution error for statement failure. #### Parameters -| Parameter | Type | -| ------ | ------ | -| `errorMessage?` | `string` | +| Parameter | Type | Description | +| ------ | ------ | ------ | +| `errorMessage?` | `string` | Human-readable error from the warehouse / SDK. Goes into `.message` for server logs only — *never* echoed to the client. Pass `clientMessage` explicitly if a sanitized text should reach the UI. | +| `errorCode?` | `string` | Structured code (e.g. "INVALID_PARAMETER_VALUE") to preserve through wrapping. Optional. Forwarded on SSE error payloads so UI can branch on it instead of substring-matching `error`. | +| `clientMessage?` | `string` | Optional client-safe replacement for `.message`. Defaults to "Query execution failed" via the `clientMessage` getter. Set this only when the upstream text is known-safe. | #### Returns diff --git a/docs/docs/api/appkit/Class.InitializationError.md b/docs/docs/api/appkit/Class.InitializationError.md index 96f135c0..0bc11943 100644 --- a/docs/docs/api/appkit/Class.InitializationError.md +++ b/docs/docs/api/appkit/Class.InitializationError.md @@ -21,6 +21,7 @@ throw new InitializationError("ServiceContext not initialized. Call ServiceConte ```ts new InitializationError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; }): InitializationError; ``` @@ -30,8 +31,9 @@ new InitializationError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | #### Returns @@ -44,6 +46,28 @@ new InitializationError(message: string, options?: { ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + ### cause? ```ts @@ -112,6 +136,27 @@ HTTP status code suggestion (can be overridden) [`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() diff --git a/docs/docs/api/appkit/Class.ServerError.md b/docs/docs/api/appkit/Class.ServerError.md index cce86cad..d2b61182 100644 --- a/docs/docs/api/appkit/Class.ServerError.md +++ b/docs/docs/api/appkit/Class.ServerError.md @@ -20,6 +20,7 @@ throw new ServerError("Server not started"); ```ts new ServerError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; }): ServerError; ``` @@ -29,8 +30,9 @@ new ServerError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | #### Returns @@ -43,6 +45,28 @@ new ServerError(message: string, options?: { ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + ### cause? ```ts @@ -111,6 +135,27 @@ HTTP status code suggestion (can be overridden) [`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() diff --git a/docs/docs/api/appkit/Class.TunnelError.md b/docs/docs/api/appkit/Class.TunnelError.md index 20cd6fb9..bfea22bd 100644 --- a/docs/docs/api/appkit/Class.TunnelError.md +++ b/docs/docs/api/appkit/Class.TunnelError.md @@ -21,6 +21,7 @@ throw new TunnelError("Failed to parse WebSocket message", { cause: parseError } ```ts new TunnelError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; }): TunnelError; ``` @@ -30,8 +31,9 @@ new TunnelError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | #### Returns @@ -44,6 +46,28 @@ new TunnelError(message: string, options?: { ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + ### cause? ```ts @@ -112,6 +136,27 @@ HTTP status code suggestion (can be overridden) [`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() diff --git a/docs/docs/api/appkit/Class.ValidationError.md b/docs/docs/api/appkit/Class.ValidationError.md index c00af5e0..f1c760cb 100644 --- a/docs/docs/api/appkit/Class.ValidationError.md +++ b/docs/docs/api/appkit/Class.ValidationError.md @@ -21,6 +21,7 @@ throw new ValidationError("maxPoolSize must be at least 1", { context: { value: ```ts new ValidationError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; }): ValidationError; ``` @@ -30,8 +31,9 @@ new ValidationError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `clientMessage?`: `string`; `context?`: `Record`\<`string`, `unknown`\>; \} | | `options.cause?` | `Error` | +| `options.clientMessage?` | `string` | | `options.context?` | `Record`\<`string`, `unknown`\> | #### Returns @@ -44,6 +46,28 @@ new ValidationError(message: string, options?: { ## Properties +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` +— `message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + ### cause? ```ts @@ -112,6 +136,27 @@ HTTP status code suggestion (can be overridden) [`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() diff --git a/packages/appkit-ui/src/js/arrow/arrow-client.test.ts b/packages/appkit-ui/src/js/arrow/arrow-client.test.ts new file mode 100644 index 00000000..8f9771f0 --- /dev/null +++ b/packages/appkit-ui/src/js/arrow/arrow-client.test.ts @@ -0,0 +1,60 @@ +import { Table, tableToIPC, Utf8, vectorFromArray } from "apache-arrow"; +import { describe, expect, test } from "vitest"; +import { ArrowClient } from "./arrow-client"; + +/** Build an Arrow IPC stream (2 cols, 2 rows) with the given field names. */ +function ipc(names: [string, string]): Uint8Array { + const table = new Table({ + [names[0]]: vectorFromArray(["a", "b"], new Utf8()), + [names[1]]: vectorFromArray(["1", "2"], new Utf8()), + }); + return tableToIPC(table, "stream"); +} + +describe("ArrowClient.processArrowBuffer column relabeling", () => { + test("relabels positional col_N to the provided manifest names", async () => { + const table = await ArrowClient.processArrowBuffer( + ipc(["col_0", "col_1"]), + ["name", "totalSpend"], + ); + expect(table.schema.fields.map((f) => f.name)).toEqual([ + "name", + "totalSpend", + ]); + // Data is preserved under the new names. + expect(table.getChild("name")?.get(0)).toBe("a"); + expect(table.getChild("totalSpend")?.get(1)).toBe("2"); + }); + + test("no-op when no names are provided", async () => { + const table = await ArrowClient.processArrowBuffer(ipc(["col_0", "col_1"])); + expect(table.schema.fields.map((f) => f.name)).toEqual(["col_0", "col_1"]); + }); + + test("no-op when the name count does not match", async () => { + const table = await ArrowClient.processArrowBuffer( + ipc(["col_0", "col_1"]), + ["only_one"], + ); + expect(table.schema.fields.map((f) => f.name)).toEqual(["col_0", "col_1"]); + }); + + test("ignores non-unique names rather than dropping a column", async () => { + const table = await ArrowClient.processArrowBuffer( + ipc(["col_0", "col_1"]), + ["dup", "dup"], + ); + expect(table.schema.fields.map((f) => f.name)).toEqual(["col_0", "col_1"]); + }); + + test("already-correct names are left as-is", async () => { + const table = await ArrowClient.processArrowBuffer( + ipc(["name", "totalSpend"]), + ["name", "totalSpend"], + ); + expect(table.schema.fields.map((f) => f.name)).toEqual([ + "name", + "totalSpend", + ]); + }); +}); diff --git a/packages/appkit-ui/src/js/arrow/arrow-client.ts b/packages/appkit-ui/src/js/arrow/arrow-client.ts index 3ee17363..4554243f 100644 --- a/packages/appkit-ui/src/js/arrow/arrow-client.ts +++ b/packages/appkit-ui/src/js/arrow/arrow-client.ts @@ -1,4 +1,4 @@ -import type { Field, Table } from "apache-arrow"; +import type { Field, Table, Vector } from "apache-arrow"; import { DATE_FIELD_PATTERNS, METADATA_DATE_PATTERNS, @@ -23,14 +23,40 @@ export class ArrowClient { * Lazily loads the Apache Arrow library on first use. * * @param buffer - The Arrow IPC format buffer + * @param columnNames - Optional real column names to relabel the schema + * with. Databricks encodes ARROW_STREAM schemas positionally (`col_0`, + * …); the analytics server sends the manifest names in a response header + * so charts can look columns up by their real name. Relabels in-place by + * rebuilding the Table around the existing vectors (no re-decode). A + * no-op when the names already match, the count differs, or the names + * aren't unique. * @returns Promise resolving to an Arrow Table */ - static async processArrowBuffer(buffer: Uint8Array): Promise { + static async processArrowBuffer( + buffer: Uint8Array, + columnNames?: string[], + ): Promise
{ try { const arrow = await getArrowModule(); // Initialize type ID sets now that Arrow is loaded await initializeTypeIdSets(); - return arrow.tableFromIPC(buffer); + const table = arrow.tableFromIPC(buffer); + + if (!columnNames || columnNames.length !== table.numCols) return table; + const fields = table.schema.fields; + const targetNames = columnNames.map((n, i) => + n && n.length > 0 ? n : fields[i].name, + ); + if (targetNames.every((n, i) => n === fields[i].name)) return table; + if (new Set(targetNames).size !== targetNames.length) return table; + + const children: Record = {}; + for (let i = 0; i < fields.length; i++) { + const child = table.getChildAt(i); + if (!child) return table; + children[targetNames[i]] = child; + } + return new arrow.Table(children); } catch (error) { throw new Error( `Failed to process Arrow buffer: ${ diff --git a/packages/appkit-ui/src/js/sse/connect-sse.ts b/packages/appkit-ui/src/js/sse/connect-sse.ts index c4fd4500..010ba3bf 100644 --- a/packages/appkit-ui/src/js/sse/connect-sse.ts +++ b/packages/appkit-ui/src/js/sse/connect-sse.ts @@ -18,7 +18,11 @@ export async function connectSSE( lastEventId: initialLastEventId = null, retryDelay = 2000, maxRetries = 3, - maxBufferSize = 1024 * 1024, // 1MB + // 1 MiB — matches the server's `streamDefaults.maxEventSize`. SSE + // carries only short JSON control messages; bulk Arrow payloads stream + // back as the raw HTTP response body of the query request, so this + // buffer never needs to hold multi-MiB attachments. + maxBufferSize = 1 * 1024 * 1024, timeout = 300000, // 5 minutes onError, } = options; diff --git a/packages/appkit-ui/src/react/charts/__tests__/types.test.ts b/packages/appkit-ui/src/react/charts/__tests__/types.test.ts index 13394dcf..d6685ce0 100644 --- a/packages/appkit-ui/src/react/charts/__tests__/types.test.ts +++ b/packages/appkit-ui/src/react/charts/__tests__/types.test.ts @@ -93,7 +93,7 @@ describe("isQueryProps", () => { const props = { queryKey: "test_query", parameters: { limit: 100 }, - format: "json" as const, + format: "json_array" as const, }; expect(isQueryProps(props as any)).toBe(true); diff --git a/packages/appkit-ui/src/react/charts/types.ts b/packages/appkit-ui/src/react/charts/types.ts index 60efb1e1..fba131ec 100644 --- a/packages/appkit-ui/src/react/charts/types.ts +++ b/packages/appkit-ui/src/react/charts/types.ts @@ -4,8 +4,22 @@ import type { Table } from "apache-arrow"; // Data Format Types // ============================================================================ -/** Supported data formats for analytics queries */ -export type DataFormat = "json" | "arrow" | "auto"; +/** + * Supported data formats for analytics queries. + * + * "json" and "arrow" are legacy aliases kept for backwards compatibility + * with appkit-ui < 0.33.0 — safe to remove once no consumer is on a + * pre-0.33.0 version. resolveFormat() normalizes them to their canonical + * equivalents before any downstream code reads the value. + */ +export type DataFormat = + | "json_array" + | "arrow_stream" + | "auto" + /** @deprecated Use "json_array". Safe to remove once no consumer is on appkit-ui < 0.33.0. */ + | "json" + /** @deprecated Use "arrow_stream". Safe to remove once no consumer is on appkit-ui < 0.33.0. */ + | "arrow"; /** Chart orientation */ export type Orientation = "vertical" | "horizontal"; @@ -89,8 +103,8 @@ export interface QueryProps extends ChartBaseProps { parameters?: Record; /** * Data format to use - * - "json": Use JSON format (smaller payloads, simpler) - * - "arrow": Use Arrow format (faster for large datasets) + * - "json_array": Use JSON format (smaller payloads, simpler) + * - "arrow_stream": Use Arrow format (faster for large datasets) * - "auto": Automatically select based on expected data size * @default "auto" */ diff --git a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts index bbadded1..4c5f1dd5 100644 --- a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts @@ -1,6 +1,7 @@ import { act, renderHook, waitFor } from "@testing-library/react"; -import { afterEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +let lastConnectArgs: any = null; let capturedCallbacks: { onMessage?: (msg: { data: string }) => void; onError?: (err: Error) => void; @@ -11,34 +12,29 @@ const mockFetchArrow = vi.fn(); const mockProcessArrowBuffer = vi.fn(); // Mock connectSSE so the hook does not attempt a real network request. -const mockConnectSSE = vi - .fn() - .mockImplementation( - (opts: { - onMessage?: (msg: { data: string }) => void; - onError?: (err: Error) => void; - signal?: AbortSignal; - }) => { - capturedCallbacks = { - onMessage: opts.onMessage, - onError: opts.onError, - signal: opts.signal, - }; - return new Promise(() => {}); - }, - ); +// Capture both the full args (used by the arrow/result/error tests) and the +// individual callbacks/signal (used by the warehouse-status and late-envelope +// tests). The hook ignores the return value. +const mockConnectSSE = vi.fn((args: any): unknown => { + lastConnectArgs = args; + capturedCallbacks = { + onMessage: args?.onMessage, + onError: args?.onError, + signal: args?.signal, + }; + return () => {}; +}); vi.mock("@/js", () => ({ + connectSSE: (...args: unknown[]) => mockConnectSSE(...(args as [any])), ArrowClient: { fetchArrow: (...args: unknown[]) => mockFetchArrow(...args), processArrowBuffer: (...args: unknown[]) => mockProcessArrowBuffer(...args), }, - connectSSE: (...args: unknown[]) => mockConnectSSE(...args), })); -// Stub useQueryHMR so we don't pull in import.meta.hot wiring. vi.mock("../use-query-hmr", () => ({ - useQueryHMR: () => {}, + useQueryHMR: vi.fn(), })); import { useAnalyticsQuery } from "../use-analytics-query"; @@ -50,9 +46,206 @@ function markAborted() { } describe("useAnalyticsQuery", () => { - afterEach(() => { - capturedCallbacks = {}; + beforeEach(() => { vi.clearAllMocks(); + lastConnectArgs = null; + capturedCallbacks = {}; + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + test("fetches ARROW_STREAM results as raw Arrow bytes directly from the query endpoint (no SSE)", async () => { + const fakeTable = { numRows: 1, schema: { fields: [] } }; + const fakeBytes = new Uint8Array([1, 2, 3]); + const fetchMock = vi.fn().mockResolvedValue({ + ok: true, + arrayBuffer: async () => fakeBytes.buffer, + headers: { get: () => null }, + }); + vi.stubGlobal("fetch", fetchMock); + mockProcessArrowBuffer.mockResolvedValueOnce(fakeTable); + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await waitFor(() => { + expect(result.current.data).toBe(fakeTable); + }); + + // ARROW_STREAM never opens an SSE stream — the bytes come straight back + // on a direct POST to the query endpoint. + expect(mockConnectSSE).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(1); + const [url, init] = fetchMock.mock.calls[0]; + expect(String(url)).toContain("/api/analytics/query/q"); + expect(init.method).toBe("POST"); + // No column-names header → decode with the raw Arrow schema names. + expect(mockProcessArrowBuffer).toHaveBeenCalledWith( + expect.any(Uint8Array), + undefined, + ); + expect(result.current.loading).toBe(false); + }); + + test("relabels ARROW_STREAM columns from the X-Appkit-Arrow-Columns header", async () => { + const fakeTable = { numRows: 1, schema: { fields: [] } }; + const names = ["name", "totalSpend"]; + const fetchMock = vi.fn().mockResolvedValue({ + ok: true, + arrayBuffer: async () => new Uint8Array([1, 2, 3]).buffer, + headers: { + get: (h: string) => + h === "X-Appkit-Arrow-Columns" + ? encodeURIComponent(JSON.stringify(names)) + : null, + }, + }); + vi.stubGlobal("fetch", fetchMock); + mockProcessArrowBuffer.mockResolvedValueOnce(fakeTable); + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await waitFor(() => { + expect(result.current.data).toBe(fakeTable); + }); + + // The parsed manifest names are handed to the decoder for relabeling. + expect(mockProcessArrowBuffer).toHaveBeenCalledWith( + expect.any(Uint8Array), + names, + ); + }); + + test("surfaces an error when the ARROW_STREAM fetch fails", async () => { + const fetchMock = vi.fn().mockRejectedValue(new Error("boom")); + vi.stubGlobal("fetch", fetchMock); + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await waitFor(() => { + expect(result.current.error).toBe( + "Unable to load data, please try again", + ); + }); + expect(result.current.loading).toBe(false); + }); + + test("surfaces a structured errorCode from an ARROW_STREAM JSON error body", async () => { + // On a pre-first-byte failure the server responds with a JSON + // `{ error, errorCode }` body; the hook exposes both so consumers can + // branch on the stable code (e.g. RESULT_TOO_LARGE_FOR_JSON_FALLBACK). + const fetchMock = vi.fn().mockResolvedValue({ + ok: false, + json: async () => ({ + error: "Result too large for JSON format", + errorCode: "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", + }), + }); + vi.stubGlobal("fetch", fetchMock); + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await waitFor(() => { + expect(result.current.error).toBe("Result too large for JSON format"); + }); + expect(result.current.errorCode).toBe("RESULT_TOO_LARGE_FOR_JSON_FALLBACK"); + expect(result.current.loading).toBe(false); + }); + + test("normalizes an empty result message (no data field) to []", async () => { + // The wire schema makes `data` optional — empty result sets may omit + // it. The hook must surface that as an explicit empty array rather + // than `undefined`, so callers can rely on `data` being either null + // (no message yet) or a value of the inferred result type. + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "JSON_ARRAY" }), + ); + + await lastConnectArgs.onMessage({ + data: JSON.stringify({ type: "result" }), + }); + + await waitFor(() => { + expect(result.current.data).toEqual([]); + }); + expect(result.current.loading).toBe(false); + expect(result.current.error).toBeNull(); + }); + + test("still handles type:result rows for JSON_ARRAY", async () => { + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "JSON_ARRAY" }), + ); + + await lastConnectArgs.onMessage({ + data: JSON.stringify({ + type: "result", + data: [{ id: 1 }, { id: 2 }], + }), + }); + + await waitFor(() => { + expect(result.current.data).toEqual([{ id: 1 }, { id: 2 }]); + }); + expect(mockProcessArrowBuffer).not.toHaveBeenCalled(); + expect(mockFetchArrow).not.toHaveBeenCalled(); + }); + + test("a malformed (non-JSON) SSE payload clears loading and surfaces an error — does not strand the hook in loading=true", async () => { + // A `JSON.parse` failure inside the SSE handler used to be swallowed + // by the outer catch with only a console.warn, leaving the hook + // permanently in `loading=true` with no error surfaced. The UI would + // spin forever. The handler now reports a user-facing error so the + // consumer can render a retry affordance. + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "JSON_ARRAY" }), + ); + + await lastConnectArgs.onMessage({ data: "not-json{" }); + + await waitFor(() => { + expect(result.current.loading).toBe(false); + }); + expect(result.current.error).toBe("Unable to load data, please try again"); + expect(result.current.data).toBeNull(); + }); + + test("a server error event carrying a structured errorCode exposes it on the hook return value", async () => { + // The SSE error broadcaster forwards an `errorCode` field for UI + // branching. The hook surfaces both the human `error` text AND the + // structured `errorCode` so consumers can branch on the stable + // identifier instead of substring-matching the sanitized human message. + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "JSON_ARRAY" }), + ); + + await lastConnectArgs.onMessage({ + data: JSON.stringify({ + type: "error", + error: "Server is at capacity, please retry", + code: "UPSTREAM_ERROR", + errorCode: "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", + }), + }); + + await waitFor(() => { + expect(result.current.error).toBe("Server is at capacity, please retry"); + }); + expect(result.current.loading).toBe(false); + expect(result.current.errorCode).toBe("RESULT_TOO_LARGE_FOR_JSON_FALLBACK"); + + errorSpy.mockRestore(); }); test("does not refetch when params object is structurally equal across renders", () => { @@ -265,25 +458,5 @@ describe("useAnalyticsQuery", () => { expect(result.current.data).toBeNull(); }); - - test("ignores late arrow envelope after the controller was aborted", async () => { - const { result } = renderHook(() => - useAnalyticsQuery("test_query", null, { format: "ARROW_STREAM" }), - ); - - await waitFor(() => expect(capturedCallbacks.signal).toBeDefined()); - - markAborted(); - - await act(async () => { - await capturedCallbacks.onMessage?.({ - data: JSON.stringify({ type: "arrow", statement_id: "stmt-123" }), - }); - }); - - expect(mockFetchArrow).not.toHaveBeenCalled(); - expect(result.current.data).toBeNull(); - expect(result.current.error).toBeNull(); - }); }); }); diff --git a/packages/appkit-ui/src/react/hooks/__tests__/use-chart-data.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/use-chart-data.test.ts index a4d99a91..686aff31 100644 --- a/packages/appkit-ui/src/react/hooks/__tests__/use-chart-data.test.ts +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-chart-data.test.ts @@ -72,7 +72,7 @@ describe("useChartData", () => { }); describe("format selection", () => { - test("uses JSON format when explicitly specified", () => { + test("uses JSON_ARRAY format when explicitly specified", () => { mockUseAnalyticsQuery.mockReturnValue({ data: [], loading: false, @@ -82,7 +82,7 @@ describe("useChartData", () => { renderHook(() => useChartData({ queryKey: "test", - format: "json", + format: "json_array", }), ); @@ -93,7 +93,7 @@ describe("useChartData", () => { ); }); - test("uses ARROW format when explicitly specified", () => { + test("uses ARROW_STREAM format when explicitly specified", () => { mockUseAnalyticsQuery.mockReturnValue({ data: [], loading: false, @@ -103,7 +103,7 @@ describe("useChartData", () => { renderHook(() => useChartData({ queryKey: "test", - format: "arrow", + format: "arrow_stream", }), ); @@ -114,7 +114,7 @@ describe("useChartData", () => { ); }); - test("auto-selects ARROW for large limit", () => { + test("auto-selects ARROW_STREAM for large limit", () => { mockUseAnalyticsQuery.mockReturnValue({ data: [], loading: false, @@ -136,7 +136,7 @@ describe("useChartData", () => { ); }); - test("auto-selects ARROW for date range queries", () => { + test("auto-selects ARROW_STREAM for date range queries", () => { mockUseAnalyticsQuery.mockReturnValue({ data: [], loading: false, @@ -205,7 +205,7 @@ describe("useChartData", () => { ); }); - test("auto-selects JSON by default when no heuristics match", () => { + test("auto-selects JSON_ARRAY by default when no heuristics match", () => { mockUseAnalyticsQuery.mockReturnValue({ data: [], loading: false, @@ -227,7 +227,7 @@ describe("useChartData", () => { ); }); - test("defaults to auto format (JSON) when format is not specified", () => { + test("defaults to auto format (JSON_ARRAY) when format is not specified", () => { mockUseAnalyticsQuery.mockReturnValue({ data: [], loading: false, @@ -353,7 +353,7 @@ describe("useChartData", () => { expect(result.current.isArrow).toBe(false); }); - test("isArrow reflects requested ARROW format when data is null", () => { + test("isArrow reflects requested ARROW_STREAM format when data is null", () => { mockUseAnalyticsQuery.mockReturnValue({ data: null, loading: true, @@ -361,13 +361,13 @@ describe("useChartData", () => { }); const { result } = renderHook(() => - useChartData({ queryKey: "test", format: "arrow" }), + useChartData({ queryKey: "test", format: "arrow_stream" }), ); expect(result.current.isArrow).toBe(true); }); - test("isArrow reflects requested JSON format when data is null", () => { + test("isArrow reflects requested JSON_ARRAY format when data is null", () => { mockUseAnalyticsQuery.mockReturnValue({ data: null, loading: true, @@ -375,7 +375,7 @@ describe("useChartData", () => { }); const { result } = renderHook(() => - useChartData({ queryKey: "test", format: "json" }), + useChartData({ queryKey: "test", format: "json_array" }), ); expect(result.current.isArrow).toBe(false); diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index 8d96b617..aa0df890 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -47,7 +47,7 @@ export interface TypedArrowTable< export interface UseAnalyticsQueryOptions< F extends AnalyticsFormat = "JSON_ARRAY", > { - /** Response format - "JSON_ARRAY" returns typed arrays, "ARROW_STREAM" returns TypedArrowTable */ + /** Response format - "JSON_ARRAY" (default) returns typed arrays, "ARROW_STREAM" uses Arrow (inline or external links) */ format?: F; /** Maximum size of serialized parameters in bytes */ @@ -92,8 +92,16 @@ export interface UseAnalyticsQueryResult { data: T | null; /** Loading state of the query */ loading: boolean; - /** Error state of the query */ + /** Error state of the query — sanitized human-readable message */ error: string | null; + /** + * Structured upstream error code when the server attaches one to the + * error payload (e.g. `RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, + * `ARROW_DELIVERY_UNSUPPORTED`, `NOT_IMPLEMENTED`). Prefer branching UI + * on this stable identifier rather than substring-matching `error`, + * which is a free-form sanitized message. + */ + errorCode: string | null; /** * Latest warehouse status emitted by the server while waiting for the SQL * warehouse to reach RUNNING. `null` until the first status event arrives; diff --git a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts index 006cdbf9..1c63ffe1 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -61,15 +61,25 @@ function getDevMode(): string { return dev ? `?dev=${dev}` : ""; } -function getArrowStreamUrl(id: string): string { - return `/api/analytics/arrow-result/${id}`; -} - const GENERIC_LOAD_ERROR = "Unable to load data, please try again"; +/** Map a fetch/SSE transport error to a user-facing message. */ +function userFacingFetchError(error: unknown): string { + if (error instanceof Error) { + if (error.name === "AbortError") { + return "Request timed out, please try again"; + } + if (error.message.includes("Failed to fetch")) { + return "Network error. Please check your connection."; + } + } + return GENERIC_LOAD_ERROR; +} + interface AnalyticsQuerySseContext { setLoading: (loading: boolean) => void; setError: (error: string | null) => void; + setErrorCode: (code: string | null) => void; setData: (data: ResultType | null) => void; setWarehouseStatus: (status: WarehouseStatus | null) => void; publishWarehouseStatus: (status: WarehouseStatus | null) => void; @@ -104,30 +114,22 @@ async function handleAnalyticsSseMessage( return; } + // JSON result. The SSE wire schema is intentionally loose (`data` is an + // optional array of unknown values), so a structural check is enough here — + // no need to ship a schema validator (zod, ~60 KB gz) to the browser just + // to read our own same-origin server's messages. Missing or non-array + // `data` normalizes to [] so `undefined` never bleeds into the hook's + // `T | null` state. if (parsed.type === "result") { ctx.setLoading(false); - ctx.setData(parsed.data as ResultType); + ctx.setData((Array.isArray(parsed.data) ? parsed.data : []) as ResultType); ctx.unpublishWarehouseStatus(); return; } - if (parsed.type === "arrow") { - try { - const arrowData = await ArrowClient.fetchArrow( - getArrowStreamUrl(parsed.statement_id as string), - ); - const table = await ArrowClient.processArrowBuffer(arrowData); - ctx.setLoading(false); - ctx.setData(table as ResultType); - ctx.unpublishWarehouseStatus(); - } catch (error) { - console.error("[useAnalyticsQuery] Failed to fetch Arrow data", error); - ctx.setLoading(false); - ctx.setError(GENERIC_LOAD_ERROR); - ctx.unpublishWarehouseStatus(); - } - return; - } + // NOTE: ARROW_STREAM no longer flows over SSE — the server streams the + // raw Arrow IPC bytes back as the query response body, handled by + // `fetchArrowDirect` instead of this SSE handler. if (parsed.type === "error" || parsed.error || parsed.code) { const errorMsg = @@ -137,16 +139,139 @@ async function handleAnalyticsSseMessage( ctx.setLoading(false); ctx.setError(errorMsg); ctx.unpublishWarehouseStatus(); + // Propagate the upstream structured code so UI consumers can branch on + // a stable identifier (e.g. format-switch on + // RESULT_TOO_LARGE_FOR_JSON_FALLBACK or ARROW_DELIVERY_UNSUPPORTED) + // instead of parsing the human-readable message. + if (typeof parsed.errorCode === "string") { + ctx.setErrorCode(parsed.errorCode); + } if (parsed.code) { console.error( `[useAnalyticsQuery] Code: ${parsed.code}, Message: ${errorMsg}`, ); } + return; } + + // Not a warehouse-status, result, or error event — surface a generic error + // rather than silently dropping an unrecognized payload. + console.error("[useAnalyticsQuery] Unrecognized SSE payload", parsed); + ctx.setLoading(false); + ctx.setError(GENERIC_LOAD_ERROR); + ctx.unpublishWarehouseStatus(); +} + +interface ArrowDirectContext { + url: string; + payload: string; + signal: AbortSignal; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; + setErrorCode: (code: string | null) => void; + setData: (data: unknown) => void; + unpublishWarehouseStatus: () => void; } /** - * Subscribe to an analytics query over SSE and returns its latest result. + * Fetch the real column names for a statement from the fallback endpoint, + * used when a very wide schema's names didn't fit in the response header. + * Returns undefined on any failure so decoding falls back to the raw Arrow + * schema names. + */ +async function fetchArrowColumns( + statementId: string, + signal: AbortSignal, +): Promise { + try { + const res = await fetch( + `/api/analytics/columns/${encodeURIComponent(statementId)}`, + { signal }, + ); + if (!res.ok) return undefined; + const body = (await res.json()) as { columns?: unknown }; + return Array.isArray(body.columns) ? (body.columns as string[]) : undefined; + } catch { + return undefined; + } +} + +/** + * Fetch an ARROW_STREAM query result as raw Arrow IPC bytes directly from + * the query endpoint (no SSE, no second /arrow-result request) and decode + * it into a Table. The server streams the bytes back as the POST response + * body; errors before the first byte arrive as a JSON `{ error, errorCode }`. + */ +async function fetchArrowDirect(ctx: ArrowDirectContext): Promise { + try { + const response = await fetch(ctx.url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: ctx.payload, + signal: ctx.signal, + }); + if (ctx.signal.aborted) return; + + if (!response.ok) { + let message = GENERIC_LOAD_ERROR; + let code: string | null = null; + try { + const body = (await response.json()) as { + error?: string; + errorCode?: string; + }; + if (body.error) message = body.error; + if (typeof body.errorCode === "string") code = body.errorCode; + } catch { + // Non-JSON error body — keep the generic message. + } + ctx.setLoading(false); + ctx.setError(message); + if (code) ctx.setErrorCode(code); + ctx.unpublishWarehouseStatus(); + return; + } + + const buffer = await response.arrayBuffer(); + if (ctx.signal.aborted) return; + // Databricks encodes ARROW_STREAM columns positionally (col_0, …); the + // server sends the real manifest names so we can relabel the decoded + // Table (charts look columns up by name). Normally inline in the + // `X-Appkit-Arrow-Columns` header; for very wide schemas the header + // carries only a statement-id reference and we fetch the names. + let columnNames: string[] | undefined; + const header = response.headers.get("X-Appkit-Arrow-Columns"); + if (header) { + try { + columnNames = JSON.parse(decodeURIComponent(header)); + } catch { + // Malformed header — fall back to the raw Arrow schema names. + } + } else { + const ref = response.headers.get("X-Appkit-Arrow-Columns-Ref"); + if (ref) { + columnNames = await fetchArrowColumns(ref, ctx.signal); + } + } + const table = await ArrowClient.processArrowBuffer( + new Uint8Array(buffer), + columnNames, + ); + ctx.setData(table); + ctx.setLoading(false); + ctx.unpublishWarehouseStatus(); + } catch (error) { + if (ctx.signal.aborted) return; + ctx.setLoading(false); + ctx.unpublishWarehouseStatus(); + ctx.setError(userFacingFetchError(error)); + } +} + +/** + * Subscribe to an analytics query and return its latest result. JSON_ARRAY + * results stream over SSE (with warehouse-readiness progress); ARROW_STREAM + * results are fetched as raw Arrow bytes directly from the query endpoint. * Integration hook between client and analytics plugin. * * The return type is automatically inferred based on the format: @@ -162,13 +287,13 @@ async function handleAnalyticsSseMessage( * @param options - Analytics query settings including format * @returns Query result state with format-appropriate data type * - * @example JSON format (default) + * @example JSON_ARRAY format (default) * ```typescript * const { data } = useAnalyticsQuery("spend_data", params); * // data: Array<{ group_key: string; cost_usd: number; ... }> | null * ``` * - * @example Arrow format + * @example ARROW_STREAM format * ```typescript * const { data } = useAnalyticsQuery("spend_data", params, { format: "ARROW_STREAM" }); * // data: TypedArrowTable<{ group_key: string; cost_usd: number; ... }> | null @@ -194,6 +319,7 @@ export function useAnalyticsQuery< const [data, setData] = useState(null); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); + const [errorCode, setErrorCode] = useState(null); const [warehouseStatus, setWarehouseStatus] = useState(null); const abortControllerRef = useRef(null); @@ -242,6 +368,7 @@ export function useAnalyticsQuery< setLoading(true); setError(null); + setErrorCode(null); setData(null); setWarehouseStatus(null); publishWarehouseStatus(null); @@ -249,9 +376,26 @@ export function useAnalyticsQuery< const abortController = new AbortController(); abortControllerRef.current = abortController; + // ARROW_STREAM: the server streams raw Arrow IPC bytes back on the query + // response body (no SSE). Fetch and decode directly. + if (format === "ARROW_STREAM") { + void fetchArrowDirect({ + url: urlSuffix, + payload, + signal: abortController.signal, + setLoading, + setError, + setErrorCode, + setData: (table) => setData(table as ResultType), + unpublishWarehouseStatus, + }); + return; + } + const sseContext: AnalyticsQuerySseContext = { setLoading, setError, + setErrorCode, setData, setWarehouseStatus, publishWarehouseStatus, @@ -270,7 +414,20 @@ export function useAnalyticsQuery< const parsed = JSON.parse(message.data) as Record; await handleAnalyticsSseMessage(parsed, sseContext); } catch (error) { + // A `JSON.parse` failure (or any other thrown error inside the + // SSE message handler) used to leave the hook permanently in + // `loading=true` with no error surfaced — the UI would just + // spin forever. Clear loading and report a user-facing error + // so the consumer can render a retry affordance. + // + // We also abort the SSE connection: if the upstream is + // emitting un-parseable frames, leaving the stream open just + // re-fires the same failure on the next message. Closing + // forces the consumer into a clean retry path. console.warn("[useAnalyticsQuery] Malformed message received", error); + setLoading(false); + setError(GENERIC_LOAD_ERROR); + abortController.abort(); } }, onError: (error) => { @@ -278,26 +435,21 @@ export function useAnalyticsQuery< setLoading(false); unpublishWarehouseStatus(); - let userMessage = GENERIC_LOAD_ERROR; if (error instanceof Error) { - if (error.name === "AbortError") { - userMessage = "Request timed out, please try again"; - } else if (error.message.includes("Failed to fetch")) { - userMessage = "Network error. Please check your connection."; - } console.error("[useAnalyticsQuery] Error", { queryKey, error: error.message, stack: error.stack, }); } - setError(userMessage); + setError(userFacingFetchError(error)); }, }); }, [ queryKey, payload, urlSuffix, + format, publishWarehouseStatus, unpublishWarehouseStatus, ]); @@ -315,5 +467,5 @@ export function useAnalyticsQuery< useQueryHMR(queryKey, start); - return { data, loading, error, warehouseStatus }; + return { data, loading, error, errorCode, warehouseStatus }; } diff --git a/packages/appkit-ui/src/react/hooks/use-chart-data.ts b/packages/appkit-ui/src/react/hooks/use-chart-data.ts index 9d858d9b..a19acd04 100644 --- a/packages/appkit-ui/src/react/hooks/use-chart-data.ts +++ b/packages/appkit-ui/src/react/hooks/use-chart-data.ts @@ -18,8 +18,8 @@ export interface UseChartDataOptions { parameters?: Record; /** * Data format preference - * - "json": Force JSON format - * - "arrow": Force Arrow format + * - "json_array": Force JSON format + * - "arrow_stream": Force Arrow format * - "auto": Auto-select based on heuristics * @default "auto" */ @@ -59,9 +59,10 @@ function resolveFormat( format: DataFormat, parameters?: Record, ): "JSON_ARRAY" | "ARROW_STREAM" { - // Explicit format selection - if (format === "json") return "JSON_ARRAY"; - if (format === "arrow") return "ARROW_STREAM"; + // Explicit format selection (legacy "json"/"arrow" accepted for back-compat + // with appkit-ui < 0.33.0 — see DataFormat in ../charts/types.ts). + if (format === "json_array" || format === "json") return "JSON_ARRAY"; + if (format === "arrow_stream" || format === "arrow") return "ARROW_STREAM"; // Auto-selection heuristics if (format === "auto") { @@ -105,7 +106,7 @@ function resolveFormat( * // Force Arrow format * const { data } = useChartData({ * queryKey: "big_query", - * format: "arrow" + * format: "arrow_stream" * }); * ``` */ diff --git a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts new file mode 100644 index 00000000..17d099e3 --- /dev/null +++ b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts @@ -0,0 +1,441 @@ +import { + Binary, + Bool, + type DataType, + DateDay, + Decimal, + DurationMicrosecond, + Field, + Float32, + Float64, + Int8, + Int16, + Int32, + Int64, + IntervalYearMonth, + List, + Map_, + Null, + Schema, + Struct, + Table, + TimestampMicrosecond, + tableToIPC, + Utf8, +} from "apache-arrow"; + +/** + * Parse a Databricks SQL type text (the value returned by the Statement + * Execution API in `ColumnInfo.type_text`) into an Apache Arrow DataType. + * + * Supports: + * - All scalar types (STRING, INT, BIGINT, DECIMAL, TIMESTAMP, etc.) + * - Parameterized scalars: DECIMAL(p,s), VARCHAR(n), CHAR(n) + * - Nested types: ARRAY, MAP, STRUCT + * - INTERVAL year-month and day-time variants + * - Backtick-quoted struct field names with embedded `` `` `` escapes + * + * Unknown or unparseable types fall back to Utf8 — empty-Table consumers + * still see a column with the right name; only the inner type is degraded. + */ +export function parseDatabricksType(typeText: string): DataType { + const parser = new TypeParser(typeText); + const result = parser.parseType(); + parser.expectEnd(); + return result; +} + +/** + * Build an empty Arrow IPC stream (base64-encoded) matching the column schema + * returned by the warehouse. Used so ARROW_STREAM responses with no rows still + * deliver a real Arrow Table to the client, preserving the hook's typed + * contract. + */ +export function buildEmptyArrowIPCBase64( + columns: Array<{ + name?: string; + type_text?: string; + type_name?: string; + }>, +): string { + const fields = columns.map((col, index) => { + const typeText = col.type_text ?? col.type_name ?? "STRING"; + let dataType: DataType; + try { + dataType = parseDatabricksType(typeText); + } catch { + dataType = new Utf8(); + } + const name = col.name && col.name.length > 0 ? col.name : `column_${index}`; + return new Field(name, dataType, true); + }); + const schema = new Schema(fields); + const table = new Table(schema); + const ipc = tableToIPC(table, "stream"); + return Buffer.from(ipc).toString("base64"); +} + +// ============================================================================ +// Recursive-descent parser +// ============================================================================ + +class TypeParser { + private readonly input: string; + private pos = 0; + + constructor(input: string) { + this.input = input; + } + + parseType(): DataType { + this.skipWs(); + + let name: string; + if (this.peek() === "`") { + name = this.consumeBacktickIdent(); + } else { + name = this.consumeIdent(); + } + const upper = name.toUpperCase(); + + this.skipWs(); + + if (upper === "INTERVAL") { + return this.parseInterval(); + } + + if (this.peek() === "(") { + this.consume("("); + const args = this.parseNumberArgs(); + this.consume(")"); + this.skipWs(); + return this.makeParameterized(upper, args); + } + + if (this.peek() === "<") { + this.consume("<"); + const result = this.makeGeneric(upper); + this.skipWs(); + this.consume(">"); + return result; + } + + return this.makeScalar(upper); + } + + expectEnd(): void { + this.skipWs(); + if (this.pos < this.input.length) { + throw new Error( + `Unexpected trailing input at position ${this.pos}: "${this.input.slice(this.pos)}"`, + ); + } + } + + // ─── Type constructors ─────────────────────────────────── + + private makeScalar(upper: string): DataType { + switch (upper) { + case "STRING": + case "VARIANT": + return new Utf8(); + case "VARCHAR": + case "CHAR": + return new Utf8(); + case "BINARY": + case "GEOGRAPHY": + case "GEOMETRY": + return new Binary(); + case "BOOLEAN": + case "BOOL": + return new Bool(); + case "TINYINT": + case "BYTE": + return new Int8(); + case "SMALLINT": + case "SHORT": + return new Int16(); + case "INT": + case "INTEGER": + return new Int32(); + case "BIGINT": + case "LONG": + return new Int64(); + case "FLOAT": + case "REAL": + return new Float32(); + case "DOUBLE": + return new Float64(); + case "DECIMAL": + case "NUMERIC": + case "DEC": + return new Decimal(0, 10, 128); + case "DATE": + return new DateDay(); + case "TIMESTAMP": + case "TIMESTAMP_LTZ": + return new TimestampMicrosecond("UTC"); + case "TIMESTAMP_NTZ": + return new TimestampMicrosecond(); + case "VOID": + case "NULL": + return new Null(); + default: + return new Utf8(); + } + } + + private makeParameterized(upper: string, args: number[]): DataType { + switch (upper) { + case "DECIMAL": + case "NUMERIC": + case "DEC": { + const precision = args[0] ?? 10; + const scale = args[1] ?? 0; + // Arrow JS Decimal constructor signature is (scale, precision, bitWidth). + return new Decimal(scale, precision, 128); + } + case "VARCHAR": + case "CHAR": + return new Utf8(); + default: + return new Utf8(); + } + } + + private makeGeneric(upper: string): DataType { + switch (upper) { + case "ARRAY": { + const inner = this.parseType(); + return new List(new Field("item", inner, true)); + } + case "MAP": { + const keyType = this.parseType(); + this.skipWs(); + this.consume(","); + this.skipWs(); + const valueType = this.parseType(); + const entriesStruct = new Struct([ + new Field("key", keyType, false), + new Field("value", valueType, true), + ]); + return new Map_(new Field("entries", entriesStruct, false), false); + } + case "STRUCT": + return this.parseStructFields(); + default: + // Unknown generic — skip to matching '>' and fall back. + this.skipBalancedAngles(); + return new Utf8(); + } + } + + private parseStructFields(): DataType { + const fields: Field[] = []; + while (true) { + this.skipWs(); + if (this.peek() === ">") break; + + let name: string; + if (this.peek() === "`") { + name = this.consumeBacktickIdent(); + } else { + name = this.consumeIdent(); + } + + this.skipWs(); + this.consume(":"); + this.skipWs(); + + const type = this.parseType(); + + // Optional `NOT NULL` and `COMMENT '...'`. Both are accepted by + // Databricks DDL and may appear in `type_text`. + this.skipWs(); + while (this.peekKeyword("NOT")) { + this.consumeIdent(); + this.skipWs(); + if (this.peekKeyword("NULL")) { + this.consumeIdent(); + } + this.skipWs(); + } + if (this.peekKeyword("COMMENT")) { + this.consumeIdent(); + this.skipWs(); + this.consumeStringLiteral(); + this.skipWs(); + } + + fields.push(new Field(name, type, true)); + + this.skipWs(); + if (this.peek() === ",") { + this.consume(","); + } else { + break; + } + } + return new Struct(fields); + } + + private parseInterval(): DataType { + // Grammar: INTERVAL [TO ] + // YEAR / MONTH variants -> IntervalYearMonth + // DAY / HOUR / MINUTE / SECOND variants -> Duration(microsecond) + const seen: string[] = []; + while (this.pos < this.input.length) { + this.skipWs(); + const c = this.peek(); + if (c === "" || c === "," || c === ">" || c === ")") break; + const word = this.consumeIdent().toUpperCase(); + seen.push(word); + } + const isYearMonth = seen.some((w) => w === "YEAR" || w === "MONTH"); + return isYearMonth ? new IntervalYearMonth() : new DurationMicrosecond(); + } + + private parseNumberArgs(): number[] { + const args: number[] = []; + while (true) { + this.skipWs(); + if (this.peek() === ")") break; + args.push(this.consumeNumber()); + this.skipWs(); + if (this.peek() === ",") { + this.consume(","); + } else { + break; + } + } + return args; + } + + // ─── Token utilities ───────────────────────────────────── + + private peek(): string { + return this.input[this.pos] ?? ""; + } + + private peekKeyword(word: string): boolean { + const slice = this.input.slice(this.pos, this.pos + word.length); + if (slice.toUpperCase() !== word.toUpperCase()) return false; + // Must be followed by a non-identifier character (boundary check). + const next = this.input[this.pos + word.length] ?? ""; + return !/[A-Za-z0-9_]/.test(next); + } + + private consume(expected: string): void { + if (this.peek() !== expected) { + throw new Error( + `Expected "${expected}" at position ${this.pos}, got "${this.peek()}" in "${this.input}"`, + ); + } + this.pos++; + } + + private skipWs(): void { + while ( + this.pos < this.input.length && + /\s/.test(this.input[this.pos] ?? "") + ) { + this.pos++; + } + } + + private consumeIdent(): string { + const start = this.pos; + while ( + this.pos < this.input.length && + /[A-Za-z0-9_]/.test(this.input[this.pos] ?? "") + ) { + this.pos++; + } + if (this.pos === start) { + throw new Error( + `Expected identifier at position ${this.pos}, got "${this.peek()}" in "${this.input}"`, + ); + } + return this.input.slice(start, this.pos); + } + + private consumeBacktickIdent(): string { + this.consume("`"); + let value = ""; + while (this.pos < this.input.length) { + if (this.input[this.pos] === "`") { + if (this.input[this.pos + 1] === "`") { + value += "`"; + this.pos += 2; + continue; + } + break; + } + value += this.input[this.pos]; + this.pos++; + } + this.consume("`"); + return value; + } + + private consumeNumber(): number { + const start = this.pos; + while ( + this.pos < this.input.length && + /[0-9]/.test(this.input[this.pos] ?? "") + ) { + this.pos++; + } + if (this.pos === start) { + throw new Error( + `Expected number at position ${this.pos}, got "${this.peek()}" in "${this.input}"`, + ); + } + return Number.parseInt(this.input.slice(start, this.pos), 10); + } + + private consumeStringLiteral(): string { + const quote = this.peek(); + if (quote !== "'" && quote !== '"') { + throw new Error( + `Expected string literal at position ${this.pos}, got "${quote}" in "${this.input}"`, + ); + } + this.pos++; + let value = ""; + while (this.pos < this.input.length) { + const c = this.input[this.pos]; + if (c === "\\") { + // Escape sequence: keep the next char verbatim. + const next = this.input[this.pos + 1]; + if (next !== undefined) { + value += next; + this.pos += 2; + continue; + } + this.pos++; + continue; + } + if (c === quote) { + this.pos++; + return value; + } + value += c; + this.pos++; + } + throw new Error(`Unterminated string literal in "${this.input}"`); + } + + private skipBalancedAngles(): void { + let depth = 1; + while (this.pos < this.input.length && depth > 0) { + const c = this.peek(); + if (c === "<") depth++; + else if (c === ">") { + depth--; + if (depth === 0) return; + } + this.pos++; + } + } +} diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 2f6a9f73..9e924635 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -12,7 +12,10 @@ import { ValidationError, } from "../../errors"; import { createLogger } from "../../logging/logger"; -import { ArrowStreamProcessor } from "../../stream/arrow-stream-processor"; +import { + ArrowStreamProcessor, + type RefreshChunkLink, +} from "../../stream/arrow-stream-processor"; import type { TelemetryProvider } from "../../telemetry"; import { type Counter, @@ -22,12 +25,49 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; +import { buildEmptyArrowIPCBase64 } from "./arrow-schema"; import { executeStatementDefaults } from "./defaults"; import { WarehousePollBackoff } from "./warehouse-poll-backoff"; import { WarehouseStatusEmitter } from "./warehouse-status-emitter"; const logger = createLogger("connectors:sql-warehouse"); +/** + * Real column names from a statement's result manifest (positional + * `column_N` fallback for any blank ones). Databricks encodes ARROW_STREAM + * schemas positionally (`col_0`, …); these names let consumers relabel an + * Arrow result to match the JSON path. Returns `undefined` when the manifest + * carries no columns. + */ +function arrowColumnNames( + response: sql.StatementResponse, +): string[] | undefined { + const cols = response.manifest?.schema?.columns; + if (!cols || cols.length === 0) return undefined; + return cols.map((c, i) => + c.name && c.name.length > 0 ? c.name : `column_${i}`, + ); +} + +/** + * Maximum size for inline Arrow IPC attachments (25 MiB decoded — the + * Databricks Statement Execution API hard cap on INLINE responses). + * + * The analytics route streams the decoded attachment straight back on the + * query response body; this cap bounds the single in-memory copy. Larger + * results fall through to `disposition: "EXTERNAL_LINKS"`, which streams + * chunk-by-chunk and is not bound by this limit. + */ +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`. + */ +const MAX_EXTERNAL_CHUNK_FOLLOWS = 10_000; + interface SQLWarehouseConfig { timeout?: number; telemetry?: TelemetryOptions; @@ -154,8 +194,6 @@ export class SQLWarehouseConnector { if (!this._arrowProcessor) { this._arrowProcessor = new ArrowStreamProcessor({ timeout: this.config.timeout || executeStatementDefaults.timeout, - maxConcurrentDownloads: - ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS, retries: ArrowStreamProcessor.DEFAULT_RETRIES, }); } @@ -274,10 +312,17 @@ export class SQLWarehouseConnector { ); break; case "SUCCEEDED": - result = this._transformDataArray(response); + result = await this._transformDataArray( + response, + workspaceClient, + signal, + ); break; case "FAILED": - throw ExecutionError.statementFailed(status.error?.message); + throw ExecutionError.statementFailed( + status.error?.message, + status.error?.error_code, + ); case "CANCELED": throw ExecutionError.canceled(); case "CLOSED": @@ -324,15 +369,7 @@ export class SQLWarehouseConnector { ); } - if (error instanceof Error && error.name === "AbortError") { - throw error; - } - if (error instanceof AppKitError) { - throw error; - } - throw ExecutionError.statementFailed( - error instanceof Error ? error.message : String(error), - ); + this._rethrowStatementError(error); } finally { // remove abort handler if (abortHandler && signal) { @@ -791,9 +828,16 @@ export class SQLWarehouseConnector { "poll.duration_ms": elapsedTime, }); span.setStatus({ code: SpanStatusCode.OK }); - return this._transformDataArray(response); + return this._transformDataArray( + response, + workspaceClient, + signal, + ); case "FAILED": - throw ExecutionError.statementFailed(status.error?.message); + throw ExecutionError.statementFailed( + status.error?.message, + status.error?.error_code, + ); case "CANCELED": throw ExecutionError.canceled(); case "CLOSED": @@ -816,15 +860,7 @@ export class SQLWarehouseConnector { }); // error logging is handled by executeStatement's catch block (gated on isAborted) - if (error instanceof Error && error.name === "AbortError") { - throw error; - } - if (error instanceof AppKitError) { - throw error; - } - throw ExecutionError.statementFailed( - error instanceof Error ? error.message : String(error), - ); + this._rethrowStatementError(error); } finally { span.end(); } @@ -833,9 +869,57 @@ export class SQLWarehouseConnector { ); } - private _transformDataArray(response: sql.StatementResponse) { + private async _transformDataArray( + response: sql.StatementResponse, + workspaceClient: WorkspaceClient, + signal?: AbortSignal, + ) { if (response.manifest?.format === "ARROW_STREAM") { - return this.updateWithArrowStatus(response); + const result = response.result as + | (sql.ResultData & { attachment?: string }) + | undefined; + + // Inline Arrow: pass the base64 IPC attachment through unmodified so + // the analytics route can stream it to the client, where the existing + // ArrowClient infrastructure decodes it into a Table. Validate size + // here to fail fast on runaway payloads. + if (result?.attachment) { + return this._validateArrowAttachment(response, result.attachment); + } + + // External links: data fetched separately via statement_id. An empty + // array is a zero-row result (some warehouses emit `external_links: []` + // 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) { + 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 + // 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); + if (hasNoRows && response.manifest?.schema?.columns) { + const synthesized = buildEmptyArrowIPCBase64( + response.manifest.schema.columns, + ); + return { + ...response, + result: { ...(result ?? {}), attachment: synthesized }, + }; + } + + // Inline data_array 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) { @@ -881,97 +965,234 @@ export class SQLWarehouseConnector { }; } - private updateWithArrowStatus(response: sql.StatementResponse): { - result: { statement_id: string; status: sql.StatementStatus }; - } { + /** + * Validate a base64 Arrow IPC attachment (INLINE ARROW_STREAM) and attach + * the real column names from the result manifest. + * + * Databricks warehouses encode the Arrow schema positionally (`col_0`, …); + * the real, aliased names live only in `manifest.schema.columns`. Rather + * than decode/re-encode the payload server-side, we carry the names on the + * result (`columnNames`) so the route can hand them to the client in a + * response header and the client relabels the decoded Table — the one + * mechanism used for both INLINE and EXTERNAL_LINKS. + */ + private _validateArrowAttachment( + response: sql.StatementResponse, + attachment: string, + ) { + // Cap the size to protect against unbounded inline payloads from + // misbehaving warehouses. `MAX_INLINE_ATTACHMENT_BYTES` tracks the API's + // ~25 MiB inline hard cap and bounds memory if a server returns a runaway + // response. + // + // Strip whitespace (rare but legal in base64) and account for trailing + // `=` padding so the byte count is exact rather than an upper bound. + const stripped = attachment.replace(/\s+/g, ""); + const padding = stripped.endsWith("==") + ? 2 + : stripped.endsWith("=") + ? 1 + : 0; + const decodedSize = Math.floor((stripped.length * 3) / 4) - padding; + if (decodedSize > MAX_INLINE_ATTACHMENT_BYTES) { + throw ExecutionError.statementFailed( + `Inline Arrow attachment exceeds maximum size (${decodedSize} > ${MAX_INLINE_ATTACHMENT_BYTES} bytes)`, + ); + } + + const columnNames = arrowColumnNames(response); + if (columnNames) { + return { + ...response, + result: { + ...(response.result as sql.ResultData & { + attachment?: string; + columnNames?: string[]; + }), + // `statement_id` 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, + columnNames, + }, + }; + } + + return response; + } + + private async updateWithArrowStatus( + response: sql.StatementResponse, + workspaceClient: WorkspaceClient, + signal?: AbortSignal, + ): Promise<{ + result: { + statement_id: string; + status: sql.StatementStatus; + columnNames?: string[]; + external_links?: sql.ExternalLink[]; + refreshChunkLink?: RefreshChunkLink; + }; + }> { + const statementId = response.statement_id as string; return { result: { - statement_id: response.statement_id as string, + statement_id: statementId, status: { state: response.status?.state, error: response.status?.error, } as sql.StatementStatus, + columnNames: arrowColumnNames(response), + // Resolve the pre-signed links for EVERY chunk in the caller's own + // execution context. Streaming these directly (see + // {@link streamExternalLinks}) avoids a second `getStatement` under + // the ambient service-principal context — which would fail for + // `.obo.sql` (user-owned) statements. + external_links: await this._resolveAllExternalLinks( + workspaceClient, + statementId, + response, + signal, + ), + // Bound to this caller's identity so the streamer can re-mint an + // expired chunk link mid-download (links live <= 15 min; a large/slow + // result can outlast the earliest ones). + refreshChunkLink: this._makeChunkLinkRefresher( + workspaceClient, + statementId, + ), }, }; } - async getArrowData( + /** + * 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 + * 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 + * would be silently truncated to its first chunk. + */ + private async _resolveAllExternalLinks( workspaceClient: WorkspaceClient, - jobId: string, + statementId: string, + response: sql.StatementResponse, signal?: AbortSignal, - ): Promise> { - const startTime = Date.now(); - - return this.telemetry.startActiveSpan( - "arrow.getData", - { - kind: SpanKind.CLIENT, - attributes: { - "db.system": "databricks", - "arrow.job_id": jobId, - }, - }, - async (span: Span) => { - try { - const response = - await workspaceClient.statementExecution.getStatement( - { statement_id: jobId }, - this._createContext(signal), - ); - - const chunks = response.result?.external_links; - const schema = response.manifest?.schema; - - if (!chunks || !schema) { - throw ExecutionError.missingData("chunks or schema"); - } - - span.setAttribute("arrow.chunk_count", chunks.length); - - const result = await this.arrowProcessor.processChunks( - chunks, - schema, - signal, - ); - - span.setAttribute("arrow.data_size_bytes", result.data.length); - span.setStatus({ code: SpanStatusCode.OK }); - - const duration = Date.now() - startTime; - this.telemetryMetrics.queryDuration.record(duration, { - operation: "arrow.getData", - status: "success", - }); + ): Promise { + const first = response.result?.external_links; + if (!first || first.length === 0) return first; + + const links: sql.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; + 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 ?? []; + if (chunkLinks.length === 0) break; + links.push(...chunkLinks); + next = this._nextChunkIndex(chunkLinks); + } + return links; + } - logger.event()?.setContext("sql-warehouse", { - arrow_data_size_bytes: result.data.length, - arrow_job_id: jobId, - }); + /** The `next_chunk_index` advertised by a chunk's links, if any. */ + private _nextChunkIndex(links: sql.ExternalLink[]): number | undefined { + for (const link of links) { + if (link.next_chunk_index != null) return link.next_chunk_index; + } + return undefined; + } - return result; - } catch (error) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: error instanceof Error ? error.message : "Unknown error", - }); - span.recordException(error as Error); + /** + * A closure that re-mints a single chunk's pre-signed link via + * `getStatementResultChunkN`, 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`. + */ + 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 duration = Date.now() - startTime; - this.telemetryMetrics.queryDuration.record(duration, { - operation: "arrow.getData", - status: "error", - }); + /** + * Stream already-resolved EXTERNAL_LINKS chunks as raw Arrow IPC bytes, one + * chunk in memory at a time. Takes the `external_links` the caller already + * has from `executeStatement` (see {@link updateWithArrowStatus}), so there + * is no extra `getStatement` round-trip and no execution-context mismatch — + * the pre-signed URLs need no auth to download. + */ + streamExternalLinks( + chunks: sql.ExternalLink[], + signal?: AbortSignal, + refresh?: RefreshChunkLink, + ): AsyncGenerator { + return this.arrowProcessor.streamChunks(chunks, signal, refresh); + } - logger.error("Failed Arrow job: %s %O", jobId, error); + /** + * Fetch just the real column names for a completed statement from its result + * manifest — no result data. Backs the `/columns/:statementId` fallback the + * client uses when a very wide schema's names exceed the response-header + * size limit. Stateless: re-derives from the warehouse, no server cache. + */ + async getColumnNames( + workspaceClient: WorkspaceClient, + jobId: string, + signal?: AbortSignal, + ): Promise { + const response = await workspaceClient.statementExecution.getStatement( + { statement_id: jobId }, + this._createContext(signal), + ); + return arrowColumnNames(response); + } - if (error instanceof AppKitError) { - throw error; - } - throw ExecutionError.statementFailed( - error instanceof Error ? error.message : String(error), - ); - } - }, + /** + * Normalize an error caught during statement execution/polling and rethrow. + * Abort and already-structured `AppKitError`s pass through untouched; + * anything else is wrapped as an `ExecutionError`, preserving the SDK's + * structured `ApiError.errorCode` (e.g. "INVALID_PARAMETER_VALUE", + * "BAD_REQUEST") so callers can branch on a stable identifier rather than + * substring-matching the message. + */ + private _rethrowStatementError(error: unknown): never { + if (error instanceof Error && error.name === "AbortError") { + throw error; + } + if (error instanceof AppKitError) { + throw error; + } + const sdkErrorCode = + error && typeof error === "object" && "errorCode" in error + ? (error as { errorCode?: unknown }).errorCode + : undefined; + throw ExecutionError.statementFailed( + error instanceof Error ? error.message : String(error), + typeof sdkErrorCode === "string" ? sdkErrorCode : undefined, ); } 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 new file mode 100644 index 00000000..e30b7315 --- /dev/null +++ b/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts @@ -0,0 +1,514 @@ +import { + Binary, + Bool, + type DataType, + DateDay, + Decimal, + DurationMicrosecond, + Float32, + Float64, + Int8, + Int16, + Int32, + Int64, + IntervalYearMonth, + List, + Map_, + Null, + Struct, + TimestampMicrosecond, + Type, + tableFromIPC, + Utf8, +} from "apache-arrow"; +import { describe, expect, test } from "vitest"; +import { buildEmptyArrowIPCBase64, parseDatabricksType } from "../arrow-schema"; + +// ============================================================================ +// Helpers +// ============================================================================ + +/** Walk the type tree and produce a stable string representation for assertions. */ +function typeSummary(t: DataType): string { + if (t instanceof Decimal) return `Decimal(${t.precision},${t.scale})`; + if (t instanceof TimestampMicrosecond) { + const tz = (t as TimestampMicrosecond & { timezone?: string }).timezone; + return tz ? `Timestamp[us,${tz}]` : "Timestamp[us]"; + } + if (t instanceof List) { + const inner = (t.children?.[0]?.type as DataType | undefined) ?? new Utf8(); + return `List<${typeSummary(inner)}>`; + } + if (t instanceof Struct) { + const inner = (t.children ?? []) + .map((f) => `${f.name}:${typeSummary(f.type as DataType)}`) + .join(","); + return `Struct<${inner}>`; + } + if (t instanceof Map_) { + const entries = + (t.children?.[0]?.type as Struct | undefined)?.children ?? []; + const k = entries[0]?.type as DataType | undefined; + const v = entries[1]?.type as DataType | undefined; + return `Map<${typeSummary(k ?? new Utf8())},${typeSummary(v ?? new Utf8())}>`; + } + // Fall back to typeId for primitives. + return Type[t.typeId] ?? t.constructor.name; +} + +// ============================================================================ +// Scalar types +// ============================================================================ + +describe("parseDatabricksType — scalars", () => { + test.each([ + ["STRING", Utf8], + ["VARIANT", Utf8], + ["BINARY", Binary], + ["GEOGRAPHY", Binary], + ["GEOMETRY", Binary], + ["BOOLEAN", Bool], + ["BOOL", Bool], + ["TINYINT", Int8], + ["BYTE", Int8], + ["SMALLINT", Int16], + ["SHORT", Int16], + ["INT", Int32], + ["INTEGER", Int32], + ["BIGINT", Int64], + ["LONG", Int64], + ["FLOAT", Float32], + ["REAL", Float32], + ["DOUBLE", Float64], + ["DATE", DateDay], + ["VOID", Null], + ["NULL", Null], + ] as const)("%s parses to expected type", (input, ctor) => { + const t = parseDatabricksType(input); + expect(t).toBeInstanceOf(ctor); + }); + + test("case-insensitive — lowercase is accepted", () => { + expect(parseDatabricksType("string")).toBeInstanceOf(Utf8); + expect(parseDatabricksType("bigint")).toBeInstanceOf(Int64); + }); + + test("TIMESTAMP defaults to UTC tz", () => { + const t = parseDatabricksType("TIMESTAMP") as TimestampMicrosecond; + expect(t).toBeInstanceOf(TimestampMicrosecond); + expect(t.timezone).toBe("UTC"); + }); + + test("TIMESTAMP_LTZ behaves like TIMESTAMP", () => { + const t = parseDatabricksType("TIMESTAMP_LTZ") as TimestampMicrosecond; + expect(t.timezone).toBe("UTC"); + }); + + test("TIMESTAMP_NTZ has no timezone", () => { + const t = parseDatabricksType("TIMESTAMP_NTZ") as TimestampMicrosecond; + expect(t).toBeInstanceOf(TimestampMicrosecond); + expect(t.timezone == null || t.timezone === "").toBe(true); + }); + + test("Unknown scalar falls back to Utf8 (degraded but doesn't throw)", () => { + expect(parseDatabricksType("SOMETHING_NEW")).toBeInstanceOf(Utf8); + }); +}); + +// ============================================================================ +// Parameterized scalars +// ============================================================================ + +describe("parseDatabricksType — parameterized scalars", () => { + test("VARCHAR(255) → Utf8 (Arrow doesn't track string length)", () => { + expect(parseDatabricksType("VARCHAR(255)")).toBeInstanceOf(Utf8); + }); + + test("CHAR(10) → Utf8", () => { + expect(parseDatabricksType("CHAR(10)")).toBeInstanceOf(Utf8); + }); + + test("DECIMAL(10,2) → Decimal(precision=10, scale=2)", () => { + const t = parseDatabricksType("DECIMAL(10,2)") as Decimal; + expect(t).toBeInstanceOf(Decimal); + expect(t.precision).toBe(10); + expect(t.scale).toBe(2); + }); + + test("DECIMAL(38,0) — max precision, no scale", () => { + const t = parseDatabricksType("DECIMAL(38,0)") as Decimal; + expect(t.precision).toBe(38); + expect(t.scale).toBe(0); + }); + + test("NUMERIC(p,s) is an alias for DECIMAL(p,s)", () => { + const t = parseDatabricksType("NUMERIC(15,4)") as Decimal; + expect(t).toBeInstanceOf(Decimal); + expect(t.precision).toBe(15); + expect(t.scale).toBe(4); + }); + + test("DEC(p,s) is an alias for DECIMAL(p,s)", () => { + const t = parseDatabricksType("DEC(7,3)") as Decimal; + expect(t.precision).toBe(7); + expect(t.scale).toBe(3); + }); + + test("DECIMAL with whitespace inside parens", () => { + const t = parseDatabricksType("DECIMAL( 10 , 2 )") as Decimal; + expect(t.precision).toBe(10); + expect(t.scale).toBe(2); + }); + + test("DECIMAL with single arg (precision only) defaults scale=0", () => { + const t = parseDatabricksType("DECIMAL(20)") as Decimal; + expect(t.precision).toBe(20); + expect(t.scale).toBe(0); + }); + + test("Bare DECIMAL falls back to default precision/scale", () => { + const t = parseDatabricksType("DECIMAL") as Decimal; + expect(t).toBeInstanceOf(Decimal); + expect(typeof t.precision).toBe("number"); + expect(typeof t.scale).toBe("number"); + }); +}); + +// ============================================================================ +// INTERVAL types +// ============================================================================ + +describe("parseDatabricksType — INTERVAL", () => { + test("INTERVAL YEAR → IntervalYearMonth", () => { + expect(parseDatabricksType("INTERVAL YEAR")).toBeInstanceOf( + IntervalYearMonth, + ); + }); + + test("INTERVAL MONTH → IntervalYearMonth", () => { + expect(parseDatabricksType("INTERVAL MONTH")).toBeInstanceOf( + IntervalYearMonth, + ); + }); + + test("INTERVAL YEAR TO MONTH → IntervalYearMonth", () => { + expect(parseDatabricksType("INTERVAL YEAR TO MONTH")).toBeInstanceOf( + IntervalYearMonth, + ); + }); + + test("INTERVAL DAY → DurationMicrosecond", () => { + expect(parseDatabricksType("INTERVAL DAY")).toBeInstanceOf( + DurationMicrosecond, + ); + }); + + test("INTERVAL DAY TO SECOND → DurationMicrosecond", () => { + expect(parseDatabricksType("INTERVAL DAY TO SECOND")).toBeInstanceOf( + DurationMicrosecond, + ); + }); + + test("INTERVAL HOUR TO MINUTE → DurationMicrosecond", () => { + expect(parseDatabricksType("INTERVAL HOUR TO MINUTE")).toBeInstanceOf( + DurationMicrosecond, + ); + }); +}); + +// ============================================================================ +// ARRAY +// ============================================================================ + +describe("parseDatabricksType — ARRAY", () => { + test("ARRAY → List", () => { + const t = parseDatabricksType("ARRAY") as List; + expect(t).toBeInstanceOf(List); + expect(t.children?.[0]?.type).toBeInstanceOf(Utf8); + }); + + test("ARRAY → List", () => { + const t = parseDatabricksType("ARRAY") as List; + expect(t.children?.[0]?.type).toBeInstanceOf(Int32); + }); + + test("ARRAY preserves precision/scale", () => { + const t = parseDatabricksType("ARRAY") as List; + const inner = t.children?.[0]?.type as Decimal; + expect(inner).toBeInstanceOf(Decimal); + expect(inner.precision).toBe(10); + expect(inner.scale).toBe(2); + }); + + test("ARRAY> — nested twice", () => { + const t = parseDatabricksType("ARRAY>") as List; + const inner1 = t.children?.[0]?.type as List; + expect(inner1).toBeInstanceOf(List); + expect(inner1.children?.[0]?.type).toBeInstanceOf(Int32); + }); + + test("ARRAY>> — three levels deep", () => { + expect( + typeSummary(parseDatabricksType("ARRAY>>")), + ).toBe("List>>"); + }); + + test("ARRAY with whitespace", () => { + const t = parseDatabricksType("ARRAY < STRING >") as List; + expect(t.children?.[0]?.type).toBeInstanceOf(Utf8); + }); +}); + +// ============================================================================ +// MAP +// ============================================================================ + +describe("parseDatabricksType — MAP", () => { + test("MAP", () => { + expect(typeSummary(parseDatabricksType("MAP"))).toBe( + "Map", + ); + }); + + test("MAP — with whitespace", () => { + expect(typeSummary(parseDatabricksType("MAP"))).toBe( + "Map", + ); + }); + + test("MAP> — value is nested", () => { + expect(typeSummary(parseDatabricksType("MAP>"))).toBe( + "Map>", + ); + }); + + test("MAP> — fully nested", () => { + expect( + typeSummary(parseDatabricksType("MAP>")), + ).toBe("Map>"); + }); +}); + +// ============================================================================ +// STRUCT +// ============================================================================ + +describe("parseDatabricksType — STRUCT", () => { + test("STRUCT", () => { + const t = parseDatabricksType("STRUCT") as Struct; + expect(t).toBeInstanceOf(Struct); + expect(t.children?.length).toBe(2); + expect(t.children?.[0]?.name).toBe("a"); + expect(t.children?.[0]?.type).toBeInstanceOf(Int32); + expect(t.children?.[1]?.name).toBe("b"); + expect(t.children?.[1]?.type).toBeInstanceOf(Utf8); + }); + + test("STRUCT with whitespace and many fields", () => { + const t = parseDatabricksType( + "STRUCT", + ) as Struct; + expect(t.children?.map((f) => f.name)).toEqual(["id", "name", "ts"]); + expect(t.children?.[0]?.type).toBeInstanceOf(Int64); + expect(t.children?.[2]?.type).toBeInstanceOf(TimestampMicrosecond); + }); + + test("STRUCT with COMMENT on a field", () => { + const t = parseDatabricksType( + "STRUCT", + ) as Struct; + expect(t.children?.length).toBe(2); + expect(t.children?.[0]?.name).toBe("id"); + expect(t.children?.[0]?.type).toBeInstanceOf(Int32); + expect(t.children?.[1]?.name).toBe("name"); + }); + + test("STRUCT with COMMENT containing escaped quote", () => { + const t = parseDatabricksType( + "STRUCT", + ) as Struct; + expect(t.children?.length).toBe(2); + expect(t.children?.[0]?.name).toBe("id"); + }); + + test("STRUCT with NOT NULL annotation on a field", () => { + const t = parseDatabricksType( + "STRUCT", + ) as Struct; + expect(t.children?.length).toBe(2); + expect(t.children?.[0]?.name).toBe("id"); + }); + + test("STRUCT with backticked field name", () => { + const t = parseDatabricksType( + "STRUCT<`weird name`:INT, normal:STRING>", + ) as Struct; + expect(t.children?.[0]?.name).toBe("weird name"); + expect(t.children?.[0]?.type).toBeInstanceOf(Int32); + }); + + test("STRUCT with backticked field name containing escaped backtick", () => { + const t = parseDatabricksType( + "STRUCT<`with``tick`:INT, other:STRING>", + ) as Struct; + expect(t.children?.[0]?.name).toBe("with`tick"); + }); + + test("STRUCT with nested STRUCT", () => { + const t = parseDatabricksType( + "STRUCT, name:STRING>", + ) as Struct; + expect(t.children?.length).toBe(2); + const nested = t.children?.[0]?.type as Struct; + expect(nested).toBeInstanceOf(Struct); + expect(nested.children?.[0]?.name).toBe("inner"); + expect(nested.children?.[0]?.type).toBeInstanceOf(Int32); + }); + + test("Empty STRUCT<>", () => { + const t = parseDatabricksType("STRUCT<>") as Struct; + expect(t).toBeInstanceOf(Struct); + expect(t.children?.length).toBe(0); + }); +}); + +// ============================================================================ +// Deep nesting / mixed types +// ============================================================================ + +describe("parseDatabricksType — deeply nested", () => { + test("MAP>>", () => { + expect( + typeSummary( + parseDatabricksType( + "MAP>>", + ), + ), + ).toBe("Map>>"); + }); + + test("ARRAY>>> — 4 levels mixed", () => { + expect( + typeSummary( + parseDatabricksType( + "ARRAY>>>", + ), + ), + ).toBe("List>>>"); + }); +}); + +// ============================================================================ +// Error / robustness behavior +// ============================================================================ + +describe("parseDatabricksType — error / robustness", () => { + test("trailing garbage throws", () => { + expect(() => parseDatabricksType("INT junk")).toThrow(); + }); + + test("unmatched < throws", () => { + expect(() => parseDatabricksType("ARRAY { + expect(() => parseDatabricksType("DECIMAL(10,2")).toThrow(); + }); + + test("empty string throws", () => { + expect(() => parseDatabricksType("")).toThrow(); + }); +}); + +// ============================================================================ +// buildEmptyArrowIPCBase64 — round-trip +// ============================================================================ + +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" }, + ]; + const b64 = buildEmptyArrowIPCBase64(columns); + const buf = Buffer.from(b64, "base64"); + const table = tableFromIPC(buf); + expect(table.numRows).toBe(0); + expect(table.numCols).toBe(5); + expect(table.schema.fields.map((f) => f.name)).toEqual([ + "user_id", + "name", + "created_at", + "balance", + "active", + ]); + expect( + (table.schema.fields[0]?.type as { bitWidth?: number }).bitWidth, + ).toBe(64); + expect(table.schema.fields[1]?.type).toBeInstanceOf(Utf8); + // After IPC round-trip Arrow JS resolves Timestamp* subclasses to a + // generic Timestamp with `unit` and `timezone`; assert structurally. + expect(table.schema.fields[2]?.type.typeId).toBe(Type.Timestamp); + expect((table.schema.fields[2]?.type as { unit?: number }).unit).toBe(2); // TimeUnit.MICROSECOND + const decimal = table.schema.fields[3]?.type as Decimal; + expect(decimal).toBeInstanceOf(Decimal); + expect(decimal.precision).toBe(10); + expect(decimal.scale).toBe(2); + expect(table.schema.fields[4]?.type).toBeInstanceOf(Bool); + }); + + 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" }, + ]; + const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); + const table = tableFromIPC(buf); + expect(table.numRows).toBe(0); + expect(table.numCols).toBe(3); + expect(table.schema.fields[0]?.type).toBeInstanceOf(List); + expect(table.schema.fields[1]?.type).toBeInstanceOf(Struct); + 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" }]; + const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); + const table = tableFromIPC(buf); + expect( + (table.schema.fields[0]?.type as { bitWidth?: number }).bitWidth, + ).toBe(64); + }); + + test("unknown type degrades to Utf8 without throwing", () => { + const columns = [ + { name: "id", type_text: "BIGINT" }, + { name: "weird", type_text: "FUTURE_TYPE_NOT_YET_SUPPORTED" }, + ]; + const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); + const table = tableFromIPC(buf); + expect( + (table.schema.fields[0]?.type as { bitWidth?: number }).bitWidth, + ).toBe(64); + expect(table.schema.fields[1]?.type).toBeInstanceOf(Utf8); + }); + + test("missing column name gets a synthesized placeholder", () => { + const columns = [{ type_text: "STRING" }, { name: "", type_text: "INT" }]; + const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); + const table = tableFromIPC(buf); + expect(table.schema.fields[0]?.name).toBe("column_0"); + expect(table.schema.fields[1]?.name).toBe("column_1"); + }); + + test("empty schema produces a valid 0-column 0-row Table", () => { + const buf = Buffer.from(buildEmptyArrowIPCBase64([]), "base64"); + const table = tableFromIPC(buf); + expect(table.numRows).toBe(0); + expect(table.numCols).toBe(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 new file mode 100644 index 00000000..f0b03c07 --- /dev/null +++ b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts @@ -0,0 +1,520 @@ +import type { sql } from "@databricks/sdk-experimental"; +import { tableFromIPC } from "apache-arrow"; +import { describe, expect, test, vi } from "vitest"; + +vi.mock("../../../telemetry", () => { + const mockMeter = { + createCounter: () => ({ add: vi.fn() }), + createHistogram: () => ({ record: vi.fn() }), + }; + return { + TelemetryManager: { + getProvider: () => ({ + startActiveSpan: vi.fn(), + getMeter: () => mockMeter, + }), + }, + SpanKind: { CLIENT: 1 }, + SpanStatusCode: { ERROR: 2 }, + }; +}); +vi.mock("../../../logging/logger", () => ({ + createLogger: () => ({ + info: vi.fn(), + debug: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + event: () => null, + }), +})); +vi.mock("../../../stream/arrow-stream-processor", () => ({ + ArrowStreamProcessor: vi.fn(), +})); + +import { SQLWarehouseConnector } from "../client"; + +function createConnector() { + return new SQLWarehouseConnector({ timeout: 30000 }); +} + +// `_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 / +// single-chunk cases; the multi-chunk tests pass a real mock. +function transform( + connector: SQLWarehouseConnector, + response: sql.StatementResponse, + workspaceClient: unknown = {}, +) { + return (connector as any)._transformDataArray(response, workspaceClient); +} + +// Real base64 Arrow IPC from a serverless warehouse returning +// `SELECT 1 AS test_col, 2 AS test_col2` with INLINE + ARROW_STREAM. +// Contains schema (two INT columns) + one record batch with values [1, 2]. +const REAL_ARROW_ATTACHMENT = + "/////7gAAAAQAAAAAAAKAAwACgAJAAQACgAAABAAAAAAAQQACAAIAAAABAAIAAAABAAAAAIAAABMAAAABAAAAMz///8QAAAAGAAAAAAAAQIUAAAAvP///yAAAAAAAAABAAAAAAkAAAB0ZXN0X2NvbDIAAAAQABQAEAAOAA8ABAAAAAgAEAAAABgAAAAgAAAAAAABAhwAAAAIAAwABAALAAgAAAAgAAAAAAAAAQAAAAAIAAAAdGVzdF9jb2wAAAAA/////7gAAAAQAAAADAAaABgAFwAEAAgADAAAACAAAAAAAQAAAAAAAAAAAAAAAAADBAAKABgADAAIAAQACgAAADwAAAAQAAAAAQAAAAAAAAAAAAAAAgAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAEAAAAAAAAAQAAAAAAAAAAEAAAAAAAAAIAAAAAAAAAAAQAAAAAAAADAAAAAAAAAAAQAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP////8AAAAA"; + +describe("SQLWarehouseConnector._transformDataArray", () => { + describe("classic warehouse (JSON_ARRAY + INLINE)", () => { + test("transforms data_array rows into named objects", async () => { + const connector = createConnector(); + // Real response shape from classic warehouse: INLINE + JSON_ARRAY + const response = { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { + format: "JSON_ARRAY", + schema: { + column_count: 2, + columns: [ + { + name: "test_col", + type_text: "INT", + type_name: "INT", + position: 0, + }, + { + name: "test_col2", + type_text: "INT", + type_name: "INT", + position: 1, + }, + ], + }, + total_row_count: 1, + truncated: false, + }, + result: { + data_array: [["1", "2"]], + }, + } as unknown as sql.StatementResponse; + + const result = await transform(connector, response); + expect(result.result.data).toEqual([{ test_col: "1", test_col2: "2" }]); + expect(result.result.data_array).toBeUndefined(); + }); + + test("parses JSON strings in STRING columns", async () => { + const connector = createConnector(); + const response = { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { + format: "JSON_ARRAY", + schema: { + columns: [ + { name: "id", type_name: "INT" }, + { name: "metadata", type_name: "STRING" }, + ], + }, + }, + result: { + data_array: [["1", '{"key":"value"}']], + }, + } as unknown as sql.StatementResponse; + + const result = await transform(connector, response); + expect(result.result.data[0].metadata).toEqual({ key: "value" }); + }); + }); + + describe("classic warehouse (EXTERNAL_LINKS + ARROW_STREAM)", () => { + test("returns statement_id for external links fetch", async () => { + const connector = createConnector(); + // Real response shape from classic warehouse: EXTERNAL_LINKS + ARROW_STREAM + const response = { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + schema: { + columns: [ + { name: "test_col", type_name: "INT" }, + { name: "test_col2", type_name: "INT" }, + ], + }, + }, + result: { + external_links: [ + { + external_link: "https://storage.example.com/chunk0", + expiration: "2026-04-15T00:00:00Z", + }, + ], + }, + } as unknown as sql.StatementResponse; + + const result = await transform(connector, response); + expect(result.result.statement_id).toBe("stmt-1"); + expect(result.result.data).toBeUndefined(); + }); + }); + + describe("serverless warehouse (INLINE + ARROW_STREAM with attachment)", () => { + 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. + const response = { + statement_id: "00000001-test-stmt", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + schema: { + column_count: 2, + columns: [ + { + name: "test_col", + type_text: "INT", + type_name: "INT", + position: 0, + }, + { + name: "test_col2", + type_text: "INT", + type_name: "INT", + position: 1, + }, + ], + total_chunk_count: 1, + chunks: [{ chunk_index: 0, row_offset: 0, row_count: 1 }], + total_row_count: 1, + }, + truncated: false, + }, + result: { + chunk_index: 0, + row_offset: 0, + row_count: 1, + attachment: REAL_ARROW_ATTACHMENT, + }, + } as unknown as sql.StatementResponse; + + const result = await transform(connector, response); + expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); + expect(result.result.data).toBeUndefined(); + // Preserves other result fields + expect(result.result.row_count).toBe(1); + // `statement_id` is projected onto the result so the route can advertise + // an `X-Appkit-Arrow-Columns-Ref` for wide inline schemas (it lives on + // the top-level response, not on `ResultData`). + expect(result.result.statement_id).toBe("00000001-test-stmt"); + }); + + test("preserves manifest and status alongside attachment", async () => { + const connector = createConnector(); + const response = { + statement_id: "00000001-test-stmt", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + schema: { + columns: [ + { name: "test_col", type_name: "INT" }, + { name: "test_col2", type_name: "INT" }, + ], + }, + }, + result: { + chunk_index: 0, + row_count: 1, + attachment: REAL_ARROW_ATTACHMENT, + }, + } as unknown as sql.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.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 + // 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", + 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: "balance", + type_text: "DECIMAL(10,2)", + type_name: "DECIMAL", + }, + ], + }, + total_row_count: 0, + }, + result: {}, + } as unknown as sql.StatementResponse; + + const transformed = await transform(connector, response); + const attachment: string = transformed.result.attachment; + expect(typeof attachment).toBe("string"); + expect(attachment.length).toBeGreaterThan(0); + + // Verify the synthesized attachment decodes into the right empty schema. + const table = tableFromIPC(Buffer.from(attachment, "base64")); + expect(table.numRows).toBe(0); + expect(table.schema.fields.map((f) => f.name)).toEqual([ + "user_id", + "name", + "balance", + ]); + }); + + test("does NOT synthesize an attachment when external_links are present", async () => { + const connector = createConnector(); + const response = { + statement_id: "stmt-ext", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + schema: { columns: [{ name: "x", type_text: "INT" }] }, + }, + result: { + external_links: [ + { external_link: "https://example.com/x", expiration: "9999" }, + ], + }, + } as unknown as sql.StatementResponse; + + const transformed = await transform(connector, response); + // External-links path returns the statement_id projection — no attachment. + expect(transformed.result.attachment).toBeUndefined(); + expect(transformed.result.statement_id).toBe("stmt-ext"); + }); + + 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 + // 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", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + schema: { + columns: [{ name: "x", type_text: "INT", type_name: "INT" }], + }, + total_row_count: 0, + }, + result: { external_links: [] }, + } as unknown as sql.StatementResponse; + + const transformed = await transform(connector, response); + const attachment: string = transformed.result.attachment; + expect(typeof attachment).toBe("string"); + const table = tableFromIPC(Buffer.from(attachment, "base64")); + expect(table.numRows).toBe(0); + expect(table.schema.fields.map((f) => f.name)).toEqual(["x"]); + }); + + test("does NOT synthesize an attachment when schema is missing", async () => { + const connector = createConnector(); + const response = { + statement_id: "stmt-no-schema", + status: { state: "SUCCEEDED" }, + manifest: { format: "ARROW_STREAM" }, + result: {}, + } as unknown as sql.StatementResponse; + + const transformed = await transform(connector, response); + // Without a schema we cannot build a Table — pass through unchanged. + expect(transformed.result?.attachment).toBeUndefined(); + }); + + test("rejects oversized attachments to bound memory", async () => { + const connector = createConnector(); + // 25 MiB decoded cap (Databricks API hard cap on INLINE) → 36 MiB of + // base64 chars decodes to ~27 MiB, comfortably above the limit. + const oversized = "A".repeat(36 * 1024 * 1024); + const response = { + statement_id: "stmt-oversized", + status: { state: "SUCCEEDED" }, + manifest: { format: "ARROW_STREAM" }, + result: { attachment: oversized }, + } as unknown as sql.StatementResponse; + + await expect(transform(connector, response)).rejects.toThrow( + /exceeds maximum size/, + ); + }); + }); + + describe("ARROW_STREAM with data_array (hypothetical inline variant)", () => { + test("transforms data_array like JSON_ARRAY path", async () => { + const connector = createConnector(); + const response = { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + schema: { + columns: [ + { name: "id", type_name: "INT" }, + { name: "value", type_name: "STRING" }, + ], + }, + }, + result: { + data_array: [ + ["1", "hello"], + ["2", "world"], + ], + }, + } as unknown as sql.StatementResponse; + + const result = await transform(connector, response); + expect(result.result.data).toEqual([ + { id: "1", value: "hello" }, + { id: "2", value: "world" }, + ]); + }); + }); + + describe("edge cases", () => { + test("returns response unchanged when no data_array, attachment, or schema", async () => { + const connector = createConnector(); + const response = { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { format: "JSON_ARRAY" }, + result: {}, + } as unknown as sql.StatementResponse; + + const result = await transform(connector, response); + expect(result).toBe(response); + }); + + test("attachment takes priority over data_array when both present", async () => { + const connector = createConnector(); + const response = { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + schema: { + columns: [ + { name: "test_col", type_name: "INT" }, + { name: "test_col2", type_name: "INT" }, + ], + }, + }, + result: { + attachment: REAL_ARROW_ATTACHMENT, + data_array: [["999", "999"]], + }, + } as unknown as sql.StatementResponse; + + const result = await transform(connector, response); + // Should pass attachment through (client decodes), not transform data_array + 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 { + return { + statement_id: "stmt-multi", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + total_chunk_count: totalChunks, + schema: { columns: [{ name: "x", type_name: "INT" }] }, + }, + result: { + external_links: [ + { + chunk_index: 0, + external_link: "https://example.com/chunk0", + next_chunk_index: 1, + }, + ], + }, + } as unknown as sql.StatementResponse; + } + + test("follows next_chunk_index to resolve every chunk's links", async () => { + const connector = createConnector(); + const getStatementResultChunkN = vi + .fn() + .mockResolvedValueOnce({ + external_links: [ + { + chunk_index: 1, + external_link: "https://example.com/chunk1", + next_chunk_index: 2, + }, + ], + }) + .mockResolvedValueOnce({ + external_links: [ + { chunk_index: 2, external_link: "https://example.com/chunk2" }, + ], + }); + const workspaceClient = { + statementExecution: { getStatementResultChunkN }, + }; + + const result = await transform( + connector, + multiChunkResponse(3), + workspaceClient, + ); + + expect(getStatementResultChunkN).toHaveBeenCalledTimes(2); + expect(getStatementResultChunkN).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ statement_id: "stmt-multi", chunk_index: 1 }), + expect.anything(), + ); + expect( + result.result.external_links.map( + (l: sql.ExternalLink) => l.external_link, + ), + ).toEqual([ + "https://example.com/chunk0", + "https://example.com/chunk1", + "https://example.com/chunk2", + ]); + }); + + test("is bounded by total_chunk_count when next_chunk_index never terminates", async () => { + const connector = createConnector(); + // Misbehaving warehouse: always advertises another chunk. + const getStatementResultChunkN = vi.fn().mockResolvedValue({ + external_links: [ + { + chunk_index: 1, + external_link: "https://example.com/loop", + next_chunk_index: 99, + }, + ], + }); + const workspaceClient = { + statementExecution: { getStatementResultChunkN }, + }; + + const result = await transform( + connector, + multiChunkResponse(2), + workspaceClient, + ); + + // Terminates (no hang) — capped at total_chunk_count fetches. + expect(getStatementResultChunkN).toHaveBeenCalledTimes(2); + expect(result.result.external_links.length).toBeGreaterThan(0); + }); + }); +}); diff --git a/packages/appkit/src/errors/base.ts b/packages/appkit/src/errors/base.ts index 50233823..7ed56d21 100644 --- a/packages/appkit/src/errors/base.ts +++ b/packages/appkit/src/errors/base.ts @@ -46,14 +46,32 @@ export abstract class AppKitError extends Error { /** Additional context for the error */ readonly context?: Record; + /** + * Client-safe error message. When set, callers serializing the error to + * a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` + * — `message` may contain raw upstream / SDK text including statement + * fragments, internal object names, and correlation IDs. + * + * Subclasses can set this in their constructor for a fixed sanitized + * string. When unset, `clientMessage` defaults to a generic per-code + * string (see the getter), and the raw `message` is kept server-side + * only. + */ + protected readonly _clientMessage?: string; + constructor( message: string, - options?: { cause?: Error; context?: Record }, + options?: { + cause?: Error; + context?: Record; + clientMessage?: string; + }, ) { super(message); this.name = this.constructor.name; this.cause = options?.cause; this.context = options?.context; + this._clientMessage = options?.clientMessage; // Maintains proper stack trace for where the error was thrown if (Error.captureStackTrace) { @@ -61,6 +79,14 @@ export abstract class AppKitError extends Error { } } + /** + * Sanitized message safe to forward to clients. Override in subclasses + * if a more specific default is appropriate. + */ + get clientMessage(): string { + return this._clientMessage ?? "An internal error occurred"; + } + /** * Convert error to JSON for logging/serialization. * Sensitive values in context are automatically redacted. diff --git a/packages/appkit/src/errors/execution.ts b/packages/appkit/src/errors/execution.ts index 42de7704..1c9acf91 100644 --- a/packages/appkit/src/errors/execution.ts +++ b/packages/appkit/src/errors/execution.ts @@ -16,20 +16,69 @@ export class ExecutionError extends AppKitError { readonly isRetryable = false; /** - * Create an execution error for statement failure + * Structured error code from the upstream source (typically the warehouse's + * `error_code` for statement-level failures, or the SDK's `ApiError.errorCode` + * for HTTP failures). Preserved through wrapping so callers can branch on a + * stable identifier without substring-matching the message. */ - static statementFailed(errorMessage?: string): ExecutionError { + readonly errorCode?: string; + + constructor( + message: string, + options?: { + cause?: Error; + context?: Record; + errorCode?: string; + clientMessage?: string; + }, + ) { + super(message, options); + this.errorCode = options?.errorCode; + } + + /** + * Execution errors default to a generic message — the raw warehouse / + * SDK text in `.message` often includes statement fragments, internal + * paths, and correlation IDs. UI code should branch on `errorCode` + * (`RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`, etc.) and not on + * the human string. + */ + override get clientMessage(): string { + return this._clientMessage ?? "Query execution failed"; + } + + /** + * Create an execution error for statement failure. + * @param errorMessage Human-readable error from the warehouse / SDK. + * Goes into `.message` for server logs only — *never* echoed to the + * client. Pass `clientMessage` explicitly if a sanitized text should + * reach the UI. + * @param errorCode Structured code (e.g. "INVALID_PARAMETER_VALUE") to + * preserve through wrapping. Optional. Forwarded on SSE error + * payloads so UI can branch on it instead of substring-matching + * `error`. + * @param clientMessage Optional client-safe replacement for `.message`. + * Defaults to "Query execution failed" via the `clientMessage` + * getter. Set this only when the upstream text is known-safe. + */ + static statementFailed( + errorMessage?: string, + errorCode?: string, + clientMessage?: string, + ): ExecutionError { const message = errorMessage ? `Statement failed: ${errorMessage}` : "Statement failed: Unknown error"; - return new ExecutionError(message); + return new ExecutionError(message, { errorCode, clientMessage }); } /** * Create an execution error for canceled operation */ static canceled(): ExecutionError { - return new ExecutionError("Statement was canceled"); + return new ExecutionError("Statement was canceled", { + clientMessage: "Query was canceled", + }); } /** @@ -38,6 +87,7 @@ export class ExecutionError extends AppKitError { static resultsClosed(): ExecutionError { return new ExecutionError( "Statement execution completed but results are no longer available (CLOSED state)", + { clientMessage: "Query results expired" }, ); } diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index c63ec094..922f8a9b 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -4,7 +4,6 @@ import type express from "express"; import pc from "picocolors"; import type { AgentAdapter, - AgentEvent, AgentRunContext, AgentToolDefinition, IAppRouter, @@ -17,7 +16,6 @@ import type { ToolProvider, } from "shared"; import { AppKitMcpClient, buildMcpHostPolicy } from "../../connectors/mcp"; -import { getWorkspaceClient } from "../../context"; import { consumeAdapterStream } from "../../core/agent/consume-adapter-stream"; import { loadAgentsFromDir } from "../../core/agent/load-agents"; import { normalizeToolResult } from "../../core/agent/normalize-result"; diff --git a/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts b/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts index 8aba4ce1..2a4e4a22 100644 --- a/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts @@ -308,7 +308,7 @@ describe("runSubAgent — depth guard", () => { * so we can drive `runSubAgent` directly against the depth guard. */ function makeRunState( - plugin: AgentsPlugin, + _plugin: AgentsPlugin, overrides: Partial<{ maxToolCalls: number; maxSubAgentDepth: number; diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 282e3cbc..4ad05cc3 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1,12 +1,13 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; import type express from "express"; -import type { - AgentToolDefinition, - IAppRouter, - PluginExecuteConfig, - SQLTypeMarker, - StreamExecutionSettings, - ToolProvider, +import { + type AgentToolDefinition, + type AnalyticsSseMessage, + type IAppRouter, + makeResultMessage, + type PluginExecuteConfig, + type SQLTypeMarker, + type StreamExecutionSettings, + type ToolProvider, } from "shared"; import { z } from "zod"; import { SQLWarehouseConnector } from "../../connectors"; @@ -22,21 +23,27 @@ import { toolsFromRegistry, } from "../../core/agent/tools/define-tool"; import { assertReadOnlySql } from "../../core/agent/tools/sql-policy"; -import { ExecutionError } from "../../errors"; +import { AppKitError, ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; -import type { - AnalyticsQueryResponse, - AnalyticsStreamMessage, - IAnalyticsConfig, - IAnalyticsQueryRequest, - WarehouseStatus, +import { + type ArrowCapability, + deliverArrowBytes, + deliverJsonResult, + type QueryExecutor, +} from "./result-delivery"; +import { + type AnalyticsQueryResponse, + type AnalyticsStreamMessage, + type IAnalyticsConfig, + type IAnalyticsQueryRequest, + normalizeAnalyticsFormat, + type WarehouseStatus, } from "./types"; -import { normalizeAnalyticsFormat } from "./types"; const logger = createLogger("analytics"); @@ -99,6 +106,16 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { private SQLClient: SQLWarehouseConnector; private queryProcessor: QueryProcessor; + /** + * In-process memo of which arrow delivery mode each warehouse supports + * (keyed by warehouse id). A standard warehouse rejects `INLINE+ARROW_STREAM` + * on every query, so once learned we skip that doomed probe; Reyden stays + * `"inline"`. Capability is a property of the warehouse, not the user, so it + * is not user-scoped. Bounded by the number of distinct warehouses a process + * talks to (effectively one). + */ + private _arrowCapability = new Map(); + constructor(config: IAnalyticsConfig) { super(config); this.config = config; @@ -111,19 +128,6 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } injectRoutes(router: IAppRouter) { - // Arrow data downloads always run as service principal and bypass the - // interceptor chain (execute/executeStream). The original query execution - // handles OBO via executeStream(); this endpoint fetches pre-computed - // results by job ID. - this.route(router, { - name: "arrow", - method: "get", - path: "/arrow-result/:jobId", - handler: async (req: express.Request, res: express.Response) => { - await this._handleArrowRoute(req, res); - }, - }); - this.route(router, { name: "query", method: "post", @@ -132,47 +136,81 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { await this._handleQueryRoute(req, res); }, }); + + // Column-names fallback for very wide Arrow schemas whose names don't fit + // in the `X-Appkit-Arrow-Columns` response header (see + // `_setArrowColumnsHeader`). The client hits this with the statement id + // from `X-Appkit-Arrow-Columns-Ref`. + this.route(router, { + name: "arrow-columns", + method: "get", + path: "/columns/:statementId", + handler: async (req: express.Request, res: express.Response) => { + await this._handleColumnsRoute(req, res); + }, + }); } /** - * Handle Arrow data download requests. - * When called via asUser(req), uses the user's Databricks credentials. + * Column-names fallback endpoint. Re-derives the real column names from the + * statement's result manifest (stateless — no server cache), for the client + * to relabel a positional Arrow schema when the names were too large for the + * response header. */ - async _handleArrowRoute( + async _handleColumnsRoute( req: express.Request, res: express.Response, ): Promise { - try { - const { jobId } = req.params; - const workspaceClient = getWorkspaceClient(); - - logger.debug("Processing Arrow job request for jobId=%s", jobId); - - const event = logger.event(req); - event?.setComponent("analytics", "getArrowData").setContext("analytics", { - job_id: jobId, - plugin: this.name, - }); - - const result = await this.getArrowData(workspaceClient, jobId); - - res.setHeader("Content-Type", "application/octet-stream"); - res.setHeader("Content-Length", result.data.length.toString()); - res.setHeader("Cache-Control", "public, max-age=3600"); + const { statementId } = req.params; + const columns = await this._resolveColumnNames(req, statementId); + if (columns && columns.length > 0) { + res.setHeader("Cache-Control", "no-store"); + res.json({ columns }); + return; + } + res.status(404).json({ + error: "Column names unavailable", + plugin: this.name, + }); + } - logger.debug( - "Sending Arrow buffer: %d bytes for job %s", - result.data.length, - jobId, - ); - res.send(Buffer.from(result.data)); - } catch (error) { - logger.error("Arrow job error: %O", error); - res.status(404).json({ - error: error instanceof Error ? error.message : "Arrow job not found", - plugin: this.name, - }); + /** + * Resolve a statement's real column names, trying the user's identity first + * (required for `.obo.sql` statements, which the service principal cannot + * `getStatement`) then falling back to the service principal (for + * SP-executed statements). Returns undefined if neither identity can read it, + * so the client falls back to the raw positional Arrow schema names. + */ + private async _resolveColumnNames( + req: express.Request, + statementId: string, + ): Promise { + const attempts: Array<() => Promise> = [ + () => this.asUser(req)._getColumnNames(statementId), + () => this._getColumnNames(statementId), + ]; + for (const attempt of attempts) { + try { + const columns = await attempt(); + if (columns && columns.length > 0) return columns; + } catch (error) { + logger.debug( + "Arrow column-names lookup attempt failed for %s: %O", + statementId, + error, + ); + } } + return undefined; + } + + /** + * Fetch column names in the current execution context. Proxied by `asUser`, + * so `getWorkspaceClient()` resolves to the user's client when invoked via + * `this.asUser(req)` and the service principal's otherwise. + */ + async _getColumnNames(statementId: string): Promise { + return this.SQLClient.getColumnNames(getWorkspaceClient(), statementId); } /** @@ -186,6 +224,19 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { const { query_key } = req.params; const { parameters, format: rawFormat = "JSON_ARRAY" } = req.body as IAnalyticsQueryRequest; + + if ( + rawFormat !== "JSON_ARRAY" && + rawFormat !== "ARROW_STREAM" && + rawFormat !== "JSON" && + rawFormat !== "ARROW" + ) { + res.status(400).json({ + error: `Invalid format: ${String(rawFormat)}. Expected "JSON_ARRAY" or "ARROW_STREAM".`, + }); + return; + } + const format = normalizeAnalyticsFormat(rawFormat); // Request-scoped logging with WideEvent tracking @@ -217,41 +268,48 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { const { query, isAsUser } = queryResult; + // ARROW_STREAM streams the raw Arrow IPC bytes back as the HTTP response + // body — no SSE, no server-side stash, no second /arrow-result request. + // INLINE attachments are piped straight through (the bytes are already in + // hand from executeStatement); a warehouse that refuses INLINE falls back + // to EXTERNAL_LINKS and streams those chunks. JSON keeps the SSE path + // below (it carries warehouse-readiness progress + cached rows). + if (format === "ARROW_STREAM") { + await this._handleArrowStreamQuery( + req, + res, + query_key, + query, + isAsUser, + parameters, + ); + return; + } + // get execution context - user-scoped if .obo.sql, otherwise service principal const executor = isAsUser ? this.asUser(req) : this; const executorKey = isAsUser ? this.resolveUserId(req) : "global"; - const queryParameters = - format === "ARROW_STREAM" - ? { - formatParameters: { - disposition: "EXTERNAL_LINKS", - format: "ARROW_STREAM", - }, - type: "arrow" as const, - } - : { - type: "result" as const, - }; - const hashedQuery = this.queryProcessor.hashQuery(query); + const cacheConfig = { + ...queryDefaults.cache, + cacheKey: [ + "analytics:query", + query_key, + JSON.stringify(parameters), + format, + hashedQuery, + executorKey, + ], + }; + // Cache/retry/timeout are scoped to the SQL execution itself (inner // `execute`) so the warehouse-readiness phase isn't subject to retries // and the generator value never leaks into the cache. const sqlConfig: PluginExecuteConfig = { ...queryDefaults, - cache: { - ...queryDefaults.cache, - cacheKey: [ - "analytics:query", - query_key, - JSON.stringify(parameters), - JSON.stringify(format), - hashedQuery, - executorKey, - ], - }, + cache: cacheConfig, }; // Outer stream: no cache/retry — `executeStream` would otherwise wrap the @@ -303,17 +361,32 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { }; } + // `execute()` reduces a thrown error to `{ status, message }`, + // dropping the rich fields (`errorCode`, `clientMessage`) the + // fallback's `ExecutionError`s carry. Capture the original here so + // we can re-throw it intact — the SSE error path + // (`StreamManager`) reads `errorCode`/`clientMessage` off it. + let originalError: unknown; const sqlResult = await executor.execute( async (sig) => { - const processedParams = - await self.queryProcessor.processQueryParams(query, parameters); - const result = await executor.query( - query, - processedParams, - queryParameters.formatParameters, - sig, - ); - return { type: queryParameters.type, ...result }; + try { + const processedParams = + await self.queryProcessor.processQueryParams(query, parameters); + // JSON_ARRAY path: tries INLINE + JSON_ARRAY and, if the + // warehouse only accepts ARROW_STREAM for INLINE, retries as + // ARROW_STREAM and decodes server-side — returning the SSE + // `result` message with plain rows. (ARROW_STREAM requests are + // handled earlier via `_handleArrowStreamQuery`.) + return await self._executeJsonArrayPath( + executor, + query, + processedParams, + sig, + ); + } catch (err) { + originalError = err; + throw err; + } }, { default: sqlConfig }, executorKey, @@ -333,6 +406,13 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { ); throw err; } + // Re-throw the original error so its structured `errorCode` (e.g. + // RESULT_TOO_LARGE_FOR_JSON_FALLBACK) and sanitized `clientMessage` + // survive to the SSE error payload. Fall back to a generic + // statement failure only if the original wasn't an AppKitError. + if (originalError instanceof AppKitError) { + throw originalError; + } const inner = msg.startsWith("Statement failed: ") ? msg.slice("Statement failed: ".length) : msg; @@ -346,6 +426,297 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { ); } + /** + * JSON_ARRAY SSE path. Delegates the disposition/format fallback to + * {@link deliverJsonResult} (INLINE JSON_ARRAY → on `needs-arrow-inline`, + * INLINE ARROW_STREAM decoded to rows) and wraps the rows in a `result` + * message. External links are never used for the JSON fallback. + */ + private async _executeJsonArrayPath( + executor: AnalyticsPlugin, + query: string, + processedParams: + | Record + | undefined, + signal?: AbortSignal, + ): Promise { + const result = await deliverJsonResult( + executor, + query, + processedParams, + signal, + ); + return makeResultMessage(result.data, { + status: result.status, + statement_id: result.statement_id, + }); + } + + /** + * Attach the real column names so the client can relabel the positional + * Arrow schema (Databricks encodes ARROW_STREAM columns as col_0, …). + * + * Small schemas ride `X-Appkit-Arrow-Columns` directly. A very wide schema + * whose URL-encoded names would blow the HTTP header size limit instead + * advertises the statement id in `X-Appkit-Arrow-Columns-Ref`, and the + * client fetches the names from `GET /columns/:statementId`. + */ + private _setArrowColumnsHeader( + res: express.Response, + columnsRef: { columnNames?: string[]; statementId?: string }, + ): void { + const names = columnsRef.columnNames; + if (!names || names.length === 0) return; + + const encoded = encodeURIComponent(JSON.stringify(names)); + if (encoded.length <= MAX_ARROW_COLUMNS_HEADER_BYTES) { + res.setHeader("X-Appkit-Arrow-Columns", encoded); + return; + } + if (columnsRef.statementId) { + res.setHeader("X-Appkit-Arrow-Columns-Ref", columnsRef.statementId); + } else { + logger.warn( + "Arrow column names exceed the header limit and no statement id is available for the fallback endpoint; client will fall back to the raw schema names", + ); + } + } + + /** + * ARROW_STREAM query handler: stream the raw Arrow IPC bytes back as the + * HTTP response body — no SSE, no server-side stash, no second + * `/arrow-result` request. + * + * The first chunk is pulled before headers are sent so a failure still + * yields a clean JSON error; once bytes are in flight a mid-stream failure + * can only abort the socket. Warehouse readiness is awaited (no SSE + * progress on this path) — a no-op for a warm warehouse, a blocking wait + * on a cold start. Runs under the user's context for `.obo.sql` queries. + */ + private async _handleArrowStreamQuery( + req: express.Request, + res: express.Response, + query_key: string, + query: string, + isAsUser: boolean, + parameters: IAnalyticsQueryRequest["parameters"], + ): Promise { + const executor = isAsUser ? this.asUser(req) : this; + const executorKey = isAsUser ? this.resolveUserId(req) : "global"; + const abortController = new AbortController(); + const onClose = () => abortController.abort(); + res.on("close", onClose); + const signal = abortController.signal; + + // Fail-fast: bound the wait for the first byte (warehouse readiness + + // execute + first chunk) so a stuck/overloaded warehouse returns a clear + // 503 instead of hanging until the client gives up. Cleared once the + // first chunk arrives — a legitimately long stream is never interrupted. + const firstByteTimeoutMs = + this.config.arrowFirstByteTimeoutMs ?? + DEFAULT_ARROW_FIRST_BYTE_TIMEOUT_MS; + let timedOut = false; + const failFast = setTimeout(() => { + timedOut = true; + abortController.abort(); + }, firstByteTimeoutMs); + + try { + // Run warehouse readiness in the SAME identity context as the query: + // for `.obo.sql`, `executor` is the `asUser(req)` proxy, so + // `getWorkspaceClient()` inside resolves to the user's client (matching + // the SSE path). Calling it bare here would auto-start the warehouse as + // the service principal even for OBO requests. + await executor._ensureArrowWarehouseReady(signal); + + const processedParams = await this.queryProcessor.processQueryParams( + query, + parameters, + ); + + const warehouseId = await getWarehouseId(); + + // Populated by `deliverArrowBytes` from the result manifest before the + // first chunk is yielded, so the header below carries the real names. + const columnsRef: { columnNames?: string[]; statementId?: string } = {}; + const bytes = deliverArrowBytes( + // Wrap the executor so the INLINE attempt runs through the interceptor + // chain (cache + retry), matching the JSON path. Only inline attachments + // are cached; EXTERNAL_LINKS carry expiring pre-signed URLs and are + // never cached (see `_arrowCachingExecutor`). + this._arrowCachingExecutor( + executor, + query_key, + query, + parameters, + executorKey, + ), + this.SQLClient, + query, + processedParams, + columnsRef, + signal, + { + // Skip the doomed INLINE probe on a warehouse already known to need + // EXTERNAL_LINKS; remember the resolved mode for next time. + capabilityHint: this._arrowCapability.get(warehouseId), + onCapabilityResolved: (capability) => + this._arrowCapability.set(warehouseId, capability), + }, + ); + const first = await bytes.next(); + // First byte in hand — stop the fail-fast clock. + clearTimeout(failFast); + + res.setHeader("Content-Type", "application/vnd.apache.arrow.stream"); + res.setHeader("Cache-Control", "no-store"); + this._setArrowColumnsHeader(res, columnsRef); + + if (!first.done) { + await writeChunk(res, first.value); + for await (const buf of bytes) { + await writeChunk(res, buf); + } + } + res.end(); + } catch (error) { + clearTimeout(failFast); + // Fail-fast timeout: the warehouse never produced a first byte. This also + // aborts the signal, so it must be handled before the generic + // `signal.aborted` branch below. Headers aren't sent yet (we time out + // before the first chunk), so a clean 503 is still possible. + if (timedOut) { + logger.warn( + "Arrow query timed out before first byte after %dms", + firstByteTimeoutMs, + ); + res.status(503).json({ + error: + "The SQL warehouse is starting or overloaded and did not respond in time. Please retry.", + errorCode: "WAREHOUSE_UNAVAILABLE", + plugin: this.name, + }); + return; + } + // Client disconnect / unmount aborts the signal (see `onClose`). That's + // routine UI behavior, not a server error — tear down quietly whether it + // fires before or after headers. Checked before the headersSent branch so + // a mid-stream disconnect doesn't spam ERROR logs / alerting. + if (signal.aborted) { + if (res.headersSent) res.destroy(); + else res.end(); + return; + } + if (res.headersSent) { + logger.error("Arrow query stream failed mid-flight: %O", error); + res.destroy(error instanceof Error ? error : new Error(String(error))); + return; + } + logger.error("Arrow query error: %O", error); + // Do not echo upstream / SDK error text — it can include statement + // fragments and correlation ids. Keep the structured code so the + // client can branch (e.g. RESULT_TOO_LARGE_FOR_JSON_FALLBACK, + // ARROW_DELIVERY_UNSUPPORTED). + const errorCode = + error instanceof ExecutionError ? error.errorCode : undefined; + res.status(500).json({ + // `clientMessage` is the sanitized, actionable text (e.g. "Re-run with + // JSON_ARRAY" for ARROW_DELIVERY_UNSUPPORTED); it never carries raw + // warehouse/SDK strings. Fall back to a generic message otherwise. + error: + error instanceof AppKitError + ? error.clientMessage + : "Unable to execute query", + errorCode, + plugin: this.name, + }); + } finally { + res.off("close", onClose); + } + } + + /** + * Await SQL warehouse readiness for the direct-binary Arrow path. There is no + * SSE progress channel here — readiness is simply awaited, bounded by the + * caller's abort signal / fail-fast timeout. Invoked via the request executor + * (`asUser(req)` for `.obo.sql`) so `getWorkspaceClient()` resolves in the + * correct identity context rather than defaulting to the service principal. + */ + async _ensureArrowWarehouseReady(signal: AbortSignal): Promise { + const workspaceClient = getWorkspaceClient(); + const warehouseId = await getWarehouseId(); + const startupTimeoutMs = + this.config.warehouseStartupTimeoutMs ?? + DEFAULT_WAREHOUSE_STARTUP_TIMEOUT_MS; + const autoStart = this.config.autoStartWarehouse ?? true; + await this.SQLClient.ensureWarehouseRunning(workspaceClient, warehouseId, { + signal, + timeoutMs: startupTimeoutMs, + autoStart, + onStatus: () => {}, + }); + } + + /** + * Wrap an executor so the `INLINE + ARROW_STREAM` attempt is served from + * (and populates) the same per-user TTL cache the JSON path uses — otherwise + * every arrow chart render is a fresh warehouse execution, unlike its JSON + * twin. Caching is deliberately scoped to inline attachments: + * + * - INLINE results carry a bounded (<=25 MiB) base64 `attachment` with the + * same lifecycle as cached JSON rows — safe to cache. + * - EXTERNAL_LINKS results carry short-lived pre-signed URLs that expire in + * minutes; caching them would serve dead links, so those pass through + * uncached (only the tiny link metadata would be cached anyway). + * + * Uses `this.cache.getOrExecute` directly rather than `this.execute()` + * because `execute()` reduces a thrown error to `{ ok:false, message }`, + * dropping the `errorCode` the capability fallback classifies on. The cache + * re-throws `AppKitError`s intact and never caches a rejection, so the + * INLINE→EXTERNAL_LINKS fallback still sees the structured rejection. + */ + private _arrowCachingExecutor( + executor: AnalyticsPlugin, + query_key: string, + query: string, + parameters: IAnalyticsQueryRequest["parameters"], + executorKey: string, + ): QueryExecutor { + const hashedQuery = this.queryProcessor.hashQuery(query); + const cache = this.cache; + const ttl = queryDefaults.cache?.ttl; + return { + query: (q, params, formatParameters, signal) => { + // Only the inline-arrow attempt is cacheable — EXTERNAL_LINKS carry + // short-lived pre-signed URLs, so those pass straight through. + if ( + formatParameters.disposition !== "INLINE" || + formatParameters.format !== "ARROW_STREAM" + ) { + return executor.query(q, params, formatParameters, signal); + } + // On a standard warehouse this throws a capability rejection — the + // cache never stores a rejection, so the fallback still sees the + // structured error. On Reyden it returns a bounded (<=25 MiB) + // attachment that caches like the JSON path's rows. The shared signal + // dedupes concurrent renders (e.g. React StrictMode double-mount). + return cache.getOrExecute( + [ + "analytics:query:arrow", + query_key, + JSON.stringify(parameters), + hashedQuery, + executorKey, + ], + (sharedSignal) => + executor.query(q, params, formatParameters, sharedSignal ?? signal), + executorKey, + { ttl, callerSignal: signal }, + ); + }, + }; + } + /** * Execute a SQL query using the current execution context. * @@ -387,17 +758,6 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { return response.result; } - /** - * Get Arrow-formatted data for a completed query job. - */ - protected async getArrowData( - workspaceClient: WorkspaceClient, - jobId: string, - signal?: AbortSignal, - ): Promise> { - return await this.SQLClient.getArrowData(workspaceClient, jobId, signal); - } - async shutdown(): Promise { this.streamManager.abortAll(); } @@ -462,6 +822,71 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } } +/** + * Write one chunk to the response honoring backpressure: if the socket + * buffer is full (`res.write` returns false), wait for `drain` before + * resolving so a slow client can't balloon Node's internal write queue and + * defeat the constant-memory goal of streaming. + * + * @internal exported for unit testing the backpressure/disconnect behavior. + */ +export function writeChunk( + res: express.Response, + bytes: Uint8Array, +): Promise { + const buf = Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength); + // If the socket is already gone, `res.write` won't return true and the + // `drain`/`close`/`error` events have already fired — so the promise below + // would never settle, wedging the for-await loop and the upstream reader. + // Reject up front instead. + if (res.destroyed || res.writableEnded) { + return Promise.reject( + new DOMException("The response stream closed", "AbortError"), + ); + } + if (res.write(buf)) return Promise.resolve(); + // Backpressured: resolve on `drain`, but also settle on `close`/`error`. A + // client that disconnects mid-backpressure never emits `drain` on the + // destroyed socket, so waiting on `drain` alone would wedge this promise — + // and with it the awaiting for-await loop and the upstream Arrow reader — + // forever. Rejecting instead unwinds the stream so its `finally` can cancel + // the reader. + return new Promise((resolve, reject) => { + const cleanup = () => { + res.off("drain", onDrain); + res.off("close", onClose); + res.off("error", onClose); + }; + const onDrain = () => { + cleanup(); + resolve(); + }; + const onClose = () => { + cleanup(); + reject(new DOMException("The response stream closed", "AbortError")); + }; + res.once("drain", onDrain); + res.once("close", onClose); + res.once("error", onClose); + }); +} + +/** + * Fail-fast ceiling on the wait for the first Arrow byte (warehouse + * readiness + execute + first chunk). Past this a stuck/overloaded warehouse + * yields a clear 503 instead of hanging. Override per plugin via + * `arrowFirstByteTimeoutMs`. + */ +const DEFAULT_ARROW_FIRST_BYTE_TIMEOUT_MS = 120_000; + +/** + * Byte ceiling for the `X-Appkit-Arrow-Columns` header value. Beyond this (a + * very wide schema) the names are served via the `/columns/:statementId` + * fallback endpoint instead of the header. Kept well under the common ~8 KiB + * per-header limit. + */ +const MAX_ARROW_COLUMNS_HEADER_BYTES = 6000; + /** * @internal */ diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts new file mode 100644 index 00000000..2a323104 --- /dev/null +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -0,0 +1,557 @@ +import type { sql } from "@databricks/sdk-experimental"; +import { type DataType, type Field, Type, tableFromIPC } from "apache-arrow"; +import type { SQLTypeMarker } from "shared"; +import { ExecutionError } from "../../errors"; +import { createLogger } from "../../logging/logger"; +import type { RefreshChunkLink } from "../../stream/arrow-stream-processor"; + +/** + * Centralized disposition/format fallback for analytics result delivery. + * + * Two warehouse capability profiles must both work, and their supported + * combinations are mutually exclusive per warehouse: + * + * - **Reyden**: `ARROW_STREAM + INLINE` (and `JSON_ARRAY + INLINE`); does NOT + * support `EXTERNAL_LINKS`. + * - **Normal SQL warehouse**: `JSON_ARRAY + INLINE` and + * `ARROW_STREAM + EXTERNAL_LINKS`; rejects `ARROW_STREAM + INLINE`. + * + * The fallback logic lives here (rather than inline in the plugin route) so it + * can be unit-tested against a fake executor without HTTP, and so the brittle + * "which rejection is this" decision is expressed once. + */ + +const logger = createLogger("analytics:delivery"); + +/** + * Minimal query surface the fallback needs. The analytics plugin (and its + * `asUser(req)` proxy) satisfy it, so the same delivery logic runs unchanged + * under service-principal and on-behalf-of identities. + */ +export interface QueryExecutor { + query( + query: string, + parameters: Record | undefined, + formatParameters: { disposition: string; format: string }, + signal?: AbortSignal, + ): Promise< + | { + attachment?: string; + data?: Record[]; + external_links?: sql.ExternalLink[]; + columnNames?: string[]; + statement_id?: string; + status?: unknown; + refreshChunkLink?: RefreshChunkLink; + } + | undefined + >; +} + +/** Streams already-resolved EXTERNAL_LINKS chunks; the connector provides it. */ +export interface ArrowChunkStreamer { + streamExternalLinks( + chunks: sql.ExternalLink[], + signal?: AbortSignal, + refresh?: RefreshChunkLink, + ): AsyncGenerator; +} + +/** The arrow delivery mode a warehouse actually supports. */ +export type ArrowCapability = "inline" | "external"; + +/** Optional per-warehouse capability memo, avoiding a doomed probe per query. */ +interface ArrowDeliveryOptions { + /** + * When `"external"`, skip the `INLINE+ARROW_STREAM` attempt and go straight + * to `EXTERNAL_LINKS` — a standard warehouse rejects INLINE arrow on every + * query, so once we've learned that, the probe is pure waste. `"inline"` or + * undefined attempts INLINE first (the common/Reyden path). + */ + capabilityHint?: ArrowCapability; + /** + * Called once the working delivery mode is known, so the caller can memoize + * it per warehouse for subsequent queries. + */ + onCapabilityResolved?: (capability: ArrowCapability) => void; +} + +/** + * How a warehouse rejected a disposition/format combination. + * + * Only the two structured error codes Databricks emits for capability + * mismatches (`INVALID_PARAMETER_VALUE`, `NOT_IMPLEMENTED`) gate a fallback — + * auth, permission, and SQL errors never match, so they propagate untouched. + * Message matching only *disambiguates* which capability is missing once the + * code has confirmed it's a capability error. + */ +type DispositionRejection = + | "needs-arrow-inline" // INLINE requires ARROW_STREAM + | "needs-json-inline" // INLINE requires JSON_ARRAY + | "external-links-unsupported" // EXTERNAL_LINKS not implemented + | null; + +/** + * Whether an error is a disposition/format *capability* rejection — the only + * class of error a fallback should react to. Gated on the two structured codes + * Databricks emits for capability mismatches (`INVALID_PARAMETER_VALUE`, + * `NOT_IMPLEMENTED`); auth, permission, and SQL errors carry other codes and + * never match. This is deliberately message-independent: a Databricks wording + * change must not turn a capability rejection into a hard 500. + */ +function isCapabilityRejection(err: unknown): boolean { + const structuredCode = + err instanceof ExecutionError ? err.errorCode : undefined; + if ( + structuredCode === "INVALID_PARAMETER_VALUE" || + structuredCode === "NOT_IMPLEMENTED" + ) { + return true; + } + const lower = ( + err instanceof Error ? err.message : String(err) + ).toLowerCase(); + return ( + lower.includes("invalid_parameter_value") || + lower.includes("not_implemented") + ); +} + +export function classifyDispositionRejection( + err: unknown, +): DispositionRejection { + // Auth / permission / SQL errors carry other codes → never fall back. + if (!isCapabilityRejection(err)) return null; + + const msg = err instanceof Error ? err.message : String(err); + const lower = msg.toLowerCase(); + + // EXTERNAL_LINKS disposition not implemented (Reyden). + if ( + /external[_\s]?links/.test(lower) && + (lower.includes("not yet implemented") || + lower.includes("not implemented") || + lower.includes("not supported")) + ) { + return "external-links-unsupported"; + } + + // Remaining cases are INLINE disposition/format mismatches. + if (!lower.includes("inline")) return null; + + if ( + /only supports\s+arrow_stream/i.test(msg) || + /must be\s+arrow_stream/i.test(msg) + ) { + return "needs-arrow-inline"; + } + if ( + /only supports\s+json_array/i.test(msg) || + /must be\s+json_array/i.test(msg) || + /arrow_stream\s+(is\s+|was\s+)?not\s+supported/i.test(msg) + ) { + return "needs-json-inline"; + } + + return null; +} + +/** Structured error: the warehouse supports no Arrow delivery mode at all. */ +export function arrowDeliveryUnsupported(): ExecutionError { + return ExecutionError.statementFailed( + "Warehouse supports neither ARROW_STREAM+INLINE nor ARROW_STREAM+EXTERNAL_LINKS", + "ARROW_DELIVERY_UNSUPPORTED", + 'This warehouse cannot return Arrow results. Re-run the query with format="JSON_ARRAY".', + ); +} + +const errMessage = (err: unknown): string => + err instanceof Error ? err.message : String(err); + +/** + * Produce the raw Arrow IPC bytes for an ARROW_STREAM query, applying the + * capability fallback: + * + * 1. `INLINE + ARROW_STREAM` — on success the base64 attachment is decoded + * once and yielded as a `Uint8Array` view (no Arrow parsing, no stash). + * 2. On a `needs-json-inline` rejection, `EXTERNAL_LINKS + ARROW_STREAM` — + * the pre-signed chunks resolved in the executor's own context are + * streamed one at a time. + * 3. If EXTERNAL_LINKS is also unsupported, a structured + * `ARROW_DELIVERY_UNSUPPORTED` error is thrown. + * + * `out` is populated with the real column names and statement id before the + * first chunk is yielded, so the caller can emit them (e.g. as a header) + * ahead of the body. + */ +export async function* deliverArrowBytes( + executor: QueryExecutor, + streamer: ArrowChunkStreamer, + query: string, + processedParams: Record | undefined, + out: { columnNames?: string[]; statementId?: string }, + signal?: AbortSignal, + opts?: ArrowDeliveryOptions, +): AsyncGenerator { + // Skip the INLINE probe when a prior query already learned this warehouse + // only serves arrow via EXTERNAL_LINKS (standard warehouses reject INLINE + // arrow on every query, so probing each time is pure waste). + if (opts?.capabilityHint !== "external") { + try { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "ARROW_STREAM" }, + signal, + ); + if (result?.attachment) { + opts?.onCapabilityResolved?.("inline"); + out.columnNames = result.columnNames; + out.statementId = result.statement_id; + // Decode base64 once; yield a view over the Buffer (no copy, no parse). + const decoded = Buffer.from(result.attachment, "base64"); + yield new Uint8Array( + decoded.buffer, + decoded.byteOffset, + decoded.byteLength, + ); + return; + } + // INLINE succeeded but returned no attachment (rare: inline data_array + // under ARROW_STREAM). Fall through to EXTERNAL_LINKS. + logger.warn( + "ARROW_STREAM INLINE returned no attachment; falling back to EXTERNAL_LINKS", + ); + } catch (err: unknown) { + if (signal?.aborted) throw err; + // Any capability-coded rejection of INLINE+ARROW_STREAM means this + // warehouse won't serve arrow inline — try EXTERNAL_LINKS. We + // intentionally do NOT require the message to match a specific phrase + // ("must be JSON_ARRAY", …): the errorCode gate already excludes auth/SQL + // errors, and relying on exact wording would turn a Databricks message + // reword into a hard 500 on every standard warehouse. Worst case for a + // genuinely bad parameter that happens to carry a capability code: one + // wasted re-execution that fails identically on EXTERNAL_LINKS and + // propagates. + if (!isCapabilityRejection(err)) throw err; + logger.warn( + "ARROW_STREAM INLINE rejected by warehouse, falling back to EXTERNAL_LINKS: %s", + errMessage(err), + ); + } + } + + let ext: Awaited>; + try { + ext = await executor.query( + query, + processedParams, + { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, + signal, + ); + } catch (err: unknown) { + if (signal?.aborted) throw err; + // Neither INLINE nor EXTERNAL_LINKS Arrow is supported — surface a clear, + // actionable error instead of a raw warehouse rejection. + if (classifyDispositionRejection(err) === "external-links-unsupported") { + throw arrowDeliveryUnsupported(); + } + throw err; + } + + if (!ext?.external_links) { + throw ExecutionError.missingData("external_links"); + } + opts?.onCapabilityResolved?.("external"); + out.columnNames = ext.columnNames; + out.statementId = ext.statement_id; + // Stream the pre-signed links resolved in THIS request's execution context + // (user creds for `.obo.sql`, service principal otherwise). Pre-signed URLs + // need no auth to download, so there is no second `getStatement` under a + // mismatched identity. `refreshChunkLink` (also bound to this context) lets + // the streamer re-mint a link that expires mid-download. + yield* streamer.streamExternalLinks( + ext.external_links, + signal, + ext.refreshChunkLink, + ); +} + +/** JSON row result plus the metadata the SSE `result` message forwards. */ +interface JsonResult { + data: Record[] | undefined; + status: unknown; + statement_id?: string; +} + +/** + * Produce JSON rows for a JSON_ARRAY query, applying the capability fallback: + * + * 1. `INLINE + JSON_ARRAY` — the native, common path. + * 2. On a `needs-arrow-inline` rejection (warehouse only accepts ARROW_STREAM + * for INLINE, e.g. Reyden), retry `INLINE + ARROW_STREAM` and decode the + * attachment to plain rows server-side so the caller's JSON contract holds. + * + * EXTERNAL_LINKS is never used for the JSON fallback. + */ +export async function deliverJsonResult( + executor: QueryExecutor, + query: string, + processedParams: Record | undefined, + signal?: AbortSignal, +): Promise { + try { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "JSON_ARRAY" }, + signal, + ); + return { + data: result?.data, + status: result?.status, + statement_id: result?.statement_id, + }; + } catch (err: unknown) { + if (signal?.aborted) throw err; + if (classifyDispositionRejection(err) !== "needs-arrow-inline") throw err; + logger.warn( + "JSON_ARRAY INLINE rejected by warehouse, retrying as ARROW_STREAM INLINE and decoding server-side: %s", + errMessage(err), + ); + } + + const arrowResult = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "ARROW_STREAM" }, + signal, + ); + if (!arrowResult?.attachment) { + throw ExecutionError.missingData("ARROW_STREAM attachment"); + } + const rows = decodeArrowAttachmentToRows( + arrowResult.attachment, + arrowResult.columnNames, + ); + return { + data: rows, + status: arrowResult.status, + statement_id: arrowResult.statement_id, + }; +} + +// ──────────────────────────────────────────────────────────────────────────── +// Server-side Arrow → JSON row materializer (JSON fallback only) +// ──────────────────────────────────────────────────────────────────────────── + +/** + * Hard caps on the server-side JSON_ARRAY fallback. The materializer builds + * every row as a plain JS object on the Node main thread (O(rows × cols) + * allocations), so a runaway result blocks the event loop and pressures GC. + * Cap on rows AND decoded bytes — either dimension can blow up independently. + */ +const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; +const JSON_ARRAY_FALLBACK_MAX_BYTES = 64 * 1024 * 1024; + +/** Render an apache-arrow Decimal cell to a fixed-point string. */ +function formatDecimalCell( + value: { toString(): string }, + scale: number, +): string { + const unscaled = value.toString(); + if (scale <= 0) return unscaled; + const negative = unscaled.startsWith("-"); + let digits = negative ? unscaled.slice(1) : unscaled; + if (digits.length <= scale) digits = digits.padStart(scale + 1, "0"); + const point = digits.length - scale; + const out = `${digits.slice(0, point)}.${digits.slice(point)}`; + return negative ? `-${out}` : out; +} + +/** + * Render an apache-arrow Timestamp cell to ISO-8601 ms precision — `Z` for + * zoned columns, no `Z` for TIMESTAMP_NTZ — matching native JSON_ARRAY. + */ +function formatTimestampCell(epochMs: number, hasTimezone: boolean): string { + const iso = new Date(epochMs).toISOString(); + return hasTimezone ? iso : iso.slice(0, -1); +} + +/** Render an apache-arrow Date cell to `yyyy-MM-dd` (UTC). */ +function formatDateCell(epochMs: number): string { + return new Date(epochMs).toISOString().slice(0, 10); +} + +/** + * Render a single scalar Arrow value to the string form the warehouse emits + * under native JSON_ARRAY. Used for both top-level cells and (recursively, via + * {@link renderNestedValue}) scalar leaves inside List/Struct/Map columns, + * keyed on the Arrow field type so a nested `DECIMAL` keeps its scale and a + * nested `TIMESTAMP` becomes ISO-8601 instead of raw epoch-ms. + */ +function formatScalarCell(value: unknown, type: DataType): string { + switch (type.typeId) { + case Type.Decimal: + return formatDecimalCell( + value as { toString(): string }, + (type as DataType & { scale: number }).scale, + ); + case Type.Timestamp: + return formatTimestampCell( + Number(value), + (type as DataType & { timezone?: string | null }).timezone != null, + ); + case Type.Date: + return formatDateCell(Number(value)); + default: + if (value instanceof Uint8Array) { + return Buffer.from(value).toString("base64"); + } + if (value instanceof Date) return value.toISOString(); + return String(value); + } +} + +/** + * Recursively render a nested Arrow value (List / Struct / Map, or a scalar + * leaf) to the plain JS shape the warehouse nests inside a JSON_ARRAY cell: + * structs and maps become objects, lists become arrays, and every scalar leaf + * is stringified per {@link formatScalarCell}. The caller `JSON.stringify`s the + * result so nested columns match the warehouse's native `data_array` exactly. + */ +function renderNestedValue(value: unknown, type: DataType): unknown { + if (value == null) return null; + switch (type.typeId) { + case Type.List: + case Type.FixedSizeList: { + const itemType = (type as DataType & { children: Field[] }).children[0] + .type; + const out: unknown[] = []; + for (const item of value as Iterable) { + out.push(renderNestedValue(item, itemType)); + } + return out; + } + case Type.Struct: { + const fields = (type as DataType & { children: Field[] }).children; + const obj: Record = {}; + for (const field of fields) { + obj[field.name] = renderNestedValue( + (value as Record)[field.name], + field.type, + ); + } + return obj; + } + case Type.Map: { + // A Map is a List>; the value child carries the type. + const entriesType = (type as DataType & { children: Field[] }).children[0] + .type; + const valueType = (entriesType as DataType & { children: Field[] }) + .children[1].type; + const obj: Record = {}; + for (const [k, v] of value as Iterable<[unknown, unknown]>) { + obj[String(k)] = renderNestedValue(v, valueType); + } + return obj; + } + default: + return formatScalarCell(value, type); + } +} + +/** Parse a STRING cell that looks like JSON into an object/array. */ +function maybeParseJsonString(value: string): unknown { + if (value && (value[0] === "{" || value[0] === "[")) { + try { + return JSON.parse(value); + } catch { + return value; + } + } + return value; +} + +/** + * Decode a base64 Arrow IPC attachment to plain row objects, matching what the + * warehouse emits natively under JSON_ARRAY so callers cannot tell which path + * served the query. + * + * The output key uses the manifest name (`columnNames[i]`) when available — + * Databricks encodes the Arrow schema positionally (col_0, …) — while the + * vector is still fetched by the Arrow field's own name. + */ +export function decodeArrowAttachmentToRows( + attachment: string, + columnNames?: string[], +): Record[] { + const decoded = Buffer.from(attachment, "base64"); + if (decoded.byteLength > JSON_ARRAY_FALLBACK_MAX_BYTES) { + throw ExecutionError.statementFailed( + `Result attachment is ${decoded.byteLength} bytes; JSON_ARRAY fallback materializer caps at ${JSON_ARRAY_FALLBACK_MAX_BYTES} bytes. Re-issue the query with format="ARROW_STREAM" to stream the full result.`, + "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", + "Result too large for JSON format. Re-run with ARROW_STREAM format.", + ); + } + const table = tableFromIPC( + new Uint8Array(decoded.buffer, decoded.byteOffset, decoded.byteLength), + ); + if (table.numRows > JSON_ARRAY_FALLBACK_MAX_ROWS) { + throw ExecutionError.statementFailed( + `Result has ${table.numRows} rows; JSON_ARRAY fallback materializer caps at ${JSON_ARRAY_FALLBACK_MAX_ROWS}. Re-issue the query with format="ARROW_STREAM" to stream the full result.`, + "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", + `Result too large for JSON format (over ${JSON_ARRAY_FALLBACK_MAX_ROWS} rows). Re-run with ARROW_STREAM format.`, + ); + } + // Resolve child vectors once (getChild walks fields on every call). The + // output key uses the manifest name; the vector is fetched by field name. + const columns = table.schema.fields.map((f, i) => { + const outName = + columnNames?.[i] && columnNames[i].length > 0 ? columnNames[i] : f.name; + return [outName, table.getChild(f.name), f.type] as const; + }); + const rows: Record[] = []; + for (let i = 0; i < table.numRows; i++) { + const row: Record = {}; + for (const [name, col, type] of columns) { + const v = col?.get(i); + if (v == null) { + row[name] = null; + continue; + } + switch (type.typeId) { + case Type.Decimal: + case Type.Timestamp: + case Type.Date: + row[name] = formatScalarCell(v, type); + continue; + // Nested columns (STRUCT / ARRAY / MAP) arrive on the native + // JSON_ARRAY wire as a JSON string with every scalar leaf stringified + // by type; render the same shape and stringify so callers cannot tell + // the fallback path from the native one. + case Type.List: + case Type.FixedSizeList: + case Type.Struct: + case Type.Map: + row[name] = JSON.stringify(renderNestedValue(v, type)); + continue; + } + if ( + typeof v === "number" || + typeof v === "bigint" || + typeof v === "boolean" + ) { + row[name] = String(v); + } else if (typeof v === "string") { + row[name] = maybeParseJsonString(v); + } else if (v instanceof Uint8Array) { + row[name] = Buffer.from(v).toString("base64"); + } else { + row[name] = formatScalarCell(v, type); + } + } + rows.push(row); + } + return rows; +} diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 08ede7bc..35677792 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -5,10 +5,21 @@ import { mockServiceContext, setupDatabricksEnv, } from "@tools/test-helpers"; +import { + DateDay, + Decimal, + makeData, + Table, + TimestampMicrosecond, + tableToIPC, + Utf8, + Vector, + vectorFromArray, +} from "apache-arrow"; import { sql } from "shared"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { ServiceContext } from "../../../context/service-context"; -import { AnalyticsPlugin, analytics } from "../analytics"; +import { AnalyticsPlugin, analytics, writeChunk } from "../analytics"; import type { IAnalyticsConfig } from "../types"; // Mock CacheManager singleton with actual caching behavior @@ -27,11 +38,7 @@ const { mockCacheStore, mockCacheInstance } = vi.hoisted(() => { set: vi.fn(), delete: vi.fn(), getOrExecute: vi.fn( - async ( - key: unknown[], - fn: (signal?: AbortSignal) => Promise, - userKey: string, - ) => { + async (key: unknown[], fn: () => Promise, userKey: string) => { const cacheKey = generateKey(key, userKey); if (store.has(cacheKey)) { return store.get(cacheKey); @@ -97,19 +104,6 @@ describe("Analytics Plugin", () => { ); }); - test("should register GET route for arrow results", () => { - const plugin = new AnalyticsPlugin(config); - const { router } = createMockRouter(); - - plugin.injectRoutes(router); - - expect(router.get).toHaveBeenCalledTimes(1); - expect(router.get).toHaveBeenCalledWith( - "/arrow-result/:jobId", - expect.any(Function), - ); - }); - test("/query/:query_key should return 400 when query_key is missing", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -642,7 +636,7 @@ describe("Analytics Plugin", () => { ); }); - test("emits warehouse_status events before the result for a STARTING warehouse", async () => { + test("/query/:query_key should pass INLINE + ARROW_STREAM format parameters when format is ARROW_STREAM", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -652,101 +646,1177 @@ describe("Analytics Plugin", () => { }); const executeMock = vi.fn().mockResolvedValue({ - result: { data: [{ id: 1, name: "test" }] }, + result: { data: [{ id: 1 }] }, }); (plugin as any).SQLClient.executeStatement = executeMock; plugin.injectRoutes(router); const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); - // Override the default RUNNING mock with a STARTING -> RUNNING sequence - // so the route streams a warehouse_status event before the result. - const warehouseGet = vi - .fn() - .mockResolvedValueOnce({ state: "STARTING" }) - .mockResolvedValueOnce({ state: "RUNNING" }); + await handler(mockReq, mockRes); + + expect(executeMock).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + statement: "SELECT * FROM test", + warehouse_id: "test-warehouse-id", + disposition: "INLINE", + format: "ARROW_STREAM", + }), + expect.any(AbortSignal), + ); + }); + + test("/query/:query_key should use INLINE + JSON_ARRAY by default when no format specified", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const executeMock = vi.fn().mockResolvedValue({ + result: { data: [{ id: 1 }] }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); const mockReq = createMockRequest({ params: { query_key: "test_query" }, body: { parameters: {} }, }); - mockReq.serviceWorkspaceClient.warehouses.get = warehouseGet; - mockReq.userWorkspaceClient.warehouses.get = warehouseGet; const mockRes = createMockResponse(); - // The connector polls every 3s between warehouse state checks; use fake - // timers so the test doesn't actually sleep. - vi.useFakeTimers(); - const handlerPromise = handler(mockReq, mockRes); - await vi.runAllTimersAsync(); - await handlerPromise; - vi.useRealTimers(); + await handler(mockReq, mockRes); - // Inspect the SSE writes: a `warehouse_status` event must precede the - // `result` event. - const eventLines = (mockRes.write as any).mock.calls - .map((call: any[]) => call[0] as string) - .filter((s: string) => s.startsWith("event: ")); - const warehouseIdx = eventLines.findIndex( - (s: string) => s === "event: warehouse_status\n", + expect(executeMock).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + disposition: "INLINE", + format: "JSON_ARRAY", + }), + expect.any(AbortSignal), ); - const resultIdx = eventLines.findIndex( - (s: string) => s === "event: result\n", + }); + + test("/query/:query_key should pass INLINE + JSON_ARRAY when format is explicitly JSON_ARRAY", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const executeMock = vi.fn().mockResolvedValue({ + result: { data: [{ id: 1 }] }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "JSON_ARRAY" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(executeMock.mock.calls[0][1]).toMatchObject({ + disposition: "INLINE", + format: "JSON_ARRAY", + }); + }); + + test("/query/:query_key falls back ARROW_STREAM INLINE→EXTERNAL_LINKS and streams the preserved links in-context", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // INLINE rejected → EXTERNAL_LINKS. The result carries the pre-signed + // `external_links` resolved in this request's context, so the route + // streams them directly — no second getStatement under the ambient + // service-principal context. + const links = [{ external_link: "https://example.com/chunk-0" }]; + const executeMock = vi + .fn() + .mockRejectedValueOnce( + new Error( + "INVALID_PARAMETER_VALUE: ARROW_STREAM not supported with INLINE disposition", + ), + ) + .mockResolvedValueOnce({ + result: { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + columnNames: ["group_key", "cost_usd"], + external_links: links, + }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + const extBytes = new Uint8Array([9, 8, 7]); + const streamExternalLinksMock = vi.fn(function* (_chunks: unknown) { + yield extBytes; + }); + (plugin as any).SQLClient.streamExternalLinks = streamExternalLinksMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // First call INLINE (rejected), second EXTERNAL_LINKS (fallback). + expect(executeMock.mock.calls[0][1]).toMatchObject({ + disposition: "INLINE", + format: "ARROW_STREAM", + }); + expect(executeMock.mock.calls[1][1]).toMatchObject({ + disposition: "EXTERNAL_LINKS", + format: "ARROW_STREAM", + }); + // Streams the pre-signed links from the execute response (no re-fetch). + expect(streamExternalLinksMock).toHaveBeenCalledTimes(1); + expect(streamExternalLinksMock.mock.calls[0][0]).toBe(links); + // Bytes stream on the response body with the Arrow content type — no + // JSON error, no missingData throw. + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Content-Type", + "application/vnd.apache.arrow.stream", ); - expect(warehouseIdx).toBeGreaterThanOrEqual(0); - expect(resultIdx).toBeGreaterThanOrEqual(0); - expect(warehouseIdx).toBeLessThan(resultIdx); + // Manifest names (from EXTERNAL_LINKS result) ride the header too. + expect(mockRes.setHeader).toHaveBeenCalledWith( + "X-Appkit-Arrow-Columns", + encodeURIComponent(JSON.stringify(["group_key", "cost_usd"])), + ); + expect(mockRes.json).not.toHaveBeenCalled(); + const written = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as Buffer, + ); + expect(written).toHaveLength(1); + expect(Array.from(written[0])).toEqual([9, 8, 7]); + }); - // The status payload should include the state field. - expect(mockRes.write).toHaveBeenCalledWith( - expect.stringMatching( - /"type":"warehouse_status".*"state":"(STARTING|RUNNING)"/, - ), + test("OBO: .obo.sql ARROW_STREAM external-links streams under the user context", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + // `.obo.sql` → isAsUser true → the route must run through asUser(req). + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: true, + }); + + const links = [{ external_link: "https://example.com/obo-chunk-0" }]; + // Fake user-context executor: INLINE rejected, EXTERNAL_LINKS resolves + // with the pre-signed links (resolved with the user's identity). + const userExecutorQuery = vi.fn(async (_q, _p, fp) => { + if (fp.disposition === "INLINE") { + throw new Error( + "INVALID_PARAMETER_VALUE: The format field must be JSON_ARRAY when the disposition field is INLINE.", + ); + } + return { + external_links: links, + columnNames: ["a"], + statement_id: "obo-stmt", + }; + }); + // Warehouse readiness must run through the user-context executor too, so + // `getWorkspaceClient()` resolves to the user (not the SP) for `.obo.sql`. + const ensureReadyMock = vi.fn().mockResolvedValue(undefined); + const asUserSpy = vi.spyOn(plugin as any, "asUser").mockReturnValue({ + query: userExecutorQuery, + _ensureArrowWarehouseReady: ensureReadyMock, + }); + + const streamExternalLinksMock = vi.fn(function* (_chunks: unknown) { + yield new Uint8Array([1, 2, 3]); + }); + (plugin as any).SQLClient.streamExternalLinks = streamExternalLinksMock; + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + // OBO requests carry the forwarded user identity; the arrow cache key + // is scoped per user, so `resolveUserId` reads this header. + headers: { + "x-forwarded-user": "user@example.com", + "x-forwarded-access-token": "user-token", + }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // The user context (asUser) executor ran the queries AND the warehouse + // readiness — not the SP `this`. + expect(asUserSpy).toHaveBeenCalledWith(mockReq); + expect(userExecutorQuery).toHaveBeenCalled(); + expect(ensureReadyMock).toHaveBeenCalled(); + // The links the user-context executor resolved are streamed directly — + // no re-fetch under a different identity. + expect(streamExternalLinksMock).toHaveBeenCalledTimes(1); + expect(streamExternalLinksMock.mock.calls[0][0]).toBe(links); + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Content-Type", + "application/vnd.apache.arrow.stream", ); + }); - expect(executeMock).toHaveBeenCalledTimes(1); + test("ARROW_STREAM: a stuck warehouse fails fast with a 503 WAREHOUSE_UNAVAILABLE", async () => { + const plugin = new AnalyticsPlugin({ + ...config, + arrowFirstByteTimeoutMs: 20, + }); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + (plugin as any).SQLClient.ensureWarehouseRunning = vi + .fn() + .mockResolvedValue(undefined); + // Warehouse never produces a first byte — only settles when aborted. + (plugin as any).SQLClient.executeStatement = vi.fn( + (_c: unknown, _i: unknown, signal?: AbortSignal) => + new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => + reject( + new DOMException("The operation was aborted.", "AbortError"), + ), + ); + }), + ); + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(503); + expect(mockRes.json).toHaveBeenCalledWith( + expect.objectContaining({ errorCode: "WAREHOUSE_UNAVAILABLE" }), + ); }); - test("should return 404 when query file is not found", async () => { + test("ARROW_STREAM: a schema too wide for the header advertises a columns-ref instead", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); - // Mock getAppQuery to return null (query not found) - (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue(null); + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + // 2000 columns → the encoded header value exceeds the size cap. + const wideNames = Array.from({ length: 2000 }, (_, i) => `column_${i}`); + (plugin as any).SQLClient.executeStatement = vi.fn().mockResolvedValue({ + result: { + attachment: Buffer.from([1, 2, 3]).toString("base64"), + columnNames: wideNames, + statement_id: "stmt-wide", + }, + }); plugin.injectRoutes(router); - const handler = getHandler("POST", "/query/:query_key"); const mockReq = createMockRequest({ - params: { query_key: "nonexistent_query" }, - body: { parameters: {} }, + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, }); const mockRes = createMockResponse(); await handler(mockReq, mockRes); - expect(mockRes.status).toHaveBeenCalledWith(404); - expect(mockRes.json).toHaveBeenCalledWith({ - error: "Query not found", + const setHeaderCalls = (mockRes.setHeader as any).mock.calls.map( + (c: any[]) => c[0], + ); + expect(setHeaderCalls).toContain("X-Appkit-Arrow-Columns-Ref"); + expect(setHeaderCalls).not.toContain("X-Appkit-Arrow-Columns"); + expect(mockRes.setHeader).toHaveBeenCalledWith( + "X-Appkit-Arrow-Columns-Ref", + "stmt-wide", + ); + }); + + test("GET /columns/:statementId returns the manifest column names", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + (plugin as any).SQLClient.getColumnNames = vi + .fn() + .mockResolvedValue(["name", "spend"]); + + plugin.injectRoutes(router); + const handler = getHandler("GET", "/columns/:statementId"); + const mockReq = createMockRequest({ + params: { statementId: "stmt-wide" }, }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.json).toHaveBeenCalledWith({ columns: ["name", "spend"] }); }); - }); - describe("toolkit()", () => { - test("produces ToolkitEntry records keyed by the plugin name", () => { - const plugin = new AnalyticsPlugin({ name: "analytics" }); - const entries = plugin.toolkit(); - expect(Object.keys(entries)).toContain("analytics.query"); - const entry = entries["analytics.query"]; - expect(entry.__toolkitRef).toBe(true); - expect(entry.pluginName).toBe("analytics"); - expect(entry.localName).toBe("query"); + test("GET /columns/:statementId falls back to the service principal when the user identity can't read it (OBO)", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + // User identity 404s (e.g. the statement was executed by the SP, which + // the user can't `getStatement`); the SP `this` resolves it. + const spGetColumns = vi.fn().mockResolvedValue(["a", "b"]); + (plugin as any).SQLClient.getColumnNames = spGetColumns; + const userGetColumnNames = vi + .fn() + .mockRejectedValue(new Error("RESOURCE_DOES_NOT_EXIST")); + const asUserSpy = vi + .spyOn(plugin as any, "asUser") + .mockReturnValue({ _getColumnNames: userGetColumnNames }); + + plugin.injectRoutes(router); + const handler = getHandler("GET", "/columns/:statementId"); + const mockReq = createMockRequest({ params: { statementId: "stmt-x" } }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // Tried the user identity first, then fell back to the service principal. + expect(asUserSpy).toHaveBeenCalledWith(mockReq); + expect(userGetColumnNames).toHaveBeenCalledWith("stmt-x"); + expect(spGetColumns).toHaveBeenCalled(); + expect(mockRes.json).toHaveBeenCalledWith({ columns: ["a", "b"] }); }); - test("respects prefix and only options", () => { - const plugin = new AnalyticsPlugin({ name: "analytics" }); - const entries = plugin.toolkit({ prefix: "", only: ["query"] }); - expect(Object.keys(entries)).toEqual(["query"]); + test("/query/:query_key falls back on a structured ExecutionError.errorCode without scanning the message", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // Properly-structured ExecutionError, as the connector now produces + // when the SDK's ApiError surfaces with errorCode set. + const { ExecutionError } = await import("../../../errors/execution"); + const structuredError = ExecutionError.statementFailed( + "ARROW_STREAM is not supported with INLINE disposition", + "INVALID_PARAMETER_VALUE", + ); + + const executeMock = vi + .fn() + .mockRejectedValueOnce(structuredError) + .mockResolvedValueOnce({ + result: { statement_id: "stmt-1", status: { state: "SUCCEEDED" } }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // Both attempts: INLINE (rejected via structured code) → EXTERNAL_LINKS. + expect(executeMock).toHaveBeenCalledTimes(2); + expect(executeMock.mock.calls[1][1]).toMatchObject({ + disposition: "EXTERNAL_LINKS", + format: "ARROW_STREAM", + }); + }); + + test("/query/:query_key falls back when error message carries a structured INVALID_PARAMETER_VALUE error_code", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // Wrapped JSON error like the SDK surfaces from a `Bad Request` HTTP + // response. Both INLINE and ARROW_STREAM appear, plus the structured code. + const wrappedJsonError = new Error( + 'Response from server (Bad Request) {"error_code":"INVALID_PARAMETER_VALUE","message":"ARROW_STREAM is not supported with INLINE disposition on this warehouse"}', + ); + const executeMock = vi + .fn() + .mockRejectedValueOnce(wrappedJsonError) + .mockResolvedValueOnce({ + result: { statement_id: "stmt-1", status: { state: "SUCCEEDED" } }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // Both attempts ran: INLINE (rejected) then EXTERNAL_LINKS (succeeded). + expect(executeMock).toHaveBeenCalledTimes(2); + expect(executeMock.mock.calls[1][1]).toMatchObject({ + disposition: "EXTERNAL_LINKS", + format: "ARROW_STREAM", + }); + }); + + test("/query/:query_key does NOT fall back on a non-capability error (auth/SQL)", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // A genuine SQL/permission error carries no capability error code + // (INVALID_PARAMETER_VALUE / NOT_IMPLEMENTED), so it must NOT trigger a + // wasted EXTERNAL_LINKS attempt — it propagates as-is. + const executeMock = vi + .fn() + .mockRejectedValue( + new Error( + 'Response from server (Bad Request) {"error_code":"TABLE_OR_VIEW_NOT_FOUND","message":"Table or view not found: foo"}', + ), + ); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // The retry interceptor may attempt the query multiple times, but every + // attempt stays on INLINE — a non-capability error never escalates to + // EXTERNAL_LINKS. + for (const call of executeMock.mock.calls) { + expect(call[1]).toMatchObject({ + disposition: "INLINE", + format: "ARROW_STREAM", + }); + } + }); + + test("/query/:query_key falls back on a capability-coded rejection regardless of exact wording", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // A capability-coded rejection with wording we don't pattern-match + // (Databricks could reword at any time). The errorCode gate alone must be + // enough to escalate to EXTERNAL_LINKS — otherwise a message reword would + // 500 every standard warehouse. + const executeMock = vi + .fn() + .mockRejectedValueOnce( + new Error( + 'Response from server (Bad Request) {"error_code":"INVALID_PARAMETER_VALUE","message":"this disposition/format combination is not available here"}', + ), + ) + .mockResolvedValueOnce({ + result: { statement_id: "stmt-1", status: { state: "SUCCEEDED" } }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(executeMock).toHaveBeenCalledTimes(2); + expect(executeMock.mock.calls[1][1]).toMatchObject({ + disposition: "EXTERNAL_LINKS", + format: "ARROW_STREAM", + }); + }); + + test("/query/:query_key should not fall back for non-format errors", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const executeMock = vi + .fn() + .mockRejectedValue(new Error("PERMISSION_DENIED: no access")); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // Only one call — non-format error is not retried with different disposition. + for (const call of executeMock.mock.calls) { + expect(call[1]).toMatchObject({ + disposition: "INLINE", + format: "ARROW_STREAM", + }); + } + }); + + test("/query/:query_key streams ARROW_STREAM INLINE bytes directly on the response body (no SSE, no stash)", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // Real base64 so the route can decode it via Buffer.from(..., "base64"). + // `columnNames` is what the connector attaches from the manifest (the + // Arrow schema itself is positional col_N). + const arrowBytes = new Uint8Array([1, 2, 3, 4, 5]); + const fakeAttachment = Buffer.from(arrowBytes).toString("base64"); + const executeMock = vi.fn().mockResolvedValue({ + result: { + attachment: fakeAttachment, + row_count: 1, + columnNames: ["name", "totalSpend"], + }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // INLINE succeeded — a single execution, no EXTERNAL_LINKS fallback. + expect(executeMock).toHaveBeenCalledTimes(1); + expect(executeMock.mock.calls[0][1]).toMatchObject({ + disposition: "INLINE", + format: "ARROW_STREAM", + }); + + // Bytes stream straight back on the response body with the Arrow + // content type — no SSE frames, no stash, no JSON error. + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Content-Type", + "application/vnd.apache.arrow.stream", + ); + // Manifest column names ride a response header so the client can relabel + // the positional Arrow schema. + expect(mockRes.setHeader).toHaveBeenCalledWith( + "X-Appkit-Arrow-Columns", + encodeURIComponent(JSON.stringify(["name", "totalSpend"])), + ); + expect(mockRes.json).not.toHaveBeenCalled(); + expect(mockRes.end).toHaveBeenCalled(); + + const written = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as Buffer, + ); + expect(written).toHaveLength(1); + expect(Buffer.isBuffer(written[0])).toBe(true); + expect(Array.from(written[0])).toEqual(Array.from(arrowBytes)); + }); + + test("/query/:query_key falls back JSON_ARRAY to ARROW_STREAM INLINE when warehouse refuses JSON_ARRAY for INLINE", async () => { + // Some serverless warehouses (the ones this PR is centrally aimed at) + // only accept ARROW_STREAM for INLINE results — JSON_ARRAY + INLINE is + // rejected outright. The caller still asked for JSON_ARRAY, so the + // server retries as ARROW_STREAM + INLINE and decodes the attachment + // back into plain row objects: the caller's contract is preserved and + // the SSE channel still carries a `result` message, not an `arrow` + // message. + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // Real base64 Arrow IPC captured from a serverless warehouse running + // `SELECT 1 AS test_col, 2 AS test_col2` (one row, two INT columns). + const REAL_ARROW_ATTACHMENT = + "/////7gAAAAQAAAAAAAKAAwACgAJAAQACgAAABAAAAAAAQQACAAIAAAABAAIAAAABAAAAAIAAABMAAAABAAAAMz///8QAAAAGAAAAAAAAQIUAAAAvP///yAAAAAAAAABAAAAAAkAAAB0ZXN0X2NvbDIAAAAQABQAEAAOAA8ABAAAAAgAEAAAABgAAAAgAAAAAAABAhwAAAAIAAwABAALAAgAAAAgAAAAAAAAAQAAAAAIAAAAdGVzdF9jb2wAAAAA/////7gAAAAQAAAADAAaABgAFwAEAAgADAAAACAAAAAAAQAAAAAAAAAAAAAAAAADBAAKABgADAAIAAQACgAAADwAAAAQAAAAAQAAAAAAAAAAAAAAAgAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAEAAAAAAAAAQAAAAAAAAAAEAAAAAAAAAIAAAAAAAAAAAQAAAAAAAADAAAAAAAAAAAQAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP////8AAAAA"; + + const executeMock = vi + .fn() + // First call: JSON_ARRAY + INLINE — warehouse rejects. + .mockRejectedValueOnce( + new Error( + 'Response from server (Bad Request) {"error_code":"INVALID_PARAMETER_VALUE","message":"Inline disposition only supports ARROW_STREAM format."}', + ), + ) + // Second call: ARROW_STREAM + INLINE — warehouse returns the bytes. + .mockResolvedValueOnce({ + result: { + attachment: REAL_ARROW_ATTACHMENT, + status: { state: "SUCCEEDED" }, + }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "JSON_ARRAY" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // Two calls: first JSON_ARRAY + INLINE (rejected), then the fallback + // ARROW_STREAM + INLINE (the warehouse's preferred shape for INLINE). + expect(executeMock).toHaveBeenCalledTimes(2); + expect(executeMock.mock.calls[0][1]).toMatchObject({ + disposition: "INLINE", + format: "JSON_ARRAY", + }); + expect(executeMock.mock.calls[1][1]).toMatchObject({ + disposition: "INLINE", + format: "ARROW_STREAM", + }); + + // The SSE wire payload must look like a JSON_ARRAY result, not an + // arrow message — the caller asked for JSON_ARRAY and the server has + // already decoded Arrow → rows. + const writeCalls = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as string, + ); + // Skip the leading warehouse_status event the route always emits; + // the terminal result/arrow/error payload is the one under test. + const payload = writeCalls.find( + (s: string) => + s.startsWith("data: ") && !s.includes("warehouse_status"), + ); + expect(payload).toBeDefined(); + expect(payload).toContain('"type":"result"'); + expect(payload).not.toContain('"type":"arrow"'); + // Real row values from the captured attachment: test_col=1, test_col2=2. + // Integer columns are coerced to strings to match what JSON_ARRAY would + // have produced for the same warehouse + same INT columns. + expect(payload).toContain('"test_col":"1"'); + expect(payload).toContain('"test_col2":"2"'); + }); + + test("/query/:query_key JSON_ARRAY fallback formats decimal, timestamp, date, and JSON-string columns to match the native JSON_ARRAY shape", async () => { + // The server-side Arrow→rows decoder (needs-arrow fallback) must + // render typed Arrow cells the same way the warehouse renders them + // under native JSON_ARRAY, or callers can tell which path served the + // query. Regression coverage for the four decode bugs: + // - DECIMAL: scale applied ("123.45"), not the raw unscaled mantissa. + // - TIMESTAMP: "yyyy-MM-dd HH:mm:ss[.SSS]", not a raw epoch number. + // - DATE: "yyyy-MM-dd", not a raw epoch number. + // - STRING holding JSON: parsed to an object, matching the JSON path. + // Build a one-row Arrow table with each type and IPC-encode it. + const decType = new Decimal(2, 10, 128); // scale=2, precision=10 + // 128-bit little-endian limbs for the signed unscaled value 12345. + const decLimbs = new Uint32Array(4); + decLimbs[0] = 12345; + const decVec = new Vector([ + makeData({ type: decType, length: 1, data: decLimbs }), + ]); + const tsVec = vectorFromArray( + [new Date(Date.UTC(2024, 5, 13, 1, 2, 3, 500))], + new TimestampMicrosecond("UTC"), + ); + // TIMESTAMP_NTZ: no timezone → ISO without the trailing Z. + const tsNtzVec = vectorFromArray( + [new Date(Date.UTC(2024, 0, 2, 3, 4, 5))], + new TimestampMicrosecond(), + ); + const dateVec = vectorFromArray( + [new Date(Date.UTC(2024, 5, 13))], + new DateDay(), + ); + const jsonVec = vectorFromArray(['{"nested":1}'], new Utf8()); + const table = new Table({ + amount: decVec, + ts: tsVec, + ts_ntz: tsNtzVec, + d: dateVec, + meta: jsonVec, + }); + const attachment = Buffer.from(tableToIPC(table, "stream")).toString( + "base64", + ); + + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + const executeMock = vi + .fn() + .mockRejectedValueOnce( + new Error( + 'Response from server (Bad Request) {"error_code":"INVALID_PARAMETER_VALUE","message":"Inline disposition only supports ARROW_STREAM format."}', + ), + ) + .mockResolvedValueOnce({ + result: { attachment, status: { state: "SUCCEEDED" } }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "JSON_ARRAY" }, + }); + const mockRes = createMockResponse(); + await handler(mockReq, mockRes); + + const writeCalls = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as string, + ); + const payload = writeCalls.find( + (s: string) => + s.startsWith("data: ") && !s.includes("warehouse_status"), + ); + expect(payload).toBeDefined(); + expect(payload).toContain('"type":"result"'); + // DECIMAL: scale applied, not the unscaled "12345". + expect(payload).toContain('"amount":"123.45"'); + // TIMESTAMP: ISO-8601 ms precision with Z (zoned), not an epoch + // number. Matches dogfood's native JSON_ARRAY rendering. + expect(payload).toContain('"ts":"2024-06-13T01:02:03.500Z"'); + // TIMESTAMP_NTZ: same ISO form without the trailing Z. + expect(payload).toContain('"ts_ntz":"2024-01-02T03:04:05.000"'); + // DATE: yyyy-MM-dd. + expect(payload).toContain('"d":"2024-06-13"'); + // STRING holding JSON parsed to a nested object. + expect(payload).toContain('"meta":{"nested":1}'); + }); + + test("/query/:query_key surfaces an error when both JSON_ARRAY + INLINE and the ARROW_STREAM retry fail", async () => { + // If the JSON_ARRAY retry path (ARROW_STREAM + INLINE) also fails — e.g. + // a downstream warehouse outage that affects both shapes — the route + // must surface the failure rather than silently dropping it. + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // First mocked call (JSON_ARRAY + INLINE) rejects with a needs-arrow + // signal; every subsequent call rejects with an unrelated failure. The + // retry interceptor may retry the second call multiple times — we only + // care that the retry path was taken and that the request ultimately + // surfaces an error rather than a successful result. + const executeMock = vi.fn().mockImplementation((_wc, opts) => { + if (opts?.disposition === "INLINE" && opts?.format === "JSON_ARRAY") { + return Promise.reject( + new Error( + 'Response from server (Bad Request) {"error_code":"INVALID_PARAMETER_VALUE","message":"Inline disposition only supports ARROW_STREAM format."}', + ), + ); + } + return Promise.reject(new Error("warehouse is down")); + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "JSON_ARRAY" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // The retry happened: at least one ARROW_STREAM + INLINE call followed + // the initial JSON_ARRAY + INLINE rejection. + const formats = executeMock.mock.calls.map((c: any[]) => c[1]); + expect( + formats.some( + (f: any) => f?.disposition === "INLINE" && f?.format === "JSON_ARRAY", + ), + ).toBe(true); + expect( + formats.some( + (f: any) => + f?.disposition === "INLINE" && f?.format === "ARROW_STREAM", + ), + ).toBe(true); + // No call should escalate to EXTERNAL_LINKS — that fallback only + // exists on the ARROW_STREAM caller path. + expect( + formats.some((f: any) => f?.disposition === "EXTERNAL_LINKS"), + ).toBe(false); + + // The SSE payload, if any was written, must NOT carry a successful + // result frame. + const writeCalls = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as string, + ); + // Skip the leading warehouse_status event the route always emits; + // the terminal result/arrow/error payload is the one under test. + const payload = writeCalls.find( + (s: string) => + s.startsWith("data: ") && !s.includes("warehouse_status"), + ); + if (payload) { + expect(payload).not.toContain('"type":"result"'); + } + }); + + test("/query/:query_key rejects unknown format values with 400", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + const executeMock = vi.fn(); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + // "CSV" is genuinely unsupported. The legacy spellings "JSON" / "ARROW" + // are *accepted* by the route (normalized to JSON_ARRAY / ARROW_STREAM + // for back-compat with appkit < 0.33.0), so they must not be used here. + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "CSV" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(400); + expect(executeMock).not.toHaveBeenCalled(); + }); + + test("/query/:query_key does not retry the fallback when the request was aborted", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const executeMock = vi.fn().mockImplementation((_wc, _opts, signal) => { + // Simulate a signal that becomes aborted before the failure surfaces — + // e.g. the client cancelled the SSE stream mid-query. Use vitest's + // getter spy rather than Object.defineProperty so we don't try to + // override the native non-configurable AbortSignal.aborted getter. + if (signal) { + vi.spyOn(signal, "aborted", "get").mockReturnValue(true); + } + return Promise.reject( + new Error( + "INVALID_PARAMETER_VALUE: ARROW_STREAM not supported with INLINE disposition", + ), + ); + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // Even though the error message would normally trigger fallback, the + // aborted signal should short-circuit and prevent a second statement. + expect(executeMock).toHaveBeenCalledTimes(1); + }); + + test("/query/:query_key does NOT fall back JSON_ARRAY when the rejection lacks a needs-arrow signal", async () => { + // A generic INVALID_PARAMETER_VALUE that doesn't mention the INLINE + // disposition could be any unrelated SQL/permission error. The classifier + // must NOT interpret it as "warehouse wants ARROW_STREAM" — falling back + // would mask the real failure. + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const executeMock = vi + .fn() + .mockRejectedValue( + new Error("INVALID_PARAMETER_VALUE: only supports ARROW_STREAM"), + ); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "JSON_ARRAY" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // All calls stay on JSON_ARRAY + INLINE — no retry path with a different + // disposition or format was taken. + for (const call of executeMock.mock.calls) { + expect(call[1]).toMatchObject({ + disposition: "INLINE", + format: "JSON_ARRAY", + }); + } + }); + + test("emits warehouse_status events before the result for a STARTING warehouse", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const executeMock = vi.fn().mockResolvedValue({ + result: { data: [{ id: 1, name: "test" }] }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + + // Override the default RUNNING mock with a STARTING -> RUNNING sequence + // so the route streams a warehouse_status event before the result. + const warehouseGet = vi + .fn() + .mockResolvedValueOnce({ state: "STARTING" }) + .mockResolvedValueOnce({ state: "RUNNING" }); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {} }, + }); + mockReq.serviceWorkspaceClient.warehouses.get = warehouseGet; + mockReq.userWorkspaceClient.warehouses.get = warehouseGet; + const mockRes = createMockResponse(); + + // The connector polls every 3s between warehouse state checks; use fake + // timers so the test doesn't actually sleep. + vi.useFakeTimers(); + const handlerPromise = handler(mockReq, mockRes); + await vi.runAllTimersAsync(); + await handlerPromise; + vi.useRealTimers(); + + // Inspect the SSE writes: a `warehouse_status` event must precede the + // `result` event. + const eventLines = (mockRes.write as any).mock.calls + .map((call: any[]) => call[0] as string) + .filter((s: string) => s.startsWith("event: ")); + const warehouseIdx = eventLines.findIndex( + (s: string) => s === "event: warehouse_status\n", + ); + const resultIdx = eventLines.findIndex( + (s: string) => s === "event: result\n", + ); + expect(warehouseIdx).toBeGreaterThanOrEqual(0); + expect(resultIdx).toBeGreaterThanOrEqual(0); + expect(warehouseIdx).toBeLessThan(resultIdx); + + // The status payload should include the state field. + expect(mockRes.write).toHaveBeenCalledWith( + expect.stringMatching( + /"type":"warehouse_status".*"state":"(STARTING|RUNNING)"/, + ), + ); + + expect(executeMock).toHaveBeenCalledTimes(1); + }); + + test("should return 404 when query file is not found", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + // Mock getAppQuery to return null (query not found) + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue(null); + + plugin.injectRoutes(router); + + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "nonexistent_query" }, + body: { parameters: {} }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(404); + expect(mockRes.json).toHaveBeenCalledWith({ + error: "Query not found", + }); + }); + }); + + describe("toolkit()", () => { + test("produces ToolkitEntry records keyed by the plugin name", () => { + const plugin = new AnalyticsPlugin({ name: "analytics" }); + const entries = plugin.toolkit(); + expect(Object.keys(entries)).toContain("analytics.query"); + const entry = entries["analytics.query"]; + expect(entry.__toolkitRef).toBe(true); + expect(entry.pluginName).toBe("analytics"); + expect(entry.localName).toBe("query"); + }); + + test("respects prefix and only options", () => { + const plugin = new AnalyticsPlugin({ name: "analytics" }); + const entries = plugin.toolkit({ prefix: "", only: ["query"] }); + expect(Object.keys(entries)).toEqual(["query"]); + }); + }); + + describe("writeChunk backpressure", () => { + // A response stub that is always backpressured (`write` → false) and lets + // the test fire lifecycle events. The shared `createMockResponse` returns a + // truthy `write`, so it never exercises the backpressure branch. + function backpressuredRes() { + const listeners: Record void>> = {}; + return { + write: vi.fn(() => false), + once: vi.fn(function (this: unknown, ev: string, h: () => void) { + listeners[ev] ??= []; + listeners[ev].push(h); + return this; + }), + off: vi.fn(function (this: unknown, ev: string, h: () => void) { + if (listeners[ev]) { + listeners[ev] = listeners[ev].filter((x) => x !== h); + } + return this; + }), + emit(ev: string) { + for (const h of [...(listeners[ev] ?? [])]) h(); + }, + listenerCount(ev: string) { + return (listeners[ev] ?? []).length; + }, + }; + } + + test("resolves once the socket drains and detaches its listeners", async () => { + const res = backpressuredRes(); + const pending = writeChunk(res as never, new Uint8Array([1, 2, 3])); + res.emit("drain"); + await expect(pending).resolves.toBeUndefined(); + expect(res.listenerCount("drain")).toBe(0); + expect(res.listenerCount("close")).toBe(0); + expect(res.listenerCount("error")).toBe(0); + }); + + test("rejects (does not hang) when the client disconnects mid-backpressure", async () => { + const res = backpressuredRes(); + const pending = writeChunk(res as never, new Uint8Array([1, 2, 3])); + // The socket closes while backpressured — `drain` will never fire. If + // writeChunk only awaited `drain`, this promise (and the stream feeding + // it) would hang forever. + res.emit("close"); + await expect(pending).rejects.toThrow(); + expect(res.listenerCount("drain")).toBe(0); + expect(res.listenerCount("close")).toBe(0); + expect(res.listenerCount("error")).toBe(0); + }); + + test("rejects up front if the socket is already closed when called", async () => { + // Socket already destroyed before writeChunk runs: `close`/`error` have + // already fired, so attaching listeners would never settle. The guard + // must reject without writing or waiting. + const res = { + destroyed: true, + writableEnded: false, + write: vi.fn(() => false), + once: vi.fn(), + off: vi.fn(), + }; + await expect( + writeChunk(res as never, new Uint8Array([1, 2, 3])), + ).rejects.toThrow(); + expect(res.write).not.toHaveBeenCalled(); + expect(res.once).not.toHaveBeenCalled(); }); }); }); 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 new file mode 100644 index 00000000..83e37b4a --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts @@ -0,0 +1,68 @@ +import { WorkspaceClient } from "@databricks/sdk-experimental"; +import { tableFromIPC } from "apache-arrow"; +import { describe, expect, test } from "vitest"; +import { SQLWarehouseConnector } from "../../../connectors"; +import { deliverArrowBytes, type QueryExecutor } from "../result-delivery"; + +/** + * Gated live integration test for the Arrow delivery fallback against a REAL + * warehouse. Skipped unless `APPKIT_INTEGRATION_WAREHOUSE_ID` is set (with + * `DATABRICKS_HOST`/`DATABRICKS_CONFIG_PROFILE` auth in the environment), so it + * never runs in CI. Run locally, e.g.: + * + * DATABRICKS_CONFIG_PROFILE=dogfood \ + * APPKIT_INTEGRATION_WAREHOUSE_ID=86e2ab2c5d30b12a \ + * pnpm exec vitest run arrow-delivery.integration + * + * Verifies the full capability fallback end-to-end: on a normal warehouse the + * INLINE attempt is rejected and EXTERNAL_LINKS chunks stream through; on + * Reyden the inline attachment streams. Column names are resolved from the + * manifest either way. + * + * NOTE: this drives the service-principal identity. Full `.obo.sql` (user) + * validation requires a user token / `x-forwarded-user` request and is + * verified manually; the wiring is covered by the mocked OBO test in + * `analytics.test.ts`. + */ +const warehouseId = process.env.APPKIT_INTEGRATION_WAREHOUSE_ID; + +describe.runIf(!!warehouseId)("arrow delivery (live warehouse)", () => { + test("ARROW_STREAM delivers valid Arrow with real column names", async () => { + const client = new WorkspaceClient({}); + const connector = new SQLWarehouseConnector({ timeout: 120_000 }); + + const executor: QueryExecutor = { + query: async (statement, _params, fp, signal) => { + const res = (await connector.executeStatement( + client, + { + statement, + warehouse_id: warehouseId as string, + wait_timeout: "50s", + on_wait_timeout: "CONTINUE", + disposition: fp.disposition as never, + format: fp.format as never, + }, + signal, + )) as { result?: unknown }; + return res.result as never; + }, + }; + + const out: { columnNames?: string[]; statementId?: string } = {}; + const chunks: Buffer[] = []; + for await (const bytes of deliverArrowBytes( + executor, + connector, + "SELECT id AS my_id, cast(id AS string) AS my_name FROM range(1000)", + undefined, + out, + )) { + chunks.push(Buffer.from(bytes)); + } + + const table = tableFromIPC(new Uint8Array(Buffer.concat(chunks))); + expect(table.numRows).toBe(1000); + expect(out.columnNames).toEqual(["my_id", "my_name"]); + }, 180_000); +}); diff --git a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts new file mode 100644 index 00000000..65ffd228 --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts @@ -0,0 +1,425 @@ +import { Table, tableToIPC, Utf8, vectorFromArray } from "apache-arrow"; +import { describe, expect, test, vi } from "vitest"; +import { ExecutionError } from "../../../errors"; +import { + type ArrowChunkStreamer, + arrowDeliveryUnsupported, + classifyDispositionRejection, + decodeArrowAttachmentToRows, + deliverArrowBytes, + deliverJsonResult, + type QueryExecutor, +} from "../result-delivery"; + +/** Build a base64 Arrow IPC stream (2 cols, given names + rows). */ +function arrowBase64(names: [string, string]): string { + const table = new Table({ + [names[0]]: vectorFromArray(["r0", "r1"], new Utf8()), + [names[1]]: vectorFromArray(["v0", "v1"], new Utf8()), + }); + return Buffer.from(tableToIPC(table, "stream")).toString("base64"); +} + +/** Fake executor whose `query` is driven by the requested disposition/format. */ +function executorFrom( + handler: (fp: { disposition: string; format: string }) => Promise, +): { + executor: QueryExecutor; + calls: { disposition: string; format: string }[]; +} { + const calls: { disposition: string; format: string }[] = []; + const executor: QueryExecutor = { + query: (_q, _p, fp) => { + calls.push(fp); + return handler(fp) as ReturnType; + }, + }; + return { executor, calls }; +} + +async function collect(gen: AsyncGenerator): Promise { + const out: Uint8Array[] = []; + for await (const b of gen) out.push(b); + return out; +} + +const reject = (code: string, message: string) => + ExecutionError.statementFailed(message, code); + +describe("classifyDispositionRejection", () => { + test("normal warehouse rejecting ARROW+INLINE → needs-json-inline", () => { + expect( + classifyDispositionRejection( + reject( + "INVALID_PARAMETER_VALUE", + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + ), + ), + ).toBe("needs-json-inline"); + }); + + test("warehouse requiring ARROW for INLINE → needs-arrow-inline", () => { + expect( + classifyDispositionRejection( + reject( + "INVALID_PARAMETER_VALUE", + "Inline disposition only supports ARROW_STREAM format.", + ), + ), + ).toBe("needs-arrow-inline"); + }); + + test("EXTERNAL_LINKS not implemented (Reyden) → external-links-unsupported", () => { + expect( + classifyDispositionRejection( + reject( + "NOT_IMPLEMENTED", + "ExternalLinks disposition is not yet implemented.", + ), + ), + ).toBe("external-links-unsupported"); + }); + + test("permission error does NOT fall back", () => { + expect( + classifyDispositionRejection( + reject("PERMISSION_DENIED", "User does not have SELECT on table t"), + ), + ).toBeNull(); + }); + + test("plain SQL error does NOT fall back", () => { + expect( + classifyDispositionRejection(new Error("Table or view not found: foo")), + ).toBeNull(); + }); + + test("capability code without an inline/external signal does NOT fall back", () => { + expect( + classifyDispositionRejection( + reject("INVALID_PARAMETER_VALUE", "Some unrelated parameter problem"), + ), + ).toBeNull(); + }); +}); + +describe("deliverArrowBytes — Reyden (ARROW + INLINE)", () => { + test("streams the inline attachment and never touches external links", async () => { + const b64 = arrowBase64(["col_0", "col_1"]); + const { executor, calls } = executorFrom(async () => ({ + attachment: b64, + columnNames: ["name", "spend"], + statement_id: "stmt-inline", + })); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () {}), + }; + const out: { columnNames?: string[]; statementId?: string } = {}; + + const chunks = await collect( + deliverArrowBytes(executor, streamer, "SELECT 1", undefined, out), + ); + + expect(calls).toEqual([{ disposition: "INLINE", format: "ARROW_STREAM" }]); + expect(streamer.streamExternalLinks).not.toHaveBeenCalled(); + expect(out.columnNames).toEqual(["name", "spend"]); + expect(out.statementId).toBe("stmt-inline"); + // Bytes are the decoded attachment. + expect( + Buffer.concat(chunks.map((c) => Buffer.from(c))).toString("base64"), + ).toBe(b64); + }); +}); + +describe("deliverArrowBytes — normal warehouse (fallback to EXTERNAL_LINKS)", () => { + test("INLINE rejected → streams the external links from the execute response", async () => { + const links = [{ external_link: "https://example.com/chunk-0" }]; + const { executor, calls } = executorFrom(async (fp) => { + if (fp.disposition === "INLINE") { + throw reject( + "INVALID_PARAMETER_VALUE", + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + ); + } + return { + external_links: links, + columnNames: ["a", "b"], + statement_id: "stmt-ext", + }; + }); + const extBytes = new Uint8Array([9, 8, 7]); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () { + yield extBytes; + }), + }; + const out: { columnNames?: string[]; statementId?: string } = {}; + + const chunks = await collect( + deliverArrowBytes(executor, streamer, "SELECT 1", undefined, out), + ); + + expect(calls).toEqual([ + { disposition: "INLINE", format: "ARROW_STREAM" }, + { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, + ]); + // Streams the links resolved in-context (no re-fetch): passed straight + // through, along with the (here absent) chunk-link refresher. + expect(streamer.streamExternalLinks).toHaveBeenCalledWith( + links, + undefined, + undefined, + ); + expect(out.columnNames).toEqual(["a", "b"]); + expect(chunks).toEqual([extBytes]); + }); + + test("both INLINE and EXTERNAL_LINKS unsupported → ARROW_DELIVERY_UNSUPPORTED", async () => { + const { executor } = executorFrom(async (fp) => { + if (fp.disposition === "INLINE") { + throw reject( + "INVALID_PARAMETER_VALUE", + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + ); + } + throw reject( + "NOT_IMPLEMENTED", + "ExternalLinks disposition is not yet implemented.", + ); + }); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () {}), + }; + + await expect( + collect(deliverArrowBytes(executor, streamer, "SELECT 1", undefined, {})), + ).rejects.toMatchObject({ errorCode: "ARROW_DELIVERY_UNSUPPORTED" }); + }); + + test("auth/permission error on INLINE propagates without a fallback", async () => { + const authErr = reject("PERMISSION_DENIED", "no access"); + const { executor, calls } = executorFrom(async () => { + throw authErr; + }); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () {}), + }; + + await expect( + collect(deliverArrowBytes(executor, streamer, "SELECT 1", undefined, {})), + ).rejects.toBe(authErr); + // Only the INLINE attempt — no EXTERNAL_LINKS fallback. + expect(calls).toEqual([{ disposition: "INLINE", format: "ARROW_STREAM" }]); + }); +}); + +describe("deliverJsonResult", () => { + test("INLINE JSON_ARRAY success returns rows directly", async () => { + const { executor, calls } = executorFrom(async () => ({ + data: [{ id: 1 }], + status: { state: "SUCCEEDED" }, + statement_id: "s1", + })); + const result = await deliverJsonResult(executor, "SELECT 1", undefined); + expect(result.data).toEqual([{ id: 1 }]); + expect(calls).toEqual([{ disposition: "INLINE", format: "JSON_ARRAY" }]); + }); + + test("needs-arrow-inline → retries as ARROW and decodes rows with manifest names", async () => { + const b64 = arrowBase64(["col_0", "col_1"]); + const { executor, calls } = executorFrom(async (fp) => { + if (fp.format === "JSON_ARRAY") { + throw reject( + "INVALID_PARAMETER_VALUE", + "Inline disposition only supports ARROW_STREAM format.", + ); + } + return { attachment: b64, columnNames: ["name", "spend"] }; + }); + const result = await deliverJsonResult(executor, "SELECT 1", undefined); + expect(calls).toEqual([ + { disposition: "INLINE", format: "JSON_ARRAY" }, + { disposition: "INLINE", format: "ARROW_STREAM" }, + ]); + // Rows carry the manifest names, not the positional col_N names. + expect(result.data).toEqual([ + { name: "r0", spend: "v0" }, + { name: "r1", spend: "v1" }, + ]); + }); +}); + +describe("decodeArrowAttachmentToRows", () => { + test("relabels positional columns from manifest names", () => { + const rows = decodeArrowAttachmentToRows(arrowBase64(["col_0", "col_1"]), [ + "name", + "spend", + ]); + expect(rows).toEqual([ + { name: "r0", spend: "v0" }, + { name: "r1", spend: "v1" }, + ]); + }); +}); + +describe("deliverArrowBytes — capability memo", () => { + test("reports 'inline' when the INLINE attempt succeeds", async () => { + const { executor } = executorFrom(async () => ({ + attachment: arrowBase64(["col_0", "col_1"]), + columnNames: ["a", "b"], + statement_id: "s", + })); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () {}), + }; + const resolved: string[] = []; + await collect( + deliverArrowBytes( + executor, + streamer, + "SELECT 1", + undefined, + {}, + undefined, + { + onCapabilityResolved: (c) => resolved.push(c), + }, + ), + ); + expect(resolved).toEqual(["inline"]); + }); + + test("reports 'external' after falling back to EXTERNAL_LINKS", async () => { + const links = [{ external_link: "https://x/0" }]; + const { executor } = executorFrom(async (fp) => { + if (fp.disposition === "INLINE") { + throw reject( + "INVALID_PARAMETER_VALUE", + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + ); + } + return { external_links: links, columnNames: ["a"], statement_id: "s" }; + }); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () { + yield new Uint8Array([1]); + }), + }; + const resolved: string[] = []; + await collect( + deliverArrowBytes( + executor, + streamer, + "SELECT 1", + undefined, + {}, + undefined, + { + onCapabilityResolved: (c) => resolved.push(c), + }, + ), + ); + expect(resolved).toEqual(["external"]); + }); + + test("capabilityHint 'external' skips the INLINE probe entirely", async () => { + const links = [{ external_link: "https://x/0" }]; + const { executor, calls } = executorFrom(async (fp) => { + if (fp.disposition === "INLINE") { + throw new Error("INLINE should not have been attempted"); + } + return { external_links: links, columnNames: ["a"], statement_id: "s" }; + }); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () { + yield new Uint8Array([1]); + }), + }; + await collect( + deliverArrowBytes( + executor, + streamer, + "SELECT 1", + undefined, + {}, + undefined, + { + capabilityHint: "external", + }, + ), + ); + // Only EXTERNAL_LINKS was attempted — no wasted INLINE probe. + expect(calls).toEqual([ + { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, + ]); + }); +}); + +describe("arrowDeliveryUnsupported", () => { + test("is a structured, client-actionable error", () => { + const err = arrowDeliveryUnsupported(); + expect(err.errorCode).toBe("ARROW_DELIVERY_UNSUPPORTED"); + expect(err.clientMessage).toMatch(/JSON_ARRAY/); + }); +}); + +// ───────────────────────────────────────────────────────────────────────── +// Captured-fixture round-trip: real bytes from a live Reyden serverless +// warehouse (id 000000000000000d), which refuses INLINE + JSON_ARRAY and only +// serves INLINE + ARROW_STREAM. Each fixture pairs the base64 Arrow IPC +// attachment with the exact `data_array` the same query returns under native +// JSON_ARRAY — the shape `decodeArrowAttachmentToRows` must reproduce so a +// JSON_ARRAY caller cannot tell the fallback path from the native one. +// Captured 2026-07-09; regenerate by re-running the query under both formats. +// ───────────────────────────────────────────────────────────────────────── +const REYDEN_SCALAR_FIXTURE = { + // amt DECIMAL(10,2), neg DECIMAL(10,2), ts TIMESTAMP, ts_ntz TIMESTAMP_NTZ, + // d DATE, s STRUCT, arr ARRAY + attachment: + "/////3gCAAAQAAAAAAAKAAwACgAJAAQACgAAABAAAAAAAQQACAAIAAAABAAIAAAABAAAAAcAAAD4AQAAuAEAAHQBAAA0AQAACAEAAGAAAAAEAAAANP7//xgAAAAMAAAAAAABDEAAAAABAAAACAAAAKT///9o////FAAAAAwAAAAAAAAHFAAAAAAAAABE/v//AgAAAAoAAAAEAAAAaXRlbQAAAAADAAAAYXJyAIz+//8gAAAADAAAAAAAAQ2MAAAAAgAAAFgAAAAMAAAABAAEAAQAAADI////FAAAAAwAAAAAAAAKIAAAAAAAAAAg////CAAAAAAAAgAHAAAARXRjL1VUQwACAAAAdHMAABAAFAAQAAAADwAEAAAACAAQAAAAFAAAAAwAAAAAAAAHFAAAAAAAAADs/v//AgAAAAoAAAADAAAAYW10AAEAAABzAAAAMP///xQAAAAMAAAAAAABCBAAAAAAAAAA1v///wAAAAABAAAAZAAAAFj///8cAAAADAAAAAAAAQogAAAAAAAAAAAABgAIAAYABgAAAAAAAgAAAAAAAAAAAAYAAAB0c19udHoAAJT///8cAAAADAAAAAAAAQooAAAAAAAAAAgADAAKAAQACAAAAAgAAAAAAAIABwAAAEV0Yy9VVEMAAgAAAHRzAADU////FAAAAAwAAAAAAAEHFAAAAAAAAADE////AgAAAAoAAAADAAAAbmVnABAAFAAQAA4ADwAEAAAACAAQAAAAHAAAAAwAAAAAAAEHHAAAAAAAAAAIAAwACAAEAAgAAAACAAAACgAAAAMAAABhbXQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP////84AgAAEAAAAAwAGgAYABcABAAIAAwAAAAgAAAAwAQAAAAAAAAAAAAAAAAAAwQACgAYAAwACAAEAAoAAAC8AAAAEAAAAAEAAAAAAAAAAAAAAAoAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAACAAAAAAAAAAAAAAAAAAAAAAAAABMAAAAAAAAAAAAAAAEAAAAAAAAAQAAAAAAAAAAQAAAAAAAAAIAAAAAAAAAAAQAAAAAAAADAAAAAAAAAABAAAAAAAAAAAAEAAAAAAAABAAAAAAAAAEABAAAAAAAACAAAAAAAAACAAQAAAAAAAAEAAAAAAAAAwAEAAAAAAAAIAAAAAAAAAAACAAAAAAAAAQAAAAAAAABAAgAAAAAAAAQAAAAAAAAAgAIAAAAAAAABAAAAAAAAAMACAAAAAAAAAQAAAAAAAAAAAwAAAAAAABAAAAAAAAAAQAMAAAAAAAABAAAAAAAAAIADAAAAAAAACAAAAAAAAADAAwAAAAAAAAEAAAAAAAAAAAQAAAAAAAAIAAAAAAAAAEAEAAAAAAAAAQAAAAAAAACABAAAAAAAACAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA5MAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAJ3///////////////////8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQPMmch+bBQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABA8yZyH5sFAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAFdHAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADkwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQPMmch+bBQAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADkwAAAAAAAAAAAAAAAAAABjAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/////AAAAAA==", + columns: ["amt", "neg", "ts", "ts_ntz", "d", "s", "arr"], + expected: { + amt: "123.45", + neg: "-0.99", + ts: "2020-01-02T03:04:05.000Z", + ts_ntz: "2020-01-02T03:04:05.000", + d: "2020-01-02", + s: '{"amt":"123.45","ts":"2020-01-02T03:04:05.000Z"}', + arr: '["123.45","0.99"]', + }, +} as const; + +const REYDEN_NESTED_FIXTURE = { + // s STRUCT with 6 scalar-leaf types, arr_dec ARRAY, arr_int + // ARRAY, m MAP, deep STRUCT>, nul + // DECIMAL (null). + attachment: + "//////gDAAAQAAAAAAAKAAwACgAJAAQACgAAABAAAAAAAQQACAAIAAAABAAIAAAABAAAAAYAAAA8AgAAzAEAAGwBAAC8AAAANAAAAAQAAADs/f//FAAAAAwAAAAAAAEHFAAAAAAAAACg/P//AgAAAAoAAAADAAAAbnVsABj+//8YAAAADAAAAAAAAQ1oAAAAAQAAAAgAAABA/f///Pz//xgAAAAMAAAAAAAADTwAAAABAAAACAAAAGD9//8c/f//FAAAAAwAAAAAAAAHFAAAAAAAAAAM/f//AgAAAAoAAAADAAAAYW10AAUAAABpbm5lcgAAAAQAAABkZWVwAAAAAJz+//8YAAAADAAAAAAAARGUAAAAAQAAAAgAAADE/f//gP3//xwAAAAMAAAAAAAADWgAAAACAAAAPAAAAAgAAADo/f//pP3//xQAAAAMAAAAAAAABxQAAAAAAAAAlP3//wIAAAAKAAAABQAAAHZhbHVlAAAA1P3//xQAAAAMAAAAAAAABQwAAAAAAAAANP7//wMAAABrZXkABwAAAGVudHJpZXMAAQAAAG0AAABI////GAAAAAwAAAAAAAEMQAAAAAEAAAAIAAAAcP7//yz+//8QAAAAGAAAAAAAAAIUAAAAYP7//yAAAAAAAAABAAAAAAQAAABpdGVtAAAAAAcAAABhcnJfaW50AKT///8YAAAADAAAAAAAAQxAAAAAAQAAAAgAAADM/v//iP7//xQAAAAMAAAAAAAABxQAAAAAAAAAeP7//wIAAAAKAAAABAAAAGl0ZW0AAAAABwAAAGFycl9kZWMAEAAUABAADgAPAAQAAAAIABAAAAAsAAAADAAAAAAAAQ1gAQAABgAAACQBAADcAAAArAAAAIAAAAA8AAAACAAAAEz///8I////HAAAAAwAAAAAAAAIGAAAAAAAAAAAAAYACAAGAAYAAAAAAAAAAQAAAGQAAAA4////HAAAAAwAAAAAAAAKKAAAAAAAAAAIAAwACgAEAAgAAAAIAAAAAAACAAcAAABFdGMvVVRDAAIAAAB0cwAAeP///xQAAAAMAAAAAAAABgwAAAAAAAAA2P///wQAAABmbGFnAAAAAKD///8YAAAADAAAAAAAAAUQAAAAAAAAAAQABAAEAAAABAAAAG5hbWUAAAAAzP///xgAAAAgAAAAAAAAAhwAAAAIAAwABAALAAgAAAAgAAAAAAAAAQAAAAABAAAAbgAAABAAFAAQAAAADwAEAAAACAAQAAAAHAAAAAwAAAAAAAAHHAAAAAAAAAAIAAwACAAEAAgAAAACAAAACgAAAAMAAABhbXQAAQAAAHMAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP/////4AwAAEAAAAAwAGgAYABcABAAIAAwAAAAgAAAAAAkAAAAAAAAAAAAAAAAAAwQACgAYAAwACAAEAAoAAABMAQAAEAAAAAEAAAAAAAAAAAAAABMAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAIAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAwAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAACAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAAAAAAIAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAEAAAAAAAAAAAAAACQAAAAAAAAAAAAAAAEAAAAAAAAAQAAAAAAAAAABAAAAAAAAAIAAAAAAAAAAEAAAAAAAAADAAAAAAAAAAAEAAAAAAAAAAAEAAAAAAAAEAAAAAAAAAEABAAAAAAAAAQAAAAAAAACAAQAAAAAAAAgAAAAAAAAAwAEAAAAAAAACAAAAAAAAAAACAAAAAAAAAQAAAAAAAABAAgAAAAAAAAEAAAAAAAAAgAIAAAAAAAABAAAAAAAAAMACAAAAAAAACAAAAAAAAAAAAwAAAAAAAAEAAAAAAAAAQAMAAAAAAAAEAAAAAAAAAIADAAAAAAAAAQAAAAAAAADAAwAAAAAAAAgAAAAAAAAAAAQAAAAAAAABAAAAAAAAAEAEAAAAAAAAIAAAAAAAAACABAAAAAAAAAEAAAAAAAAAwAQAAAAAAAAIAAAAAAAAAAAFAAAAAAAAAQAAAAAAAABABQAAAAAAAAwAAAAAAAAAgAUAAAAAAAABAAAAAAAAAMAFAAAAAAAACAAAAAAAAAAABgAAAAAAAAEAAAAAAAAAQAYAAAAAAAABAAAAAAAAAIAGAAAAAAAADAAAAAAAAADABgAAAAAAAAQAAAAAAAAAAAcAAAAAAAABAAAAAAAAAEAHAAAAAAAAIAAAAAAAAACABwAAAAAAAAEAAAAAAAAAwAcAAAAAAAABAAAAAAAAAAAIAAAAAAAAAQAAAAAAAABACAAAAAAAABAAAAAAAAAAgAgAAAAAAAABAAAAAAAAAMAIAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAJYAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAKgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAaGkAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEDzJnIfmwUAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAV0cAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAJYAAAAAAAAAAAAAAAAAAAD6AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAMAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAAAAAgAAAAMAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABrMWsyAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAJYAAAAAAAAAAAAAAAAAAAD6AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA/wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP8AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADnAwAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD/////AAAAAA==", + columns: ["s", "arr_dec", "arr_int", "m", "deep", "nul"], + expected: { + s: '{"amt":"1.50","n":"42","name":"hi","flag":"true","ts":"2020-01-02T03:04:05.000Z","d":"2020-01-02"}', + arr_dec: '["1.50","2.50"]', + arr_int: '["1","2","3"]', + m: '{"k1":"1.50","k2":"2.50"}', + deep: '{"inner":{"amt":"9.99"}}', + nul: null, + }, +} as const; + +describe("decodeArrowAttachmentToRows — captured Reyden fixtures", () => { + test("scalar decimal/timestamp/timestamp_ntz/date + one-level nested match native JSON_ARRAY", () => { + const rows = decodeArrowAttachmentToRows( + REYDEN_SCALAR_FIXTURE.attachment, + REYDEN_SCALAR_FIXTURE.columns as unknown as string[], + ); + expect(rows).toEqual([REYDEN_SCALAR_FIXTURE.expected]); + }); + + test("nested struct/list/map leaves are formatted by type, not corrupted to unscaled ints or epoch-ms", () => { + const rows = decodeArrowAttachmentToRows( + REYDEN_NESTED_FIXTURE.attachment, + REYDEN_NESTED_FIXTURE.columns as unknown as string[], + ); + expect(rows).toEqual([REYDEN_NESTED_FIXTURE.expected]); + }); +}); diff --git a/packages/appkit/src/plugins/analytics/types.ts b/packages/appkit/src/plugins/analytics/types.ts index 9e3dd1d2..95c987d1 100644 --- a/packages/appkit/src/plugins/analytics/types.ts +++ b/packages/appkit/src/plugins/analytics/types.ts @@ -15,6 +15,14 @@ export interface IAnalyticsConfig extends BasePluginConfig { * `ConfigurationError`. */ autoStartWarehouse?: boolean; + /** + * Fail-fast ceiling (ms) for an `ARROW_STREAM` query to produce its first + * byte (warehouse readiness + execute + first chunk). Past this, a stuck or + * overloaded warehouse returns a `503` (`WAREHOUSE_UNAVAILABLE`) instead of + * hanging until the client disconnects. Defaults to 2 min. Once the first + * byte arrives the stream is not time-bounded. + */ + arrowFirstByteTimeoutMs?: number; } /** diff --git a/packages/appkit/src/stream/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index d3118e7e..94eb8257 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -4,243 +4,223 @@ import { createLogger } from "../logging/logger"; const logger = createLogger("stream:arrow"); -type ResultManifest = sql.ResultManifest; 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 + * finish draining can hit an expired link; this fetches a fresh one for the + * given chunk index. Created in the request's identity context (it captures the + * caller's workspace client), so it stays valid even though streaming runs + * outside that context. Returns `undefined` if the link can't be re-resolved. + */ +export type RefreshChunkLink = ( + chunkIndex: number, + signal?: AbortSignal, +) => Promise; + interface ArrowStreamOptions { - maxConcurrentDownloads: number; + /** Idle timeout (ms) for a single chunk: no progress → abort the download. */ timeout: number; + /** Attempts to establish a chunk's response before any bytes are yielded. */ retries: number; } -/** - * Result from zero-copy Arrow chunk processing. - * Contains raw IPC bytes without server-side parsing. - */ -interface ArrowRawResult { - /** Concatenated raw Arrow IPC bytes */ - data: Uint8Array; - /** Schema from Databricks manifest (not parsed from Arrow) */ - schema: ResultManifest["schema"]; -} - const BACKOFF_MULTIPLIER = 1000; +/** + * Streams Arrow IPC bytes for a completed statement's EXTERNAL_LINKS chunks + * without ever buffering a whole chunk — each chunk's response body is piped + * through as it arrives, so peak memory is a single network read rather than + * a full chunk (let alone the full result). No Arrow parsing on the server; + * the client parses the concatenated IPC bytes. + */ export class ArrowStreamProcessor { - static readonly DEFAULT_MAX_CONCURRENT_DOWNLOADS = 5; static readonly DEFAULT_TIMEOUT = 30000; static readonly DEFAULT_RETRIES = 3; - constructor( - private options: ArrowStreamOptions = { - maxConcurrentDownloads: - ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS, - timeout: ArrowStreamProcessor.DEFAULT_TIMEOUT, - retries: ArrowStreamProcessor.DEFAULT_RETRIES, - }, - ) { + private options: ArrowStreamOptions; + + constructor(options?: Partial) { this.options = { - maxConcurrentDownloads: - options.maxConcurrentDownloads ?? - ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS, - timeout: options.timeout ?? ArrowStreamProcessor.DEFAULT_TIMEOUT, - retries: options.retries ?? ArrowStreamProcessor.DEFAULT_RETRIES, + timeout: options?.timeout ?? ArrowStreamProcessor.DEFAULT_TIMEOUT, + retries: options?.retries ?? ArrowStreamProcessor.DEFAULT_RETRIES, }; } /** - * Process Arrow chunks using zero-copy proxy pattern. + * Stream Arrow chunks in array order, piping each chunk's response body. * - * Downloads raw IPC bytes from external links and concatenates them - * without parsing into Arrow Tables on the server. This reduces: - * - Memory usage by ~50% (no parsed Table representation) - * - CPU usage (no tableFromIPC/tableToIPC calls) - * - * The client is responsible for parsing the IPC bytes. - * - * @param chunks - External links to Arrow IPC data - * @param schema - Schema from Databricks manifest - * @param signal - Optional abort signal - * @returns Raw concatenated IPC bytes with schema + * Yields network-sized pieces as they arrive (not whole chunks), so peak + * memory is one read buffer. Chunks are fetched sequentially; the + * concatenation of everything yielded is byte-identical to the raw Arrow + * result, so the client parses it exactly as a buffered response. */ - async processChunks( + async *streamChunks( chunks: ExternalLink[], - schema: ResultManifest["schema"], signal?: AbortSignal, - ): Promise { + refresh?: RefreshChunkLink, + ): AsyncGenerator { if (chunks.length === 0) { throw ValidationError.missingField("chunks"); } - const buffers = await this.downloadChunksRaw(chunks, signal); - const data = this.concatenateBuffers(buffers); - - return { data, schema }; - } - - /** - * Download all chunks as raw bytes with concurrency control. - */ - private async downloadChunksRaw( - chunks: ExternalLink[], - signal?: AbortSignal, - ): Promise { - const semaphore = new Semaphore(this.options.maxConcurrentDownloads); - - const downloadPromises = chunks.map(async (chunk) => { - await semaphore.acquire(); - try { - return await this.downloadChunkRaw(chunk, signal); - } finally { - semaphore.release(); - } - }); - - return Promise.all(downloadPromises); + for (const chunk of chunks) { + yield* this.streamChunkBody(chunk, signal, refresh); + } } /** - * Download a single chunk as raw bytes with retry logic. + * Pipe one chunk's response body. + * + * Retry applies only while *establishing* the response (connection failure + * or non-2xx) — once bytes have been yielded downstream we cannot re-fetch, + * so a mid-body failure propagates and the caller aborts the response. An + * idle timeout, reset before every read, aborts a stalled download. */ - private async downloadChunkRaw( + private async *streamChunkBody( chunk: ExternalLink, signal?: AbortSignal, - ): Promise { - let lastError: Error | null = null; + refresh?: RefreshChunkLink, + ): AsyncGenerator { + let externalLink = chunk.external_link; + if (!externalLink) { + // A missing link cannot be fixed by retrying — fail loudly. + throw ExecutionError.statementFailed( + `External link missing for chunk ${chunk.chunk_index}`, + ); + } + + let response: Response | undefined; + let controller: AbortController | undefined; + let onOuterAbort: (() => void) | undefined; + let lastError: unknown; for (let attempt = 0; attempt < this.options.retries; attempt++) { - const timeoutController = new AbortController(); - const timeoutId = setTimeout(() => { - timeoutController.abort(); - }, this.options.timeout); + if (signal?.aborted) throw ExecutionError.canceled(); - const combinedSignal = signal - ? this.combineAbortSignals(signal, timeoutController.signal) - : timeoutController.signal; + const attemptController = new AbortController(); + const listener = () => attemptController.abort(); + signal?.addEventListener("abort", listener, { once: true }); + const timer = setTimeout( + () => attemptController.abort(), + this.options.timeout, + ); try { - const externalLink = chunk.external_link; - if (!externalLink) { - logger.error("External link is required for chunk: %O", chunk); - continue; - } - - const response = await fetch(externalLink, { - signal: combinedSignal, + const r = await fetch(externalLink, { + signal: attemptController.signal, }); - - if (!response.ok) { + clearTimeout(timer); + if (!r.ok) { throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunk_index}: ${response.status} ${response.statusText}`, + `Failed to download chunk ${chunk.chunk_index}: ${r.status} ${r.statusText}`, ); } - - const arrayBuffer = await response.arrayBuffer(); - return new Uint8Array(arrayBuffer); + // Keep this attempt's controller alive to drive the body read + idle + // timeout below. + response = r; + controller = attemptController; + onOuterAbort = listener; + break; } catch (error) { - lastError = error as Error; - - if (timeoutController.signal.aborted) { - lastError = new Error( - `Chunk ${chunk.chunk_index} download timed out after ${this.options.timeout}ms`, - ); - } - - if (signal?.aborted) { - throw ExecutionError.canceled(); - } - + clearTimeout(timer); + signal?.removeEventListener("abort", listener); + lastError = error; + if (signal?.aborted) throw ExecutionError.canceled(); if (attempt < this.options.retries - 1) { - await this.delay(2 ** attempt * BACKOFF_MULTIPLIER); + await this.delay(2 ** attempt * BACKOFF_MULTIPLIER, signal); + // Pre-signed URLs expire (<= 15 min). Before retrying, re-mint this + // 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) { + try { + const fresh = await refresh(chunk.chunk_index, signal); + if (fresh?.external_link) externalLink = fresh.external_link; + } 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, + refreshError, + ); + } + } } - } finally { - clearTimeout(timeoutId); } } - throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunk_index} after ${this.options.retries} attempts: ${lastError?.message}`, - ); - } - - /** - * Concatenate multiple Uint8Array buffers into a single buffer. - * Pre-allocates the result array for efficiency. - */ - private concatenateBuffers(buffers: Uint8Array[]): Uint8Array { - if (buffers.length === 0) { - throw ValidationError.missingField("buffers"); - } - - if (buffers.length === 1) { - return buffers[0]; + if (!response || !controller) { + throw ExecutionError.statementFailed( + `Failed to download chunk ${chunk.chunk_index} after ${this.options.retries} attempts: ${ + lastError instanceof Error ? lastError.message : String(lastError) + }`, + ); } - const totalLength = buffers.reduce((sum, buf) => sum + buf.length, 0); - const result = new Uint8Array(totalLength); - - let offset = 0; - for (const buffer of buffers) { - result.set(buffer, offset); - offset += buffer.length; - } - - return result; - } - - /** - * Combines multiple AbortSignals into one. - * The combined signal aborts when any of the input signals abort. - */ - private combineAbortSignals(...signals: AbortSignal[]): AbortSignal { - const controller = new AbortController(); - - for (const signal of signals) { - if (signal.aborted) { - controller.abort(); - return controller.signal; + const body = response.body; + let idleTimer: ReturnType | undefined; + try { + if (!body) return; // empty chunk (0 rows) + const reader = body.getReader(); + try { + while (true) { + // The idle timeout must cover ONLY the upstream read, never the + // downstream `yield`. Arm it right before `reader.read()` and clear + // it the instant the read resolves — if it stayed armed across the + // yield, a slow client backpressuring `writeChunk` would look like an + // upstream stall and abort a perfectly healthy download. + idleTimer = setTimeout( + () => controller.abort(), + this.options.timeout, + ); + let done: boolean; + let value: Uint8Array | undefined; + try { + ({ done, value } = await reader.read()); + } finally { + clearTimeout(idleTimer); + } + if (done) break; + if (value && value.byteLength > 0) yield value; + } + } finally { + // Release the stream on early exit (downstream abort / error). + reader.cancel().catch(() => {}); } - signal.addEventListener("abort", () => controller.abort(), { - once: true, - }); + } catch (error) { + if (signal?.aborted) throw ExecutionError.canceled(); + logger.error( + "Failed streaming chunk %s body: %O", + chunk.chunk_index, + error, + ); + throw error instanceof ExecutionError + ? error + : ExecutionError.statementFailed( + `Failed streaming chunk ${chunk.chunk_index}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } finally { + clearTimeout(idleTimer); + if (onOuterAbort) signal?.removeEventListener("abort", onOuterAbort); } - - return controller.signal; - } - - private delay(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); - } -} - -class Semaphore { - private permits: number; - private waiting: (() => void)[] = []; - - constructor(permits: number) { - this.permits = permits; } - async acquire(): Promise { - if (this.permits > 0) { - this.permits--; - return; - } - - return new Promise((resolve) => { - this.waiting.push(resolve); + private delay(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(); + }, ms); + const onAbort = () => { + clearTimeout(timer); + reject(ExecutionError.canceled()); + }; + signal?.addEventListener("abort", onAbort, { once: true }); }); } - - release(): void { - if (this.waiting.length > 0) { - const next = this.waiting.shift(); - - if (next) { - next(); - } - } else { - this.permits++; - } - } } diff --git a/packages/appkit/src/stream/defaults.ts b/packages/appkit/src/stream/defaults.ts index 304a6311..c8863b4f 100644 --- a/packages/appkit/src/stream/defaults.ts +++ b/packages/appkit/src/stream/defaults.ts @@ -1,6 +1,10 @@ export const streamDefaults = { bufferSize: 100, - maxEventSize: 1024 * 1024, // 1MB + // 1 MiB. SSE carries only short JSON control messages — JSON_ARRAY result + // rows (already row-size-bounded by the warehouse) plus warehouse-readiness + // and error events. ARROW_STREAM never uses SSE: the raw Arrow IPC bytes + // stream back on the query response body (`_handleArrowStreamQuery`). + maxEventSize: 1 * 1024 * 1024, bufferTTL: 10 * 60 * 1000, // 10 minutes cleanupInterval: 5 * 60 * 1000, // 5 minutes maxPersistentBuffers: 10000, // 10000 buffers diff --git a/packages/appkit/src/stream/sse-writer.ts b/packages/appkit/src/stream/sse-writer.ts index b73cdb58..2b49d1de 100644 --- a/packages/appkit/src/stream/sse-writer.ts +++ b/packages/appkit/src/stream/sse-writer.ts @@ -39,12 +39,14 @@ export class SSEWriter { eventId: string, error: string, code: SSEErrorCode = SSEErrorCode.INTERNAL_ERROR, + errorCode?: string, ): void { if (res.writableEnded) return; const errorData: SSEError = { error, code, + ...(errorCode ? { errorCode } : {}), }; res.write(`id: ${eventId}\n`); diff --git a/packages/appkit/src/stream/stream-manager.ts b/packages/appkit/src/stream/stream-manager.ts index 14741d64..cf28ead5 100644 --- a/packages/appkit/src/stream/stream-manager.ts +++ b/packages/appkit/src/stream/stream-manager.ts @@ -1,6 +1,8 @@ import { randomUUID } from "node:crypto"; import { context } from "@opentelemetry/api"; import type { IAppResponse, StreamConfig } from "shared"; +import { AppKitError } from "../errors/base"; +import { ExecutionError } from "../errors/execution"; import { createLogger } from "../logging/logger"; import { EventRingBuffer } from "./buffers"; import { streamDefaults } from "./defaults"; @@ -296,8 +298,21 @@ export class StreamManager { // cleanup if no clients are connected this._cleanupStream(streamEntry); } catch (error) { - const errorMsg = + // Two distinct messages: a *raw* one for server-side logs (full + // detail, statement fragments, correlation IDs) and a *client* + // one for the SSE payload (sanitized, stable, safe to render in + // a UI). Mixing them leaks upstream wording to anyone connected + // to the stream — see CWE-209. + const rawMsg = error instanceof Error ? error.message : "Internal server error"; + const clientMsg = + error instanceof AppKitError + ? error.clientMessage + : "Internal server error"; + // Upstream structured code (e.g. RESULT_TOO_LARGE_FOR_JSON_FALLBACK, + // NOT_IMPLEMENTED). UI should branch on this, not on `error`. + const upstreamCode = + error instanceof ExecutionError ? error.errorCode : undefined; const errorEventId = randomUUID(); const errorCode = this._categorizeError(error); @@ -312,16 +327,23 @@ export class StreamManager { } logger.error( - "Stream execution failed: %s (code=%s)", - errorMsg, + "Stream execution failed: %s (code=%s upstreamCode=%s)", + rawMsg, errorCode, + upstreamCode ?? "n/a", ); + const payload: Record = { + error: clientMsg, + code: errorCode, + }; + if (upstreamCode) payload.errorCode = upstreamCode; + // buffer error event streamEntry.eventBuffer.add({ id: errorEventId, type: "error", - data: JSON.stringify({ error: errorMsg, code: errorCode }), + data: JSON.stringify(payload), timestamp: Date.now(), }); @@ -329,9 +351,10 @@ export class StreamManager { this._broadcastErrorToClients( streamEntry, errorEventId, - errorMsg, + clientMsg, errorCode, true, + upstreamCode, ); streamEntry.isCompleted = true; this._clearGraceTimer(streamEntry); @@ -387,10 +410,17 @@ export class StreamManager { errorMessage: string, errorCode: SSEErrorCode, closeClients: boolean = false, + upstreamCode?: string, ): void { for (const client of streamEntry.clients) { if (!client.writableEnded) { - this.sseWriter.writeError(client, eventId, errorMessage, errorCode); + this.sseWriter.writeError( + client, + eventId, + errorMessage, + errorCode, + upstreamCode, + ); if (closeClients) { client.end(); } 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 8a95fbdf..eb367929 100644 --- a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts +++ b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts @@ -1,11 +1,17 @@ -import type { sql } from "@databricks/sdk-experimental"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { ArrowStreamProcessor } from "../arrow-stream-processor"; -type ResultSchema = sql.ResultManifest["schema"]; +/** A ReadableStream that emits the given pieces in order, then closes. */ +function streamOf(...pieces: Uint8Array[]): ReadableStream { + return new ReadableStream({ + start(controller) { + for (const piece of pieces) controller.enqueue(piece); + controller.close(); + }, + }); +} -// Helper to create mock chunks -function createMockChunks(count: number) { +function mockChunks(count: number) { return Array.from({ length: count }, (_, i) => ({ chunk_index: i, external_link: `https://example.com/chunk-${i}`, @@ -14,46 +20,33 @@ function createMockChunks(count: number) { })); } -// Helper to create mock schema -function createMockSchema(): ResultSchema { - return { - columns: [ - { name: "id", type_name: "INT" }, - { name: "name", type_name: "STRING" }, - ], - } as ResultSchema; +/** Drain a byte generator to a flat number[]. */ +async function drain( + gen: AsyncGenerator, +): Promise { + const out: number[] = []; + for await (const piece of gen) out.push(...piece); + return out; } -describe("ArrowStreamProcessor", () => { +describe("ArrowStreamProcessor.streamChunks", () => { let processor: ArrowStreamProcessor; let originalFetch: typeof globalThis.fetch; beforeEach(() => { - processor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 3, - timeout: 5000, - retries: 3, - }); - + processor = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); originalFetch = globalThis.fetch; - - // Default mock: successful fetch returning raw bytes - globalThis.fetch = vi.fn().mockImplementation(async (url: string) => { - // Extract chunk index from URL to create unique data - const match = url.match(/chunk-(\d+)/); - const chunkIndex = match ? parseInt(match[1], 10) : 0; - - // Return raw bytes (simulating Arrow IPC data) + // Default: each chunk's body is a single piece [100 + index]. + globalThis.fetch = vi.fn(async (input: string | URL | Request) => { + const match = String(input).match(/chunk-(\d+)/); + const index = match ? Number.parseInt(match[1], 10) : 0; return { ok: true, status: 200, statusText: "OK", - arrayBuffer: async () => - new Uint8Array([chunkIndex + 1, chunkIndex + 2, chunkIndex + 3]) - .buffer, - }; + body: streamOf(new Uint8Array([100 + index])), + } as unknown as Response; }); - vi.clearAllMocks(); }); @@ -62,484 +55,211 @@ describe("ArrowStreamProcessor", () => { vi.restoreAllMocks(); }); - describe("constructor", () => { - test("should use default options when not provided", () => { - const defaultProcessor = new ArrowStreamProcessor(); - expect(defaultProcessor).toBeDefined(); - }); - - test("should accept custom options", () => { - const customProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 10, - timeout: 60000, - retries: 5, - }); - expect(customProcessor).toBeDefined(); - }); - - test("should use defaults for missing option properties", () => { - const partialProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 2, - } as any); - expect(partialProcessor).toBeDefined(); - }); + test("throws when no chunks are provided", async () => { + await expect(drain(processor.streamChunks([]))).rejects.toThrow(); }); - describe("processChunks", () => { - test("should throw error when no chunks provided", async () => { - await expect( - processor.processChunks([], createMockSchema()), - ).rejects.toThrow("Missing required field: chunks"); - }); - - test("should process single chunk successfully", async () => { - const chunks = createMockChunks(1); - const schema = createMockSchema(); - - const result = await processor.processChunks(chunks, schema); - - expect(result).toHaveProperty("data"); - expect(result).toHaveProperty("schema", schema); - expect(result.data).toBeInstanceOf(Uint8Array); - expect(globalThis.fetch).toHaveBeenCalledTimes(1); - }); - - test("should process multiple chunks successfully", async () => { - const chunks = createMockChunks(5); - const schema = createMockSchema(); - - const result = await processor.processChunks(chunks, schema); - - expect(result.data).toBeInstanceOf(Uint8Array); - expect(result.schema).toBe(schema); - expect(globalThis.fetch).toHaveBeenCalledTimes(5); - }); - - test("should concatenate raw bytes from multiple chunks", async () => { - const chunks = createMockChunks(3); - const schema = createMockSchema(); - - const result = await processor.processChunks(chunks, schema); - - // Each chunk returns 3 bytes: [chunkIndex+1, chunkIndex+2, chunkIndex+3] - // Chunk 0: [1, 2, 3], Chunk 1: [2, 3, 4], Chunk 2: [3, 4, 5] - expect(result.data).toEqual(new Uint8Array([1, 2, 3, 2, 3, 4, 3, 4, 5])); - }); - - test("should return single chunk data without modification", async () => { - const chunks = createMockChunks(1); - const schema = createMockSchema(); - - const result = await processor.processChunks(chunks, schema); - - // Single chunk returns [1, 2, 3] - expect(result.data).toEqual(new Uint8Array([1, 2, 3])); - }); - - test("should pass abort signal to fetch", async () => { - const chunks = createMockChunks(1); - const schema = createMockSchema(); - const abortController = new AbortController(); - - await processor.processChunks(chunks, schema, abortController.signal); - - expect(globalThis.fetch).toHaveBeenCalledWith( - expect.any(String), - expect.objectContaining({ - signal: expect.any(AbortSignal), - }), - ); - }); + test("streams chunks in array order", async () => { + const bytes = await drain(processor.streamChunks(mockChunks(3))); + expect(bytes).toEqual([100, 101, 102]); }); - describe("concurrent downloads", () => { - test("should limit concurrent downloads with semaphore", async () => { - const maxConcurrent = 2; - const limitedProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: maxConcurrent, - timeout: 5000, - retries: 1, - }); - - let currentConcurrent = 0; - let maxObservedConcurrent = 0; - - globalThis.fetch = vi.fn().mockImplementation(async () => { - currentConcurrent++; - maxObservedConcurrent = Math.max( - maxObservedConcurrent, - currentConcurrent, - ); - - // Simulate network delay - await new Promise((resolve) => setTimeout(resolve, 10)); - - currentConcurrent--; - - return { - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; - }); - - const chunks = createMockChunks(10); - await limitedProcessor.processChunks(chunks, createMockSchema()); - - expect(maxObservedConcurrent).toBeLessThanOrEqual(maxConcurrent); - expect(globalThis.fetch).toHaveBeenCalledTimes(10); - }); - }); - - describe("retry logic", () => { - test("should retry on fetch failure", async () => { - let attempts = 0; - globalThis.fetch = vi.fn().mockImplementation(async () => { - attempts++; - if (attempts < 3) { - throw new Error("Network error"); - } - return { - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; - }); - - const chunks = createMockChunks(1); - const result = await processor.processChunks(chunks, createMockSchema()); - - expect(attempts).toBe(3); - expect(result.data).toBeDefined(); - }); - - test("should throw after exhausting retries", async () => { - globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network error")); - - const chunks = createMockChunks(1); - - await expect( - processor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/Failed to download chunk 0 after 3 attempts/); - - expect(globalThis.fetch).toHaveBeenCalledTimes(3); - }); - - test("should retry on non-ok response", async () => { - let attempts = 0; - globalThis.fetch = vi.fn().mockImplementation(async () => { - attempts++; - if (attempts < 2) { - return { - ok: false, - status: 500, - statusText: "Internal Server Error", - }; - } - return { + test("pipes a chunk body piece-by-piece (no whole-chunk buffering)", async () => { + globalThis.fetch = vi.fn( + async () => + ({ ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; - }); - - const chunks = createMockChunks(1); - const result = await processor.processChunks(chunks, createMockSchema()); - - expect(attempts).toBe(2); - expect(result.data).toBeDefined(); - }); - - test("should use exponential backoff between retries", async () => { - vi.useFakeTimers(); - - let attempts = 0; - globalThis.fetch = vi.fn().mockImplementation(async () => { - attempts++; - if (attempts < 3) { - throw new Error("Network error"); - } - return { - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; - }); - - const chunks = createMockChunks(1); - const promise = processor.processChunks(chunks, createMockSchema()); - - // First attempt - immediate - await vi.advanceTimersByTimeAsync(0); - expect(attempts).toBe(1); - - // First retry after 1000ms (2^0 * 1000) - await vi.advanceTimersByTimeAsync(1000); - expect(attempts).toBe(2); - - // Second retry after 2000ms (2^1 * 1000) - await vi.advanceTimersByTimeAsync(2000); - expect(attempts).toBe(3); - - await promise; - vi.useRealTimers(); - }); - }); - - describe("timeout handling", () => { - test("should timeout slow requests", async () => { - const shortTimeoutProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 1, - timeout: 50, - retries: 1, - }); - - globalThis.fetch = vi.fn().mockImplementation( - (_url: string, options?: { signal?: AbortSignal }) => - new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - resolve({ - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }); - }, 5000); // Much longer than timeout - - // Listen for abort (from timeout) - options?.signal?.addEventListener("abort", () => { - clearTimeout(timeout); - reject(new DOMException("Aborted", "AbortError")); - }); - }), - ); - - const chunks = createMockChunks(1); - - // The processor should timeout and reject - await expect( - shortTimeoutProcessor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/timed out|Failed to download/); - }); - - test("should clear timeout after successful fetch", async () => { - const clearTimeoutSpy = vi.spyOn(globalThis, "clearTimeout"); - - const chunks = createMockChunks(1); - await processor.processChunks(chunks, createMockSchema()); - - expect(clearTimeoutSpy).toHaveBeenCalled(); - clearTimeoutSpy.mockRestore(); - }); + status: 200, + body: streamOf(new Uint8Array([1, 2]), new Uint8Array([3, 4, 5])), + }) as unknown as Response, + ); + // Collect per-yield pieces (not flattened) to prove the body streams in + // multiple reads rather than arriving as one buffer. + const pieces: number[][] = []; + for await (const piece of processor.streamChunks(mockChunks(1))) { + pieces.push([...piece]); + } + expect(pieces).toEqual([ + [1, 2], + [3, 4, 5], + ]); }); - describe("abort signal handling", () => { - test("should abort immediately if signal already aborted", async () => { - const abortController = new AbortController(); - abortController.abort(); - - // Mock fetch to check if it receives an aborted signal - globalThis.fetch = vi - .fn() - .mockImplementation( - async (_url: string, options?: { signal?: AbortSignal }) => { - // If signal is already aborted, throw - if (options?.signal?.aborted) { - throw new DOMException("Aborted", "AbortError"); - } - return { - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; + test("does not abort a healthy download when the consumer stalls between reads (backpressure)", async () => { + // Body emits [1] immediately, then nothing until the test enqueues more. + // It errors if the attempt's abort signal fires — mirroring a real fetch + // being aborted. If the idle timer stayed armed across the downstream + // `yield`, a slow consumer would trip it and kill this healthy stream. + let bodyController!: ReadableStreamDefaultController; + globalThis.fetch = vi.fn( + async (_input: string | URL | Request, init?: RequestInit) => { + const body = new ReadableStream({ + start(c) { + bodyController = c; + c.enqueue(new Uint8Array([1])); }, + }); + init?.signal?.addEventListener("abort", () => + bodyController.error(new Error("idle-aborted")), ); + return { ok: true, status: 200, body } as unknown as Response; + }, + ); - const chunks = createMockChunks(1); - - await expect( - processor.processChunks( - chunks, - createMockSchema(), - abortController.signal, - ), - ).rejects.toThrow("Statement was canceled"); - }); + const p = new ArrowStreamProcessor({ timeout: 50, retries: 1 }); + const gen = p.streamChunks(mockChunks(1)); - test("should abort in-flight requests when signal fires", async () => { - const abortController = new AbortController(); - let fetchStarted = false; - - globalThis.fetch = vi - .fn() - .mockImplementation( - async (_url: string, options?: { signal?: AbortSignal }) => { - fetchStarted = true; - - // Simulate slow request that checks signal - return new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - resolve({ - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }); - }, 1000); - - options?.signal?.addEventListener("abort", () => { - clearTimeout(timeout); - reject(new DOMException("Aborted", "AbortError")); - }); - }); - }, - ); + const first = await gen.next(); + expect([...(first.value as Uint8Array)]).toEqual([1]); - const chunks = createMockChunks(1); - const promise = processor.processChunks( - chunks, - createMockSchema(), - abortController.signal, - ); + // Consumer stalls well past the 50ms idle timeout while suspended at yield. + await new Promise((r) => setTimeout(r, 120)); - // Wait for fetch to start, then abort - await vi.waitFor(() => expect(fetchStarted).toBe(true)); - abortController.abort(); + // The download is still healthy — the rest streams through. + bodyController.enqueue(new Uint8Array([2])); + bodyController.close(); - await expect(promise).rejects.toThrow("Statement was canceled"); - }); + const second = await gen.next(); + expect(second.done).toBe(false); + expect([...(second.value as Uint8Array)]).toEqual([2]); + expect((await gen.next()).done).toBe(true); }); - describe("buffer concatenation", () => { - test("should concatenate buffers from multiple chunks", async () => { - // Custom fetch that returns distinct byte patterns - globalThis.fetch = vi.fn().mockImplementation(async (url: string) => { - const match = url.match(/chunk-(\d+)/); - const index = match ? parseInt(match[1], 10) : 0; - // Each chunk returns a unique byte pattern - return { - ok: true, - arrayBuffer: async () => new Uint8Array([100 + index]).buffer, - }; - }); - - const chunks = createMockChunks(3); - const result = await processor.processChunks(chunks, createMockSchema()); + test("downloads lazily — one fetch per pulled chunk, not all upfront", async () => { + const fetchMock = globalThis.fetch as ReturnType; + const iterator = processor.streamChunks(mockChunks(3)); - // Should be [100, 101, 102] - expect(result.data).toEqual(new Uint8Array([100, 101, 102])); - }); + // Draining the first chunk's body fetches exactly one chunk. + await iterator.next(); // first body piece of chunk 0 + await iterator.next(); // chunk 0 done → advances to chunk 1's fetch + expect(fetchMock).toHaveBeenCalledTimes(2); + }); - test("should return single buffer without unnecessary allocation", async () => { - globalThis.fetch = vi.fn().mockResolvedValue({ + test("retries establishing the response, then streams the body", async () => { + const fetchMock = vi + .fn() + .mockRejectedValueOnce(new Error("connection reset")) + .mockResolvedValueOnce({ ok: true, - arrayBuffer: async () => new Uint8Array([42, 43, 44]).buffer, - }); - - const chunks = createMockChunks(1); - const result = await processor.processChunks(chunks, createMockSchema()); - - // Single chunk should return as-is - expect(result.data).toEqual(new Uint8Array([42, 43, 44])); - }); + status: 200, + body: streamOf(new Uint8Array([7, 8, 9])), + } as unknown as Response); + globalThis.fetch = fetchMock; + + const bytes = await drain(processor.streamChunks(mockChunks(1))); + + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(bytes).toEqual([7, 8, 9]); + }, 10000); + + test("throws after exhausting retries on a non-2xx response", async () => { + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 1 }); + globalThis.fetch = vi.fn( + async () => + ({ ok: false, status: 500, statusText: "Server Error" }) as Response, + ); + await expect(drain(p.streamChunks(mockChunks(1)))).rejects.toThrow( + /chunk 0/, + ); }); - describe("missing external_link handling", () => { - test("should log error and continue retrying on missing external_link", async () => { - // Create chunk without external_link - const chunks = [ - { chunk_index: 0, external_link: undefined }, - { chunk_index: 1, external_link: "https://example.com/chunk-1" }, - ] as any; - - const consoleSpy = vi - .spyOn(console, "error") - .mockImplementation(() => {}); - - // This will fail because chunk 0 has no external_link and retries will be exhausted - await expect( - processor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(); - - // Logger uses util.format, so the message is pre-formatted - expect(consoleSpy).toHaveBeenCalledWith( - "[appkit:stream:arrow]", - expect.stringContaining("External link is required for chunk:"), - ); - - consoleSpy.mockRestore(); - }); + test("throws immediately when a chunk has no external_link", async () => { + const chunks = [{ chunk_index: 0 }] as any; + await expect(drain(processor.streamChunks(chunks))).rejects.toThrow( + /External link missing/, + ); }); - describe("error messages", () => { - test("should include chunk index in error message", async () => { - globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network error")); - - const chunks = createMockChunks(1); - - await expect( - processor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/chunk 0/); - }); + test("aborts during retry backoff instead of waiting it out", async () => { + const controller = new AbortController(); + // Every attempt fails to connect, forcing the retry backoff between them. + const fetchMock = vi.fn().mockRejectedValue(new Error("connection reset")); + globalThis.fetch = fetchMock; - test("should include HTTP status in error message for failed responses", async () => { - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: false, - status: 403, - statusText: "Forbidden", - }); + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const drained = drain(p.streamChunks(mockChunks(1), controller.signal)); - const singleRetryProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 1, - timeout: 5000, - retries: 1, - }); + // Let the first attempt fail and enter the (~1s) backoff, then abort. + await new Promise((r) => setTimeout(r, 50)); + controller.abort(); - const chunks = createMockChunks(1); - - await expect( - singleRetryProcessor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/403 Forbidden/); - }); + await expect(drained).rejects.toThrow(); + // Aborting mid-backoff cancels immediately — no second fetch attempt. + expect(fetchMock).toHaveBeenCalledTimes(1); }); -}); - -describe("Semaphore (via ArrowStreamProcessor)", () => { - let originalFetch: typeof globalThis.fetch; - beforeEach(() => { - originalFetch = globalThis.fetch; - vi.clearAllMocks(); + test("an already-aborted signal cancels before fetching", async () => { + const controller = new AbortController(); + controller.abort(); + const fetchMock = globalThis.fetch as ReturnType; + await expect( + drain(processor.streamChunks(mockChunks(1), controller.signal)), + ).rejects.toThrow(); + expect(fetchMock).not.toHaveBeenCalled(); }); - afterEach(() => { - globalThis.fetch = originalFetch; - }); - - test("should properly queue and release permits", async () => { - const processor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 1, - timeout: 5000, - retries: 1, - }); - - const order: number[] = []; - - globalThis.fetch = vi.fn().mockImplementation(async (url: string) => { - const match = url.match(/chunk-(\d+)/); - const index = match ? parseInt(match[1], 10) : 0; - order.push(index); - - // Simulate varying response times - await new Promise((resolve) => setTimeout(resolve, 5)); - - return { + test("re-resolves an expired chunk link before retrying, then streams the fresh one", async () => { + // First attempt hits an expired (403) link; the refresher mints a fresh + // URL that the retry succeeds against. + const fetchMock = vi + .fn() + .mockResolvedValueOnce({ + ok: false, + status: 403, + statusText: "Forbidden", + } as Response) + .mockResolvedValueOnce({ ok: true, - arrayBuffer: async () => new Uint8Array([index + 1]).buffer, - }; - }); + status: 200, + body: streamOf(new Uint8Array([42])), + } as unknown as Response); + globalThis.fetch = fetchMock; + + const refresh = vi.fn(async (chunkIndex: number) => ({ + chunk_index: chunkIndex, + external_link: "https://example.com/fresh-link", + })); + + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const bytes = await drain( + p.streamChunks(mockChunks(1), undefined, refresh), + ); + + expect(refresh).toHaveBeenCalledWith(0, undefined); + // The retry fetched the fresh URL, not the stale one. + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(String(fetchMock.mock.calls[1][0])).toBe( + "https://example.com/fresh-link", + ); + expect(bytes).toEqual([42]); + }, 10000); + + test("survives a refresher that throws — keeps retrying the current link", async () => { + const fetchMock = vi + .fn() + .mockResolvedValueOnce({ + ok: false, + status: 403, + statusText: "Forbidden", + } as Response) + .mockResolvedValueOnce({ + ok: true, + status: 200, + body: streamOf(new Uint8Array([7])), + } as unknown as Response); + globalThis.fetch = fetchMock; - const chunks = [ - { chunk_index: 0, external_link: "https://example.com/chunk-0" }, - { chunk_index: 1, external_link: "https://example.com/chunk-1" }, - { chunk_index: 2, external_link: "https://example.com/chunk-2" }, - ]; + const refresh = vi.fn(async () => { + throw new Error("re-resolve failed"); + }); - await processor.processChunks(chunks as any, { columns: [] }); + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const bytes = await drain( + p.streamChunks(mockChunks(1), undefined, refresh), + ); - // With concurrency of 1, they should complete in order - expect(order).toHaveLength(3); - expect(globalThis.fetch).toHaveBeenCalledTimes(3); - }); + expect(refresh).toHaveBeenCalled(); + // Refresh failure is swallowed; the retry still runs against the original. + expect(bytes).toEqual([7]); + }, 10000); }); diff --git a/packages/appkit/src/stream/tests/stream-registry.test.ts b/packages/appkit/src/stream/tests/stream-registry.test.ts index d3f70e95..efb88c88 100644 --- a/packages/appkit/src/stream/tests/stream-registry.test.ts +++ b/packages/appkit/src/stream/tests/stream-registry.test.ts @@ -374,7 +374,7 @@ describe("StreamRegistry", () => { expect(dataCall).toBeDefined(); const payload = JSON.parse( - (dataCall![0] as string).replace("data: ", "").trim(), + (dataCall?.[0] as string).replace("data: ", "").trim(), ); expect(payload).toEqual({ error: "Stream evicted", diff --git a/packages/appkit/src/stream/tests/stream.test.ts b/packages/appkit/src/stream/tests/stream.test.ts index 1de49b5f..e111283c 100644 --- a/packages/appkit/src/stream/tests/stream.test.ts +++ b/packages/appkit/src/stream/tests/stream.test.ts @@ -410,21 +410,30 @@ describe("StreamManager", () => { }); describe("error handling", () => { - test("should send error event when handler throws", async () => { + test("should send a sanitized error event when handler throws — raw error text is kept server-side only", async () => { + // The SSE error broadcaster never echoes a bare Error's `.message` + // to the wire: that path can carry statement fragments, internal + // object names, and correlation IDs (CWE-209). Non-AppKitError + // throws collapse to "Internal server error"; AppKitErrors carry + // their `clientMessage` (sanitized) and `errorCode` (structured). const { mockRes, events } = createMockResponse(); async function* generator() { yield { type: "start" }; - throw new Error("Test error"); + throw new Error("Test error with /internal/path and corrId=abc-123"); } await streamManager.stream(mockRes as any, generator); expect(events).toContain('data: {"type":"start"}\n\n'); expect(events).toContain("event: error\n"); - expect(events).toContain( - 'data: {"error":"Test error","code":"INTERNAL_ERROR"}\n\n', - ); + const dataLines = events.filter((e) => e.startsWith("data: ")); + const errorLine = dataLines.find((e) => e.includes('"error"')); + expect(errorLine).toContain('"error":"Internal server error"'); + expect(errorLine).toContain('"code":"INTERNAL_ERROR"'); + // The raw Error.message must not appear anywhere on the wire. + expect(events.join("")).not.toContain("/internal/path"); + expect(events.join("")).not.toContain("corrId=abc-123"); expect(mockRes.end).toHaveBeenCalled(); }); @@ -690,10 +699,15 @@ describe("StreamManager", () => { await streamManager.stream(mockRes2 as any, generator2, { streamId }); - const replayedError = events2.some((e) => - e.includes("Something went wrong"), + // The buffered error event must be replayed — sanitized, not raw. + // We don't pin the exact wording (that's a separate test) but we + // do require the error frame to be present. + const replayedError = events2.some( + (e) => e.startsWith("data:") && e.includes('"error"'), ); expect(replayedError).toBe(true); + // And the raw thrown message must not have leaked. + expect(events2.join("")).not.toContain("Something went wrong"); }); test("should detect buffer overflow on completed stream and close connection", async () => { diff --git a/packages/appkit/src/stream/types.ts b/packages/appkit/src/stream/types.ts index e9a6bbc0..78b14e35 100644 --- a/packages/appkit/src/stream/types.ts +++ b/packages/appkit/src/stream/types.ts @@ -25,6 +25,12 @@ export type SSEErrorCode = (typeof SSEErrorCode)[keyof typeof SSEErrorCode]; export interface SSEError { error: string; code: SSEErrorCode; + /** + * Upstream-domain structured code (e.g. `RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, + * `NOT_IMPLEMENTED`). UI code should branch on this instead of parsing + * the human-readable `error` string. + */ + errorCode?: string; } export interface BufferedEvent { diff --git a/packages/appkit/src/type-generator/query-registry.ts b/packages/appkit/src/type-generator/query-registry.ts index 1844720f..fb239487 100644 --- a/packages/appkit/src/type-generator/query-registry.ts +++ b/packages/appkit/src/type-generator/query-registry.ts @@ -1,6 +1,7 @@ import fs from "node:fs/promises"; import path from "node:path"; import { WorkspaceClient } from "@databricks/sdk-experimental"; +import { tableFromIPC } from "apache-arrow"; import pc from "picocolors"; import { createLogger } from "../logging/logger"; import { CACHE_VERSION, hashSQL, loadCache, saveCache } from "./cache"; @@ -147,18 +148,69 @@ function formatParametersType(sql: string): string { : "Record"; } +/** + * 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. + * + * IMPORTANT: a DESCRIBE QUERY response is itself a result *table* with rows + * shaped like `(col_name, data_type, comment)` describing the user query's + * output schema. We must read those rows — NOT `table.schema.fields`, which + * would describe DESCRIBE QUERY's own output (`col_name`, `data_type`, + * `comment`) and yield bogus types for every query. + */ +function columnsFromArrowAttachment( + attachment: string, +): Array<{ name: string; type_name: string; comment: string | undefined }> { + const buf = Buffer.from(attachment, "base64"); + const table = tableFromIPC(buf); + return table.toArray().map((row) => { + const obj = row.toJSON() as { + col_name?: unknown; + data_type?: unknown; + comment?: unknown; + }; + return { + name: typeof obj.col_name === "string" ? obj.col_name : "", + type_name: + typeof obj.data_type === "string" + ? obj.data_type.toUpperCase() + : "STRING", + comment: + typeof obj.comment === "string" && obj.comment !== "" + ? obj.comment + : undefined, + }; + }); +} + export function convertToQueryType( result: DatabricksStatementExecutionResponse, sql: string, queryName: string, ): { type: string; hasResults: boolean } { const dataRows = result.result?.data_array || []; - const columns = dataRows.map((row) => ({ + let columns = dataRows.map((row) => ({ name: row[0] || "", type_name: row[1]?.toUpperCase() || "STRING", comment: row[2] || undefined, })); + // Fallback: serverless warehouses return ARROW_STREAM format with an inline + // base64 attachment instead of data_array. 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"); + try { + columns = columnsFromArrowAttachment(result.result.attachment); + } catch (err) { + logger.warn( + "Failed to decode Arrow IPC attachment: %s", + err instanceof Error ? err.message : String(err), + ); + } + } + const paramsType = formatParametersType(sql); // generate result fields with JSDoc @@ -748,10 +800,11 @@ export async function generateQueriesFromDescribe( ); logger.debug( - "DESCRIBE result for %s: state=%s, rows=%d", + "DESCRIBE result for %s: state=%s, rows=%d, hasAttachment=%s", queryName, result.status.state, result.result?.data_array?.length ?? 0, + !!result.result?.attachment, ); if (result.status.state === "FAILED") { 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 1e37c802..1f8156a2 100644 --- a/packages/appkit/src/type-generator/tests/query-registry.test.ts +++ b/packages/appkit/src/type-generator/tests/query-registry.test.ts @@ -1,4 +1,24 @@ -import { describe, expect, test } from "vitest"; +import { Table, tableToIPC, vectorFromArray } from "apache-arrow"; +import { describe, expect, test, vi } from "vitest"; + +const { mockLoggerWarn, mockLoggerDebug } = vi.hoisted(() => ({ + mockLoggerWarn: vi.fn(), + mockLoggerDebug: vi.fn(), +})); +vi.mock("../../logging/logger", () => ({ + createLogger: vi.fn(() => ({ + debug: mockLoggerDebug, + info: vi.fn(), + warn: mockLoggerWarn, + error: vi.fn(), + event: vi.fn(() => ({ + set: vi.fn().mockReturnThis(), + setComponent: vi.fn().mockReturnThis(), + setContext: vi.fn().mockReturnThis(), + })), + })), +})); + import { convertToQueryType, defaultForType, @@ -13,6 +33,20 @@ import { } from "../query-registry"; import type { DatabricksStatementExecutionResponse } from "../types"; +// Build a base64 Arrow IPC payload that mimics a DESCRIBE QUERY response — +// a result *table* with columns (col_name, data_type, comment) describing +// the user query's output schema. +function describeQueryAttachment( + rows: Array<{ col_name: string; data_type: string; comment: string | null }>, +): string { + const table = new Table({ + col_name: vectorFromArray(rows.map((r) => r.col_name)), + data_type: vectorFromArray(rows.map((r) => r.data_type)), + comment: vectorFromArray(rows.map((r) => r.comment ?? "")), + }); + return Buffer.from(tableToIPC(table, "stream")).toString("base64"); +} + describe("normalizeTypeName", () => { test("returns simple types unchanged", () => { expect(normalizeTypeName("STRING")).toBe("STRING"); @@ -391,6 +425,107 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` ); expect(hasResults).toBe(false); }); + + describe("ARROW_STREAM attachment fallback (serverless warehouses)", () => { + test("decodes column metadata from Arrow IPC data rows, not schema fields", () => { + // Critical regression test: it would be a bug to read + // `table.schema.fields` here, which would generate types like + // { col_name: string; data_type: string; comment: string } for every + // query (those are DESCRIBE QUERY's own output columns). We must read + // the data rows. + const attachment = describeQueryAttachment([ + { col_name: "user_id", data_type: "BIGINT", comment: null }, + { col_name: "name", data_type: "STRING", comment: "display name" }, + { col_name: "active", data_type: "BOOLEAN", comment: null }, + ]); + const response: DatabricksStatementExecutionResponse = { + statement_id: "test-arrow", + status: { state: "SUCCEEDED" }, + result: { attachment }, + }; + + const { type, hasResults } = convertToQueryType( + response, + "SELECT user_id, name, active FROM users", + "users", + ); + + expect(hasResults).toBe(true); + // Real query columns appear in the generated type: + expect(type).toContain("user_id: number"); + expect(type).toContain("name: string"); + expect(type).toContain("active: boolean"); + // Column comments survive: + expect(type).toContain("/** display name"); + // The DESCRIBE QUERY metadata column names must NOT leak as user types: + expect(type).not.toContain("col_name: string"); + expect(type).not.toContain("data_type: string"); + }); + + test("normalizes lowercase data_type values to uppercase", () => { + const attachment = describeQueryAttachment([ + { col_name: "id", data_type: "int", comment: null }, + ]); + const response: DatabricksStatementExecutionResponse = { + statement_id: "test-arrow", + status: { state: "SUCCEEDED" }, + result: { attachment }, + }; + + const { type } = convertToQueryType(response, "SELECT 1", "test"); + expect(type).toContain("@sqlType INT"); + expect(type).toContain("id: number"); + }); + + test("prefers data_array 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", + status: { state: "SUCCEEDED" }, + result: { + data_array: [["from_data_array", "INT", null]], + attachment, + }, + }; + + const { type } = convertToQueryType(response, "SELECT 1", "test"); + expect(type).toContain("from_data_array: number"); + expect(type).not.toContain("from_arrow"); + }); + + test("logs a warning and yields the unknown-result fallback on malformed attachment", () => { + mockLoggerWarn.mockClear(); + const response: DatabricksStatementExecutionResponse = { + statement_id: "test-bad", + status: { state: "SUCCEEDED" }, + result: { attachment: "not-valid-arrow-ipc" }, + }; + + const { hasResults, type } = convertToQueryType( + response, + "SELECT 1", + "test", + ); + + // No columns extracted → unknown-result type, hasResults false. + expect(hasResults).toBe(false); + expect(type).toContain("unknown"); + // None of DESCRIBE QUERY's metadata column names should leak in as + // user-facing type fields — that would mean the parser swallowed + // the failure and produced bogus columns instead. + expect(type).not.toContain("col_name"); + expect(type).not.toContain("data_type"); + + // The warning must fire so a regression that silently produces empty + // types (no telemetry signal) fails this test. + expect(mockLoggerWarn).toHaveBeenCalledWith( + expect.stringContaining("Failed to decode Arrow IPC attachment"), + expect.any(String), + ); + }); + }); }); describe("inferParameterTypes", () => { diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index 9829729a..d036e0db 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -4,4 +4,5 @@ export * from "./execute"; export * from "./genie"; export * from "./plugin"; export * from "./sql"; +export * from "./sse/analytics"; export * from "./tunnel"; diff --git a/packages/shared/src/sse/analytics.test.ts b/packages/shared/src/sse/analytics.test.ts new file mode 100644 index 00000000..cdc8a70e --- /dev/null +++ b/packages/shared/src/sse/analytics.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, test } from "vitest"; +import { AnalyticsSseMessage, makeResultMessage } from "./analytics"; + +describe("AnalyticsSseMessage schema", () => { + test("accepts a result message with rows", () => { + const parsed = AnalyticsSseMessage.parse({ + type: "result", + data: [{ id: 1, name: "alice" }], + }); + expect(parsed.type).toBe("result"); + }); + + test("accepts a result message with no data (empty result)", () => { + expect(() => AnalyticsSseMessage.parse({ type: "result" })).not.toThrow(); + }); + + test("rejects a retired arrow message — ARROW_STREAM no longer uses SSE", () => { + // Arrow bytes now stream on the query response body, not the SSE channel, + // so there is no `arrow` message type to parse. + expect(() => + AnalyticsSseMessage.parse({ type: "arrow", statement_id: "stmt-1" }), + ).toThrow(); + }); + + test("rejects an unknown type", () => { + expect(() => + AnalyticsSseMessage.parse({ type: "unknown_kind", foo: "bar" }), + ).toThrow(); + }); + + test("safeParse returns success: false for malformed payloads", () => { + const r = AnalyticsSseMessage.safeParse({ type: "arrow" }); + expect(r.success).toBe(false); + }); +}); + +describe("typed builder", () => { + test("makeResultMessage roundtrips through the schema", () => { + const msg = makeResultMessage([{ id: 1 }], { statement_id: "s-1" }); + expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); + }); +}); diff --git a/packages/shared/src/sse/analytics.ts b/packages/shared/src/sse/analytics.ts new file mode 100644 index 00000000..41022672 --- /dev/null +++ b/packages/shared/src/sse/analytics.ts @@ -0,0 +1,78 @@ +import { z } from "zod"; + +/** + * Wire protocol for analytics SSE messages emitted by `/api/analytics/query`. + * + * The SSE channel carries only the JSON_ARRAY path (warehouse-readiness + * events + a `result` message of rows). ARROW_STREAM does NOT use SSE — + * the server streams the raw Arrow IPC bytes back on the query response body + * (`_handleArrowStreamQuery`) and the client reads them directly + * (`fetchArrowDirect`), so there is no `arrow` message type here. + * + * These schemas are the single source of truth for the JSON contract between + * the server (`AnalyticsPlugin._handleQueryRoute`) and the client + * (`useAnalyticsQuery`). Both sides validate with the same schema: + * + * - Server uses the typed builder (`makeResultMessage`) to construct messages + * with compile-time guarantees that all required fields are present. + * - Client calls `AnalyticsSseMessage.safeParse(JSON.parse(event.data))` to + * fail loudly on a malformed payload instead of silently treating an + * undefined field as data. + */ + +/** Successful row-shaped result (JSON_ARRAY format, or empty results). */ +export const AnalyticsResultMessage = z.object({ + type: z.literal("result"), + // `data` is intentionally `z.array(z.unknown())` rather than a deep + // row schema. Validating every row's keys for shape costs O(rows × cols) + // CPU and main-thread blocking time on the *client* (the schema is + // also reused for `safeParse` in `useAnalyticsQuery`); for large JSON + // results that pushes hundreds of ms to seconds of TBT into the + // render pipeline. The server constructs `data` via the typed + // `makeResultMessage` builder, so the per-row shape is enforced at + // the source, not at validation time. The TS-level interface below + // narrows `data` to `Record[]` for callers. + data: z.array(z.unknown()).optional(), + // Status is opaque metadata forwarded from the warehouse — keep it as + // `unknown` so we don't bake the SDK's detailed shape into the contract. + status: z.unknown().optional(), + statement_id: z.string().optional(), +}); + +/** + * TS-level shape of a successful row-shaped result message. + * + * **Kept in sync by hand** with `AnalyticsResultMessage` above. The Zod + * schema is intentionally loose (`z.array(z.unknown())`) to keep client + * validation cheap; this interface narrows `data` to + * `Record[]` so consumers don't have to cast at every + * call site. If you add a field to the Zod schema, add it here too. + */ +export interface AnalyticsResultMessage { + type: "result"; + data?: Record[]; + status?: unknown; + statement_id?: string; +} + +/** + * Every message the analytics SSE stream may emit. Currently only the + * row-shaped `result` message (JSON_ARRAY path); `warehouse_status` and + * `error` events are handled off-schema by the client. ARROW_STREAM never + * uses SSE. + */ +export const AnalyticsSseMessage = AnalyticsResultMessage; +export type AnalyticsSseMessage = z.infer; + +// ──────────────────────────────────────────────────────────────────────────── +// Typed builder — call from the server route handler. The compiler enforces +// that every required field is supplied, and the return type narrows so +// downstream code (executeStream / SSE writer) keeps full type information. +// ──────────────────────────────────────────────────────────────────────────── + +export function makeResultMessage( + data: Record[] | undefined, + extras: { status?: unknown; statement_id?: string } = {}, +): AnalyticsResultMessage { + return { type: "result", data, ...extras }; +} diff --git a/tools/test-helpers.ts b/tools/test-helpers.ts index f15630d7..9161a0f6 100644 --- a/tools/test-helpers.ts +++ b/tools/test-helpers.ts @@ -140,9 +140,16 @@ export function createMockResponse() { const eventListeners: Record void>> = {}; const res = { + // Flips to true once headers/body have gone out — mirrors Express so + // streaming handlers can branch between a JSON error (pre-headers) and + // aborting the socket (mid-stream). + headersSent: false, status: vi.fn().mockReturnThis(), json: vi.fn().mockReturnThis(), - send: vi.fn().mockReturnThis(), + send: vi.fn(function (this: any) { + this.headersSent = true; + return this; + }), sendStatus: vi.fn().mockReturnThis(), end: vi.fn(function (this: any) { this.writableEnded = true; @@ -154,9 +161,16 @@ export function createMockResponse() { } return this; }), - write: vi.fn().mockReturnThis(), - setHeader: vi.fn().mockReturnThis(), + write: vi.fn(function (this: any) { + this.headersSent = true; + return this; + }), + setHeader: vi.fn(function (this: any) { + this.headersSent = true; + return this; + }), flushHeaders: vi.fn().mockReturnThis(), + destroy: vi.fn().mockReturnThis(), on: vi.fn(function ( this: any, event: string, @@ -168,6 +182,18 @@ export function createMockResponse() { eventListeners[event].push(handler); return this; }), + off: vi.fn(function ( + this: any, + event: string, + handler: (...args: any[]) => void, + ) { + if (eventListeners[event]) { + eventListeners[event] = eventListeners[event].filter( + (h) => h !== handler, + ); + } + return this; + }), writableEnded: false, }; return res;