From 840a67b058d6cd53d519509a8d511d14bed8bf9b Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Wed, 29 Apr 2026 15:10:30 +0000 Subject: [PATCH 01/25] refactor: rename AnalyticsFormat values to API enum names MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Renames the client-side analytics format model from "JSON"/"ARROW" to "JSON_ARRAY"/"ARROW_STREAM" to match the Statement Execution API enum verbatim — no more local-name to API-name translation. Pure mechanical rename. No behavior change. Internal type values only; the lowercase user-facing values passed to useChartData ("json", "arrow", "auto") are unchanged. Carved out of #256 (#327 is layer 1, this is layer 2). The actual inline-Arrow-IPC + warehouse-fallback fix sits on top of this in layer 3. Note: this is a breaking change for any direct consumer of useAnalyticsQuery passing explicit format: "JSON" or "ARROW" — they will need to update to "JSON_ARRAY" / "ARROW_STREAM". Consumers using useChartData (lowercase "json"/"arrow"/"auto") are unaffected. Co-authored-by: Isaac --- .../hooks/__tests__/use-chart-data.test.ts | 16 +++++++-------- packages/appkit-ui/src/react/hooks/types.ts | 8 ++++---- .../src/react/hooks/use-analytics-query.ts | 10 +++++----- .../src/react/hooks/use-chart-data.ts | 20 +++++++++---------- .../appkit/src/plugins/analytics/analytics.ts | 4 ++-- .../appkit/src/plugins/analytics/types.ts | 2 +- 6 files changed, 30 insertions(+), 30 deletions(-) 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 3d5e96f1..a4d99a91 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 @@ -89,7 +89,7 @@ describe("useChartData", () => { expect(mockUseAnalyticsQuery).toHaveBeenCalledWith( "test", undefined, - expect.objectContaining({ format: "JSON" }), + expect.objectContaining({ format: "JSON_ARRAY" }), ); }); @@ -110,7 +110,7 @@ describe("useChartData", () => { expect(mockUseAnalyticsQuery).toHaveBeenCalledWith( "test", undefined, - expect.objectContaining({ format: "ARROW" }), + expect.objectContaining({ format: "ARROW_STREAM" }), ); }); @@ -132,7 +132,7 @@ describe("useChartData", () => { expect(mockUseAnalyticsQuery).toHaveBeenCalledWith( "test", { limit: 1000 }, - expect.objectContaining({ format: "ARROW" }), + expect.objectContaining({ format: "ARROW_STREAM" }), ); }); @@ -157,7 +157,7 @@ describe("useChartData", () => { expect(mockUseAnalyticsQuery).toHaveBeenCalledWith( "test", expect.objectContaining({ startDate: "2025-01-01" }), - expect.objectContaining({ format: "ARROW" }), + expect.objectContaining({ format: "ARROW_STREAM" }), ); }); @@ -179,7 +179,7 @@ describe("useChartData", () => { expect(mockUseAnalyticsQuery).toHaveBeenCalledWith( "test", expect.anything(), - expect.objectContaining({ format: "JSON" }), + expect.objectContaining({ format: "JSON_ARRAY" }), ); }); @@ -201,7 +201,7 @@ describe("useChartData", () => { expect(mockUseAnalyticsQuery).toHaveBeenCalledWith( "test", expect.anything(), - expect.objectContaining({ format: "ARROW" }), + expect.objectContaining({ format: "ARROW_STREAM" }), ); }); @@ -223,7 +223,7 @@ describe("useChartData", () => { expect(mockUseAnalyticsQuery).toHaveBeenCalledWith( "test", { limit: 100 }, - expect.objectContaining({ format: "JSON" }), + expect.objectContaining({ format: "JSON_ARRAY" }), ); }); @@ -243,7 +243,7 @@ describe("useChartData", () => { expect(mockUseAnalyticsQuery).toHaveBeenCalledWith( "test", undefined, - expect.objectContaining({ format: "JSON" }), + expect.objectContaining({ format: "JSON_ARRAY" }), ); }); }); diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index 03e943e2..e5e178c9 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -5,7 +5,7 @@ import type { Table } from "apache-arrow"; // ============================================================================ /** Supported response formats for analytics queries */ -export type AnalyticsFormat = "JSON" | "ARROW"; +export type AnalyticsFormat = "JSON_ARRAY" | "ARROW_STREAM"; /** * Typed Arrow Table - preserves row type information for type inference. @@ -32,8 +32,8 @@ export interface TypedArrowTable< // ============================================================================ /** Options for configuring an analytics SSE query */ -export interface UseAnalyticsQueryOptions { - /** Response format - "JSON" returns typed arrays, "ARROW" returns TypedArrowTable */ +export interface UseAnalyticsQueryOptions { + /** Response format - "JSON_ARRAY" returns typed arrays, "ARROW_STREAM" returns TypedArrowTable */ format?: F; /** Maximum size of serialized parameters in bytes */ @@ -120,7 +120,7 @@ export type InferResultByFormat< T, K, F extends AnalyticsFormat, -> = F extends "ARROW" ? TypedArrowTable> : InferResult; +> = F extends "ARROW_STREAM" ? TypedArrowTable> : InferResult; /** * Infers parameters type from QueryRegistry[K]["parameters"] 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 24e03ea3..0bd0b2f0 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -27,8 +27,8 @@ function getArrowStreamUrl(id: string) { * Integration hook between client and analytics plugin. * * The return type is automatically inferred based on the format: - * - `format: "JSON"` (default): Returns typed array from QueryRegistry - * - `format: "ARROW"`: Returns TypedArrowTable with row type preserved + * - `format: "JSON_ARRAY"` (default): Returns typed array from QueryRegistry + * - `format: "ARROW_STREAM"`: Returns TypedArrowTable with row type preserved * * Note: User context execution is determined by query file naming: * - `queryKey.obo.sql`: Executes as user (OBO = on-behalf-of / user delegation) @@ -47,20 +47,20 @@ function getArrowStreamUrl(id: string) { * * @example Arrow format * ```typescript - * const { data } = useAnalyticsQuery("spend_data", params, { format: "ARROW" }); + * const { data } = useAnalyticsQuery("spend_data", params, { format: "ARROW_STREAM" }); * // data: TypedArrowTable<{ group_key: string; cost_usd: number; ... }> | null * ``` */ export function useAnalyticsQuery< T = unknown, K extends QueryKey = QueryKey, - F extends AnalyticsFormat = "JSON", + F extends AnalyticsFormat = "JSON_ARRAY", >( queryKey: K, parameters?: InferParams | null, options: UseAnalyticsQueryOptions = {} as UseAnalyticsQueryOptions, ): UseAnalyticsQueryResult> { - const format = options?.format ?? "JSON"; + const format = options?.format ?? "JSON_ARRAY"; const maxParametersSize = options?.maxParametersSize ?? 100 * 1024; const autoStart = options?.autoStart ?? true; 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 d8d0bd38..a90481a2 100644 --- a/packages/appkit-ui/src/react/hooks/use-chart-data.ts +++ b/packages/appkit-ui/src/react/hooks/use-chart-data.ts @@ -50,32 +50,32 @@ export interface UseChartDataResult { function resolveFormat( format: DataFormat, parameters?: Record, -): "JSON" | "ARROW" { +): "JSON_ARRAY" | "ARROW_STREAM" { // Explicit format selection - if (format === "json") return "JSON"; - if (format === "arrow") return "ARROW"; + if (format === "json") return "JSON_ARRAY"; + if (format === "arrow") return "ARROW_STREAM"; // Auto-selection heuristics if (format === "auto") { // Check for explicit hint in parameters - if (parameters?._preferArrow === true) return "ARROW"; - if (parameters?._preferJson === true) return "JSON"; + if (parameters?._preferArrow === true) return "ARROW_STREAM"; + if (parameters?._preferJson === true) return "JSON_ARRAY"; // Check limit parameter as data size hint const limit = parameters?.limit; if (typeof limit === "number" && limit > ARROW_THRESHOLD) { - return "ARROW"; + return "ARROW_STREAM"; } // Check for date range queries (often large) if (parameters?.startDate && parameters?.endDate) { - return "ARROW"; + return "ARROW_STREAM"; } - return "JSON"; + return "JSON_ARRAY"; } - return "JSON"; + return "JSON_ARRAY"; } // ============================================================================ @@ -110,7 +110,7 @@ export function useChartData(options: UseChartDataOptions): UseChartDataResult { [format, parameters], ); - const isArrowFormat = resolvedFormat === "ARROW"; + const isArrowFormat = resolvedFormat === "ARROW_STREAM"; // Fetch data using the analytics query hook const { diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index fdcb16b4..564b4cfb 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -128,7 +128,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { res: express.Response, ): Promise { const { query_key } = req.params; - const { parameters, format = "JSON" } = req.body as IAnalyticsQueryRequest; + const { parameters, format = "JSON_ARRAY" } = req.body as IAnalyticsQueryRequest; // Request-scoped logging with WideEvent tracking logger.debug(req, "Executing query: %s (format=%s)", query_key, format); @@ -164,7 +164,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { const executorKey = isAsUser ? this.resolveUserId(req) : "global"; const queryParameters = - format === "ARROW" + format === "ARROW_STREAM" ? { formatParameters: { disposition: "EXTERNAL_LINKS", diff --git a/packages/appkit/src/plugins/analytics/types.ts b/packages/appkit/src/plugins/analytics/types.ts index c58b6ecf..c0e72fdb 100644 --- a/packages/appkit/src/plugins/analytics/types.ts +++ b/packages/appkit/src/plugins/analytics/types.ts @@ -4,7 +4,7 @@ export interface IAnalyticsConfig extends BasePluginConfig { timeout?: number; } -export type AnalyticsFormat = "JSON" | "ARROW"; +export type AnalyticsFormat = "JSON_ARRAY" | "ARROW_STREAM"; export interface IAnalyticsQueryRequest { parameters?: Record; format?: AnalyticsFormat; From 09392bbc23d7adaaa3683cb4ce877d7149980c61 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Mon, 11 May 2026 16:11:14 +0000 Subject: [PATCH 02/25] refactor(analytics): accept legacy "JSON"/"ARROW" format aliases Widen AnalyticsFormat to also include the pre-rename "JSON" and "ARROW" spellings, both marked @deprecated with a JSDoc note describing the removal condition (no consumer on appkit/appkit-ui < 0.33.0). Add a normalizeAnalyticsFormat helper and call it at the analytics route handler entry point so all downstream code (cache key, format branching, formatParameters) continues to operate on the canonical "JSON_ARRAY" | "ARROW_STREAM" values. InferResultByFormat is widened to also match "ARROW" so callers passing the legacy spelling still get TypedArrowTable<...> inferred. This lifts the breaking-change carve-out from the rename, so callers of useAnalyticsQuery({ format: "JSON" | "ARROW" }) keep working with only an IDE deprecation hint. Signed-off-by: James Broadhead --- packages/appkit-ui/src/react/hooks/types.ts | 20 +++++++++-- .../appkit/src/plugins/analytics/analytics.ts | 5 ++- .../appkit/src/plugins/analytics/types.ts | 33 ++++++++++++++++++- 3 files changed, 53 insertions(+), 5 deletions(-) diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index e5e178c9..f775113f 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -4,8 +4,20 @@ import type { Table } from "apache-arrow"; // Data Format Types // ============================================================================ -/** Supported response formats for analytics queries */ -export type AnalyticsFormat = "JSON_ARRAY" | "ARROW_STREAM"; +/** + * Supported response formats for analytics queries. + * + * "JSON" and "ARROW" are legacy aliases kept for backwards compatibility + * with appkit/appkit-ui < 0.33.0 — safe to remove once no consumer is on + * a pre-0.33.0 version. + */ +export type AnalyticsFormat = + | "JSON_ARRAY" + | "ARROW_STREAM" + /** @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"; /** * Typed Arrow Table - preserves row type information for type inference. @@ -120,7 +132,9 @@ export type InferResultByFormat< T, K, F extends AnalyticsFormat, -> = F extends "ARROW_STREAM" ? TypedArrowTable> : InferResult; +> = F extends "ARROW_STREAM" | "ARROW" + ? TypedArrowTable> + : InferResult; /** * Infers parameters type from QueryRegistry[K]["parameters"] diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 564b4cfb..bfee250f 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -24,6 +24,7 @@ import type { PluginManifest } from "../../registry"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; +import { normalizeAnalyticsFormat } from "./types"; import type { AnalyticsQueryResponse, IAnalyticsConfig, @@ -128,7 +129,9 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { res: express.Response, ): Promise { const { query_key } = req.params; - const { parameters, format = "JSON_ARRAY" } = req.body as IAnalyticsQueryRequest; + const { parameters, format: rawFormat = "JSON_ARRAY" } = + req.body as IAnalyticsQueryRequest; + const format = normalizeAnalyticsFormat(rawFormat); // Request-scoped logging with WideEvent tracking logger.debug(req, "Executing query: %s (format=%s)", query_key, format); diff --git a/packages/appkit/src/plugins/analytics/types.ts b/packages/appkit/src/plugins/analytics/types.ts index c0e72fdb..2a84e32e 100644 --- a/packages/appkit/src/plugins/analytics/types.ts +++ b/packages/appkit/src/plugins/analytics/types.ts @@ -4,7 +4,38 @@ export interface IAnalyticsConfig extends BasePluginConfig { timeout?: number; } -export type AnalyticsFormat = "JSON_ARRAY" | "ARROW_STREAM"; +/** + * Supported response formats for analytics queries. + * + * "JSON" and "ARROW" are legacy aliases kept for backwards compatibility + * with appkit/appkit-ui < 0.33.0 — safe to remove once no consumer is on + * a pre-0.33.0 version. The route handler normalizes them to their + * canonical equivalents before any downstream code reads the value. + */ +export type AnalyticsFormat = + | "JSON_ARRAY" + | "ARROW_STREAM" + /** @deprecated Use "JSON_ARRAY". Safe to remove once no consumer is on appkit < 0.33.0. */ + | "JSON" + /** @deprecated Use "ARROW_STREAM". Safe to remove once no consumer is on appkit < 0.33.0. */ + | "ARROW"; + +/** Canonical (post-normalization) analytics format values. */ +export type CanonicalAnalyticsFormat = "JSON_ARRAY" | "ARROW_STREAM"; + +/** + * Map a (possibly legacy) AnalyticsFormat to its canonical form. + * Legacy values come from appkit/appkit-ui < 0.33.0 and can be removed + * along with the deprecated aliases once no such consumer remains. + */ +export function normalizeAnalyticsFormat( + f: AnalyticsFormat, +): CanonicalAnalyticsFormat { + if (f === "JSON") return "JSON_ARRAY"; + if (f === "ARROW") return "ARROW_STREAM"; + return f; +} + export interface IAnalyticsQueryRequest { parameters?: Record; format?: AnalyticsFormat; From 08c5486c886b6558f33b79ddd9e2a69cac9d541d Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Wed, 29 Apr 2026 15:13:38 +0000 Subject: [PATCH 03/25] feat: decode inline Arrow IPC + warehouse-compat fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Serverless warehouses return ARROW_STREAM + INLINE results as base64 Arrow IPC in result.attachment rather than result.data_array. The previous code path discarded inline data for any ARROW_STREAM response (designed for EXTERNAL_LINKS), so these warehouses silently returned empty results. This commit makes the analytics plugin work across classic and serverless warehouses by handling both dispositions for ARROW_STREAM, decoding inline Arrow IPC attachments server-side, and falling back to JSON_ARRAY when a warehouse rejects ARROW_STREAM + INLINE. Changes - Inline Arrow IPC decoding (new arrow-schema.ts) via apache-arrow's tableFromIPC, producing the same row-object shape as JSON_ARRAY regardless of warehouse backend. apache-arrow@21.1.0 added as a server dep. - Format fallback: ARROW_STREAM + INLINE requests automatically fall back to JSON_ARRAY if a classic warehouse rejects them. Explicit format requests are respected without fallback. - Zod-validated SSE wire protocol for /api/analytics/query (shared schema between server and client; malformed payloads surface a clear error instead of silent undefined). - Default remains JSON_ARRAY for compatibility. Stack: layer 3 of 3 carved from #256. - #327 — coverage backfill (layer 1) - #328 — AnalyticsFormat rename to API enum names (layer 2) - (this PR) — the actual fix Fixes #242 Co-authored-by: Isaac --- docs/docs/api/appkit/Class.ExecutionError.md | 30 +- packages/appkit-ui/src/js/sse/connect-sse.ts | 5 +- .../src/react/charts/__tests__/types.test.ts | 2 +- packages/appkit-ui/src/react/charts/types.ts | 6 +- .../__tests__/use-analytics-query.test.ts | 143 +++++ .../hooks/__tests__/use-chart-data.test.ts | 24 +- packages/appkit-ui/src/react/hooks/types.ts | 6 +- .../src/react/hooks/use-analytics-query.ts | 98 +++- .../src/react/hooks/use-chart-data.ts | 10 +- packages/appkit/package.json | 4 +- .../connectors/sql-warehouse/arrow-schema.ts | 441 +++++++++++++++ .../src/connectors/sql-warehouse/client.ts | 114 +++- .../sql-warehouse/tests/arrow-schema.test.ts | 514 ++++++++++++++++++ .../sql-warehouse/tests/client.test.ts | 382 +++++++++++++ packages/appkit/src/errors/execution.ts | 32 +- .../appkit/src/plugins/analytics/analytics.ts | 204 +++++-- .../plugins/analytics/tests/analytics.test.ts | 449 +++++++++++++++ packages/appkit/src/stream/defaults.ts | 8 +- .../src/type-generator/query-registry.ts | 97 +++- .../tests/query-registry.test.ts | 137 ++++- packages/appkit/src/type-generator/types.ts | 2 + packages/shared/package.json | 3 +- packages/shared/src/index.ts | 1 + packages/shared/src/sse/analytics.test.ts | 87 +++ packages/shared/src/sse/analytics.ts | 95 ++++ pnpm-lock.yaml | 53 +- 26 files changed, 2825 insertions(+), 122 deletions(-) create mode 100644 packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts create mode 100644 packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts create mode 100644 packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts create mode 100644 packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts create mode 100644 packages/shared/src/sse/analytics.test.ts create mode 100644 packages/shared/src/sse/analytics.ts diff --git a/docs/docs/api/appkit/Class.ExecutionError.md b/docs/docs/api/appkit/Class.ExecutionError.md index 75886c4d..eacfdc23 100644 --- a/docs/docs/api/appkit/Class.ExecutionError.md +++ b/docs/docs/api/appkit/Class.ExecutionError.md @@ -22,6 +22,7 @@ throw new ExecutionError("Statement was canceled"); new ExecutionError(message: string, options?: { cause?: Error; context?: Record; + errorCode?: string; }): ExecutionError; ``` @@ -30,15 +31,16 @@ new ExecutionError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; \} | +| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; `errorCode?`: `string`; \} | | `options.cause?` | `Error` | | `options.context?` | `Record`\<`string`, `unknown`\> | +| `options.errorCode?` | `string` | #### Returns `ExecutionError` -#### Inherited from +#### Overrides [`AppKitError`](Class.AppKitError.md).[`constructor`](Class.AppKitError.md#constructor) @@ -86,6 +88,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 @@ -202,16 +217,17 @@ Create an execution error for closed/expired results ### statementFailed() ```ts -static statementFailed(errorMessage?: string): ExecutionError; +static statementFailed(errorMessage?: string, errorCode?: 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. | +| `errorCode?` | `string` | Structured code (e.g. "INVALID_PARAMETER_VALUE") to preserve through wrapping. Optional. | #### Returns diff --git a/packages/appkit-ui/src/js/sse/connect-sse.ts b/packages/appkit-ui/src/js/sse/connect-sse.ts index c4fd4500..089ddf57 100644 --- a/packages/appkit-ui/src/js/sse/connect-sse.ts +++ b/packages/appkit-ui/src/js/sse/connect-sse.ts @@ -18,7 +18,10 @@ export async function connectSSE( lastEventId: initialLastEventId = null, retryDelay = 2000, maxRetries = 3, - maxBufferSize = 1024 * 1024, // 1MB + // 8 MiB — sized to receive inline Arrow IPC attachments from + // ARROW_STREAM analytics responses; matches the server's stream + // `maxEventSize`. Most events are well under 1 MiB in practice. + maxBufferSize = 8 * 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 65804a74..ec8a15dc 100644 --- a/packages/appkit-ui/src/react/charts/types.ts +++ b/packages/appkit-ui/src/react/charts/types.ts @@ -5,7 +5,7 @@ import type { Table } from "apache-arrow"; // ============================================================================ /** Supported data formats for analytics queries */ -export type DataFormat = "json" | "arrow" | "auto"; +export type DataFormat = "json_array" | "arrow_stream" | "auto"; /** Chart orientation */ export type Orientation = "vertical" | "horizontal"; @@ -77,8 +77,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 new file mode 100644 index 00000000..65de7d10 --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts @@ -0,0 +1,143 @@ +import { renderHook, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, test, vi } from "vitest"; + +// Capture the onMessage handler so tests can drive SSE messages directly. +let lastConnectArgs: any = null; +const mockProcessArrowBuffer = vi.fn(); +const mockFetchArrow = vi.fn(); + +vi.mock("@/js", () => ({ + connectSSE: vi.fn((args: any) => { + lastConnectArgs = args; + return () => {}; + }), + ArrowClient: { + fetchArrow: (...args: unknown[]) => mockFetchArrow(...args), + processArrowBuffer: (...args: unknown[]) => mockProcessArrowBuffer(...args), + }, +})); + +// useQueryHMR is a no-op shim for tests; mock to avoid HMR side effects. +vi.mock("../use-query-hmr", () => ({ + useQueryHMR: vi.fn(), +})); + +import { useAnalyticsQuery } from "../use-analytics-query"; + +describe("useAnalyticsQuery", () => { + beforeEach(() => { + vi.clearAllMocks(); + lastConnectArgs = null; + }); + + test("decodes arrow_inline base64 attachment via ArrowClient.processArrowBuffer", async () => { + const fakeTable = { numRows: 1, schema: { fields: [] } }; + mockProcessArrowBuffer.mockResolvedValueOnce(fakeTable); + + // 'AQID' decodes to bytes [1, 2, 3]. + const base64 = "AQID"; + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + // Drive the SSE onMessage handler with an arrow_inline payload. + await lastConnectArgs.onMessage({ + data: JSON.stringify({ type: "arrow_inline", attachment: base64 }), + }); + + await waitFor(() => { + expect(result.current.data).toBe(fakeTable); + }); + + expect(mockProcessArrowBuffer).toHaveBeenCalledTimes(1); + const passedBuffer = mockProcessArrowBuffer.mock.calls[0][0] as Uint8Array; + expect(passedBuffer).toBeInstanceOf(Uint8Array); + expect(Array.from(passedBuffer)).toEqual([1, 2, 3]); + // Inline path must NOT trigger a network fetch. + expect(mockFetchArrow).not.toHaveBeenCalled(); + }); + + test("surfaces an error when arrow_inline decode fails", async () => { + mockProcessArrowBuffer.mockRejectedValueOnce(new Error("bad ipc")); + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await lastConnectArgs.onMessage({ + data: JSON.stringify({ type: "arrow_inline", attachment: "AQID" }), + }); + + await waitFor(() => { + expect(result.current.error).toBe( + "Unable to load data, please try again", + ); + }); + expect(result.current.loading).toBe(false); + }); + + test("rejects arrow_inline with missing/empty/non-string attachment without crashing atob", async () => { + const cases: Array = [undefined, null, "", 123, { foo: "bar" }]; + + for (const attachment of cases) { + mockProcessArrowBuffer.mockClear(); + const { result, unmount } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await lastConnectArgs.onMessage({ + data: JSON.stringify({ type: "arrow_inline", attachment }), + }); + + await waitFor(() => { + expect(result.current.error).toBe( + "Unable to load data, please try again", + ); + }); + // Critically: must NOT call processArrowBuffer (or atob) on the bad input. + expect(mockProcessArrowBuffer).not.toHaveBeenCalled(); + + unmount(); + } + }); + + test("rejects oversized arrow_inline attachment without allocating a huge buffer", async () => { + // Base64 string that would decode to ~9 MiB (>8 MiB cap). The hook + // should reject before calling decodeBase64 / processArrowBuffer. + const oversized = "A".repeat(13 * 1024 * 1024); + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await lastConnectArgs.onMessage({ + data: JSON.stringify({ type: "arrow_inline", attachment: oversized }), + }); + + await waitFor(() => { + expect(result.current.error).toBe( + "Unable to load data, please try again", + ); + }); + expect(mockProcessArrowBuffer).not.toHaveBeenCalled(); + }); + + 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(); + }); +}); 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 f775113f..314638ec 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -44,8 +44,10 @@ export interface TypedArrowTable< // ============================================================================ /** Options for configuring an analytics SSE query */ -export interface UseAnalyticsQueryOptions { - /** Response format - "JSON_ARRAY" returns typed arrays, "ARROW_STREAM" returns TypedArrowTable */ +export interface UseAnalyticsQueryOptions< + F extends AnalyticsFormat = "JSON_ARRAY", +> { + /** 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 */ 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 0bd0b2f0..5d18a2ee 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -1,4 +1,5 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { AnalyticsSseMessage } from "shared"; import { ArrowClient, connectSSE } from "@/js"; import type { AnalyticsFormat, @@ -22,6 +23,29 @@ function getArrowStreamUrl(id: string) { return `/api/analytics/arrow-result/${id}`; } +/** + * Client-side defensive cap on inline Arrow IPC attachments (8 MiB decoded). + * Mirrors the server's MAX_INLINE_ATTACHMENT_BYTES so a misconfigured proxy + * (or a future server bug) can't push us into allocating an unbounded + * Uint8Array and hanging the browser. + * + * REMOVE THIS GUARD if PR #320 (stash + serve via /arrow-result) lands — + * that proposal eliminates the arrow_inline SSE path entirely, so bulk + * bytes flow over HTTP where the browser handles backpressure natively + * and Content-Length is exposed up-front. + */ +const MAX_INLINE_ATTACHMENT_BYTES = 8 * 1024 * 1024; + +/** Decode a base64 string into a Uint8Array suitable for Arrow IPC parsing. */ +function decodeBase64(b64: string): Uint8Array { + const binary = atob(b64); + const bytes = new Uint8Array(binary.length); + for (let i = 0; i < binary.length; i++) { + bytes[i] = binary.charCodeAt(i); + } + return bytes; +} + /** * Subscribe to an analytics query over SSE and returns its latest result. * Integration hook between client and analytics plugin. @@ -39,13 +63,13 @@ function getArrowStreamUrl(id: string) { * @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 @@ -120,20 +144,28 @@ export function useAnalyticsQuery< signal: abortController.signal, onMessage: async (message) => { try { - const parsed = JSON.parse(message.data); + const rawParsed = JSON.parse(message.data); + + // The error/code branch below predates the SSE wire schema and + // can fire for messages that don't match any AnalyticsSseMessage + // variant (e.g. server-side error events from executeStream). + // Try schema validation first; if it fails, fall through to the + // generic error/code handling below. + const validated = AnalyticsSseMessage.safeParse(rawParsed); + const msg = validated.success ? validated.data : null; // success - JSON format - if (parsed.type === "result") { + if (msg?.type === "result") { setLoading(false); - setData(parsed.data as ResultType); + setData(msg.data as ResultType); return; } - // success - Arrow format - if (parsed.type === "arrow") { + // success - Arrow format (external links: fetch from server) + if (msg?.type === "arrow") { try { const arrowData = await ArrowClient.fetchArrow( - getArrowStreamUrl(parsed.statement_id), + getArrowStreamUrl(msg.statement_id), ); const table = await ArrowClient.processArrowBuffer(arrowData); setLoading(false); @@ -151,6 +183,44 @@ export function useAnalyticsQuery< } } + // success - Arrow format (inline: decode base64 IPC payload locally) + if (msg?.type === "arrow_inline") { + // Schema already enforced non-empty string; just check size. + // base64 length L decodes to ~L*3/4 bytes; reject before + // allocating a multi-MiB Uint8Array. + const decodedSize = Math.ceil((msg.attachment.length * 3) / 4); + if (decodedSize > MAX_INLINE_ATTACHMENT_BYTES) { + console.error( + "[useAnalyticsQuery] arrow_inline attachment exceeds %d bytes (got %d)", + MAX_INLINE_ATTACHMENT_BYTES, + decodedSize, + ); + setLoading(false); + setError("Unable to load data, please try again"); + return; + } + try { + const buffer = decodeBase64(msg.attachment); + const table = await ArrowClient.processArrowBuffer(buffer); + setLoading(false); + setData(table as ResultType); + return; + } catch (error) { + console.error( + "[useAnalyticsQuery] Failed to decode inline Arrow data", + error, + ); + setLoading(false); + setError("Unable to load data, please try again"); + return; + } + } + + // The schema didn't match — fall through to error/code handling + // below for legacy error events or surface a malformed-payload + // error if no error fields are present. + const parsed = rawParsed; + // error if (parsed.type === "error" || parsed.error || parsed.code) { const errorMsg = @@ -166,6 +236,18 @@ export function useAnalyticsQuery< } return; } + + // The payload matched neither AnalyticsSseMessage nor an error + // event — surface a generic error rather than silently dropping it. + if (!validated.success) { + console.error( + "[useAnalyticsQuery] Malformed SSE payload", + validated.error.flatten(), + ); + setLoading(false); + setError("Unable to load data, please try again"); + return; + } } catch (error) { console.warn("[useAnalyticsQuery] Malformed message received", error); } 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 a90481a2..ec4b2d4e 100644 --- a/packages/appkit-ui/src/react/hooks/use-chart-data.ts +++ b/packages/appkit-ui/src/react/hooks/use-chart-data.ts @@ -17,8 +17,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" */ @@ -52,8 +52,8 @@ function resolveFormat( parameters?: Record, ): "JSON_ARRAY" | "ARROW_STREAM" { // Explicit format selection - if (format === "json") return "JSON_ARRAY"; - if (format === "arrow") return "ARROW_STREAM"; + if (format === "json_array") return "JSON_ARRAY"; + if (format === "arrow_stream") return "ARROW_STREAM"; // Auto-selection heuristics if (format === "auto") { @@ -97,7 +97,7 @@ function resolveFormat( * // Force Arrow format * const { data } = useChartData({ * queryKey: "big_query", - * format: "arrow" + * format: "arrow_stream" * }); * ``` */ diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 026065fa..5b9429c3 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -73,6 +73,7 @@ "@opentelemetry/sdk-trace-base": "2.6.0", "@opentelemetry/semantic-conventions": "1.38.0", "@types/semver": "7.7.1", + "apache-arrow": "21.1.0", "dotenv": "16.6.1", "express": "4.22.0", "get-port": "7.2.0", @@ -83,8 +84,7 @@ "semver": "7.7.3", "shared": "workspace:*", "vite": "npm:rolldown-vite@7.1.14", - "ws": "8.18.3", - "zod": "4.3.6" + "ws": "8.18.3" }, "devDependencies": { "@opentelemetry/context-async-hooks": "2.6.1", 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 d0a1c181..4ae43941 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -21,10 +21,24 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; +import { buildEmptyArrowIPCBase64 } from "./arrow-schema"; import { executeStatementDefaults } from "./defaults"; const logger = createLogger("connectors:sql-warehouse"); +/** + * Maximum size for inline Arrow IPC attachments (8 MiB decoded). + * Aligned with `streamDefaults.maxEventSize` so anything that would exceed + * the SSE event cap fails here with a clear error rather than a confusing + * "Buffer size exceeded" downstream. Larger results should use + * `disposition: "EXTERNAL_LINKS"`, which the analytics fallback handles. + * + * RAISE TO 25 MiB (Databricks API hard cap on INLINE) if PR #320 (stash + + * serve via /arrow-result) lands — that proposal moves bulk bytes off SSE + * onto HTTP, so the SSE event-size constraint no longer applies here. + */ +const MAX_INLINE_ATTACHMENT_BYTES = 8 * 1024 * 1024; + interface SQLWarehouseConfig { timeout?: number; telemetry?: TelemetryOptions; @@ -196,7 +210,10 @@ export class SQLWarehouseConnector { result = this._transformDataArray(response); 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": @@ -236,18 +253,22 @@ export class SQLWarehouseConnector { code: SpanStatusCode.ERROR, message: error instanceof Error ? error.message : String(error), }); - - logger.error( - "Statement execution failed: %s", - error instanceof Error ? error.message : String(error), - ); } if (error instanceof AppKitError) { throw error; } + // Preserve the SDK's structured ApiError.errorCode (e.g. + // "INVALID_PARAMETER_VALUE", "BAD_REQUEST") through the wrap so + // callers can branch on a stable identifier rather than + // substring-matching the message. + 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, ); } finally { // remove abort handler @@ -360,7 +381,10 @@ export class SQLWarehouseConnector { span.setStatus({ code: SpanStatusCode.OK }); return this._transformDataArray(response); 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": @@ -382,12 +406,16 @@ export class SQLWarehouseConnector { message: error instanceof Error ? error.message : String(error), }); - // error logging is handled by executeStatement's catch block (gated on isAborted) 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, ); } finally { span.end(); @@ -399,7 +427,40 @@ export class SQLWarehouseConnector { private _transformDataArray(response: sql.StatementResponse) { 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. + if (result?.external_links) { + return this.updateWithArrowStatus(response); + } + + // 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. + if (!result?.data_array && 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) { @@ -445,6 +506,41 @@ export class SQLWarehouseConnector { }; } + /** + * Validate (but do not decode) a base64 Arrow IPC attachment. + * Some serverless warehouses return inline results as Arrow IPC in + * `result.attachment`. We pass the base64 string through to the client, + * which decodes it into an Arrow Table via the existing ArrowClient + * infrastructure. This keeps the wire contract for ARROW_STREAM + * consistent (client always receives an Arrow Table) and avoids + * decode/re-encode work on the server. + */ + private _validateArrowAttachment( + response: sql.StatementResponse, + attachment: string, + ) { + // Cap the size to protect against unbounded inline payloads from + // misbehaving warehouses. 64 MiB is well above the typical inline limit + // (~25 MiB hard cap on the API) but 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)`, + ); + } + return response; + } + private updateWithArrowStatus(response: sql.StatementResponse): { result: { statement_id: string; status: sql.StatementStatus }; } { 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..c7f73c98 --- /dev/null +++ b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts @@ -0,0 +1,382 @@ +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 }); +} + +// 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", () => { + 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 = (connector as any)._transformDataArray(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", () => { + 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 = (connector as any)._transformDataArray(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", () => { + 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 = (connector as any)._transformDataArray(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", () => { + 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 = (connector as any)._transformDataArray(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); + }); + + test("preserves manifest and status alongside attachment", () => { + 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 = (connector as any)._transformDataArray(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", () => { + 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 = (connector as any)._transformDataArray(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", () => { + 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 = (connector as any)._transformDataArray(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("does NOT synthesize an attachment when schema is missing", () => { + 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 = (connector as any)._transformDataArray(response); + // Without a schema we cannot build a Table — pass through unchanged. + expect(transformed.result?.attachment).toBeUndefined(); + }); + + test("rejects oversized attachments to bound memory", () => { + const connector = createConnector(); + // 8 MiB decoded cap → ~12 MiB of base64 chars decodes to >8 MiB. + const oversized = "A".repeat(12 * 1024 * 1024); + const response = { + statement_id: "stmt-oversized", + status: { state: "SUCCEEDED" }, + manifest: { format: "ARROW_STREAM" }, + result: { attachment: oversized }, + } as unknown as sql.StatementResponse; + + expect(() => (connector as any)._transformDataArray(response)).toThrow( + /exceeds maximum size/, + ); + }); + }); + + describe("ARROW_STREAM with data_array (hypothetical inline variant)", () => { + test("transforms data_array like JSON_ARRAY path", () => { + 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 = (connector as any)._transformDataArray(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", () => { + const connector = createConnector(); + const response = { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { format: "JSON_ARRAY" }, + result: {}, + } as unknown as sql.StatementResponse; + + const result = (connector as any)._transformDataArray(response); + expect(result).toBe(response); + }); + + test("attachment takes priority over data_array when both present", () => { + 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 = (connector as any)._transformDataArray(response); + // Should pass attachment through (client decodes), not transform data_array + expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); + expect(result.result.data).toBeUndefined(); + }); + }); +}); diff --git a/packages/appkit/src/errors/execution.ts b/packages/appkit/src/errors/execution.ts index 42de7704..1e6d1f5f 100644 --- a/packages/appkit/src/errors/execution.ts +++ b/packages/appkit/src/errors/execution.ts @@ -16,13 +16,39 @@ 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; + }, + ) { + super(message, options); + this.errorCode = options?.errorCode; + } + + /** + * Create an execution error for statement failure. + * @param errorMessage Human-readable error from the warehouse / SDK. + * @param errorCode Structured code (e.g. "INVALID_PARAMETER_VALUE") to + * preserve through wrapping. Optional. + */ + static statementFailed( + errorMessage?: string, + errorCode?: string, + ): ExecutionError { const message = errorMessage ? `Statement failed: ${errorMessage}` : "Statement failed: Unknown error"; - return new ExecutionError(message); + return new ExecutionError(message, { errorCode }); } /** diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index bfee250f..b34d3dc1 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1,12 +1,16 @@ 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, + makeArrowInlineMessage, + makeArrowMessage, + makeResultMessage, + type PluginExecuteConfig, + type SQLTypeMarker, + type StreamExecutionSettings, + type ToolProvider, } from "shared"; import { z } from "zod"; import { SQLWarehouseConnector } from "../../connectors"; @@ -18,6 +22,7 @@ import { toolsFromRegistry, } from "../../core/agent/tools/define-tool"; import { assertReadOnlySql } from "../../core/agent/tools/sql-policy"; +import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; @@ -26,6 +31,7 @@ import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; import { normalizeAnalyticsFormat } from "./types"; import type { + AnalyticsFormat, AnalyticsQueryResponse, IAnalyticsConfig, IAnalyticsQueryRequest, @@ -131,6 +137,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 @@ -166,34 +185,33 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { 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", - } - : { - type: "result", - }; - const hashedQuery = this.queryProcessor.hashQuery(query); + // ARROW_STREAM may resolve to EXTERNAL_LINKS, which returns pre-signed URLs + // that typically expire ~15 minutes after issue. Cap the cache TTL well + // under that for ARROW_STREAM so we never hand out dead URLs from cache, + // while still benefiting from caching INLINE attachment responses (and + // EXTERNAL_LINKS responses inside their valid window). + const cacheTtl = + format === "ARROW_STREAM" + ? Math.min(queryDefaults.cache?.ttl ?? 600, 600) + : queryDefaults.cache?.ttl; + const cacheConfig = { + ...queryDefaults.cache, + ttl: cacheTtl, + cacheKey: [ + "analytics:query", + query_key, + JSON.stringify(parameters), + format, + hashedQuery, + executorKey, + ], + }; + const defaultConfig: PluginExecuteConfig = { ...queryDefaults, - cache: { - ...queryDefaults.cache, - cacheKey: [ - "analytics:query", - query_key, - JSON.stringify(parameters), - JSON.stringify(format), - hashedQuery, - executorKey, - ], - }, + cache: cacheConfig, }; const streamExecutionSettings: StreamExecutionSettings = { @@ -208,20 +226,94 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { parameters, ); - const result = await executor.query( + return this._executeWithFormatFallback( + executor, query, processedParams, - queryParameters.formatParameters, + format, signal, ); - - return { type: queryParameters.type, ...result }; }, streamExecutionSettings, executorKey, ); } + /** + * Execute a query with automatic disposition fallback for ARROW_STREAM. + * + * - JSON_ARRAY: always uses INLINE disposition, no fallback. + * - ARROW_STREAM: tries INLINE first, falls back to EXTERNAL_LINKS. + * This handles warehouses that only support one disposition. + */ + private async _executeWithFormatFallback( + executor: AnalyticsPlugin, + query: string, + processedParams: + | Record + | undefined, + requestedFormat: AnalyticsFormat, + signal?: AbortSignal, + ): Promise { + if (requestedFormat === "JSON_ARRAY") { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "JSON_ARRAY" }, + signal, + ); + return makeResultMessage(result?.data, { + status: result?.status, + statement_id: result?.statement_id, + }); + } + + // ARROW_STREAM: try INLINE first, fall back to EXTERNAL_LINKS. + try { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "ARROW_STREAM" }, + signal, + ); + // INLINE responses with an Arrow IPC attachment are forwarded as base64 + // for the client to decode into an Arrow Table. Anything else (rare: + // data_array under ARROW_STREAM, or an empty result) falls back to the + // generic "result" payload. + if (result?.attachment) { + return makeArrowInlineMessage(result.attachment); + } + return makeResultMessage(result?.data, { + status: result?.status, + statement_id: result?.statement_id, + }); + } catch (err: unknown) { + // If the request was aborted, do not retry — the signal is dead and + // a second statement would be billed but never read. + if (signal?.aborted) { + throw err; + } + + if (!_isInlineArrowUnsupported(err)) { + throw err; + } + + const msg = err instanceof Error ? err.message : String(err); + logger.warn( + "ARROW_STREAM INLINE rejected by warehouse, falling back to EXTERNAL_LINKS: %s", + msg, + ); + } + + const result = await executor.query( + query, + processedParams, + { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, + signal, + ); + return makeArrowMessage(result.statement_id, { status: result.status }); + } + /** * Execute a SQL query using the current execution context. * @@ -338,6 +430,48 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } } +/** + * Determine whether a warehouse error indicates that ARROW_STREAM + INLINE + * is unsupported, vs an unrelated SQL/permission error. + * + * Preferred path: read the structured `errorCode` we now propagate from the + * SDK's `ApiError.errorCode` and the warehouse's `status.error.error_code` + * through `ExecutionError`. This is stable across error-message wording + * changes. + * + * Substring backstop: if the upstream error didn't surface a code (legacy + * SDK builds, or errors thrown outside the connector's wrap path), fall + * back to requiring both INLINE and ARROW_STREAM keywords in the message + * plus a marker phrase. The pair-requirement avoids matching unrelated SQL + * errors that happen to mention one of the words (e.g. a column named + * `INLINE_USERS`). + */ +function _isInlineArrowUnsupported(err: unknown): boolean { + const structuredCode = + err instanceof ExecutionError ? err.errorCode : undefined; + if ( + structuredCode === "INVALID_PARAMETER_VALUE" || + structuredCode === "NOT_IMPLEMENTED" + ) { + // Structured code already tells us the warehouse rejected the request. + // Require keyword pairing to confirm it's the disposition/format combo + // (vs an INVALID_PARAMETER_VALUE for something else entirely). + const msg = err instanceof Error ? err.message : String(err); + return msg.includes("INLINE") && msg.includes("ARROW_STREAM"); + } + + // Backstop for errors without a structured code. + const msg = err instanceof Error ? err.message : String(err); + if (!msg.includes("INLINE") || !msg.includes("ARROW_STREAM")) { + return false; + } + return ( + msg.includes("not supported") || + msg.includes("INVALID_PARAMETER_VALUE") || + msg.includes("NOT_IMPLEMENTED") + ); +} + /** * @internal */ diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index eb06ea95..89bfac7d 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -581,6 +581,455 @@ describe("Analytics Plugin", () => { ); }); + 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(); + + (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: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + 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: {} }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(executeMock).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + disposition: "INLINE", + format: "JSON_ARRAY", + }), + expect.any(AbortSignal), + ); + }); + + 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 should fall back ARROW_STREAM from INLINE to EXTERNAL_LINKS when warehouse rejects INLINE", 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() + .mockRejectedValueOnce( + new Error( + "INVALID_PARAMETER_VALUE: ARROW_STREAM not supported with INLINE disposition", + ), + ) + .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); + + // First call: INLINE (rejected) + expect(executeMock.mock.calls[0][1]).toMatchObject({ + disposition: "INLINE", + format: "ARROW_STREAM", + }); + // Second call: EXTERNAL_LINKS (fallback) + expect(executeMock.mock.calls[1][1]).toMatchObject({ + disposition: "EXTERNAL_LINKS", + format: "ARROW_STREAM", + }); + }); + + 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 when only one of INLINE/ARROW_STREAM appears in the error", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + // Realistic non-format error that mentions just one of the keywords — + // e.g. an unrelated INVALID_PARAMETER_VALUE about a different param. + const executeMock = vi + .fn() + .mockRejectedValue( + new Error( + 'Response from server (Bad Request) {"error_code":"INVALID_PARAMETER_VALUE","message":"INLINE is not a valid value for parameter `mode`"}', + ), + ); + (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 the + // analytics plugin must never escalate to EXTERNAL_LINKS for an error + // that doesn't actually indicate a format/disposition rejection. + for (const call of executeMock.mock.calls) { + expect(call[1]).toMatchObject({ + disposition: "INLINE", + 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 emits arrow_inline SSE event when ARROW_STREAM INLINE returns an attachment", 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 fakeAttachment = "BASE64_ARROW_IPC_BYTES"; + const executeMock = vi.fn().mockResolvedValue({ + result: { attachment: fakeAttachment, row_count: 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(); + + await handler(mockReq, mockRes); + + // The route should not fall back to EXTERNAL_LINKS — INLINE succeeded. + expect(executeMock).toHaveBeenCalledTimes(1); + expect(executeMock.mock.calls[0][1]).toMatchObject({ + disposition: "INLINE", + format: "ARROW_STREAM", + }); + // SSE payload should use the new arrow_inline message type. + const writeCalls = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as string, + ); + const payload = writeCalls.find((s: string) => s.startsWith("data: ")); + expect(payload).toBeDefined(); + expect(payload).toContain('"type":"arrow_inline"'); + expect(payload).toContain(`"attachment":"${fakeAttachment}"`); + }); + + 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"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "JSON" }, + }); + 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 should not fall back 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() + .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 use JSON_ARRAY + INLINE — explicit JSON_ARRAY, no fallback. + for (const call of executeMock.mock.calls) { + expect(call[1]).toMatchObject({ + disposition: "INLINE", + format: "JSON_ARRAY", + }); + } + }); + test("should return 404 when query file is not found", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); diff --git a/packages/appkit/src/stream/defaults.ts b/packages/appkit/src/stream/defaults.ts index c8fc9159..9212ebca 100644 --- a/packages/appkit/src/stream/defaults.ts +++ b/packages/appkit/src/stream/defaults.ts @@ -1,6 +1,12 @@ export const streamDefaults = { bufferSize: 100, - maxEventSize: 1024 * 1024, // 1MB + // 8 MiB. Sized to fit base64-encoded inline Arrow IPC attachments from + // serverless warehouses (analytics queries typically return well under 1 MiB, + // but ARROW_STREAM + INLINE can carry up to ~25 MiB per the Databricks API). + // The connector enforces the same cap (`MAX_INLINE_ATTACHMENT_BYTES`) so + // anything that would exceed this fails fast at the connector with a clear + // error rather than a confusing SSE buffer-exceeded. + maxEventSize: 8 * 1024 * 1024, bufferTTL: 10 * 60 * 1000, // 10 minutes cleanupInterval: 5 * 60 * 1000, // 5 minutes maxPersistentBuffers: 10000, // 10000 buffers diff --git a/packages/appkit/src/type-generator/query-registry.ts b/packages/appkit/src/type-generator/query-registry.ts index 196690c2..63c531d1 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"; @@ -129,18 +130,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 @@ -386,10 +438,42 @@ export async function generateQueriesFromDescribe( sqlHash, cleanedSql, }: (typeof uncachedQueries)[number]): Promise => { - const result = (await client.statementExecution.executeStatement({ - statement: `DESCRIBE QUERY ${cleanedSql}`, - warehouse_id: warehouseId, - })) as DatabricksStatementExecutionResponse; + // Prefer JSON_ARRAY + INLINE so `data_array` parsing works directly. + // Some serverless warehouses reject this combination — fall back to + // ARROW_STREAM + INLINE (still inline, just a different format) and + // let `convertToQueryType` decode the inline attachment. Forcing + // INLINE on the retry avoids EXTERNAL_LINKS, which would silently + // produce empty `data_array` and degrade types to `unknown`. + let result: DatabricksStatementExecutionResponse; + try { + result = (await client.statementExecution.executeStatement({ + statement: `DESCRIBE QUERY ${cleanedSql}`, + warehouse_id: warehouseId, + format: "JSON_ARRAY", + disposition: "INLINE", + })) as DatabricksStatementExecutionResponse; + } catch (err: unknown) { + const msg = err instanceof Error ? err.message : String(err); + const looksLikeFormatRejection = + msg.includes("JSON_ARRAY") && + (msg.includes("not supported") || + msg.includes("INVALID_PARAMETER_VALUE") || + msg.includes("NOT_IMPLEMENTED")); + if (looksLikeFormatRejection) { + logger.debug( + "Warehouse rejected JSON_ARRAY+INLINE for %s, retrying with ARROW_STREAM+INLINE", + queryName, + ); + result = (await client.statementExecution.executeStatement({ + statement: `DESCRIBE QUERY ${cleanedSql}`, + warehouse_id: warehouseId, + format: "ARROW_STREAM", + disposition: "INLINE", + })) as DatabricksStatementExecutionResponse; + } else { + throw err; + } + } completed++; spinner.update( @@ -397,10 +481,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 8d46f98e..63a5636b 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, @@ -11,6 +31,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"); @@ -346,6 +380,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/appkit/src/type-generator/types.ts b/packages/appkit/src/type-generator/types.ts index 5af43591..9a591f51 100644 --- a/packages/appkit/src/type-generator/types.ts +++ b/packages/appkit/src/type-generator/types.ts @@ -12,6 +12,8 @@ export interface DatabricksStatementExecutionResponse { }; result?: { data_array?: (string | null)[][]; + /** Base64-encoded Arrow IPC bytes (returned by serverless warehouses using ARROW_STREAM format) */ + attachment?: string; }; } diff --git a/packages/shared/package.json b/packages/shared/package.json index 27d268ca..bff3a542 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -40,6 +40,7 @@ "ajv": "8.17.1", "ajv-formats": "3.0.1", "@clack/prompts": "1.0.1", - "commander": "12.1.0" + "commander": "12.1.0", + "zod": "3.23.8" } } 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..5abeb83e --- /dev/null +++ b/packages/shared/src/sse/analytics.test.ts @@ -0,0 +1,87 @@ +import { describe, expect, test } from "vitest"; +import { + AnalyticsSseMessage, + makeArrowInlineMessage, + makeArrowMessage, + 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("accepts an arrow message with statement_id", () => { + const parsed = AnalyticsSseMessage.parse({ + type: "arrow", + statement_id: "stmt-1", + }); + expect(parsed.type).toBe("arrow"); + }); + + test("rejects an arrow message with empty statement_id", () => { + expect(() => + AnalyticsSseMessage.parse({ type: "arrow", statement_id: "" }), + ).toThrow(); + }); + + test("rejects an arrow message with no statement_id", () => { + expect(() => AnalyticsSseMessage.parse({ type: "arrow" })).toThrow(); + }); + + test("accepts an arrow_inline message with non-empty attachment", () => { + const parsed = AnalyticsSseMessage.parse({ + type: "arrow_inline", + attachment: "AQID", + }); + expect(parsed.type).toBe("arrow_inline"); + }); + + test("rejects an arrow_inline message with empty attachment", () => { + expect(() => + AnalyticsSseMessage.parse({ type: "arrow_inline", attachment: "" }), + ).toThrow(); + }); + + test("rejects an arrow_inline message with non-string attachment", () => { + expect(() => + AnalyticsSseMessage.parse({ type: "arrow_inline", attachment: 123 }), + ).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_inline" }); + expect(r.success).toBe(false); + }); +}); + +describe("typed builders", () => { + test("makeResultMessage roundtrips through the schema", () => { + const msg = makeResultMessage([{ id: 1 }], { statement_id: "s-1" }); + expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); + }); + + test("makeArrowMessage roundtrips through the schema", () => { + const msg = makeArrowMessage("stmt-2"); + expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); + }); + + test("makeArrowInlineMessage roundtrips through the schema", () => { + const msg = makeArrowInlineMessage("AQID"); + 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..1f136af2 --- /dev/null +++ b/packages/shared/src/sse/analytics.ts @@ -0,0 +1,95 @@ +import { z } from "zod"; + +/** + * Wire protocol for analytics SSE messages emitted by `/api/analytics/query`. + * + * These schemas are the single source of truth for the contract between the + * server (`AnalyticsPlugin._handleQueryRoute`) and the client + * (`useAnalyticsQuery`). Both sides validate with the same schema: + * + * - Server uses the typed builders (`makeResultMessage`, `makeArrowMessage`, + * `makeArrowInlineMessage`) to construct messages with compile-time + * guarantees that all required fields are present. + * - Client calls `AnalyticsSseMessage.parse(JSON.parse(event.data))` to fail + * loudly on a malformed payload instead of silently treating an undefined + * field as data. + * + * Adding a new message variant requires a schema update here, which keeps + * server and client in lockstep. + */ + +/** Successful row-shaped result (JSON_ARRAY format, or empty results). */ +export const AnalyticsResultMessage = z.object({ + type: z.literal("result"), + // zod 4 requires both key and value type for z.record(); zod 3 took + // value only. Using the explicit two-arg form keeps the schema valid + // under whichever zod major resolves at install time. + data: z.array(z.record(z.string(), 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(), +}); +export type AnalyticsResultMessage = z.infer; + +/** + * ARROW_STREAM result delivered via /arrow-result/:jobId — used for + * EXTERNAL_LINKS responses (statement_id from the warehouse) and, if PR #320 + * lands, also for INLINE responses (synthetic `inline-` prefixed id from + * the server-side stash). + */ +export const AnalyticsArrowMessage = z.object({ + type: z.literal("arrow"), + statement_id: z.string().min(1), + status: z.unknown().optional(), +}); +export type AnalyticsArrowMessage = z.infer; + +/** + * ARROW_STREAM + INLINE result with the base64-encoded Arrow IPC bytes + * embedded in the SSE message. The client decodes locally via + * `ArrowClient.processArrowBuffer`. + * + * Note: this variant goes away if the proposal in PR #320 lands. + */ +export const AnalyticsArrowInlineMessage = z.object({ + type: z.literal("arrow_inline"), + attachment: z.string().min(1), +}); +export type AnalyticsArrowInlineMessage = z.infer< + typeof AnalyticsArrowInlineMessage +>; + +/** Discriminated union of every message the analytics SSE stream may emit. */ +export const AnalyticsSseMessage = z.discriminatedUnion("type", [ + AnalyticsResultMessage, + AnalyticsArrowMessage, + AnalyticsArrowInlineMessage, +]); +export type AnalyticsSseMessage = z.infer; + +// ──────────────────────────────────────────────────────────────────────────── +// Typed builders — 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 }; +} + +export function makeArrowMessage( + statement_id: string, + extras: { status?: unknown } = {}, +): AnalyticsArrowMessage { + return { type: "arrow", statement_id, ...extras }; +} + +export function makeArrowInlineMessage( + attachment: string, +): AnalyticsArrowInlineMessage { + return { type: "arrow_inline", attachment }; +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4db8fbe3..63efe770 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -299,6 +299,9 @@ importers: '@types/semver': specifier: 7.7.1 version: 7.7.1 + apache-arrow: + specifier: 21.1.0 + version: 21.1.0 dotenv: specifier: 16.6.1 version: 16.6.1 @@ -332,9 +335,6 @@ importers: ws: specifier: 8.18.3 version: 8.18.3(bufferutil@4.0.9) - zod: - specifier: 4.3.6 - version: 4.3.6 devDependencies: '@opentelemetry/context-async-hooks': specifier: 2.6.1 @@ -554,6 +554,9 @@ importers: commander: specifier: 12.1.0 version: 12.1.0 + zod: + specifier: 3.23.8 + version: 3.23.8 devDependencies: '@types/express': specifier: 4.17.23 @@ -11943,12 +11946,12 @@ packages: peerDependencies: zod: ^3.25.0 || ^4.0.0 + zod@3.23.8: + resolution: {integrity: sha512-XBx9AXhXktjUqnepgTiE5flcKIYWi/rme0Eaj+5Y0lftuGBq+jyRu/md4WnuxqgP1ubdpNCsYEYPxrzVHD8d6g==} + zod@4.1.13: resolution: {integrity: sha512-AvvthqfqrAhNH9dnfmrfKzX5upOdjUVJYFqNSlkmGf64gRaTzlPwz99IHYnVs28qYAybvAlBV+H7pn0saFY4Ig==} - zod@4.3.6: - resolution: {integrity: sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==} - zrender@6.0.0: resolution: {integrity: sha512-41dFXEEXuJpNecuUQq6JlbybmnHaqqpGlbH1yxnA5V9MMP4SbohSVZsJIwz+zdjQXSSlR1Vc34EgH1zxyTDvhg==} @@ -11960,33 +11963,33 @@ packages: snapshots: - '@ai-sdk/gateway@2.0.21(zod@4.3.6)': + '@ai-sdk/gateway@2.0.21(zod@4.1.13)': dependencies: '@ai-sdk/provider': 2.0.0 - '@ai-sdk/provider-utils': 3.0.19(zod@4.3.6) + '@ai-sdk/provider-utils': 3.0.19(zod@4.1.13) '@vercel/oidc': 3.0.5 - zod: 4.3.6 + zod: 4.1.13 - '@ai-sdk/provider-utils@3.0.19(zod@4.3.6)': + '@ai-sdk/provider-utils@3.0.19(zod@4.1.13)': dependencies: '@ai-sdk/provider': 2.0.0 '@standard-schema/spec': 1.1.0 eventsource-parser: 3.0.6 - zod: 4.3.6 + zod: 4.1.13 '@ai-sdk/provider@2.0.0': dependencies: json-schema: 0.4.0 - '@ai-sdk/react@2.0.115(react@19.2.0)(zod@4.3.6)': + '@ai-sdk/react@2.0.115(react@19.2.0)(zod@4.1.13)': dependencies: - '@ai-sdk/provider-utils': 3.0.19(zod@4.3.6) - ai: 5.0.113(zod@4.3.6) + '@ai-sdk/provider-utils': 3.0.19(zod@4.1.13) + ai: 5.0.113(zod@4.1.13) react: 19.2.0 swr: 2.3.8(react@19.2.0) throttleit: 2.1.0 optionalDependencies: - zod: 4.3.6 + zod: 4.1.13 '@algolia/abtesting@1.12.0': dependencies: @@ -13647,14 +13650,14 @@ snapshots: '@docsearch/react@4.3.2(@algolia/client-search@5.46.0)(@types/react@19.2.7)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)(search-insights@2.17.3)': dependencies: - '@ai-sdk/react': 2.0.115(react@19.2.0)(zod@4.3.6) + '@ai-sdk/react': 2.0.115(react@19.2.0)(zod@4.1.13) '@algolia/autocomplete-core': 1.19.2(@algolia/client-search@5.46.0)(algoliasearch@5.46.0)(search-insights@2.17.3) '@docsearch/core': 4.3.1(@types/react@19.2.7)(react-dom@19.2.0(react@19.2.0))(react@19.2.0) '@docsearch/css': 4.3.2 - ai: 5.0.113(zod@4.3.6) + ai: 5.0.113(zod@4.1.13) algoliasearch: 5.46.0 marked: 16.4.2 - zod: 4.3.6 + zod: 4.1.13 optionalDependencies: '@types/react': 19.2.7 react: 19.2.0 @@ -17807,13 +17810,13 @@ snapshots: clean-stack: 2.2.0 indent-string: 4.0.0 - ai@5.0.113(zod@4.3.6): + ai@5.0.113(zod@4.1.13): dependencies: - '@ai-sdk/gateway': 2.0.21(zod@4.3.6) + '@ai-sdk/gateway': 2.0.21(zod@4.1.13) '@ai-sdk/provider': 2.0.0 - '@ai-sdk/provider-utils': 3.0.19(zod@4.3.6) + '@ai-sdk/provider-utils': 3.0.19(zod@4.1.13) '@opentelemetry/api': 1.9.0 - zod: 4.3.6 + zod: 4.1.13 ajv-formats@2.1.1(ajv@8.17.1): optionalDependencies: @@ -21050,7 +21053,7 @@ snapshots: typescript: 5.9.3 unbash: 2.2.0 yaml: 2.8.2 - zod: 4.3.6 + zod: 4.1.13 langium@3.3.1: dependencies: @@ -25319,9 +25322,9 @@ snapshots: dependencies: zod: 4.1.13 - zod@4.1.13: {} + zod@3.23.8: {} - zod@4.3.6: {} + zod@4.1.13: {} zrender@6.0.0: dependencies: From 2ef0c65dd0bdd25d021c49887ac55dc44cced135 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 08:52:28 +0000 Subject: [PATCH 04/25] fix: address ACE multi-model review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Six issues surfaced by GPT 5.4 xhigh + Gemini 3.1 Pro parallel review followed by an adversarial debate round (reviewer: GPT, critic: Gemini, meta: Claude Opus). 1. Raise SSE event-size cap from 8 MiB to 12 MiB on both server (streamDefaults.maxEventSize) and client (connectSSE.maxBufferSize). The inline Arrow attachment cap (MAX_INLINE_ATTACHMENT_BYTES) stays at 8 MiB *decoded*; base64 encoding + JSON + SSE framing inflate that to ~10.6 MiB on the wire, so 12 MiB leaves enough headroom for legal 8-MiB-decoded payloads to traverse the buffer. 2. Empty `data_array: []` is truthy, so zero-row ARROW_STREAM responses skipped empty-table synthesis and fell through to the JSON row transform — callers requesting Arrow got [] JSON rows. Length-check explicitly. 3. The arrow-fix commit dropped lowercase legacy "json" / "arrow" from DataFormat / resolveFormat(), silently breaking existing useChartData callers passing those spellings. Restore them as @deprecated aliases on the DataFormat union; resolveFormat() normalizes them to the canonical "JSON_ARRAY" / "ARROW_STREAM" return values. 4. The JSON_ARRAY -> ARROW_STREAM retry in DESCRIBE QUERY only fired on thrown exceptions. Some warehouses signal the rejection as `status.state === "FAILED"` instead. Extract the rejection-matcher helper and retry on both paths before degrading the typegen result to `unknown`. 5. analytics.test.ts:946 asserted `format: "JSON"` returns 400, but the route now accepts "JSON" as a legacy alias (normalized to JSON_ARRAY). Use a truly unsupported value ("CSV") so the test still exercises the malformed-format path. 6. Restore `zod: 4.3.6` to @databricks/appkit dependencies. main has it; the rebase conflict-resolution accepted the branch's older deps list which lacked it. appkit imports `zod` directly from several files (analytics.ts, agent tools, tests). Co-authored-by: Isaac Signed-off-by: James Broadhead --- packages/appkit-ui/src/js/sse/connect-sse.ts | 10 +++-- packages/appkit-ui/src/react/charts/types.ts | 18 ++++++++- .../src/react/hooks/use-chart-data.ts | 7 ++-- packages/appkit/package.json | 3 +- .../src/connectors/sql-warehouse/client.ts | 9 ++++- .../plugins/analytics/tests/analytics.test.ts | 5 ++- packages/appkit/src/stream/defaults.ts | 13 +++---- .../src/type-generator/query-registry.ts | 38 +++++++++++++++---- 8 files changed, 77 insertions(+), 26 deletions(-) diff --git a/packages/appkit-ui/src/js/sse/connect-sse.ts b/packages/appkit-ui/src/js/sse/connect-sse.ts index 089ddf57..5057bc1d 100644 --- a/packages/appkit-ui/src/js/sse/connect-sse.ts +++ b/packages/appkit-ui/src/js/sse/connect-sse.ts @@ -18,10 +18,12 @@ export async function connectSSE( lastEventId: initialLastEventId = null, retryDelay = 2000, maxRetries = 3, - // 8 MiB — sized to receive inline Arrow IPC attachments from - // ARROW_STREAM analytics responses; matches the server's stream - // `maxEventSize`. Most events are well under 1 MiB in practice. - maxBufferSize = 8 * 1024 * 1024, + // 12 MiB — matches the server's stream `maxEventSize`. Sized to + // receive inline Arrow IPC attachments from ARROW_STREAM analytics + // responses: the connector caps decoded payloads at 8 MiB, which + // inflates to ~10.6 MiB once base64-encoded and wrapped in JSON + SSE + // framing. Most events are well under 1 MiB in practice. + maxBufferSize = 12 * 1024 * 1024, timeout = 300000, // 5 minutes onError, } = options; diff --git a/packages/appkit-ui/src/react/charts/types.ts b/packages/appkit-ui/src/react/charts/types.ts index ec8a15dc..ef738aad 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_array" | "arrow_stream" | "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"; 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 ec4b2d4e..64b6e167 100644 --- a/packages/appkit-ui/src/react/hooks/use-chart-data.ts +++ b/packages/appkit-ui/src/react/hooks/use-chart-data.ts @@ -51,9 +51,10 @@ function resolveFormat( format: DataFormat, parameters?: Record, ): "JSON_ARRAY" | "ARROW_STREAM" { - // Explicit format selection - if (format === "json_array") return "JSON_ARRAY"; - if (format === "arrow_stream") 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") { diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 5b9429c3..4846f88a 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -84,7 +84,8 @@ "semver": "7.7.3", "shared": "workspace:*", "vite": "npm:rolldown-vite@7.1.14", - "ws": "8.18.3" + "ws": "8.18.3", + "zod": "4.3.6" }, "devDependencies": { "@opentelemetry/context-async-hooks": "2.6.1", diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 4ae43941..d76ac590 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -447,7 +447,14 @@ export class SQLWarehouseConnector { // 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. - if (!result?.data_array && response.manifest?.schema?.columns) { + // 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, ); diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 89bfac7d..6884e2bb 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -941,9 +941,12 @@ describe("Analytics Plugin", () => { 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: "JSON" }, + body: { parameters: {}, format: "CSV" }, }); const mockRes = createMockResponse(); diff --git a/packages/appkit/src/stream/defaults.ts b/packages/appkit/src/stream/defaults.ts index 9212ebca..7cde49c4 100644 --- a/packages/appkit/src/stream/defaults.ts +++ b/packages/appkit/src/stream/defaults.ts @@ -1,12 +1,11 @@ export const streamDefaults = { bufferSize: 100, - // 8 MiB. Sized to fit base64-encoded inline Arrow IPC attachments from - // serverless warehouses (analytics queries typically return well under 1 MiB, - // but ARROW_STREAM + INLINE can carry up to ~25 MiB per the Databricks API). - // The connector enforces the same cap (`MAX_INLINE_ATTACHMENT_BYTES`) so - // anything that would exceed this fails fast at the connector with a clear - // error rather than a confusing SSE buffer-exceeded. - maxEventSize: 8 * 1024 * 1024, + // 12 MiB. Headroom for base64-encoded inline Arrow IPC attachments: the + // connector caps the *decoded* attachment at 8 MiB (MAX_INLINE_ATTACHMENT_BYTES), + // which inflates to ~10.6 MiB once base64-encoded and is then wrapped in JSON + + // SSE framing. 12 MiB leaves enough room for that overhead so legal 8-MiB-decoded + // attachments do not trip the stream-manager cap before the connector-level cap. + maxEventSize: 12 * 1024 * 1024, bufferTTL: 10 * 60 * 1000, // 10 minutes cleanupInterval: 5 * 60 * 1000, // 5 minutes maxPersistentBuffers: 10000, // 10000 buffers diff --git a/packages/appkit/src/type-generator/query-registry.ts b/packages/appkit/src/type-generator/query-registry.ts index 63c531d1..9bbeb01e 100644 --- a/packages/appkit/src/type-generator/query-registry.ts +++ b/packages/appkit/src/type-generator/query-registry.ts @@ -431,6 +431,16 @@ export async function generateQueriesFromDescribe( `Describing ${total} ${total === 1 ? "query" : "queries"} (0/${total})`, ); + // Some serverless warehouses reject JSON_ARRAY+INLINE for DESCRIBE — and + // they signal the rejection two different ways: either as a thrown error, + // or as a `status.state === "FAILED"` response. Both paths funnel through + // this matcher so we can retry with ARROW_STREAM+INLINE consistently. + const looksLikeFormatRejection = (msg: string): boolean => + msg.includes("JSON_ARRAY") && + (msg.includes("not supported") || + msg.includes("INVALID_PARAMETER_VALUE") || + msg.includes("NOT_IMPLEMENTED")); + const describeOne = async ({ index, queryName, @@ -454,14 +464,9 @@ export async function generateQueriesFromDescribe( })) as DatabricksStatementExecutionResponse; } catch (err: unknown) { const msg = err instanceof Error ? err.message : String(err); - const looksLikeFormatRejection = - msg.includes("JSON_ARRAY") && - (msg.includes("not supported") || - msg.includes("INVALID_PARAMETER_VALUE") || - msg.includes("NOT_IMPLEMENTED")); - if (looksLikeFormatRejection) { + if (looksLikeFormatRejection(msg)) { logger.debug( - "Warehouse rejected JSON_ARRAY+INLINE for %s, retrying with ARROW_STREAM+INLINE", + "Warehouse rejected JSON_ARRAY+INLINE for %s (thrown), retrying with ARROW_STREAM+INLINE", queryName, ); result = (await client.statementExecution.executeStatement({ @@ -475,6 +480,25 @@ export async function generateQueriesFromDescribe( } } + // Some warehouses surface the format rejection as `status.state === + // "FAILED"` instead of throwing. Detect that shape and retry with + // ARROW_STREAM before we degrade the type to `unknown`. + if ( + result.status.state === "FAILED" && + looksLikeFormatRejection(result.status.error?.message ?? "") + ) { + logger.debug( + "Warehouse rejected JSON_ARRAY+INLINE for %s (state=FAILED), retrying with ARROW_STREAM+INLINE", + queryName, + ); + result = (await client.statementExecution.executeStatement({ + statement: `DESCRIBE QUERY ${cleanedSql}`, + warehouse_id: warehouseId, + format: "ARROW_STREAM", + disposition: "INLINE", + })) as DatabricksStatementExecutionResponse; + } + completed++; spinner.update( `Describing ${total} ${total === 1 ? "query" : "queries"} (${completed}/${total})`, From 2b22d5600460568e299e3b45c363018762a845be Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 08:57:00 +0000 Subject: [PATCH 05/25] chore(shared): align zod with appkit's 4.3.6 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The original commit added zod@3.23.8 to shared for the new SSE wire protocol schema. With zod restored on appkit at 4.3.6 (matching main), the workspace now had two different zod majors resolving in different packages — a latent peer-dep / type-incompatibility foot-gun even though the schema itself was already cross-major-compatible. Bump shared's zod to 4.3.6 so the whole workspace lands on one major. The schema's two-arg `z.record(z.string(), z.unknown())` form is the zod 4 spelling, so no functional change is needed; drop the now-stale "keeps it valid under either major" comment. Signed-off-by: James Broadhead --- packages/shared/package.json | 2 +- packages/shared/src/sse/analytics.ts | 3 --- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/packages/shared/package.json b/packages/shared/package.json index bff3a542..542f7a96 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -41,6 +41,6 @@ "ajv-formats": "3.0.1", "@clack/prompts": "1.0.1", "commander": "12.1.0", - "zod": "3.23.8" + "zod": "4.3.6" } } diff --git a/packages/shared/src/sse/analytics.ts b/packages/shared/src/sse/analytics.ts index 1f136af2..a7f68d54 100644 --- a/packages/shared/src/sse/analytics.ts +++ b/packages/shared/src/sse/analytics.ts @@ -21,9 +21,6 @@ import { z } from "zod"; /** Successful row-shaped result (JSON_ARRAY format, or empty results). */ export const AnalyticsResultMessage = z.object({ type: z.literal("result"), - // zod 4 requires both key and value type for z.record(); zod 3 took - // value only. Using the explicit two-arg form keeps the schema valid - // under whichever zod major resolves at install time. data: z.array(z.record(z.string(), 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. From 90ecd8a72ea4a3c681fc2bece5a63c46463dff40 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 09:48:42 +0000 Subject: [PATCH 06/25] chore: regenerate pnpm-lock.yaml after zod restoration Restoring zod@4.3.6 to appkit and bumping shared's zod from 3.23.8 to 4.3.6 left the lockfile out of sync with package.json, breaking CI's pnpm install --frozen-lockfile step on every job. Regenerate the lockfile so both specifier entries match the manifests. Signed-off-by: James Broadhead --- pnpm-lock.yaml | 53 ++++++++++++++++++++++++++------------------------ 1 file changed, 28 insertions(+), 25 deletions(-) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 63efe770..d7d75516 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -335,6 +335,9 @@ importers: ws: specifier: 8.18.3 version: 8.18.3(bufferutil@4.0.9) + zod: + specifier: 4.3.6 + version: 4.3.6 devDependencies: '@opentelemetry/context-async-hooks': specifier: 2.6.1 @@ -555,8 +558,8 @@ importers: specifier: 12.1.0 version: 12.1.0 zod: - specifier: 3.23.8 - version: 3.23.8 + specifier: 4.3.6 + version: 4.3.6 devDependencies: '@types/express': specifier: 4.17.23 @@ -5567,7 +5570,7 @@ packages: basic-ftp@5.0.5: resolution: {integrity: sha512-4Bcg1P8xhUuqcii/S0Z9wiHIrQVPMermM1any+MX5GeGD7faD3/msQUDGLol9wOcz4/jbg/WJnGqoJF6LiBdtg==} engines: {node: '>=10.0.0'} - deprecated: Security vulnerability fixed in 5.2.0, please upgrade + deprecated: Security vulnerability fixed in 5.2.1, please upgrade batch@0.6.1: resolution: {integrity: sha512-x+VAiMRL6UPkx+kudNvxTl6hB2XNNCG2r+7wixVfIYwu/2HKRXimwQyaumLjMveWvT2Hkd/cAJw+QBMfJ/EKVw==} @@ -11946,12 +11949,12 @@ packages: peerDependencies: zod: ^3.25.0 || ^4.0.0 - zod@3.23.8: - resolution: {integrity: sha512-XBx9AXhXktjUqnepgTiE5flcKIYWi/rme0Eaj+5Y0lftuGBq+jyRu/md4WnuxqgP1ubdpNCsYEYPxrzVHD8d6g==} - zod@4.1.13: resolution: {integrity: sha512-AvvthqfqrAhNH9dnfmrfKzX5upOdjUVJYFqNSlkmGf64gRaTzlPwz99IHYnVs28qYAybvAlBV+H7pn0saFY4Ig==} + zod@4.3.6: + resolution: {integrity: sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==} + zrender@6.0.0: resolution: {integrity: sha512-41dFXEEXuJpNecuUQq6JlbybmnHaqqpGlbH1yxnA5V9MMP4SbohSVZsJIwz+zdjQXSSlR1Vc34EgH1zxyTDvhg==} @@ -11963,33 +11966,33 @@ packages: snapshots: - '@ai-sdk/gateway@2.0.21(zod@4.1.13)': + '@ai-sdk/gateway@2.0.21(zod@4.3.6)': dependencies: '@ai-sdk/provider': 2.0.0 - '@ai-sdk/provider-utils': 3.0.19(zod@4.1.13) + '@ai-sdk/provider-utils': 3.0.19(zod@4.3.6) '@vercel/oidc': 3.0.5 - zod: 4.1.13 + zod: 4.3.6 - '@ai-sdk/provider-utils@3.0.19(zod@4.1.13)': + '@ai-sdk/provider-utils@3.0.19(zod@4.3.6)': dependencies: '@ai-sdk/provider': 2.0.0 '@standard-schema/spec': 1.1.0 eventsource-parser: 3.0.6 - zod: 4.1.13 + zod: 4.3.6 '@ai-sdk/provider@2.0.0': dependencies: json-schema: 0.4.0 - '@ai-sdk/react@2.0.115(react@19.2.0)(zod@4.1.13)': + '@ai-sdk/react@2.0.115(react@19.2.0)(zod@4.3.6)': dependencies: - '@ai-sdk/provider-utils': 3.0.19(zod@4.1.13) - ai: 5.0.113(zod@4.1.13) + '@ai-sdk/provider-utils': 3.0.19(zod@4.3.6) + ai: 5.0.113(zod@4.3.6) react: 19.2.0 swr: 2.3.8(react@19.2.0) throttleit: 2.1.0 optionalDependencies: - zod: 4.1.13 + zod: 4.3.6 '@algolia/abtesting@1.12.0': dependencies: @@ -13650,14 +13653,14 @@ snapshots: '@docsearch/react@4.3.2(@algolia/client-search@5.46.0)(@types/react@19.2.7)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)(search-insights@2.17.3)': dependencies: - '@ai-sdk/react': 2.0.115(react@19.2.0)(zod@4.1.13) + '@ai-sdk/react': 2.0.115(react@19.2.0)(zod@4.3.6) '@algolia/autocomplete-core': 1.19.2(@algolia/client-search@5.46.0)(algoliasearch@5.46.0)(search-insights@2.17.3) '@docsearch/core': 4.3.1(@types/react@19.2.7)(react-dom@19.2.0(react@19.2.0))(react@19.2.0) '@docsearch/css': 4.3.2 - ai: 5.0.113(zod@4.1.13) + ai: 5.0.113(zod@4.3.6) algoliasearch: 5.46.0 marked: 16.4.2 - zod: 4.1.13 + zod: 4.3.6 optionalDependencies: '@types/react': 19.2.7 react: 19.2.0 @@ -17810,13 +17813,13 @@ snapshots: clean-stack: 2.2.0 indent-string: 4.0.0 - ai@5.0.113(zod@4.1.13): + ai@5.0.113(zod@4.3.6): dependencies: - '@ai-sdk/gateway': 2.0.21(zod@4.1.13) + '@ai-sdk/gateway': 2.0.21(zod@4.3.6) '@ai-sdk/provider': 2.0.0 - '@ai-sdk/provider-utils': 3.0.19(zod@4.1.13) + '@ai-sdk/provider-utils': 3.0.19(zod@4.3.6) '@opentelemetry/api': 1.9.0 - zod: 4.1.13 + zod: 4.3.6 ajv-formats@2.1.1(ajv@8.17.1): optionalDependencies: @@ -21053,7 +21056,7 @@ snapshots: typescript: 5.9.3 unbash: 2.2.0 yaml: 2.8.2 - zod: 4.1.13 + zod: 4.3.6 langium@3.3.1: dependencies: @@ -25322,10 +25325,10 @@ snapshots: dependencies: zod: 4.1.13 - zod@3.23.8: {} - zod@4.1.13: {} + zod@4.3.6: {} + zrender@6.0.0: dependencies: tslib: 2.3.0 From 3d540096433c20b70d19039978c9e121a45af948 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 09:54:51 +0000 Subject: [PATCH 07/25] style: consolidate normalizeAnalyticsFormat into the types import block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The merge with main left two separate import statements for "./types" — one for the type-only specifiers and a duplicate value import of normalizeAnalyticsFormat. Biome rejected this as both an organize-imports failure and a noRedeclare error. Merge them into a single mixed type/value import. Signed-off-by: James Broadhead --- packages/appkit/src/plugins/analytics/analytics.ts | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 8ca41086..6e83fadd 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -29,14 +29,13 @@ import type { PluginManifest } from "../../registry"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; -import { normalizeAnalyticsFormat } from "./types"; -import type { - AnalyticsFormat, - AnalyticsQueryResponse, - IAnalyticsConfig, - IAnalyticsQueryRequest, +import { + type AnalyticsFormat, + type AnalyticsQueryResponse, + type IAnalyticsConfig, + type IAnalyticsQueryRequest, + normalizeAnalyticsFormat, } from "./types"; -import { normalizeAnalyticsFormat } from "./types"; const logger = createLogger("analytics"); From e6c2aae3825792d8f8552b9ce49ba8c8bcefa898 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 10:02:13 +0000 Subject: [PATCH 08/25] fix: restore logger.error in executeStatement catch block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The merge resolution in client.ts dropped the logger.error call from the executeStatement catch block — main has it, our pre-merge branch had it, the resolved version lost it. Without that line the "error log redaction" tests fail because the connector no longer surfaces the failure message to the log spy. Restore the call. Test plan: the two sql-warehouse.test.ts redaction tests pass locally; behavior matches the comment "executeStatement's catch ... is the single point that logs (gated on isAborted)". Signed-off-by: James Broadhead --- packages/appkit/src/connectors/sql-warehouse/client.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index d76ac590..f011a9d9 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -253,6 +253,11 @@ export class SQLWarehouseConnector { code: SpanStatusCode.ERROR, message: error instanceof Error ? error.message : String(error), }); + + logger.error( + "Statement execution failed: %s", + error instanceof Error ? error.message : String(error), + ); } if (error instanceof AppKitError) { From a7434f6bce07750b27d9175af290fd6b7541b928 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 15:43:15 +0000 Subject: [PATCH 09/25] refactor(analytics): stash inline Arrow server-side, drop arrow_inline SSE message Address Mario's design feedback: SSE is for short control messages, not bulk binary. Inline Arrow IPC payloads from serverless warehouses no longer ride the SSE channel as base64; they are stashed server-side and fetched out-of-band through the existing /arrow-result/:jobId endpoint with the canonical application/vnd.apache.arrow.stream content-type. Wire protocol - Discriminated union shrinks from three variants to two: the arrow_inline message type is gone. Both INLINE and EXTERNAL_LINKS ARROW_STREAM responses now flow as a single `arrow` message whose statement_id discriminates dispatch: warehouse-issued ids hit the warehouse path, synthetic "inline-" ids hit the stash. The client sees one path. Server - New InlineArrowStash: TTL'd (10 min), bounded-memory (256 MiB), drain-on-read, per-user-keyed map of decoded Arrow IPC bytes. Stash key is the request's user id (or "global" for SP contexts) and is symmetric between put and take. - AnalyticsPlugin holds one stash instance and uses it in two places: - _executeWithFormatFallback decodes result.attachment once, puts the bytes in the stash, and emits an arrow message with the synthetic id. Bulk bytes never traverse SSE. - _handleArrowRoute prefix-dispatches on the jobId: "inline-" drains the stash and serves with application/vnd.apache.arrow.stream + a no-store cache header; other ids fall through to the existing warehouse-fetch path unchanged. - Connector's MAX_INLINE_ATTACHMENT_BYTES raised from 8 MiB to 25 MiB (the Databricks API hard cap on INLINE) since the SSE event-size budget no longer constrains it. Client - useAnalyticsQuery loses the arrow_inline branch and the local base64 decoder. Both inline and external-links responses fetch through /api/analytics/arrow-result/:id; the prefix branch lives server-side. - The dead client-side MAX_INLINE_ATTACHMENT_BYTES guard goes away. SSE buffers - streamDefaults.maxEventSize: 12 MiB -> 1 MiB - connectSSE.maxBufferSize: 12 MiB -> 1 MiB SSE now carries only short JSON control messages (result rows, arrow envelope with statement id, error frames). Multi-MiB caps are no longer needed and would mask buffer regressions. Tests - New InlineArrowStash unit tests (TTL eviction, max-bytes LRU, drain- on-read, per-user scoping). - Reworked the route's "emits arrow_inline" test into a stash + arrow- message assertion: the SSE payload must not contain the base64 bytes or the arrow_inline type literal, and the decoded bytes must be in the stash keyed by the same synthetic id. - New /arrow-result tests cover the inline path: success drain, 410 on unknown id, 410 on user mismatch. - Client tests rewritten to assert both warehouse and inline-prefixed ids fetch through the same /arrow-result URL with no local decoding. - Shared schema tests assert the retired arrow_inline type no longer parses. - The /arrow-result content-type for warehouse hits stays application/ octet-stream (no behavior change there). Signed-off-by: James Broadhead --- packages/appkit-ui/src/js/sse/connect-sse.ts | 11 +- .../__tests__/use-analytics-query.test.ts | 105 +++++++------- .../src/react/hooks/use-analytics-query.ts | 62 +-------- .../src/connectors/sql-warehouse/client.ts | 17 ++- .../sql-warehouse/tests/client.test.ts | 5 +- .../appkit/src/plugins/analytics/analytics.ts | 119 ++++++++++++++-- .../plugins/analytics/inline-arrow-stash.ts | 129 ++++++++++++++++++ .../plugins/analytics/tests/analytics.test.ts | 105 +++++++++++++- .../tests/inline-arrow-stash.test.ts | 99 ++++++++++++++ packages/appkit/src/stream/defaults.ts | 14 +- packages/shared/src/sse/analytics.test.ts | 42 +++--- packages/shared/src/sse/analytics.ts | 45 +++--- 12 files changed, 555 insertions(+), 198 deletions(-) create mode 100644 packages/appkit/src/plugins/analytics/inline-arrow-stash.ts create mode 100644 packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts diff --git a/packages/appkit-ui/src/js/sse/connect-sse.ts b/packages/appkit-ui/src/js/sse/connect-sse.ts index 5057bc1d..13d9053d 100644 --- a/packages/appkit-ui/src/js/sse/connect-sse.ts +++ b/packages/appkit-ui/src/js/sse/connect-sse.ts @@ -18,12 +18,11 @@ export async function connectSSE( lastEventId: initialLastEventId = null, retryDelay = 2000, maxRetries = 3, - // 12 MiB — matches the server's stream `maxEventSize`. Sized to - // receive inline Arrow IPC attachments from ARROW_STREAM analytics - // responses: the connector caps decoded payloads at 8 MiB, which - // inflates to ~10.6 MiB once base64-encoded and wrapped in JSON + SSE - // framing. Most events are well under 1 MiB in practice. - maxBufferSize = 12 * 1024 * 1024, + // 1 MiB — matches the server's `streamDefaults.maxEventSize`. SSE + // carries only short JSON control messages; bulk Arrow payloads flow + // over plain HTTP via `/api/analytics/arrow-result/:jobId`, 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/hooks/__tests__/use-analytics-query.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts index 65de7d10..c2705476 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 @@ -30,89 +30,70 @@ describe("useAnalyticsQuery", () => { lastConnectArgs = null; }); - test("decodes arrow_inline base64 attachment via ArrowClient.processArrowBuffer", async () => { + test("fetches an arrow message (warehouse statement id) via /arrow-result", async () => { const fakeTable = { numRows: 1, schema: { fields: [] } }; + const fakeBytes = new Uint8Array([1, 2, 3]); + mockFetchArrow.mockResolvedValueOnce(fakeBytes); mockProcessArrowBuffer.mockResolvedValueOnce(fakeTable); - // 'AQID' decodes to bytes [1, 2, 3]. - const base64 = "AQID"; - const { result } = renderHook(() => useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), ); - // Drive the SSE onMessage handler with an arrow_inline payload. await lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow_inline", attachment: base64 }), + data: JSON.stringify({ type: "arrow", statement_id: "stmt-warehouse-1" }), }); await waitFor(() => { expect(result.current.data).toBe(fakeTable); }); - expect(mockProcessArrowBuffer).toHaveBeenCalledTimes(1); - const passedBuffer = mockProcessArrowBuffer.mock.calls[0][0] as Uint8Array; - expect(passedBuffer).toBeInstanceOf(Uint8Array); - expect(Array.from(passedBuffer)).toEqual([1, 2, 3]); - // Inline path must NOT trigger a network fetch. - expect(mockFetchArrow).not.toHaveBeenCalled(); + expect(mockFetchArrow).toHaveBeenCalledTimes(1); + expect(mockFetchArrow).toHaveBeenCalledWith( + "/api/analytics/arrow-result/stmt-warehouse-1", + ); + expect(mockProcessArrowBuffer).toHaveBeenCalledWith(fakeBytes); }); - test("surfaces an error when arrow_inline decode fails", async () => { - mockProcessArrowBuffer.mockRejectedValueOnce(new Error("bad ipc")); + test("fetches an arrow message with synthetic inline- id through the same /arrow-result path", async () => { + // The client must treat inline and external-links responses uniformly — + // it never decodes base64 locally. The /arrow-result route on the + // server is the only place that knows which path the bytes came from. + const fakeTable = { numRows: 1, schema: { fields: [] } }; + const fakeBytes = new Uint8Array([1, 2, 3, 4, 5]); + mockFetchArrow.mockResolvedValueOnce(fakeBytes); + mockProcessArrowBuffer.mockResolvedValueOnce(fakeTable); const { result } = renderHook(() => useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), ); await lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow_inline", attachment: "AQID" }), + data: JSON.stringify({ + type: "arrow", + statement_id: "inline-abc-xyz", + }), }); await waitFor(() => { - expect(result.current.error).toBe( - "Unable to load data, please try again", - ); + expect(result.current.data).toBe(fakeTable); }); - expect(result.current.loading).toBe(false); - }); - - test("rejects arrow_inline with missing/empty/non-string attachment without crashing atob", async () => { - const cases: Array = [undefined, null, "", 123, { foo: "bar" }]; - - for (const attachment of cases) { - mockProcessArrowBuffer.mockClear(); - const { result, unmount } = renderHook(() => - useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), - ); - - await lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow_inline", attachment }), - }); - await waitFor(() => { - expect(result.current.error).toBe( - "Unable to load data, please try again", - ); - }); - // Critically: must NOT call processArrowBuffer (or atob) on the bad input. - expect(mockProcessArrowBuffer).not.toHaveBeenCalled(); - - unmount(); - } + expect(mockFetchArrow).toHaveBeenCalledTimes(1); + expect(mockFetchArrow).toHaveBeenCalledWith( + "/api/analytics/arrow-result/inline-abc-xyz", + ); }); - test("rejects oversized arrow_inline attachment without allocating a huge buffer", async () => { - // Base64 string that would decode to ~9 MiB (>8 MiB cap). The hook - // should reject before calling decodeBase64 / processArrowBuffer. - const oversized = "A".repeat(13 * 1024 * 1024); + test("surfaces an error when the arrow fetch fails", async () => { + mockFetchArrow.mockRejectedValueOnce(new Error("network")); const { result } = renderHook(() => useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), ); await lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow_inline", attachment: oversized }), + data: JSON.stringify({ type: "arrow", statement_id: "stmt-1" }), }); await waitFor(() => { @@ -120,7 +101,34 @@ describe("useAnalyticsQuery", () => { "Unable to load data, please try again", ); }); + expect(result.current.loading).toBe(false); + }); + + test("rejects the retired arrow_inline message type as schema-invalid", async () => { + // arrow_inline was the prior wire shape. The discriminated union no + // longer accepts it, so it falls through to the generic error/code + // branch — but critically, it must NEVER trigger ArrowClient calls. + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await lastConnectArgs.onMessage({ + data: JSON.stringify({ type: "arrow_inline", attachment: "AQID" }), + }); + + // Whatever the hook surfaces (error or noop), it must not have tried to + // decode the payload locally. + await waitFor(() => { + // Either an error is set or loading completed without data — both are + // acceptable, but processArrowBuffer must never run on a base64 input. + expect( + result.current.loading || + result.current.error || + result.current.data === null, + ).toBeTruthy(); + }); expect(mockProcessArrowBuffer).not.toHaveBeenCalled(); + expect(mockFetchArrow).not.toHaveBeenCalled(); }); test("still handles type:result rows for JSON_ARRAY", async () => { @@ -139,5 +147,6 @@ describe("useAnalyticsQuery", () => { expect(result.current.data).toEqual([{ id: 1 }, { id: 2 }]); }); expect(mockProcessArrowBuffer).not.toHaveBeenCalled(); + expect(mockFetchArrow).not.toHaveBeenCalled(); }); }); 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 5d18a2ee..a192adc2 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -23,29 +23,6 @@ function getArrowStreamUrl(id: string) { return `/api/analytics/arrow-result/${id}`; } -/** - * Client-side defensive cap on inline Arrow IPC attachments (8 MiB decoded). - * Mirrors the server's MAX_INLINE_ATTACHMENT_BYTES so a misconfigured proxy - * (or a future server bug) can't push us into allocating an unbounded - * Uint8Array and hanging the browser. - * - * REMOVE THIS GUARD if PR #320 (stash + serve via /arrow-result) lands — - * that proposal eliminates the arrow_inline SSE path entirely, so bulk - * bytes flow over HTTP where the browser handles backpressure natively - * and Content-Length is exposed up-front. - */ -const MAX_INLINE_ATTACHMENT_BYTES = 8 * 1024 * 1024; - -/** Decode a base64 string into a Uint8Array suitable for Arrow IPC parsing. */ -function decodeBase64(b64: string): Uint8Array { - const binary = atob(b64); - const bytes = new Uint8Array(binary.length); - for (let i = 0; i < binary.length; i++) { - bytes[i] = binary.charCodeAt(i); - } - return bytes; -} - /** * Subscribe to an analytics query over SSE and returns its latest result. * Integration hook between client and analytics plugin. @@ -161,7 +138,11 @@ export function useAnalyticsQuery< return; } - // success - Arrow format (external links: fetch from server) + // success - Arrow format. Both INLINE (server-stashed, + // statement_id prefixed with "inline-") and EXTERNAL_LINKS + // (warehouse statement_id) flow through this single branch — the + // /arrow-result route dispatches based on the id prefix so the + // client doesn't need to know which path the bytes came from. if (msg?.type === "arrow") { try { const arrowData = await ArrowClient.fetchArrow( @@ -183,39 +164,6 @@ export function useAnalyticsQuery< } } - // success - Arrow format (inline: decode base64 IPC payload locally) - if (msg?.type === "arrow_inline") { - // Schema already enforced non-empty string; just check size. - // base64 length L decodes to ~L*3/4 bytes; reject before - // allocating a multi-MiB Uint8Array. - const decodedSize = Math.ceil((msg.attachment.length * 3) / 4); - if (decodedSize > MAX_INLINE_ATTACHMENT_BYTES) { - console.error( - "[useAnalyticsQuery] arrow_inline attachment exceeds %d bytes (got %d)", - MAX_INLINE_ATTACHMENT_BYTES, - decodedSize, - ); - setLoading(false); - setError("Unable to load data, please try again"); - return; - } - try { - const buffer = decodeBase64(msg.attachment); - const table = await ArrowClient.processArrowBuffer(buffer); - setLoading(false); - setData(table as ResultType); - return; - } catch (error) { - console.error( - "[useAnalyticsQuery] Failed to decode inline Arrow data", - error, - ); - setLoading(false); - setError("Unable to load data, please try again"); - return; - } - } - // The schema didn't match — fall through to error/code handling // below for legacy error events or surface a malformed-payload // error if no error fields are present. diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index f011a9d9..a0016d7b 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -27,17 +27,16 @@ import { executeStatementDefaults } from "./defaults"; const logger = createLogger("connectors:sql-warehouse"); /** - * Maximum size for inline Arrow IPC attachments (8 MiB decoded). - * Aligned with `streamDefaults.maxEventSize` so anything that would exceed - * the SSE event cap fails here with a clear error rather than a confusing - * "Buffer size exceeded" downstream. Larger results should use - * `disposition: "EXTERNAL_LINKS"`, which the analytics fallback handles. + * Maximum size for inline Arrow IPC attachments (25 MiB decoded — the + * Databricks Statement Execution API hard cap on INLINE responses). * - * RAISE TO 25 MiB (Databricks API hard cap on INLINE) if PR #320 (stash + - * serve via /arrow-result) lands — that proposal moves bulk bytes off SSE - * onto HTTP, so the SSE event-size constraint no longer applies here. + * Bulk Arrow payloads no longer traverse SSE — the analytics route stashes + * them via `InlineArrowStash` and the client fetches over HTTP — so this + * cap is bounded by the upstream API rather than our event-size budget. + * Larger results still fall through to `disposition: "EXTERNAL_LINKS"`, + * handled by the analytics format-fallback. */ -const MAX_INLINE_ATTACHMENT_BYTES = 8 * 1024 * 1024; +const MAX_INLINE_ATTACHMENT_BYTES = 25 * 1024 * 1024; interface SQLWarehouseConfig { timeout?: number; diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts index c7f73c98..c6780f3f 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts @@ -293,8 +293,9 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("rejects oversized attachments to bound memory", () => { const connector = createConnector(); - // 8 MiB decoded cap → ~12 MiB of base64 chars decodes to >8 MiB. - const oversized = "A".repeat(12 * 1024 * 1024); + // 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" }, diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 6e83fadd..b99fb9a2 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -4,7 +4,6 @@ import { type AgentToolDefinition, type AnalyticsSseMessage, type IAppRouter, - makeArrowInlineMessage, makeArrowMessage, makeResultMessage, type PluginExecuteConfig, @@ -27,6 +26,7 @@ import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; import { queryDefaults } from "./defaults"; +import { InlineArrowStash } from "./inline-arrow-stash"; import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; import { @@ -50,6 +50,17 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { private SQLClient: SQLWarehouseConnector; private queryProcessor: QueryProcessor; + /** + * Server-side stash for inline Arrow IPC payloads. + * + * INLINE ARROW_STREAM responses do not ride the SSE control channel — + * the route puts the decoded bytes here and emits an `arrow` SSE + * message with a synthetic `inline-` id, and the client fetches + * the bytes through the existing `/arrow-result/:jobId` endpoint with + * a real binary content-type. + */ + protected inlineArrowStash: InlineArrowStash = new InlineArrowStash(); + constructor(config: IAnalyticsConfig) { super(config); this.config = config; @@ -87,24 +98,60 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { /** * Handle Arrow data download requests. - * When called via asUser(req), uses the user's Databricks credentials. + * + * Two id shapes are supported: + * - `inline-`: bytes were stashed server-side by the query route. + * Drain the stash, serve directly with the canonical Arrow content + * type. No warehouse round-trip. + * - any other id: a warehouse-issued statement id. Fetch the Arrow + * stream from the warehouse via the SDK; serve the bytes. + * + * When called via asUser(req), uses the user's Databricks credentials + * for the warehouse path. The inline path is user-scoped at the stash + * layer instead. */ async _handleArrowRoute( req: express.Request, res: express.Response, ): Promise { + const { jobId } = req.params; + const event = logger.event(req); + event?.setComponent("analytics", "getArrowData").setContext("analytics", { + job_id: jobId, + plugin: this.name, + }); + + if (jobId.startsWith("inline-")) { + const userKey = this._stashUserKey(req); + const bytes = this.inlineArrowStash.take(jobId, userKey); + if (!bytes) { + // Already drained, expired, or never belonged to this user. 410 + // distinguishes this from "warehouse statement id not found" (404) + // so the client can surface a useful error. + logger.debug("Inline Arrow stash miss for jobId=%s", jobId); + res.status(410).json({ + error: "Inline Arrow result expired or unknown", + plugin: this.name, + }); + return; + } + logger.debug( + "Serving inline Arrow buffer: %d bytes for jobId=%s", + bytes.length, + jobId, + ); + res.setHeader("Content-Type", "application/vnd.apache.arrow.stream"); + res.setHeader("Content-Length", bytes.length.toString()); + // Inline payloads are single-use and short-lived; no public caching. + res.setHeader("Cache-Control", "no-store"); + res.send(Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength)); + return; + } + 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"); @@ -126,6 +173,28 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } } + /** + * Stash key used at put-time (in `_handleQueryRoute`) and take-time + * (in `_handleArrowRoute`). Centralized so the two sides cannot drift. + * + * Returns the user id when an `x-forwarded-user` header is present, + * otherwise `"global"` for service-principal contexts (no user header). + * Both queries from the same request resolve to the same key, so the + * subsequent /arrow-result fetch reliably hits the entry stashed + * during the SSE query. + * + * `resolveUserId` throws when no header is present — catch and degrade + * to "global" rather than letting that failure mode bubble through the + * route handler. + */ + protected _stashUserKey(req: express.Request): string { + try { + return this.resolveUserId(req) || "global"; + } catch { + return "global"; + } + } + /** * Handle SQL query execution requests. * When called via asUser(req), uses the user's Databricks credentials. @@ -184,6 +253,11 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { // 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"; + // Stash key is always per-request user (never "global"), independent + // of the executor's cache scope. Inline Arrow payloads are single-use + // and short-lived — there is no benefit to sharing them across users, + // and per-user scoping is defense in depth on top of unguessable ids. + const stashUserKey = this._stashUserKey(req); const hashedQuery = this.queryProcessor.hashQuery(query); @@ -231,6 +305,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { query, processedParams, format, + stashUserKey, signal, ); }, @@ -245,6 +320,10 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { * - JSON_ARRAY: always uses INLINE disposition, no fallback. * - ARROW_STREAM: tries INLINE first, falls back to EXTERNAL_LINKS. * This handles warehouses that only support one disposition. + * + * INLINE attachments are decoded once and put on the plugin's + * `inlineArrowStash`; the SSE message carries the synthetic stash id so + * the client fetches the bytes out-of-band via `/arrow-result/`. */ private async _executeWithFormatFallback( executor: AnalyticsPlugin, @@ -253,6 +332,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { | Record | undefined, requestedFormat: AnalyticsFormat, + stashUserKey: string, signal?: AbortSignal, ): Promise { if (requestedFormat === "JSON_ARRAY") { @@ -276,12 +356,21 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { { disposition: "INLINE", format: "ARROW_STREAM" }, signal, ); - // INLINE responses with an Arrow IPC attachment are forwarded as base64 - // for the client to decode into an Arrow Table. Anything else (rare: - // data_array under ARROW_STREAM, or an empty result) falls back to the - // generic "result" payload. + // INLINE responses with an Arrow IPC attachment go through the + // stash-and-serve path: decode the base64 once, hold the bytes + // server-side, emit a synthetic statement id. The client fetches via + // /arrow-result so multi-MiB Arrow blobs never traverse SSE. if (result?.attachment) { - return makeArrowInlineMessage(result.attachment); + const decoded = Buffer.from(result.attachment, "base64"); + const inlineId = this.inlineArrowStash.put( + stashUserKey, + new Uint8Array( + decoded.buffer, + decoded.byteOffset, + decoded.byteLength, + ), + ); + return makeArrowMessage(inlineId, { status: result.status }); } return makeResultMessage(result?.data, { status: result?.status, diff --git a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts new file mode 100644 index 00000000..0fd1a776 --- /dev/null +++ b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts @@ -0,0 +1,129 @@ +import { randomUUID } from "node:crypto"; + +/** + * Server-side stash for inline Arrow IPC payloads. + * + * When a warehouse returns ARROW_STREAM + INLINE results, the bytes are + * stashed here and a synthetic "inline-" job id is emitted on the + * SSE control channel. The client then fetches the bytes out-of-band via + * `/arrow-result/`, which drains the stash and serves the payload as + * `application/vnd.apache.arrow.stream`. + * + * Keeps multi-MiB Arrow blobs off SSE, lets the existing /arrow-result + * pipeline handle both inline and EXTERNAL_LINKS results uniformly, and + * delivers the bytes with a real binary content-type instead of base64 + * inside JSON inside SSE framing. + * + * Properties: + * - **Drain-on-read**: a successful `take()` removes the entry. There is + * no replay path — a lost client connection means the bytes are gone. + * - **TTL bounded**: entries past their expiry are evicted on every + * `put()` and `take()`. No background timer. + * - **Per-user keyed**: `take()` only returns bytes if the requesting + * user matches the user that originally put them. Defense in depth on + * top of unguessable ids. + * - **Memory bounded**: total stashed bytes are capped. `put()` evicts + * the oldest entries first when the cap would be exceeded. + */ +interface InlineArrowStashOptions { + /** Entries older than this are dropped on the next gc tick. */ + ttlMs?: number; + /** Soft cap on total bytes held. Oldest entries are evicted to fit. */ + maxBytes?: number; + /** Test seam: override the synthetic-id generator. */ + idGenerator?: () => string; + /** Test seam: override the clock. */ + now?: () => number; +} + +interface StashEntry { + userId: string; + bytes: Uint8Array; + expiresAt: number; + insertedAt: number; +} + +export class InlineArrowStash { + private entries = new Map(); + private totalBytes = 0; + private readonly ttlMs: number; + private readonly maxBytes: number; + private readonly idGenerator: () => string; + private readonly now: () => number; + + constructor(opts: InlineArrowStashOptions = {}) { + this.ttlMs = opts.ttlMs ?? 10 * 60 * 1000; + this.maxBytes = opts.maxBytes ?? 256 * 1024 * 1024; + this.idGenerator = opts.idGenerator ?? randomUUID; + this.now = opts.now ?? Date.now; + } + + /** Stash a payload and return its synthetic job id. */ + put(userId: string, bytes: Uint8Array): string { + if (bytes.length > this.maxBytes) { + throw new Error( + `Inline Arrow payload (${bytes.length} bytes) exceeds stash maxBytes (${this.maxBytes})`, + ); + } + this.gc(); + this.evictUntilFits(bytes.length); + const id = `inline-${this.idGenerator()}`; + const now = this.now(); + this.entries.set(id, { + userId, + bytes, + expiresAt: now + this.ttlMs, + insertedAt: now, + }); + this.totalBytes += bytes.length; + return id; + } + + /** + * Drain a payload from the stash. Returns `undefined` if the id is + * unknown, expired, or belongs to a different user. + */ + take(id: string, userId: string): Uint8Array | undefined { + this.gc(); + const entry = this.entries.get(id); + if (!entry) return undefined; + if (entry.userId !== userId) return undefined; + this.entries.delete(id); + this.totalBytes -= entry.bytes.length; + return entry.bytes; + } + + /** Inspection helpers (primarily for tests). */ + size(): number { + return this.totalBytes; + } + count(): number { + return this.entries.size; + } + + /** Drop all entries (used in plugin shutdown). */ + clear(): void { + this.entries.clear(); + this.totalBytes = 0; + } + + private gc(): void { + const now = this.now(); + for (const [id, entry] of this.entries) { + if (entry.expiresAt <= now) { + this.entries.delete(id); + this.totalBytes -= entry.bytes.length; + } + } + } + + private evictUntilFits(incoming: number): void { + if (this.totalBytes + incoming <= this.maxBytes) return; + // Insertion order in a Map mirrors FIFO, so the first key is the oldest. + for (const [id, entry] of this.entries) { + if (this.totalBytes + incoming <= this.maxBytes) return; + this.entries.delete(id); + this.totalBytes -= entry.bytes.length; + } + } +} diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 6884e2bb..f3dfd8e4 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -106,6 +106,87 @@ describe("Analytics Plugin", () => { ); }); + test("/arrow-result/inline-* drains the stash and serves bytes as application/vnd.apache.arrow.stream", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const arrowBytes = new Uint8Array([0xff, 0xfe, 0xfd, 0xfc]); + const id = (plugin as any).inlineArrowStash.put("global", arrowBytes); + expect(id.startsWith("inline-")).toBe(true); + + const handler = getHandler("GET", "/arrow-result/:jobId"); + const mockReq = createMockRequest({ params: { jobId: id } }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Content-Type", + "application/vnd.apache.arrow.stream", + ); + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Content-Length", + String(arrowBytes.length), + ); + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Cache-Control", + "no-store", + ); + expect(mockRes.send).toHaveBeenCalledTimes(1); + const sentBuf = (mockRes.send as any).mock.calls[0][0] as Buffer; + expect(Buffer.isBuffer(sentBuf)).toBe(true); + expect(Array.from(sentBuf)).toEqual(Array.from(arrowBytes)); + + // Drain-on-read: a second fetch must return 410, not the bytes again. + const secondRes = createMockResponse(); + await handler(mockReq, secondRes); + expect(secondRes.status).toHaveBeenCalledWith(410); + }); + + test("/arrow-result/inline-* returns 410 when the stash entry never existed", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const handler = getHandler("GET", "/arrow-result/:jobId"); + const mockReq = createMockRequest({ + params: { jobId: "inline-does-not-exist" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(410); + expect(mockRes.json).toHaveBeenCalledWith( + expect.objectContaining({ + error: expect.stringMatching(/expired or unknown/), + }), + ); + }); + + test("/arrow-result/inline-* returns 410 when the stash entry belongs to a different user", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + // Stash entry keyed to user-a, but the request resolves to "global" + // (no x-forwarded-user header) — keys differ, take must return + // nothing, and the entry stays put (single-user view). + const bytes = new Uint8Array([1, 2, 3]); + const id = (plugin as any).inlineArrowStash.put("user-a", bytes); + + const handler = getHandler("GET", "/arrow-result/:jobId"); + const mockReq = createMockRequest({ params: { jobId: id } }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(410); + // The entry must still be there for the real owner. + expect((plugin as any).inlineArrowStash.take(id, "user-a")).toBeDefined(); + }); + test("/query/:query_key should return 400 when query_key is missing", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -889,7 +970,7 @@ describe("Analytics Plugin", () => { } }); - test("/query/:query_key emits arrow_inline SSE event when ARROW_STREAM INLINE returns an attachment", async () => { + test("/query/:query_key stashes ARROW_STREAM INLINE bytes and emits an arrow message with a synthetic inline- id", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -898,7 +979,9 @@ describe("Analytics Plugin", () => { isAsUser: false, }); - const fakeAttachment = "BASE64_ARROW_IPC_BYTES"; + // Real base64 so the route can decode it via Buffer.from(..., "base64"). + 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 }, }); @@ -921,14 +1004,26 @@ describe("Analytics Plugin", () => { disposition: "INLINE", format: "ARROW_STREAM", }); - // SSE payload should use the new arrow_inline message type. + // SSE payload: unified `arrow` message with an inline- prefixed id. + // The base64 attachment must NOT appear on the SSE channel. const writeCalls = (mockRes.write as any).mock.calls.map( (c: any[]) => c[0] as string, ); const payload = writeCalls.find((s: string) => s.startsWith("data: ")); expect(payload).toBeDefined(); - expect(payload).toContain('"type":"arrow_inline"'); - expect(payload).toContain(`"attachment":"${fakeAttachment}"`); + expect(payload).toContain('"type":"arrow"'); + expect(payload).toMatch(/"statement_id":"inline-[^"]+"/); + expect(payload).not.toContain("arrow_inline"); + expect(payload).not.toContain(fakeAttachment); + + // The decoded bytes should be in the stash, keyed by the same + // synthetic id; a subsequent /arrow-result fetch will drain them. + const idMatch = payload?.match(/"statement_id":"(inline-[^"]+)"/); + expect(idMatch).not.toBeNull(); + const inlineId = idMatch![1]; + const stashed = (plugin as any).inlineArrowStash.take(inlineId, "global"); + expect(stashed).toBeDefined(); + expect(Array.from(stashed)).toEqual(Array.from(arrowBytes)); }); test("/query/:query_key rejects unknown format values with 400", async () => { diff --git a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts new file mode 100644 index 00000000..87543199 --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts @@ -0,0 +1,99 @@ +import { describe, expect, test } from "vitest"; +import { InlineArrowStash } from "../inline-arrow-stash"; + +function bytes(n: number): Uint8Array { + return new Uint8Array(n); +} + +describe("InlineArrowStash", () => { + test("put returns an inline-prefixed synthetic id", () => { + const stash = new InlineArrowStash({ idGenerator: () => "abc" }); + const id = stash.put("user-1", bytes(100)); + expect(id).toBe("inline-abc"); + }); + + test("take drains the entry", () => { + const stash = new InlineArrowStash(); + const id = stash.put("user-1", bytes(100)); + expect(stash.count()).toBe(1); + expect(stash.size()).toBe(100); + + const got = stash.take(id, "user-1"); + expect(got).toBeDefined(); + expect(got!.length).toBe(100); + expect(stash.count()).toBe(0); + expect(stash.size()).toBe(0); + // Drain-on-read: second take returns undefined. + expect(stash.take(id, "user-1")).toBeUndefined(); + }); + + test("take returns undefined for unknown id", () => { + const stash = new InlineArrowStash(); + expect(stash.take("inline-nope", "user-1")).toBeUndefined(); + }); + + test("take returns undefined when userId does not match", () => { + const stash = new InlineArrowStash(); + const id = stash.put("user-1", bytes(100)); + expect(stash.take(id, "user-2")).toBeUndefined(); + // Entry is still there for the right user. + expect(stash.take(id, "user-1")).toBeDefined(); + }); + + test("entries past TTL are evicted on next gc tick", () => { + let clock = 0; + const stash = new InlineArrowStash({ ttlMs: 1000, now: () => clock }); + const id = stash.put("user-1", bytes(50)); + clock = 999; + expect(stash.take(id, "user-1")).toBeDefined(); + + const id2 = stash.put("user-1", bytes(50)); + clock = 2000; + // Bump the clock past TTL and trigger gc via another put. + stash.put("user-2", bytes(10)); + expect(stash.take(id2, "user-1")).toBeUndefined(); + }); + + test("put evicts oldest entries to fit when maxBytes is exceeded", () => { + let seq = 0; + const stash = new InlineArrowStash({ + maxBytes: 200, + idGenerator: () => String(seq++), + }); + const a = stash.put("user-1", bytes(80)); + const b = stash.put("user-1", bytes(80)); + expect(stash.size()).toBe(160); + + // This 80-byte entry pushes total to 240; evicts `a` first. + const c = stash.put("user-1", bytes(80)); + expect(stash.size()).toBe(160); + expect(stash.take(a, "user-1")).toBeUndefined(); + expect(stash.take(b, "user-1")).toBeDefined(); + expect(stash.take(c, "user-1")).toBeDefined(); + }); + + test("put rejects a single payload larger than maxBytes", () => { + const stash = new InlineArrowStash({ maxBytes: 100 }); + expect(() => stash.put("user-1", bytes(200))).toThrow( + /exceeds stash maxBytes/, + ); + }); + + test("synthetic ids are unique across puts", () => { + const stash = new InlineArrowStash(); + const a = stash.put("user-1", bytes(10)); + const b = stash.put("user-1", bytes(10)); + expect(a).not.toBe(b); + expect(a.startsWith("inline-")).toBe(true); + expect(b.startsWith("inline-")).toBe(true); + }); + + test("clear drops every entry", () => { + const stash = new InlineArrowStash(); + stash.put("user-1", bytes(10)); + stash.put("user-2", bytes(20)); + stash.clear(); + expect(stash.count()).toBe(0); + expect(stash.size()).toBe(0); + }); +}); diff --git a/packages/appkit/src/stream/defaults.ts b/packages/appkit/src/stream/defaults.ts index 7cde49c4..5cb822ef 100644 --- a/packages/appkit/src/stream/defaults.ts +++ b/packages/appkit/src/stream/defaults.ts @@ -1,11 +1,13 @@ export const streamDefaults = { bufferSize: 100, - // 12 MiB. Headroom for base64-encoded inline Arrow IPC attachments: the - // connector caps the *decoded* attachment at 8 MiB (MAX_INLINE_ATTACHMENT_BYTES), - // which inflates to ~10.6 MiB once base64-encoded and is then wrapped in JSON + - // SSE framing. 12 MiB leaves enough room for that overhead so legal 8-MiB-decoded - // attachments do not trip the stream-manager cap before the connector-level cap. - maxEventSize: 12 * 1024 * 1024, + // 1 MiB. SSE is used only for short JSON control messages — JSON_ARRAY + // result rows (already row-size-bounded by the warehouse) and the small + // `arrow` envelope (statement id + status) for ARROW_STREAM. Bulk Arrow + // payloads do not traverse SSE; they are fetched over HTTP via + // `/api/analytics/arrow-result/:jobId`, which dispatches to the warehouse + // (EXTERNAL_LINKS) or the server-side `InlineArrowStash` (INLINE) based + // on the id prefix. + maxEventSize: 1 * 1024 * 1024, bufferTTL: 10 * 60 * 1000, // 10 minutes cleanupInterval: 5 * 60 * 1000, // 5 minutes maxPersistentBuffers: 10000, // 10000 buffers diff --git a/packages/shared/src/sse/analytics.test.ts b/packages/shared/src/sse/analytics.test.ts index 5abeb83e..f66437c3 100644 --- a/packages/shared/src/sse/analytics.test.ts +++ b/packages/shared/src/sse/analytics.test.ts @@ -1,7 +1,6 @@ import { describe, expect, test } from "vitest"; import { AnalyticsSseMessage, - makeArrowInlineMessage, makeArrowMessage, makeResultMessage, } from "./analytics"; @@ -19,7 +18,7 @@ describe("AnalyticsSseMessage schema", () => { expect(() => AnalyticsSseMessage.parse({ type: "result" })).not.toThrow(); }); - test("accepts an arrow message with statement_id", () => { + test("accepts an arrow message with warehouse statement_id", () => { const parsed = AnalyticsSseMessage.parse({ type: "arrow", statement_id: "stmt-1", @@ -27,6 +26,18 @@ describe("AnalyticsSseMessage schema", () => { expect(parsed.type).toBe("arrow"); }); + test("accepts an arrow message with synthetic inline- id", () => { + // Inline Arrow payloads are stashed server-side and surfaced through the + // same `arrow` message variant — the `inline-` prefix tells the + // /arrow-result handler to drain the stash instead of hitting the + // warehouse. The schema must accept both id shapes transparently. + const parsed = AnalyticsSseMessage.parse({ + type: "arrow", + statement_id: "inline-abc-123", + }); + expect(parsed.statement_id).toBe("inline-abc-123"); + }); + test("rejects an arrow message with empty statement_id", () => { expect(() => AnalyticsSseMessage.parse({ type: "arrow", statement_id: "" }), @@ -37,23 +48,12 @@ describe("AnalyticsSseMessage schema", () => { expect(() => AnalyticsSseMessage.parse({ type: "arrow" })).toThrow(); }); - test("accepts an arrow_inline message with non-empty attachment", () => { - const parsed = AnalyticsSseMessage.parse({ - type: "arrow_inline", - attachment: "AQID", - }); - expect(parsed.type).toBe("arrow_inline"); - }); - - test("rejects an arrow_inline message with empty attachment", () => { - expect(() => - AnalyticsSseMessage.parse({ type: "arrow_inline", attachment: "" }), - ).toThrow(); - }); - - test("rejects an arrow_inline message with non-string attachment", () => { + test("rejects the retired arrow_inline message type", () => { + // arrow_inline was the prior wire shape (base64 payload on the SSE + // channel). The current protocol routes all Arrow payloads through + // /arrow-result; the type must no longer parse. expect(() => - AnalyticsSseMessage.parse({ type: "arrow_inline", attachment: 123 }), + AnalyticsSseMessage.parse({ type: "arrow_inline", attachment: "AQID" }), ).toThrow(); }); @@ -64,7 +64,7 @@ describe("AnalyticsSseMessage schema", () => { }); test("safeParse returns success: false for malformed payloads", () => { - const r = AnalyticsSseMessage.safeParse({ type: "arrow_inline" }); + const r = AnalyticsSseMessage.safeParse({ type: "arrow" }); expect(r.success).toBe(false); }); }); @@ -80,8 +80,8 @@ describe("typed builders", () => { expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); }); - test("makeArrowInlineMessage roundtrips through the schema", () => { - const msg = makeArrowInlineMessage("AQID"); + test("makeArrowMessage accepts synthetic inline- ids", () => { + const msg = makeArrowMessage("inline-some-uuid"); expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); }); }); diff --git a/packages/shared/src/sse/analytics.ts b/packages/shared/src/sse/analytics.ts index a7f68d54..f37d38a5 100644 --- a/packages/shared/src/sse/analytics.ts +++ b/packages/shared/src/sse/analytics.ts @@ -7,13 +7,19 @@ import { z } from "zod"; * server (`AnalyticsPlugin._handleQueryRoute`) and the client * (`useAnalyticsQuery`). Both sides validate with the same schema: * - * - Server uses the typed builders (`makeResultMessage`, `makeArrowMessage`, - * `makeArrowInlineMessage`) to construct messages with compile-time - * guarantees that all required fields are present. + * - Server uses the typed builders (`makeResultMessage`, `makeArrowMessage`) + * to construct messages with compile-time guarantees that all required + * fields are present. * - Client calls `AnalyticsSseMessage.parse(JSON.parse(event.data))` to fail * loudly on a malformed payload instead of silently treating an undefined * field as data. * + * Arrow payloads — inline or external-links — never traverse the SSE control + * channel; both flow through `/api/analytics/arrow-result/:jobId` and are + * differentiated by an `inline-` prefix on the job id (see + * `InlineArrowStash`). The wire shape from the client's perspective is + * therefore uniform: an `arrow` message carries an id, the client fetches. + * * Adding a new message variant requires a schema update here, which keeps * server and client in lockstep. */ @@ -30,10 +36,13 @@ export const AnalyticsResultMessage = z.object({ export type AnalyticsResultMessage = z.infer; /** - * ARROW_STREAM result delivered via /arrow-result/:jobId — used for - * EXTERNAL_LINKS responses (statement_id from the warehouse) and, if PR #320 - * lands, also for INLINE responses (synthetic `inline-` prefixed id from - * the server-side stash). + * ARROW_STREAM result delivered via /arrow-result/:jobId. The id is either: + * - the warehouse-issued `statement_id` for EXTERNAL_LINKS responses, or + * - a synthetic `inline-` id pointing at the server-side + * `InlineArrowStash` for INLINE responses. + * + * Both shapes are fetched the same way; the prefix tells the route handler + * which path to take. */ export const AnalyticsArrowMessage = z.object({ type: z.literal("arrow"), @@ -42,26 +51,10 @@ export const AnalyticsArrowMessage = z.object({ }); export type AnalyticsArrowMessage = z.infer; -/** - * ARROW_STREAM + INLINE result with the base64-encoded Arrow IPC bytes - * embedded in the SSE message. The client decodes locally via - * `ArrowClient.processArrowBuffer`. - * - * Note: this variant goes away if the proposal in PR #320 lands. - */ -export const AnalyticsArrowInlineMessage = z.object({ - type: z.literal("arrow_inline"), - attachment: z.string().min(1), -}); -export type AnalyticsArrowInlineMessage = z.infer< - typeof AnalyticsArrowInlineMessage ->; - /** Discriminated union of every message the analytics SSE stream may emit. */ export const AnalyticsSseMessage = z.discriminatedUnion("type", [ AnalyticsResultMessage, AnalyticsArrowMessage, - AnalyticsArrowInlineMessage, ]); export type AnalyticsSseMessage = z.infer; @@ -84,9 +77,3 @@ export function makeArrowMessage( ): AnalyticsArrowMessage { return { type: "arrow", statement_id, ...extras }; } - -export function makeArrowInlineMessage( - attachment: string, -): AnalyticsArrowInlineMessage { - return { type: "arrow_inline", attachment }; -} From f34e18ec23e375a7309d603c26af2e68e7f95d26 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 15:59:09 +0000 Subject: [PATCH 10/25] fix: address ACE multi-model review on the inline-stash redesign Four findings surfaced by the GPT pass on the reworked PR: 1. ARROW_STREAM cache replay returned drained inline-* ids (HIGH). The previous code capped the cache TTL at 10 min for ARROW_STREAM, which made sense for EXTERNAL_LINKS pre-signed URLs that expire in ~15 min but is broken for inline ids: the stash drains on the first /arrow-result fetch, so any cache hit replays an id whose bytes are gone and reliably 410s. Bypass cache entirely for ARROW_STREAM (TTL = 0); JSON_ARRAY responses still cache normally. 2. Stash evict-on-fit invalidated already-issued ids (MEDIUM). The earlier `evictUntilFits` dropped the oldest entries when a new payload would push total bytes past `maxBytes`, but those oldest entries had ids that were already in flight to clients. Replace eviction with rejection: `put()` now returns `string | null` and the caller falls back to EXTERNAL_LINKS when the stash is full. Every id we hand out stays valid until naturally drained or expired. 3. Aborted stream still decoded + stashed (MEDIUM). If the client cancels the SSE between query completion and stash write, we still decoded the base64 attachment and held the bytes until TTL eviction. Re-check `signal.aborted` before decode/put so canceled streams exit cleanly. 4. Empty result message wrote `undefined` to the hook's state (LOW). The wire schema makes `data` optional; an empty result set may omit it. Normalize the missing case to `[]` so consumers can rely on `data` being either `null` (no message yet) or a value of the inferred result type. Also documents the process-local-memory constraint on the stash in its docstring: a `GET /arrow-result/inline-*` that lands on a different replica than the original SSE request will 410. Multi-replica deployments need sticky sessions or a shared external store, neither in scope for this PR. Tests: - `inline-arrow-stash`: replaced the eviction test with a rejection test that asserts `put()` returns null when the stash is full and that previously-issued ids remain takeable. - `useAnalyticsQuery`: new test asserts an empty result message normalizes to []. Signed-off-by: James Broadhead --- .../__tests__/use-analytics-query.test.ts | 20 ++++++++ .../src/react/hooks/use-analytics-query.ts | 7 ++- .../appkit/src/plugins/analytics/analytics.ts | 47 +++++++++++++------ .../plugins/analytics/inline-arrow-stash.ts | 44 +++++++++++------ .../tests/inline-arrow-stash.test.ts | 16 ++++--- 5 files changed, 97 insertions(+), 37 deletions(-) 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 c2705476..1e748346 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 @@ -131,6 +131,26 @@ describe("useAnalyticsQuery", () => { expect(mockFetchArrow).not.toHaveBeenCalled(); }); + 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" }), 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 a192adc2..88419313 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -131,10 +131,13 @@ export function useAnalyticsQuery< const validated = AnalyticsSseMessage.safeParse(rawParsed); const msg = validated.success ? validated.data : null; - // success - JSON format + // success - JSON format. The wire schema makes `data` optional + // (e.g. an empty result set may omit it), so normalize the + // missing case to an explicit empty array rather than letting + // `undefined` bleed into the hook's `T | null` state. if (msg?.type === "result") { setLoading(false); - setData(msg.data as ResultType); + setData((msg.data ?? []) as ResultType); return; } diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index b99fb9a2..d38077d8 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -261,15 +261,17 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { const hashedQuery = this.queryProcessor.hashQuery(query); - // ARROW_STREAM may resolve to EXTERNAL_LINKS, which returns pre-signed URLs - // that typically expire ~15 minutes after issue. Cap the cache TTL well - // under that for ARROW_STREAM so we never hand out dead URLs from cache, - // while still benefiting from caching INLINE attachment responses (and - // EXTERNAL_LINKS responses inside their valid window). - const cacheTtl = - format === "ARROW_STREAM" - ? Math.min(queryDefaults.cache?.ttl ?? 600, 600) - : queryDefaults.cache?.ttl; + // ARROW_STREAM responses reference ephemeral resources that cannot be + // safely replayed from cache: + // - EXTERNAL_LINKS pre-signed URLs expire ~15 min after issue, and + // the warehouse rotates them per execution. + // - INLINE responses point at a synthetic `inline-` job id + // backed by `InlineArrowStash`, which drains on the first + // /arrow-result fetch. A cache hit would replay an id whose bytes + // are already gone and reliably 410 the client. + // So we bypass cache for ARROW_STREAM and let every request execute + // a fresh statement. JSON_ARRAY responses still cache normally. + const cacheTtl = format === "ARROW_STREAM" ? 0 : queryDefaults.cache?.ttl; const cacheConfig = { ...queryDefaults.cache, ttl: cacheTtl, @@ -361,6 +363,12 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { // server-side, emit a synthetic statement id. The client fetches via // /arrow-result so multi-MiB Arrow blobs never traverse SSE. if (result?.attachment) { + // If the client has already disconnected, the SSE write would be + // dropped anyway — skip the decode + stash so the bytes do not + // linger in memory until TTL eviction. + if (signal?.aborted) { + throw ExecutionError.canceled(); + } const decoded = Buffer.from(result.attachment, "base64"); const inlineId = this.inlineArrowStash.put( stashUserKey, @@ -370,12 +378,23 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { decoded.byteLength, ), ); - return makeArrowMessage(inlineId, { status: result.status }); + if (inlineId === null) { + // Stash is full — every id we have already handed out must + // stay valid, so the stash refuses new entries rather than + // evicting in-flight ones. Fall back to EXTERNAL_LINKS for + // this request so the client still gets its result. + logger.warn( + "Inline Arrow stash full, falling back to EXTERNAL_LINKS for the current query", + ); + } else { + return makeArrowMessage(inlineId, { status: result.status }); + } + } else { + return makeResultMessage(result?.data, { + status: result?.status, + statement_id: result?.statement_id, + }); } - return makeResultMessage(result?.data, { - status: result?.status, - statement_id: result?.statement_id, - }); } catch (err: unknown) { // If the request was aborted, do not retry — the signal is dead and // a second statement would be billed but never read. diff --git a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts index 0fd1a776..2eda5ecb 100644 --- a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts +++ b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts @@ -22,8 +22,19 @@ import { randomUUID } from "node:crypto"; * - **Per-user keyed**: `take()` only returns bytes if the requesting * user matches the user that originally put them. Defense in depth on * top of unguessable ids. - * - **Memory bounded**: total stashed bytes are capped. `put()` evicts - * the oldest entries first when the cap would be exceeded. + * - **Memory bounded with rejection**: total stashed bytes are capped. + * When `put()` cannot fit a payload without exceeding the cap it + * returns `null` rather than evicting older entries — every issued id + * stays valid until it is drained, expires, or the process exits. + * Callers are expected to fall back to a different delivery path (e.g. + * EXTERNAL_LINKS) when `put()` rejects. + * + * Caveat (multi-replica deployments): this stash is process-local. A + * subsequent `GET /arrow-result/inline-*` that lands on a different + * replica than the one that stashed the bytes will 410. Deployments + * that run more than one replica need sticky sessions (route both + * requests in the same logical session to the same replica) or a + * shared external store, neither of which is in scope here. */ interface InlineArrowStashOptions { /** Entries older than this are dropped on the next gc tick. */ @@ -58,15 +69,28 @@ export class InlineArrowStash { this.now = opts.now ?? Date.now; } - /** Stash a payload and return its synthetic job id. */ - put(userId: string, bytes: Uint8Array): string { + /** + * Stash a payload and return its synthetic job id, or `null` when the + * stash cannot accept it without evicting older entries. The caller is + * expected to fall back to an out-of-band delivery path (e.g. + * EXTERNAL_LINKS) when the return value is `null`. + * + * Single payloads that exceed `maxBytes` outright throw so the caller + * sees the misconfiguration loudly instead of degrading silently every + * time. + */ + put(userId: string, bytes: Uint8Array): string | null { if (bytes.length > this.maxBytes) { throw new Error( `Inline Arrow payload (${bytes.length} bytes) exceeds stash maxBytes (${this.maxBytes})`, ); } this.gc(); - this.evictUntilFits(bytes.length); + if (this.totalBytes + bytes.length > this.maxBytes) { + // Refuse rather than evicting: every id we have already issued must + // remain valid until naturally drained or expired. + return null; + } const id = `inline-${this.idGenerator()}`; const now = this.now(); this.entries.set(id, { @@ -116,14 +140,4 @@ export class InlineArrowStash { } } } - - private evictUntilFits(incoming: number): void { - if (this.totalBytes + incoming <= this.maxBytes) return; - // Insertion order in a Map mirrors FIFO, so the first key is the oldest. - for (const [id, entry] of this.entries) { - if (this.totalBytes + incoming <= this.maxBytes) return; - this.entries.delete(id); - this.totalBytes -= entry.bytes.length; - } - } } diff --git a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts index 87543199..60e48218 100644 --- a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts @@ -54,7 +54,7 @@ describe("InlineArrowStash", () => { expect(stash.take(id2, "user-1")).toBeUndefined(); }); - test("put evicts oldest entries to fit when maxBytes is exceeded", () => { + test("put returns null when adding the payload would exceed maxBytes, leaving existing entries intact", () => { let seq = 0; const stash = new InlineArrowStash({ maxBytes: 200, @@ -62,17 +62,21 @@ describe("InlineArrowStash", () => { }); const a = stash.put("user-1", bytes(80)); const b = stash.put("user-1", bytes(80)); + expect(a).not.toBeNull(); + expect(b).not.toBeNull(); expect(stash.size()).toBe(160); - // This 80-byte entry pushes total to 240; evicts `a` first. + // This third 80-byte entry would push total to 240 (>200). It must + // be rejected, and both prior entries must survive — every id we have + // already handed out stays valid until drained or expired. const c = stash.put("user-1", bytes(80)); + expect(c).toBeNull(); expect(stash.size()).toBe(160); - expect(stash.take(a, "user-1")).toBeUndefined(); - expect(stash.take(b, "user-1")).toBeDefined(); - expect(stash.take(c, "user-1")).toBeDefined(); + expect(stash.take(a as string, "user-1")).toBeDefined(); + expect(stash.take(b as string, "user-1")).toBeDefined(); }); - test("put rejects a single payload larger than maxBytes", () => { + test("put throws for a single payload larger than maxBytes (caller misconfiguration)", () => { const stash = new InlineArrowStash({ maxBytes: 100 }); expect(() => stash.put("user-1", bytes(200))).toThrow( /exceeds stash maxBytes/, From 09f801f48872373ba371bb9afb9ef7b554ca8f05 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 15:59:36 +0000 Subject: [PATCH 11/25] docs(stash): correct maxBytes comment after switch to reject-on-full Signed-off-by: James Broadhead --- packages/appkit/src/plugins/analytics/inline-arrow-stash.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts index 2eda5ecb..3ae98330 100644 --- a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts +++ b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts @@ -39,7 +39,11 @@ import { randomUUID } from "node:crypto"; interface InlineArrowStashOptions { /** Entries older than this are dropped on the next gc tick. */ ttlMs?: number; - /** Soft cap on total bytes held. Oldest entries are evicted to fit. */ + /** + * Hard cap on total bytes held. `put()` rejects (returns `null`) once + * the cap would be exceeded; entries already in the stash are not + * evicted to fit new ones. + */ maxBytes?: number; /** Test seam: override the synthetic-id generator. */ idGenerator?: () => string; From 698a26419c3ff514c77d71a7983227a8a8338b77 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Tue, 12 May 2026 16:12:54 +0000 Subject: [PATCH 12/25] style: drop unused imports and tidy stash test types - agents.ts had two unused imports that biome's noUnusedImports rule flags as errors in CI. Drop them; behavior unchanged. - inline-arrow-stash.test.ts: introduce a mustPut() helper that asserts the non-null contract for successful puts, so the new `put(): string | null` return type does not poison every downstream take() call with a string-vs-string|null TS error. - Minor formatter touch-ups picked up by biome --write. Signed-off-by: James Broadhead --- packages/appkit/src/plugins/agents/agents.ts | 2 - .../plugins/agents/tests/dos-limits.test.ts | 2 +- .../plugins/analytics/tests/analytics.test.ts | 2 +- .../tests/inline-arrow-stash.test.ts | 47 ++++++++++++------- .../src/stream/tests/stream-registry.test.ts | 2 +- 5 files changed, 33 insertions(+), 22 deletions(-) diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 3c20d616..87a46d34 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, @@ -16,7 +15,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 e2bbcbe9..a0c64e57 100644 --- a/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dos-limits.test.ts @@ -307,7 +307,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/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index f3dfd8e4..b4522eab 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -1020,7 +1020,7 @@ describe("Analytics Plugin", () => { // synthetic id; a subsequent /arrow-result fetch will drain them. const idMatch = payload?.match(/"statement_id":"(inline-[^"]+)"/); expect(idMatch).not.toBeNull(); - const inlineId = idMatch![1]; + const inlineId = idMatch?.[1]; const stashed = (plugin as any).inlineArrowStash.take(inlineId, "global"); expect(stashed).toBeDefined(); expect(Array.from(stashed)).toEqual(Array.from(arrowBytes)); diff --git a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts index 60e48218..9bd59813 100644 --- a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts @@ -5,22 +5,37 @@ function bytes(n: number): Uint8Array { return new Uint8Array(n); } +// `put()` returns `string | null` — it rejects with null when the stash is +// full. Every test below that exercises a successful put narrows via this +// helper so the non-null contract is explicit at the call site. +function mustPut( + stash: InlineArrowStash, + userId: string, + b: Uint8Array, +): string { + const id = stash.put(userId, b); + if (id === null) { + throw new Error("test setup: stash unexpectedly rejected put"); + } + return id; +} + describe("InlineArrowStash", () => { test("put returns an inline-prefixed synthetic id", () => { const stash = new InlineArrowStash({ idGenerator: () => "abc" }); - const id = stash.put("user-1", bytes(100)); + const id = mustPut(stash, "user-1", bytes(100)); expect(id).toBe("inline-abc"); }); test("take drains the entry", () => { const stash = new InlineArrowStash(); - const id = stash.put("user-1", bytes(100)); + const id = mustPut(stash, "user-1", bytes(100)); expect(stash.count()).toBe(1); expect(stash.size()).toBe(100); const got = stash.take(id, "user-1"); expect(got).toBeDefined(); - expect(got!.length).toBe(100); + expect(got?.length).toBe(100); expect(stash.count()).toBe(0); expect(stash.size()).toBe(0); // Drain-on-read: second take returns undefined. @@ -34,7 +49,7 @@ describe("InlineArrowStash", () => { test("take returns undefined when userId does not match", () => { const stash = new InlineArrowStash(); - const id = stash.put("user-1", bytes(100)); + const id = mustPut(stash, "user-1", bytes(100)); expect(stash.take(id, "user-2")).toBeUndefined(); // Entry is still there for the right user. expect(stash.take(id, "user-1")).toBeDefined(); @@ -43,14 +58,14 @@ describe("InlineArrowStash", () => { test("entries past TTL are evicted on next gc tick", () => { let clock = 0; const stash = new InlineArrowStash({ ttlMs: 1000, now: () => clock }); - const id = stash.put("user-1", bytes(50)); + const id = mustPut(stash, "user-1", bytes(50)); clock = 999; expect(stash.take(id, "user-1")).toBeDefined(); - const id2 = stash.put("user-1", bytes(50)); + const id2 = mustPut(stash, "user-1", bytes(50)); clock = 2000; // Bump the clock past TTL and trigger gc via another put. - stash.put("user-2", bytes(10)); + mustPut(stash, "user-2", bytes(10)); expect(stash.take(id2, "user-1")).toBeUndefined(); }); @@ -60,10 +75,8 @@ describe("InlineArrowStash", () => { maxBytes: 200, idGenerator: () => String(seq++), }); - const a = stash.put("user-1", bytes(80)); - const b = stash.put("user-1", bytes(80)); - expect(a).not.toBeNull(); - expect(b).not.toBeNull(); + const a = mustPut(stash, "user-1", bytes(80)); + const b = mustPut(stash, "user-1", bytes(80)); expect(stash.size()).toBe(160); // This third 80-byte entry would push total to 240 (>200). It must @@ -72,8 +85,8 @@ describe("InlineArrowStash", () => { const c = stash.put("user-1", bytes(80)); expect(c).toBeNull(); expect(stash.size()).toBe(160); - expect(stash.take(a as string, "user-1")).toBeDefined(); - expect(stash.take(b as string, "user-1")).toBeDefined(); + expect(stash.take(a, "user-1")).toBeDefined(); + expect(stash.take(b, "user-1")).toBeDefined(); }); test("put throws for a single payload larger than maxBytes (caller misconfiguration)", () => { @@ -85,8 +98,8 @@ describe("InlineArrowStash", () => { test("synthetic ids are unique across puts", () => { const stash = new InlineArrowStash(); - const a = stash.put("user-1", bytes(10)); - const b = stash.put("user-1", bytes(10)); + const a = mustPut(stash, "user-1", bytes(10)); + const b = mustPut(stash, "user-1", bytes(10)); expect(a).not.toBe(b); expect(a.startsWith("inline-")).toBe(true); expect(b.startsWith("inline-")).toBe(true); @@ -94,8 +107,8 @@ describe("InlineArrowStash", () => { test("clear drops every entry", () => { const stash = new InlineArrowStash(); - stash.put("user-1", bytes(10)); - stash.put("user-2", bytes(20)); + mustPut(stash, "user-1", bytes(10)); + mustPut(stash, "user-2", bytes(20)); stash.clear(); expect(stash.count()).toBe(0); expect(stash.size()).toBe(0); 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", From 8f0e31c4df00dba260259170adafbbe8b3fbf024 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Fri, 15 May 2026 15:25:50 +0000 Subject: [PATCH 13/25] test(analytics): cover stash-full fallback to EXTERNAL_LINKS When the inline Arrow stash refuses a new entry (put returns null), the route must retry the statement with EXTERNAL_LINKS instead of emitting a useless inline- id. The stash itself was unit-tested already; this adds the integration test through the /query route. Signed-off-by: James Broadhead --- .../plugins/analytics/tests/analytics.test.ts | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index b4522eab..fe64f2f5 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -1026,6 +1026,72 @@ describe("Analytics Plugin", () => { expect(Array.from(stashed)).toEqual(Array.from(arrowBytes)); }); + test("/query/:query_key falls back to EXTERNAL_LINKS when the inline stash is full", async () => { + // When the stash refuses a new entry (put returns null), the route + // must not strand the client with a useless inline- id. It retries + // the same statement with EXTERNAL_LINKS so the warehouse hands + // back a real, fetchable statement id and the client gets bytes. + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const fakeAttachment = Buffer.from(new Uint8Array([1, 2, 3])).toString( + "base64", + ); + const executeMock = vi + .fn() + // First call: INLINE succeeds and produces an attachment. + .mockResolvedValueOnce({ + result: { attachment: fakeAttachment, row_count: 1 }, + }) + // Second call: EXTERNAL_LINKS — returns a real warehouse id. + .mockResolvedValueOnce({ + result: { + statement_id: "stmt-warehouse-real", + status: { state: "SUCCEEDED" }, + }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + // Force the stash to reject the put — simulates capacity exhaustion. + vi.spyOn((plugin as any).inlineArrowStash, "put").mockReturnValue(null); + + 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[0][1]).toMatchObject({ + disposition: "INLINE", + format: "ARROW_STREAM", + }); + expect(executeMock.mock.calls[1][1]).toMatchObject({ + disposition: "EXTERNAL_LINKS", + format: "ARROW_STREAM", + }); + + // SSE payload carries the real warehouse statement id, not an inline- id. + const writeCalls = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as string, + ); + const payload = writeCalls.find((s: string) => s.startsWith("data: ")); + expect(payload).toBeDefined(); + expect(payload).toContain('"type":"arrow"'); + expect(payload).toContain('"statement_id":"stmt-warehouse-real"'); + expect(payload).not.toMatch(/"statement_id":"inline-/); + }); + test("/query/:query_key rejects unknown format values with 400", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); From 0f5022fb47de5f469e3a912654008914df032cd2 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Fri, 15 May 2026 16:05:19 +0000 Subject: [PATCH 14/25] feat(appkit): retry JSON_ARRAY as ARROW_STREAM on inline-arrow-only warehouses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Some serverless warehouse variants only support ARROW_STREAM for the INLINE disposition — JSON_ARRAY + INLINE is rejected with 'Inline disposition only supports ARROW_STREAM format.' Before this change, every default useAnalyticsQuery call against such a warehouse failed. The plugin now classifies inline rejections into 'needs-arrow' vs 'needs-json' signals. For a JSON_ARRAY caller hitting needs-arrow, the plugin retries as ARROW_STREAM + INLINE and decodes the Arrow IPC attachment back to plain row objects server-side, keeping the caller's JSON_ARRAY contract intact (scalar values stringified to match the warehouse's native JSON_ARRAY shape). The existing ARROW_STREAM + INLINE → EXTERNAL_LINKS path now uses the same classifier with the 'needs-json' signal. Matching is case-insensitive and handles the real warehouse error wordings rather than the case-sensitive 'INLINE' + 'ARROW_STREAM' substring pair the old heuristic required, which never matched the actual wire errors. Verified live against three e2-dogfood warehouses: one that refuses JSON_ARRAY + INLINE, one other serverless, and one classic — all three now produce identical JSON row output for the same SELECT. Signed-off-by: James Broadhead --- .../appkit/src/plugins/analytics/analytics.ts | 190 ++++++++++++++---- .../plugins/analytics/tests/analytics.test.ts | 158 ++++++++++++++- 2 files changed, 304 insertions(+), 44 deletions(-) diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index e61e5be4..b3a268eb 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1,4 +1,5 @@ import type { WorkspaceClient } from "@databricks/sdk-experimental"; +import { tableFromIPC } from "apache-arrow"; import type express from "express"; import { type AgentToolDefinition, @@ -317,15 +318,21 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } /** - * Execute a query with automatic disposition fallback for ARROW_STREAM. + * Execute a query with automatic disposition/format fallback. * - * - JSON_ARRAY: always uses INLINE disposition, no fallback. - * - ARROW_STREAM: tries INLINE first, falls back to EXTERNAL_LINKS. - * This handles warehouses that only support one disposition. + * - **JSON_ARRAY** first tries `INLINE + JSON_ARRAY`. If the warehouse + * only supports `ARROW_STREAM` for `INLINE` (some serverless variants), + * retries as `INLINE + ARROW_STREAM`, decodes the Arrow IPC attachment + * server-side, and returns plain row objects — the caller's + * `JSON_ARRAY` contract is preserved. + * - **ARROW_STREAM** first tries `INLINE + ARROW_STREAM`. If the + * warehouse refuses (most classic + some serverless variants), or the + * inline stash is full, falls back to `EXTERNAL_LINKS + ARROW_STREAM`. * - * INLINE attachments are decoded once and put on the plugin's - * `inlineArrowStash`; the SSE message carries the synthetic stash id so - * the client fetches the bytes out-of-band via `/arrow-result/`. + * INLINE Arrow attachments under the ARROW_STREAM path are decoded once + * and put on the plugin's `inlineArrowStash`; the SSE message carries the + * synthetic stash id so the client fetches the bytes out-of-band via + * `/arrow-result/`. */ private async _executeWithFormatFallback( executor: AnalyticsPlugin, @@ -338,15 +345,44 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { signal?: AbortSignal, ): Promise { if (requestedFormat === "JSON_ARRAY") { - const result = await executor.query( + try { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "JSON_ARRAY" }, + signal, + ); + return makeResultMessage(result?.data, { + status: result?.status, + statement_id: result?.statement_id, + }); + } catch (err: unknown) { + if (signal?.aborted) throw err; + if (_classifyInlineRejection(err) !== "needs-arrow") throw err; + + const msg = err instanceof Error ? err.message : String(err); + logger.warn( + "JSON_ARRAY INLINE rejected by warehouse, retrying as ARROW_STREAM INLINE and decoding server-side: %s", + msg, + ); + } + + // Retry as ARROW_STREAM + INLINE so the warehouse will accept the + // request, then decode the Arrow IPC attachment to plain row + // objects so the caller still gets JSON_ARRAY-shaped data. + const arrowResult = await executor.query( query, processedParams, - { disposition: "INLINE", format: "JSON_ARRAY" }, + { disposition: "INLINE", format: "ARROW_STREAM" }, signal, ); - return makeResultMessage(result?.data, { - status: result?.status, - statement_id: result?.statement_id, + if (!arrowResult?.attachment) { + throw ExecutionError.missingData("ARROW_STREAM attachment"); + } + const rows = decodeArrowAttachmentToRows(arrowResult.attachment); + return makeResultMessage(rows, { + status: arrowResult.status, + statement_id: arrowResult.statement_id, }); } @@ -402,7 +438,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { throw err; } - if (!_isInlineArrowUnsupported(err)) { + if (_classifyInlineRejection(err) !== "needs-json") { throw err; } @@ -539,45 +575,115 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } /** - * Determine whether a warehouse error indicates that ARROW_STREAM + INLINE - * is unsupported, vs an unrelated SQL/permission error. + * Decode a base64 Arrow IPC attachment to plain row objects. * - * Preferred path: read the structured `errorCode` we now propagate from the - * SDK's `ApiError.errorCode` and the warehouse's `status.error.error_code` - * through `ExecutionError`. This is stable across error-message wording - * changes. + * Used by the JSON_ARRAY fallback path when a warehouse refuses + * `JSON_ARRAY + INLINE` and we have to satisfy the request via + * `ARROW_STREAM + INLINE` — the bytes come back as Arrow IPC but the + * caller's contract is JSON-shaped rows, so we convert server-side. * - * Substring backstop: if the upstream error didn't surface a code (legacy - * SDK builds, or errors thrown outside the connector's wrap path), fall - * back to requiring both INLINE and ARROW_STREAM keywords in the message - * plus a marker phrase. The pair-requirement avoids matching unrelated SQL - * errors that happen to mention one of the words (e.g. a column named - * `INLINE_USERS`). + * Scalar values are stringified to match what the warehouse itself emits + * for INT/BIGINT/etc. columns under the JSON_ARRAY format (everything in + * `result.data_array` is a string on the wire) — so callers see the same + * row shape regardless of which path the bytes took. BigInts get the same + * stringification treatment (also necessary for JSON-serializability). */ -function _isInlineArrowUnsupported(err: unknown): boolean { +function decodeArrowAttachmentToRows( + attachment: string, +): Record[] { + const decoded = Buffer.from(attachment, "base64"); + const table = tableFromIPC( + new Uint8Array(decoded.buffer, decoded.byteOffset, decoded.byteLength), + ); + const colNames = table.schema.fields.map((f) => f.name); + const rows: Record[] = []; + for (let i = 0; i < table.numRows; i++) { + const row: Record = {}; + for (const name of colNames) { + const col = table.getChild(name); + const v = col?.get(i); + if (v == null) { + row[name] = null; + } else if ( + typeof v === "number" || + typeof v === "bigint" || + typeof v === "boolean" + ) { + row[name] = String(v); + } else if (typeof v === "string") { + row[name] = v; + } else { + // Nested types (List, Struct, Map) — leave as-is. The JSON_ARRAY + // wire format renders these as JSON strings server-side, but that + // serialization isn't exposed to us here. Round-tripping through + // JSON.stringify would mismatch, so pass the typed value through. + row[name] = v; + } + } + rows.push(row); + } + return rows; +} + +/** + * Classify a warehouse rejection of an INLINE statement. + * + * Two distinct rejection modes are observed in the wild: + * + * - **needs-arrow**: warehouse refuses `JSON_ARRAY + INLINE`, only accepts + * `ARROW_STREAM + INLINE`. Example message: + * `"Inline disposition only supports ARROW_STREAM format."` + * Action: retry as `ARROW_STREAM + INLINE` and decode server-side. + * + * - **needs-json**: warehouse refuses `ARROW_STREAM + INLINE`, only accepts + * `JSON_ARRAY + INLINE`. Examples: + * `"The format field must be JSON_ARRAY when the disposition field is INLINE."` + * `"ARROW_STREAM not supported with INLINE disposition"` + * `"ExternalLinks disposition is not yet implemented."` (same family — + * the warehouse rejected the disposition/format combo we sent). + * Action: retry as `ARROW_STREAM + EXTERNAL_LINKS`. + * + * The structured `errorCode` (INVALID_PARAMETER_VALUE / NOT_IMPLEMENTED) + * gates the classification so unrelated SQL errors don't trigger a retry. + * Message matching is case-insensitive — warehouses are inconsistent about + * casing of "Inline"/"INLINE". + */ +type InlineRejection = "needs-arrow" | "needs-json" | null; + +function _classifyInlineRejection(err: unknown): InlineRejection { + const msg = err instanceof Error ? err.message : String(err); + const lower = msg.toLowerCase(); + const structuredCode = err instanceof ExecutionError ? err.errorCode : undefined; - if ( + const hasCode = structuredCode === "INVALID_PARAMETER_VALUE" || - structuredCode === "NOT_IMPLEMENTED" + structuredCode === "NOT_IMPLEMENTED" || + lower.includes("invalid_parameter_value") || + lower.includes("not_implemented"); + if (!hasCode) return null; + + // Must mention the inline disposition to count as a disposition-rejection. + if (!lower.includes("inline")) return null; + + // "needs-arrow": warehouse only supports ARROW_STREAM for INLINE. + if ( + /only supports\s+arrow_stream/i.test(msg) || + /must be\s+arrow_stream/i.test(msg) ) { - // Structured code already tells us the warehouse rejected the request. - // Require keyword pairing to confirm it's the disposition/format combo - // (vs an INVALID_PARAMETER_VALUE for something else entirely). - const msg = err instanceof Error ? err.message : String(err); - return msg.includes("INLINE") && msg.includes("ARROW_STREAM"); + return "needs-arrow"; } - // Backstop for errors without a structured code. - const msg = err instanceof Error ? err.message : String(err); - if (!msg.includes("INLINE") || !msg.includes("ARROW_STREAM")) { - return false; + // "needs-json": warehouse only supports JSON_ARRAY for 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"; } - return ( - msg.includes("not supported") || - msg.includes("INVALID_PARAMETER_VALUE") || - msg.includes("NOT_IMPLEMENTED") - ); + + return null; } /** diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index fe64f2f5..0683c44c 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -1092,6 +1092,155 @@ describe("Analytics Plugin", () => { expect(payload).not.toMatch(/"statement_id":"inline-/); }); + 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, + ); + const payload = writeCalls.find((s: string) => s.startsWith("data: ")); + 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 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, + ); + const payload = writeCalls.find((s: string) => s.startsWith("data: ")); + 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(); @@ -1158,7 +1307,11 @@ describe("Analytics Plugin", () => { expect(executeMock).toHaveBeenCalledTimes(1); }); - test("/query/:query_key should not fall back when format is explicitly JSON_ARRAY", async () => { + 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(); @@ -1185,7 +1338,8 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); - // All calls use JSON_ARRAY + INLINE — explicit JSON_ARRAY, no fallback. + // 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", From 4afce1f6d688d20152f8808a802345680ed61678 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Thu, 21 May 2026 16:47:32 +0000 Subject: [PATCH 15/25] fix(appkit): address Xavier v3 review on inline-Arrow path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses Mario's v3 review on #329 — 1 critical + 3 high findings: - 🔴 critical: stash-full double-execution + divergent result. The ARROW_STREAM INLINE fallback used to discard the already-decoded bytes when the stash was at cap and re-run the same statement on EXTERNAL_LINKS — that path billed the warehouse twice and could return a divergent result for non-deterministic SQL. The stash now has a second pool ("overflow") sized at `maxOverflowBytes` (default same as `maxBytes`). When the regular pool is full, decoded bytes spill into overflow rather than being thrown away. Only when both pools are at cap does the route surface `INLINE_ARROW_STASH_EXHAUSTED` — single execution, never silent double-billing. Memory is bounded above by `maxBytes + maxOverflowBytes`. - 🟠 high: warehouse / SDK error text reflected verbatim to clients (CWE-209). Added `clientMessage` to AppKitError with a sanitized default per subclass, and a stable `errorCode` field on the SSE error payload. Stream-manager and the /arrow-result 410/404 responses now emit `clientMessage` (sanitized) over the wire; raw upstream wording stays in server logs only. UI can branch on `errorCode` (e.g. `INLINE_ARROW_STASH_EXHAUSTED`) instead of substring-matching the human string. - 🟠 high: JSON_ARRAY fallback shape divergence + heavy materialization. `decodeArrowAttachmentToRows` now `JSON.stringify`s nested values (List / Struct / Map) to match what the warehouse emits natively under JSON_ARRAY — apache-arrow's typed values expose `toJSON()`, so the output shape matches. Added a hard 100k-row cap on the fallback materializer; past the cap surfaces `RESULT_TOO_LARGE_FOR_JSON_FALLBACK` with guidance to re-issue with ARROW_STREAM, rather than blocking the Node event loop for hundreds of ms. - 🟠 high: client `safeParse` O(rows × keys) Zod validation. Loosened `AnalyticsResultMessage.data` to `z.array(z.unknown())` — the per-row shape is already enforced at the source by the typed `makeResultMessage` builder, so the deep client-side validation was pure cost. TS-level type narrows back to `Record[]` for callers. Also a medium from the review: - `JSON.parse` failure in `useAnalyticsQuery` no longer strands the hook in `loading=true` — the outer catch now clears loading and surfaces a user-facing error. Co-authored-by: Isaac Signed-off-by: James Broadhead --- .../__tests__/use-analytics-query.test.ts | 47 ++++++++ .../src/react/hooks/use-analytics-query.ts | 7 ++ packages/appkit/src/errors/base.ts | 28 ++++- packages/appkit/src/errors/execution.ts | 49 +++++++- .../appkit/src/plugins/analytics/analytics.ts | 75 ++++++++---- .../plugins/analytics/inline-arrow-stash.ts | 96 +++++++++++---- .../plugins/analytics/tests/analytics.test.ts | 109 +++++++++++++----- .../tests/inline-arrow-stash.test.ts | 47 ++++++-- packages/appkit/src/stream/sse-writer.ts | 2 + packages/appkit/src/stream/stream-manager.ts | 42 ++++++- .../appkit/src/stream/tests/stream.test.ts | 28 +++-- packages/appkit/src/stream/types.ts | 6 + packages/shared/src/sse/analytics.ts | 20 +++- 13 files changed, 455 insertions(+), 101 deletions(-) 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 a3ec3190..078c2aad 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 @@ -165,6 +165,53 @@ describe("useAnalyticsQuery", () => { 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 surfaces it through the error path", async () => { + // The SSE error broadcaster forwards an `errorCode` field for + // UI branching (e.g. INLINE_ARROW_STASH_EXHAUSTED). The hook reports + // the human `error` text; downstream code can read `errorCode` from + // the parsed payload if needed via console.error. + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + const { result } = renderHook(() => + useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + ); + + await lastConnectArgs.onMessage({ + data: JSON.stringify({ + type: "error", + error: "Server is at capacity, please retry", + code: "UPSTREAM_ERROR", + errorCode: "INLINE_ARROW_STASH_EXHAUSTED", + }), + }); + + await waitFor(() => { + expect(result.current.error).toBe("Server is at capacity, please retry"); + }); + expect(result.current.loading).toBe(false); + + errorSpy.mockRestore(); + }); + test("does not refetch when params object is structurally equal across renders", () => { const { rerender } = renderHook( ({ limit }: { limit: number }) => 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 a93da6f5..ed6ec602 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -258,7 +258,14 @@ export function useAnalyticsQuery< return; } } 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. console.warn("[useAnalyticsQuery] Malformed message received", error); + setLoading(false); + setError("Unable to load data, please try again"); } }, onError: (error) => { 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 1e6d1f5f..3c9df7fe 100644 --- a/packages/appkit/src/errors/execution.ts +++ b/packages/appkit/src/errors/execution.ts @@ -29,33 +29,56 @@ export class ExecutionError extends AppKitError { 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` + * (`INLINE_ARROW_STASH_EXHAUSTED`, `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. + * 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, { errorCode }); + 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", + }); } /** @@ -64,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" }, ); } @@ -84,4 +108,23 @@ export class ExecutionError extends AppKitError { context: { dataType }, }); } + + /** + * Create an execution error for the inline Arrow stash being unable to + * accept a payload (both regular and overflow pools at cap). + * + * The route deliberately surfaces this rather than silently re-running + * the statement on EXTERNAL_LINKS — a second execution can be billed + * and return divergent results for non-deterministic SQL. Operators + * should tune stash capacity or back off load when this fires. + */ + static stashExhausted(): ExecutionError { + return new ExecutionError( + "Inline Arrow stash exhausted; retry shortly or increase stash capacity", + { + errorCode: "INLINE_ARROW_STASH_EXHAUSTED", + clientMessage: "Server is at capacity, please retry", + }, + ); + } } diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index b3a268eb..e9ea58a5 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -167,8 +167,14 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { res.send(Buffer.from(result.data)); } catch (error) { logger.error("Arrow job error: %O", error); + // Do not echo upstream / SDK error text to the client — it can + // include statement fragments, internal object names, and + // correlation IDs. The full detail stays in the server log above. + const errorCode = + error instanceof ExecutionError ? error.errorCode : undefined; res.status(404).json({ - error: error instanceof Error ? error.message : "Arrow job not found", + error: "Arrow result unavailable", + errorCode, plugin: this.name, }); } @@ -326,8 +332,10 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { * server-side, and returns plain row objects — the caller's * `JSON_ARRAY` contract is preserved. * - **ARROW_STREAM** first tries `INLINE + ARROW_STREAM`. If the - * warehouse refuses (most classic + some serverless variants), or the - * inline stash is full, falls back to `EXTERNAL_LINKS + ARROW_STREAM`. + * warehouse refuses (most classic + some serverless variants), falls + * back to `EXTERNAL_LINKS + ARROW_STREAM`. When the regular stash is + * full, the already-decoded bytes spill into the stash's overflow + * pool — never thrown away to trigger a second warehouse round-trip. * * INLINE Arrow attachments under the ARROW_STREAM path are decoded once * and put on the plugin's `inlineArrowStash`; the SSE message carries the @@ -415,16 +423,20 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { ), ); if (inlineId === null) { - // Stash is full — every id we have already handed out must - // stay valid, so the stash refuses new entries rather than - // evicting in-flight ones. Fall back to EXTERNAL_LINKS for - // this request so the client still gets its result. + // Both regular and overflow pools are at cap. Throw with a + // clear signal rather than re-running the same statement on + // EXTERNAL_LINKS: the bytes we already decoded are gone, the + // warehouse has been billed, and a second execution could + // return a divergent result for non-deterministic SQL + // (CURRENT_TIMESTAMP, RAND, late-arriving rows). The right + // recovery is upstream backpressure or capacity tuning, not + // silently double-billing. logger.warn( - "Inline Arrow stash full, falling back to EXTERNAL_LINKS for the current query", + "Inline Arrow stash exhausted (regular + overflow); rejecting", ); - } else { - return makeArrowMessage(inlineId, { status: result.status }); + throw ExecutionError.stashExhausted(); } + return makeArrowMessage(inlineId, { status: result.status }); } else { return makeResultMessage(result?.data, { status: result?.status, @@ -574,6 +586,15 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } } +/** + * Hard cap on rows produced by the server-side JSON_ARRAY fallback. The + * fallback materializes every row into a plain JS object on the Node + * main thread (O(rows × cols) allocations), so a runaway result would + * block the event loop for hundreds of ms and pressure GC. Past this + * cap, surface a clear error instead of silently degrading throughput. + */ +const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; + /** * Decode a base64 Arrow IPC attachment to plain row objects. * @@ -582,11 +603,19 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { * `ARROW_STREAM + INLINE` — the bytes come back as Arrow IPC but the * caller's contract is JSON-shaped rows, so we convert server-side. * - * Scalar values are stringified to match what the warehouse itself emits - * for INT/BIGINT/etc. columns under the JSON_ARRAY format (everything in - * `result.data_array` is a string on the wire) — so callers see the same - * row shape regardless of which path the bytes took. BigInts get the same - * stringification treatment (also necessary for JSON-serializability). + * Shape rules (chosen to match what the warehouse emits natively under + * JSON_ARRAY, so callers cannot tell which path served their query): + * - **Scalars (number / bigint / boolean)**: stringified. JSON_ARRAY + * wire format emits everything in `result.data_array` as strings; + * bigint stringification is also necessary for JSON-serializability. + * - **Strings**: passthrough. + * - **Nested (List / Struct / Map)**: JSON-stringified. The warehouse + * emits nested values as JSON-encoded strings under JSON_ARRAY; + * apache-arrow's `Vector.get()` returns a typed wrapper whose + * `toJSON()` yields plain JS values, so `JSON.stringify` produces + * the same string shape the warehouse would have produced natively. + * This eliminates the previous correctness gap where callers saw an + * Arrow vector wrapper instead of a JSON string. */ function decodeArrowAttachmentToRows( attachment: string, @@ -595,6 +624,13 @@ function decodeArrowAttachmentToRows( 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.`, + ); + } const colNames = table.schema.fields.map((f) => f.name); const rows: Record[] = []; for (let i = 0; i < table.numRows; i++) { @@ -613,11 +649,10 @@ function decodeArrowAttachmentToRows( } else if (typeof v === "string") { row[name] = v; } else { - // Nested types (List, Struct, Map) — leave as-is. The JSON_ARRAY - // wire format renders these as JSON strings server-side, but that - // serialization isn't exposed to us here. Round-tripping through - // JSON.stringify would mismatch, so pass the typed value through. - row[name] = v; + // Nested types (List, Struct, Map). Native JSON_ARRAY emits + // these as JSON-encoded strings; apache-arrow's typed values + // expose `toJSON()` so `JSON.stringify` yields the same shape. + row[name] = JSON.stringify(v); } } rows.push(row); diff --git a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts index 3ae98330..c4a5e422 100644 --- a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts +++ b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts @@ -22,12 +22,16 @@ import { randomUUID } from "node:crypto"; * - **Per-user keyed**: `take()` only returns bytes if the requesting * user matches the user that originally put them. Defense in depth on * top of unguessable ids. - * - **Memory bounded with rejection**: total stashed bytes are capped. - * When `put()` cannot fit a payload without exceeding the cap it - * returns `null` rather than evicting older entries — every issued id - * stays valid until it is drained, expires, or the process exits. - * Callers are expected to fall back to a different delivery path (e.g. - * EXTERNAL_LINKS) when `put()` rejects. + * - **Memory bounded with overflow slot**: total stashed bytes are capped + * at `maxBytes`. When `put()` cannot fit a payload, it spills into a + * separate overflow pool capped at `maxOverflowBytes` (default same as + * `maxBytes`). Overflow entries behave identically to regular ones + * except they do not count against the regular cap — they are bytes + * the caller has *already paid to decode for this single request*, so + * throwing them away would force a second warehouse round-trip. Only + * when both pools are full does `put()` return `null` and the caller + * has to fall back to a different delivery path. Memory is bounded + * above by `maxBytes + maxOverflowBytes`. * * Caveat (multi-replica deployments): this stash is process-local. A * subsequent `GET /arrow-result/inline-*` that lands on a different @@ -40,11 +44,19 @@ interface InlineArrowStashOptions { /** Entries older than this are dropped on the next gc tick. */ ttlMs?: number; /** - * Hard cap on total bytes held. `put()` rejects (returns `null`) once - * the cap would be exceeded; entries already in the stash are not - * evicted to fit new ones. + * Hard cap on total bytes held in the regular pool. `put()` spills to + * the overflow pool when this cap would be exceeded; entries already + * in the stash are not evicted to fit new ones. */ maxBytes?: number; + /** + * Hard cap on total bytes held in the overflow pool. Overflow holds + * bytes that have already been decoded for an in-flight request — its + * purpose is to avoid throwing those bytes away and double-billing + * the warehouse. Defaults to `maxBytes`. `put()` returns `null` only + * when both regular and overflow pools are at cap. + */ + maxOverflowBytes?: number; /** Test seam: override the synthetic-id generator. */ idGenerator?: () => string; /** Test seam: override the clock. */ @@ -56,43 +68,61 @@ interface StashEntry { bytes: Uint8Array; expiresAt: number; insertedAt: number; + /** True for entries in the overflow pool (do not count against maxBytes). */ + overflow: boolean; } export class InlineArrowStash { private entries = new Map(); private totalBytes = 0; + private overflowBytes = 0; private readonly ttlMs: number; private readonly maxBytes: number; + private readonly maxOverflowBytes: number; private readonly idGenerator: () => string; private readonly now: () => number; constructor(opts: InlineArrowStashOptions = {}) { this.ttlMs = opts.ttlMs ?? 10 * 60 * 1000; this.maxBytes = opts.maxBytes ?? 256 * 1024 * 1024; + this.maxOverflowBytes = opts.maxOverflowBytes ?? this.maxBytes; this.idGenerator = opts.idGenerator ?? randomUUID; this.now = opts.now ?? Date.now; } /** - * Stash a payload and return its synthetic job id, or `null` when the - * stash cannot accept it without evicting older entries. The caller is - * expected to fall back to an out-of-band delivery path (e.g. - * EXTERNAL_LINKS) when the return value is `null`. + * Stash a payload and return its synthetic job id. + * + * Tries the regular pool first; if it would overflow, spills into the + * overflow pool (sized at `maxOverflowBytes`). Returns `null` only when + * both pools are at cap — the caller then has no choice but to fall + * back to a different delivery path (e.g. EXTERNAL_LINKS). + * + * The overflow pool exists because the caller has *already decoded* + * these bytes for an in-flight request: throwing them away would force + * a second warehouse round-trip (extra latency, double billing, and + * potentially divergent results for non-deterministic SQL). Holding + * them transiently in a bounded overflow region — they are drained + * single-use on the next `/arrow-result` fetch — is strictly safer + * than re-execution. * - * Single payloads that exceed `maxBytes` outright throw so the caller - * sees the misconfiguration loudly instead of degrading silently every - * time. + * Single payloads that exceed `maxBytes + maxOverflowBytes` outright + * throw so the caller sees the misconfiguration loudly. */ put(userId: string, bytes: Uint8Array): string | null { - if (bytes.length > this.maxBytes) { + const totalCap = this.maxBytes + this.maxOverflowBytes; + if (bytes.length > totalCap) { throw new Error( - `Inline Arrow payload (${bytes.length} bytes) exceeds stash maxBytes (${this.maxBytes})`, + `Inline Arrow payload (${bytes.length} bytes) exceeds stash capacity (${totalCap})`, ); } this.gc(); - if (this.totalBytes + bytes.length > this.maxBytes) { - // Refuse rather than evicting: every id we have already issued must - // remain valid until naturally drained or expired. + const fitsRegular = this.totalBytes + bytes.length <= this.maxBytes; + const fitsOverflow = + !fitsRegular && + this.overflowBytes + bytes.length <= this.maxOverflowBytes; + if (!fitsRegular && !fitsOverflow) { + // Both pools are full — refuse rather than evicting any issued id. return null; } const id = `inline-${this.idGenerator()}`; @@ -102,8 +132,13 @@ export class InlineArrowStash { bytes, expiresAt: now + this.ttlMs, insertedAt: now, + overflow: !fitsRegular, }); - this.totalBytes += bytes.length; + if (fitsRegular) { + this.totalBytes += bytes.length; + } else { + this.overflowBytes += bytes.length; + } return id; } @@ -117,7 +152,11 @@ export class InlineArrowStash { if (!entry) return undefined; if (entry.userId !== userId) return undefined; this.entries.delete(id); - this.totalBytes -= entry.bytes.length; + if (entry.overflow) { + this.overflowBytes -= entry.bytes.length; + } else { + this.totalBytes -= entry.bytes.length; + } return entry.bytes; } @@ -125,6 +164,10 @@ export class InlineArrowStash { size(): number { return this.totalBytes; } + /** Bytes currently held in the overflow pool. */ + overflowSize(): number { + return this.overflowBytes; + } count(): number { return this.entries.size; } @@ -133,6 +176,7 @@ export class InlineArrowStash { clear(): void { this.entries.clear(); this.totalBytes = 0; + this.overflowBytes = 0; } private gc(): void { @@ -140,7 +184,11 @@ export class InlineArrowStash { for (const [id, entry] of this.entries) { if (entry.expiresAt <= now) { this.entries.delete(id); - this.totalBytes -= entry.bytes.length; + if (entry.overflow) { + this.overflowBytes -= entry.bytes.length; + } else { + this.totalBytes -= entry.bytes.length; + } } } } diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 0683c44c..fb6ecbf9 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -9,6 +9,7 @@ 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 { InlineArrowStash } from "../inline-arrow-stash"; import type { IAnalyticsConfig } from "../types"; // Mock CacheManager singleton with actual caching behavior @@ -1026,11 +1027,14 @@ describe("Analytics Plugin", () => { expect(Array.from(stashed)).toEqual(Array.from(arrowBytes)); }); - test("/query/:query_key falls back to EXTERNAL_LINKS when the inline stash is full", async () => { - // When the stash refuses a new entry (put returns null), the route - // must not strand the client with a useless inline- id. It retries - // the same statement with EXTERNAL_LINKS so the warehouse hands - // back a real, fetchable statement id and the client gets bytes. + test("/query/:query_key spills the already-decoded bytes into the stash overflow pool when the regular pool is full — single execution, no double-billing", async () => { + // The regular stash put may refuse new entries when at cap. We must + // NOT respond by re-executing the same statement with EXTERNAL_LINKS: + // the warehouse has already been billed, the bytes are already + // decoded server-side, and a second execution can return a divergent + // result for non-deterministic SQL (CURRENT_TIMESTAMP, RAND, late + // rows). The stash's overflow pool absorbs the bytes so the client + // still gets an inline- id pointing at the original result. const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -1042,23 +1046,19 @@ describe("Analytics Plugin", () => { const fakeAttachment = Buffer.from(new Uint8Array([1, 2, 3])).toString( "base64", ); - const executeMock = vi - .fn() - // First call: INLINE succeeds and produces an attachment. - .mockResolvedValueOnce({ - result: { attachment: fakeAttachment, row_count: 1 }, - }) - // Second call: EXTERNAL_LINKS — returns a real warehouse id. - .mockResolvedValueOnce({ - result: { - statement_id: "stmt-warehouse-real", - status: { state: "SUCCEEDED" }, - }, - }); + const executeMock = vi.fn().mockResolvedValueOnce({ + result: { attachment: fakeAttachment, row_count: 1 }, + }); (plugin as any).SQLClient.executeStatement = executeMock; - // Force the stash to reject the put — simulates capacity exhaustion. - vi.spyOn((plugin as any).inlineArrowStash, "put").mockReturnValue(null); + // Force the regular pool to be at cap so the put spills into overflow. + // We construct a tiny stash and pre-fill the regular pool. + const tinyStash = new InlineArrowStash({ + maxBytes: 4, + maxOverflowBytes: 64, + }); + tinyStash.put("filler", new Uint8Array([9, 9, 9, 9])); + (plugin as any).inlineArrowStash = tinyStash; plugin.injectRoutes(router); @@ -1071,25 +1071,78 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); - expect(executeMock).toHaveBeenCalledTimes(2); + // Single execution: no EXTERNAL_LINKS retry. + expect(executeMock).toHaveBeenCalledTimes(1); 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", - }); - // SSE payload carries the real warehouse statement id, not an inline- id. + // SSE payload carries an inline- id pointing at the overflow entry. const writeCalls = (mockRes.write as any).mock.calls.map( (c: any[]) => c[0] as string, ); const payload = writeCalls.find((s: string) => s.startsWith("data: ")); expect(payload).toBeDefined(); expect(payload).toContain('"type":"arrow"'); - expect(payload).toContain('"statement_id":"stmt-warehouse-real"'); - expect(payload).not.toMatch(/"statement_id":"inline-/); + expect(payload).toMatch(/"statement_id":"inline-[^"]+"/); + + // Bytes landed in the overflow pool, regular pool size unchanged. + expect(tinyStash.overflowSize()).toBe(3); + expect(tinyStash.size()).toBe(4); + }); + + test("/query/:query_key surfaces a stable error when both regular and overflow pools are exhausted — never silently double-bills", async () => { + // When even the overflow pool cannot fit the payload, the route + // surfaces INLINE_ARROW_STASH_EXHAUSTED instead of re-running the + // statement on EXTERNAL_LINKS. The previous behavior (silent retry) + // double-billed the warehouse and could return divergent results. + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const fakeAttachment = Buffer.from(new Uint8Array([1, 2, 3])).toString( + "base64", + ); + const executeMock = vi.fn().mockResolvedValueOnce({ + result: { attachment: fakeAttachment, row_count: 1 }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + // Force both put paths to refuse: spy returns null unconditionally. + vi.spyOn((plugin as any).inlineArrowStash, "put").mockReturnValue(null); + + 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); + + // Single execution: never re-runs on EXTERNAL_LINKS. + expect(executeMock).toHaveBeenCalledTimes(1); + + // SSE error payload carries the stable errorCode for UI branching. + // The writer emits separate lines (`id:`, `event: error`, `data: ...`) + // — join them so we can match across the framing. + const writeCalls = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as string, + ); + const errorFrame = writeCalls + .filter((s: string) => s.startsWith("data: ")) + .join("\n"); + expect(errorFrame).toContain("INLINE_ARROW_STASH_EXHAUSTED"); + // Client-safe message reaches the wire; raw upstream text does not. + expect(errorFrame).toContain("Server is at capacity"); + expect(errorFrame).not.toContain("Inline Arrow stash exhausted"); }); test("/query/:query_key falls back JSON_ARRAY to ARROW_STREAM INLINE when warehouse refuses JSON_ARRAY for INLINE", async () => { diff --git a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts index 9bd59813..79e1a405 100644 --- a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts @@ -69,30 +69,57 @@ describe("InlineArrowStash", () => { expect(stash.take(id2, "user-1")).toBeUndefined(); }); - test("put returns null when adding the payload would exceed maxBytes, leaving existing entries intact", () => { + test("put spills into the overflow pool when the regular pool is at cap — every issued id remains valid", () => { let seq = 0; const stash = new InlineArrowStash({ maxBytes: 200, + maxOverflowBytes: 200, idGenerator: () => String(seq++), }); const a = mustPut(stash, "user-1", bytes(80)); const b = mustPut(stash, "user-1", bytes(80)); expect(stash.size()).toBe(160); + expect(stash.overflowSize()).toBe(0); - // This third 80-byte entry would push total to 240 (>200). It must - // be rejected, and both prior entries must survive — every id we have - // already handed out stays valid until drained or expired. - const c = stash.put("user-1", bytes(80)); - expect(c).toBeNull(); + // The third 80-byte entry would push the regular pool to 240 (>200), + // so it spills into overflow. Both prior entries must survive. + const c = mustPut(stash, "user-1", bytes(80)); expect(stash.size()).toBe(160); + expect(stash.overflowSize()).toBe(80); + expect(stash.take(a, "user-1")).toBeDefined(); + expect(stash.take(b, "user-1")).toBeDefined(); + expect(stash.take(c, "user-1")).toBeDefined(); + // After draining the overflow entry, the counter reflects it. + expect(stash.overflowSize()).toBe(0); + }); + + test("put returns null only when both regular and overflow pools are full", () => { + let seq = 0; + const stash = new InlineArrowStash({ + maxBytes: 100, + maxOverflowBytes: 100, + idGenerator: () => String(seq++), + }); + const a = mustPut(stash, "user-1", bytes(100)); // fills regular + const b = mustPut(stash, "user-1", bytes(100)); // fills overflow + expect(stash.size()).toBe(100); + expect(stash.overflowSize()).toBe(100); + + // Both pools at cap — refuse rather than evict. + const c = stash.put("user-1", bytes(50)); + expect(c).toBeNull(); expect(stash.take(a, "user-1")).toBeDefined(); expect(stash.take(b, "user-1")).toBeDefined(); }); - test("put throws for a single payload larger than maxBytes (caller misconfiguration)", () => { - const stash = new InlineArrowStash({ maxBytes: 100 }); - expect(() => stash.put("user-1", bytes(200))).toThrow( - /exceeds stash maxBytes/, + test("put throws when a single payload would not fit even with overflow (caller misconfiguration)", () => { + const stash = new InlineArrowStash({ + maxBytes: 100, + maxOverflowBytes: 100, + }); + // 300 > maxBytes + maxOverflowBytes (200), so this can never fit. + expect(() => stash.put("user-1", bytes(300))).toThrow( + /exceeds stash capacity/, ); }); 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 901e0b46..2cacbcb9 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"; @@ -285,8 +287,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. INLINE_ARROW_STASH_EXHAUSTED, + // 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); @@ -295,17 +310,24 @@ export class StreamManager { logger.info("Stream aborted by client (code=%s)", errorCode); } else { 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(), }); @@ -313,9 +335,10 @@ export class StreamManager { this._broadcastErrorToClients( streamEntry, errorEventId, - errorMsg, + clientMsg, errorCode, true, + upstreamCode, ); streamEntry.isCompleted = true; } @@ -370,10 +393,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/stream.test.ts b/packages/appkit/src/stream/tests/stream.test.ts index c18d0f67..a9924203 100644 --- a/packages/appkit/src/stream/tests/stream.test.ts +++ b/packages/appkit/src/stream/tests/stream.test.ts @@ -233,21 +233,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(); }); @@ -513,10 +522,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 bb6f65f6..ba0d4915 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. `INLINE_ARROW_STASH_EXHAUSTED`, + * `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/shared/src/sse/analytics.ts b/packages/shared/src/sse/analytics.ts index f37d38a5..31aaf511 100644 --- a/packages/shared/src/sse/analytics.ts +++ b/packages/shared/src/sse/analytics.ts @@ -27,13 +27,29 @@ import { z } from "zod"; /** Successful row-shaped result (JSON_ARRAY format, or empty results). */ export const AnalyticsResultMessage = z.object({ type: z.literal("result"), - data: z.array(z.record(z.string(), z.unknown())).optional(), + // `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. + 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(), }); -export type AnalyticsResultMessage = z.infer; +// The TS-level shape narrows `data` to `Record[]` to +// match the typed `makeResultMessage` builder — the Zod schema is +// intentionally looser at runtime for performance (see comment above). +export type AnalyticsResultMessage = Omit< + z.infer, + "data" +> & { + data?: Record[]; +}; /** * ARROW_STREAM result delivered via /arrow-result/:jobId. The id is either: From 28fdb928907d82c110afed6287835da2efd96a99 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Thu, 21 May 2026 21:39:11 +0000 Subject: [PATCH 16/25] fix(appkit): address Xavier v4 self-review on inline-Arrow path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second iteration on #329 after running Xavier multi-model review against 4afce1f6. Eleven findings (1 critical, 4 high, 6 medium), all addressed: 🔴 critical — `decodeArrowAttachmentToRows` regression: bare `JSON.stringify` on Arrow nested values throws on `bigint` (the spec doesn't serialize it) and produces broken shapes for `Uint8Array` / `Date`. Added `nestedJsonReplacer` that stringifies bigint, base64s `Uint8Array`, and ISO-strings `Date`. Top-level `Date` / `Uint8Array` columns get the same treatment in their own branches so we never reach the nested fallback for them. 🟠 high — single-payload throw threshold was `maxBytes + maxOverflowBytes`, but a payload can only land in *one* pool (pools do not split). Fixed to `Math.max(maxBytes, maxOverflowBytes)`. 🟠 high — overflow pool inherited the regular 10-min TTL, which kept transient bytes around for the full duration of the regular pool's lifetime and amplified cross-user memory pressure in multi-tenant deployments. Split into a separate `overflowTtlMs` (default 30s) — overflow exists only to bridge the immediate `/arrow-result` fetch. 🟠 high — overflow defaulted to the same size as `maxBytes`, doubling the worst-case memory headroom. Dropped default to `maxBytes / 4`; overflow only needs to absorb transient spillover, not hold long-term state. Combined with the shorter TTL, the practical memory delta vs pre-overflow code is small. 🟠 high — JSON_ARRAY fallback row cap ignored cell size (10 rows × 10MB cells = OOM under the 100k-row cap). Added a 64MB byte cap on the decoded attachment, checked before `tableFromIPC` so we never even allocate the table for an oversize payload. 🟢 medium — `gc()` ran O(N) on every `put`. Throttled to `gcMinIntervalMs` (default 5s). Tests that drive the clock past TTL on a sub-throttle scale opt out via `gcMinIntervalMs: 0`. 🟢 medium — `table.getChild(name)` was being called for every cell, walking the schema fields on every lookup. Hoisted into a `[name, vector]` array resolved once before the row loop. Real win at 100k × 50 columns. 🟢 medium — `useAnalyticsQuery` dropped the new structured `errorCode` from the SSE error payload, so consumers could not actually branch on the stable identifier the previous commit added. Added `errorCode: string | null` to `UseAnalyticsQueryResult` and plumbed it through from `parsed.errorCode`. 🟢 medium — outer `catch` in the SSE handler used to clear loading on a parse failure but leave the stream open, re-firing the catch on every subsequent malformed message. Now also calls `abortController.abort()` so the consumer can retry cleanly. Validation: 2,097 appkit + 292 appkit-ui tests pass; tsc clean on appkit / appkit-ui / shared; biome clean on changed files. Co-authored-by: Isaac Signed-off-by: James Broadhead --- .../__tests__/use-analytics-query.test.ts | 10 ++- packages/appkit-ui/src/react/hooks/types.ts | 10 ++- .../src/react/hooks/use-analytics-query.ts | 18 ++++- .../appkit/src/plugins/analytics/analytics.ts | 76 ++++++++++++----- .../plugins/analytics/inline-arrow-stash.ts | 81 ++++++++++++++----- .../tests/inline-arrow-stash.test.ts | 37 ++++++++- 6 files changed, 180 insertions(+), 52 deletions(-) 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 078c2aad..ba4c962b 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 @@ -184,11 +184,12 @@ describe("useAnalyticsQuery", () => { expect(result.current.data).toBeNull(); }); - test("a server error event carrying a structured errorCode surfaces it through the error path", async () => { + 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 (e.g. INLINE_ARROW_STASH_EXHAUSTED). The hook reports - // the human `error` text; downstream code can read `errorCode` from - // the parsed payload if needed via console.error. + // UI branching (e.g. INLINE_ARROW_STASH_EXHAUSTED). 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(() => @@ -208,6 +209,7 @@ describe("useAnalyticsQuery", () => { expect(result.current.error).toBe("Server is at capacity, please retry"); }); expect(result.current.loading).toBe(false); + expect(result.current.errorCode).toBe("INLINE_ARROW_STASH_EXHAUSTED"); errorSpy.mockRestore(); }); diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index 7c249bf0..a603e10a 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -63,8 +63,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 + * SSE error payload (e.g. `INLINE_ARROW_STASH_EXHAUSTED`, + * `RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`). Prefer + * branching UI on this stable identifier rather than substring- + * matching `error`, which is a free-form sanitized message. + */ + errorCode: string | null; } /** 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 ed6ec602..a3d2cf30 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -121,6 +121,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 abortControllerRef = useRef(null); if (!queryKey || queryKey.trim().length === 0) { @@ -168,6 +169,7 @@ export function useAnalyticsQuery< setLoading(true); setError(null); + setErrorCode(null); setData(null); const abortController = new AbortController(); @@ -237,6 +239,14 @@ export function useAnalyticsQuery< setLoading(false); setError(errorMsg); + // Propagate the upstream structured code so UI consumers + // can branch on a stable identifier (e.g. retry on + // INLINE_ARROW_STASH_EXHAUSTED, format-switch on + // RESULT_TOO_LARGE_FOR_JSON_FALLBACK) instead of parsing + // the human-readable message. + if (typeof parsed.errorCode === "string") { + setErrorCode(parsed.errorCode); + } if (parsed.code) { console.error( @@ -263,9 +273,15 @@ export function useAnalyticsQuery< // `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("Unable to load data, please try again"); + abortController.abort(); } }, onError: (error) => { @@ -305,5 +321,5 @@ export function useAnalyticsQuery< // Enable HMR for query updates in dev mode useQueryHMR(queryKey, start); - return { data, loading, error }; + return { data, loading, error, errorCode }; } diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index e9ea58a5..cf13a16d 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -587,13 +587,33 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } /** - * Hard cap on rows produced by the server-side JSON_ARRAY fallback. The - * fallback materializes every row into a plain JS object on the Node - * main thread (O(rows × cols) allocations), so a runaway result would - * block the event loop for hundreds of ms and pressure GC. Past this - * cap, surface a clear error instead of silently degrading throughput. + * 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. We cap two ways — rows AND decoded bytes — + * because either dimension can blow up independently (100k×1B rows is + * fine, 10 rows × 10MB cells is not). */ const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; +const JSON_ARRAY_FALLBACK_MAX_BYTES = 64 * 1024 * 1024; + +/** + * Replacer used by the nested-value `JSON.stringify` path. Arrow IPC + * carries column values whose JS representations include `bigint` + * (which the spec'd `JSON.stringify` cannot serialize and throws on), + * `Uint8Array` (binary buffers — would serialize as `{"0":...,"1":...}` + * which is wrong), and `Date` (would lose timezone fidelity inside a + * nested struct). Coerce each to a string form that matches what the + * warehouse itself emits under native JSON_ARRAY. + */ +function nestedJsonReplacer(_key: string, value: unknown): unknown { + if (typeof value === "bigint") return value.toString(); + if (value instanceof Uint8Array) { + return Buffer.from(value).toString("base64"); + } + if (value instanceof Date) return value.toISOString(); + return value; +} /** * Decode a base64 Arrow IPC attachment to plain row objects. @@ -608,19 +628,25 @@ const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; * - **Scalars (number / bigint / boolean)**: stringified. JSON_ARRAY * wire format emits everything in `result.data_array` as strings; * bigint stringification is also necessary for JSON-serializability. - * - **Strings**: passthrough. - * - **Nested (List / Struct / Map)**: JSON-stringified. The warehouse - * emits nested values as JSON-encoded strings under JSON_ARRAY; - * apache-arrow's `Vector.get()` returns a typed wrapper whose - * `toJSON()` yields plain JS values, so `JSON.stringify` produces - * the same string shape the warehouse would have produced natively. - * This eliminates the previous correctness gap where callers saw an - * Arrow vector wrapper instead of a JSON string. + * - **Strings / Date / Uint8Array**: stringified consistently — + * `Date` to ISO string, `Uint8Array` to base64. + * - **Nested (List / Struct / Map)**: `JSON.stringify` with the + * `nestedJsonReplacer` so nested bigints / Dates / binary buffers + * serialize sanely instead of throwing or producing object-as-map + * output. Apache-arrow's typed values expose `toJSON()` so the + * resulting string matches what the warehouse would have emitted. */ function decodeArrowAttachmentToRows( attachment: 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), ); @@ -631,12 +657,17 @@ function decodeArrowAttachmentToRows( `Result too large for JSON format (over ${JSON_ARRAY_FALLBACK_MAX_ROWS} rows). Re-run with ARROW_STREAM format.`, ); } - const colNames = table.schema.fields.map((f) => f.name); + // Resolve the child vectors once. `table.getChild(name)` walks the + // schema fields on every call; without this hoist we'd pay that cost + // O(rows × cols) times. At 100k × 50 columns that's millions of + // redundant lookups on the event loop. + const columns = table.schema.fields.map( + (f) => [f.name, table.getChild(f.name)] as const, + ); const rows: Record[] = []; for (let i = 0; i < table.numRows; i++) { const row: Record = {}; - for (const name of colNames) { - const col = table.getChild(name); + for (const [name, col] of columns) { const v = col?.get(i); if (v == null) { row[name] = null; @@ -648,11 +679,16 @@ function decodeArrowAttachmentToRows( row[name] = String(v); } else if (typeof v === "string") { row[name] = v; + } else if (v instanceof Date) { + row[name] = v.toISOString(); + } else if (v instanceof Uint8Array) { + row[name] = Buffer.from(v).toString("base64"); } else { - // Nested types (List, Struct, Map). Native JSON_ARRAY emits - // these as JSON-encoded strings; apache-arrow's typed values - // expose `toJSON()` so `JSON.stringify` yields the same shape. - row[name] = JSON.stringify(v); + // Nested types (List, Struct, Map). Use the replacer so nested + // bigint / Date / Uint8Array values serialize predictably — + // bare `JSON.stringify` throws on bigint and emits broken + // shapes for binary buffers. + row[name] = JSON.stringify(v, nestedJsonReplacer); } } rows.push(row); diff --git a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts index c4a5e422..31a91dfd 100644 --- a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts +++ b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts @@ -24,14 +24,16 @@ import { randomUUID } from "node:crypto"; * top of unguessable ids. * - **Memory bounded with overflow slot**: total stashed bytes are capped * at `maxBytes`. When `put()` cannot fit a payload, it spills into a - * separate overflow pool capped at `maxOverflowBytes` (default same as - * `maxBytes`). Overflow entries behave identically to regular ones - * except they do not count against the regular cap — they are bytes - * the caller has *already paid to decode for this single request*, so - * throwing them away would force a second warehouse round-trip. Only - * when both pools are full does `put()` return `null` and the caller - * has to fall back to a different delivery path. Memory is bounded - * above by `maxBytes + maxOverflowBytes`. + * separate overflow pool capped at `maxOverflowBytes` (default + * `maxBytes / 4` — kept small because overflow only needs to bridge + * the immediate `/arrow-result` fetch). Overflow entries behave like + * regular ones except (a) they do not count against the regular cap + * and (b) they expire on a much shorter TTL (`overflowTtlMs`, + * default 30s) — they exist to absorb already-decoded bytes for a + * single request, not to hold them long-term. Only when both pools + * are full does `put()` return `null` and the caller has to fall + * back to a different delivery path. Memory is bounded above by + * `maxBytes + maxOverflowBytes`. * * Caveat (multi-replica deployments): this stash is process-local. A * subsequent `GET /arrow-result/inline-*` that lands on a different @@ -41,8 +43,17 @@ import { randomUUID } from "node:crypto"; * shared external store, neither of which is in scope here. */ interface InlineArrowStashOptions { - /** Entries older than this are dropped on the next gc tick. */ + /** Regular-pool entries older than this are dropped on the next gc tick. */ ttlMs?: number; + /** + * Overflow-pool entries older than this are dropped on the next gc + * tick. Defaults to 30s — overflow exists solely to bridge the + * immediate `/arrow-result` fetch that follows the SSE `arrow` + * message, so it should drain on the order of seconds, not minutes. + * A short TTL bounds the cross-user memory pressure that overflow + * can sustain in a multi-tenant deployment. + */ + overflowTtlMs?: number; /** * Hard cap on total bytes held in the regular pool. `put()` spills to * the overflow pool when this cap would be exceeded; entries already @@ -53,10 +64,18 @@ interface InlineArrowStashOptions { * Hard cap on total bytes held in the overflow pool. Overflow holds * bytes that have already been decoded for an in-flight request — its * purpose is to avoid throwing those bytes away and double-billing - * the warehouse. Defaults to `maxBytes`. `put()` returns `null` only - * when both regular and overflow pools are at cap. + * the warehouse. Defaults to `maxBytes / 4` (kept small because the + * pool only needs to absorb transient spillover, not hold long-term + * state). `put()` returns `null` only when both regular and overflow + * pools are at cap. */ maxOverflowBytes?: number; + /** + * Minimum interval between gc passes. `gc()` is O(N) in entry count, + * so on hot paths we skip when the previous pass was recent enough. + * Defaults to 5s. Set to 0 to disable throttling (test seam). + */ + gcMinIntervalMs?: number; /** Test seam: override the synthetic-id generator. */ idGenerator?: () => string; /** Test seam: override the clock. */ @@ -76,16 +95,22 @@ export class InlineArrowStash { private entries = new Map(); private totalBytes = 0; private overflowBytes = 0; + private lastGcAt = 0; private readonly ttlMs: number; + private readonly overflowTtlMs: number; private readonly maxBytes: number; private readonly maxOverflowBytes: number; + private readonly gcMinIntervalMs: number; private readonly idGenerator: () => string; private readonly now: () => number; constructor(opts: InlineArrowStashOptions = {}) { this.ttlMs = opts.ttlMs ?? 10 * 60 * 1000; + this.overflowTtlMs = opts.overflowTtlMs ?? 30 * 1000; this.maxBytes = opts.maxBytes ?? 256 * 1024 * 1024; - this.maxOverflowBytes = opts.maxOverflowBytes ?? this.maxBytes; + this.maxOverflowBytes = + opts.maxOverflowBytes ?? Math.floor(this.maxBytes / 4); + this.gcMinIntervalMs = opts.gcMinIntervalMs ?? 5000; this.idGenerator = opts.idGenerator ?? randomUUID; this.now = opts.now ?? Date.now; } @@ -106,14 +131,17 @@ export class InlineArrowStash { * single-use on the next `/arrow-result` fetch — is strictly safer * than re-execution. * - * Single payloads that exceed `maxBytes + maxOverflowBytes` outright - * throw so the caller sees the misconfiguration loudly. + * A single payload can only land in one pool — the pools are not + * split across — so the largest payload we can accept is + * `Math.max(maxBytes, maxOverflowBytes)`. Exceeding that throws + * synchronously so the caller sees the misconfiguration loudly + * rather than burning a warehouse round-trip and then getting `null`. */ put(userId: string, bytes: Uint8Array): string | null { - const totalCap = this.maxBytes + this.maxOverflowBytes; - if (bytes.length > totalCap) { + const largestSlot = Math.max(this.maxBytes, this.maxOverflowBytes); + if (bytes.length > largestSlot) { throw new Error( - `Inline Arrow payload (${bytes.length} bytes) exceeds stash capacity (${totalCap})`, + `Inline Arrow payload (${bytes.length} bytes) exceeds largest stash slot (${largestSlot}); cannot fit in either pool`, ); } this.gc(); @@ -127,17 +155,18 @@ export class InlineArrowStash { } const id = `inline-${this.idGenerator()}`; const now = this.now(); + const overflow = !fitsRegular; this.entries.set(id, { userId, bytes, - expiresAt: now + this.ttlMs, + expiresAt: now + (overflow ? this.overflowTtlMs : this.ttlMs), insertedAt: now, - overflow: !fitsRegular, + overflow, }); - if (fitsRegular) { - this.totalBytes += bytes.length; - } else { + if (overflow) { this.overflowBytes += bytes.length; + } else { + this.totalBytes += bytes.length; } return id; } @@ -181,6 +210,14 @@ export class InlineArrowStash { private gc(): void { const now = this.now(); + if (now - this.lastGcAt < this.gcMinIntervalMs) { + // Skip the O(N) sweep — recent enough to assume nothing + // significant has expired since the last pass. Worst case an + // entry lingers an extra `gcMinIntervalMs` past its TTL, which + // is negligible relative to either pool's intended lifetime. + return; + } + this.lastGcAt = now; for (const [id, entry] of this.entries) { if (entry.expiresAt <= now) { this.entries.delete(id); diff --git a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts index 79e1a405..31d9b01d 100644 --- a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts @@ -57,7 +57,14 @@ describe("InlineArrowStash", () => { test("entries past TTL are evicted on next gc tick", () => { let clock = 0; - const stash = new InlineArrowStash({ ttlMs: 1000, now: () => clock }); + // gcMinIntervalMs: 0 disables gc throttling so the test can drive + // the clock past TTL on a sub-throttle scale without the gc pass + // being skipped. + const stash = new InlineArrowStash({ + ttlMs: 1000, + gcMinIntervalMs: 0, + now: () => clock, + }); const id = mustPut(stash, "user-1", bytes(50)); clock = 999; expect(stash.take(id, "user-1")).toBeDefined(); @@ -112,17 +119,39 @@ describe("InlineArrowStash", () => { expect(stash.take(b, "user-1")).toBeDefined(); }); - test("put throws when a single payload would not fit even with overflow (caller misconfiguration)", () => { + test("put throws when a single payload would not fit in the largest pool (caller misconfiguration)", () => { const stash = new InlineArrowStash({ maxBytes: 100, maxOverflowBytes: 100, }); - // 300 > maxBytes + maxOverflowBytes (200), so this can never fit. + // 300 > max(maxBytes, maxOverflowBytes) (100); pools don't split, + // so no individual put can ever succeed. expect(() => stash.put("user-1", bytes(300))).toThrow( - /exceeds stash capacity/, + /exceeds largest stash slot/, ); }); + test("overflow entries expire on a shorter TTL than regular entries", () => { + let clock = 0; + const stash = new InlineArrowStash({ + maxBytes: 80, + maxOverflowBytes: 80, + ttlMs: 600_000, + overflowTtlMs: 30_000, + gcMinIntervalMs: 0, + now: () => clock, + }); + const reg = mustPut(stash, "user-1", bytes(80)); // fills regular + const ovf = mustPut(stash, "user-1", bytes(80)); // spills to overflow + expect(stash.size()).toBe(80); + expect(stash.overflowSize()).toBe(80); + + // 45 s in: overflow expired, regular still alive. + clock = 45_000; + expect(stash.take(ovf, "user-1")).toBeUndefined(); + expect(stash.take(reg, "user-1")).toBeDefined(); + }); + test("synthetic ids are unique across puts", () => { const stash = new InlineArrowStash(); const a = mustPut(stash, "user-1", bytes(10)); From f69e57794a1126cdaf4c6bb58c3076cef377d8e1 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Thu, 21 May 2026 22:08:02 +0000 Subject: [PATCH 17/25] refactor(appkit): stash telemetry, type cleanup, format-path extraction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three follow-ups from self-review on the inline-Arrow path: - **Telemetry**: `InlineArrowStash.put()` now returns `{ id, pool }` so callers can label outcomes without re-introspecting the stash. The analytics plugin emits two instruments (lazy-init since `this.telemetry` isn't bound at construct time): - `analytics.inline_arrow_stash.put.count{result}` — counter with label "regular" | "overflow" | "exhausted" - `analytics.inline_arrow_stash.put.bytes{result}` — histogram of accepted payload sizes Operators can now drive capacity-tuning dashboards off `result="overflow"` spill rate and `result="exhausted"` rejection rate without inferring state from response codes. - **Type cleanup**: `AnalyticsResultMessage` had a Zod-inferred type wrapped in `Omit<…, "data"> & { data?: ... }` to narrow `data` to `Record[]` for callers. Replaced with a hand-written interface alongside the Zod schema, with an explicit "keep in sync" comment. Same runtime behavior, dramatically easier to read. - **Format-path extraction**: `_executeWithFormatFallback` was a 120-line function handling two distinct format paths inline. Split into: - `_executeJsonArrayPath` — JSON_ARRAY with ARROW_STREAM retry - `_executeArrowStreamPath` — ARROW_STREAM with EXTERNAL_LINKS fallback - `_stashAndEmitInline` — the cap/overflow/exhaustion logic that both paths share, pulled into one named method `_executeWithFormatFallback` is now a thin dispatcher. Each method's failure modes are documented in its docstring so the contract is visible without reading the body. Validation: 2,098 appkit + 292 appkit-ui tests pass; tsc clean on appkit / appkit-ui / shared; biome clean on changed files. Co-authored-by: Isaac Signed-off-by: James Broadhead --- .../appkit/src/plugins/analytics/analytics.ts | 279 +++++++++++++----- .../plugins/analytics/inline-arrow-stash.ts | 14 +- .../plugins/analytics/tests/analytics.test.ts | 4 +- .../tests/inline-arrow-stash.test.ts | 24 +- packages/shared/src/sse/analytics.ts | 26 +- 5 files changed, 250 insertions(+), 97 deletions(-) diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index cf13a16d..5ba121d1 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -26,6 +26,7 @@ import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; +import type { Counter, Histogram } from "../../telemetry"; import { queryDefaults } from "./defaults"; import { InlineArrowStash } from "./inline-arrow-stash"; import manifest from "./manifest.json"; @@ -62,6 +63,22 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { */ protected inlineArrowStash: InlineArrowStash = new InlineArrowStash(); + /** + * Stash telemetry — created lazily on first stash operation because + * `this.telemetry` may not be bound at constructor time (see + * `Plugin.attachContext`). Cached via `_stashMetrics` after the first + * call; subsequent puts hit the cached counter/histogram instances. + * + * Counter labels: `result` ∈ {"regular", "overflow", "exhausted"} — + * one counter with a label is friendlier to dashboards than three + * separate metrics, and aggregates trivially in Prometheus / OTel. + */ + private _stashMetrics?: { + putCount: Counter; + putBytes: Histogram; + }; + private _stashMetricsAttempted = false; + constructor(config: IAnalyticsConfig) { super(config); this.config = config; @@ -180,6 +197,51 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } } + /** + * Lazy accessor for stash telemetry instruments. Returns `undefined` + * when telemetry isn't bound (factory-constructed plugins before + * `attachContext`, no-op test paths) — callers should branch on the + * return value rather than relying on it being present. + * + * Records once per call site: + * - `analytics.inline_arrow_stash.put.count{result}` — counter with + * label "regular" | "overflow" | "exhausted" + * - `analytics.inline_arrow_stash.put.bytes` — histogram of accepted + * payload sizes (label same as counter; exhausted puts record 0) + */ + private _getStashMetrics() { + if (this._stashMetricsAttempted) return this._stashMetrics; + this._stashMetricsAttempted = true; + if (!this.telemetry) return undefined; + const meter = this.telemetry.getMeter(); + this._stashMetrics = { + putCount: meter.createCounter("analytics.inline_arrow_stash.put.count", { + description: + "Count of inline-Arrow stash put attempts, labeled by pool outcome (regular | overflow | exhausted).", + unit: "1", + }), + putBytes: meter.createHistogram( + "analytics.inline_arrow_stash.put.bytes", + { + description: + "Sizes of accepted inline-Arrow stash payloads. Aggregate by `result` label for utilization analysis.", + unit: "By", + }, + ), + }; + return this._stashMetrics; + } + + private _recordStashOutcome( + result: "regular" | "overflow" | "exhausted", + bytes: number, + ): void { + const m = this._getStashMetrics(); + if (!m) return; + m.putCount.add(1, { result }); + m.putBytes.record(result === "exhausted" ? 0 : bytes, { result }); + } + /** * Stash key used at put-time (in `_handleQueryRoute`) and take-time * (in `_handleArrowRoute`). Centralized so the two sides cannot drift. @@ -352,49 +414,108 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { stashUserKey: string, signal?: AbortSignal, ): Promise { - if (requestedFormat === "JSON_ARRAY") { - try { - const result = await executor.query( + return requestedFormat === "JSON_ARRAY" + ? this._executeJsonArrayPath(executor, query, processedParams, signal) + : this._executeArrowStreamPath( + executor, query, processedParams, - { disposition: "INLINE", format: "JSON_ARRAY" }, + stashUserKey, signal, ); - return makeResultMessage(result?.data, { - status: result?.status, - statement_id: result?.statement_id, - }); - } catch (err: unknown) { - if (signal?.aborted) throw err; - if (_classifyInlineRejection(err) !== "needs-arrow") throw err; - - const msg = err instanceof Error ? err.message : String(err); - logger.warn( - "JSON_ARRAY INLINE rejected by warehouse, retrying as ARROW_STREAM INLINE and decoding server-side: %s", - msg, - ); - } + } - // Retry as ARROW_STREAM + INLINE so the warehouse will accept the - // request, then decode the Arrow IPC attachment to plain row - // objects so the caller still gets JSON_ARRAY-shaped data. - const arrowResult = await executor.query( + /** + * JSON_ARRAY path. Tries `INLINE + JSON_ARRAY` first; on a structured + * `needs-arrow` rejection (warehouse only accepts ARROW_STREAM under + * INLINE), retries as `INLINE + ARROW_STREAM` and decodes the Arrow + * IPC attachment server-side so the caller's JSON_ARRAY contract is + * preserved. + * + * Failure modes surfaced to the caller: + * - Unrelated warehouse errors (anything not classified as + * `needs-arrow`) propagate untouched. + * - The retry can throw `ExecutionError.missingData("ARROW_STREAM + * attachment")` if the warehouse returned no attachment. + * - The decode itself can throw `RESULT_TOO_LARGE_FOR_JSON_FALLBACK` + * when the attachment exceeds the row or byte cap. + */ + private async _executeJsonArrayPath( + executor: AnalyticsPlugin, + query: string, + processedParams: + | Record + | undefined, + signal?: AbortSignal, + ): Promise { + try { + const result = await executor.query( query, processedParams, - { disposition: "INLINE", format: "ARROW_STREAM" }, + { disposition: "INLINE", format: "JSON_ARRAY" }, signal, ); - if (!arrowResult?.attachment) { - throw ExecutionError.missingData("ARROW_STREAM attachment"); - } - const rows = decodeArrowAttachmentToRows(arrowResult.attachment); - return makeResultMessage(rows, { - status: arrowResult.status, - statement_id: arrowResult.statement_id, + return makeResultMessage(result?.data, { + status: result?.status, + statement_id: result?.statement_id, }); + } catch (err: unknown) { + if (signal?.aborted) throw err; + if (_classifyInlineRejection(err) !== "needs-arrow") throw err; + + const msg = err instanceof Error ? err.message : String(err); + logger.warn( + "JSON_ARRAY INLINE rejected by warehouse, retrying as ARROW_STREAM INLINE and decoding server-side: %s", + msg, + ); } - // ARROW_STREAM: try INLINE first, fall back to EXTERNAL_LINKS. + // Retry as ARROW_STREAM + INLINE so the warehouse will accept the + // request, then decode the Arrow IPC attachment to plain row + // objects so the caller still gets JSON_ARRAY-shaped data. + 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); + return makeResultMessage(rows, { + status: arrowResult.status, + statement_id: arrowResult.statement_id, + }); + } + + /** + * ARROW_STREAM path. Tries `INLINE + ARROW_STREAM` first; on a + * structured `needs-json` rejection (warehouse refuses ARROW_STREAM + * under INLINE), falls back to `EXTERNAL_LINKS + ARROW_STREAM`. When + * INLINE succeeds with an Arrow attachment, the decoded bytes go + * through `inlineArrowStash` and the client fetches them via + * `/arrow-result/inline-` — never the SSE channel. + * + * Failure modes surfaced to the caller: + * - Unrelated warehouse errors (not classified as `needs-json`) + * propagate untouched. + * - `ExecutionError.canceled()` if the client disconnects between + * the INLINE response and the stash put. + * - `ExecutionError.stashExhausted()` if both stash pools are full — + * we never silently re-run on EXTERNAL_LINKS because that would + * double-bill the warehouse and risk a divergent result for + * non-deterministic SQL. + */ + private async _executeArrowStreamPath( + executor: AnalyticsPlugin, + query: string, + processedParams: + | Record + | undefined, + stashUserKey: string, + signal?: AbortSignal, + ): Promise { try { const result = await executor.query( query, @@ -402,57 +523,21 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { { disposition: "INLINE", format: "ARROW_STREAM" }, signal, ); - // INLINE responses with an Arrow IPC attachment go through the - // stash-and-serve path: decode the base64 once, hold the bytes - // server-side, emit a synthetic statement id. The client fetches via - // /arrow-result so multi-MiB Arrow blobs never traverse SSE. if (result?.attachment) { - // If the client has already disconnected, the SSE write would be - // dropped anyway — skip the decode + stash so the bytes do not - // linger in memory until TTL eviction. - if (signal?.aborted) { - throw ExecutionError.canceled(); - } - const decoded = Buffer.from(result.attachment, "base64"); - const inlineId = this.inlineArrowStash.put( - stashUserKey, - new Uint8Array( - decoded.buffer, - decoded.byteOffset, - decoded.byteLength, - ), - ); - if (inlineId === null) { - // Both regular and overflow pools are at cap. Throw with a - // clear signal rather than re-running the same statement on - // EXTERNAL_LINKS: the bytes we already decoded are gone, the - // warehouse has been billed, and a second execution could - // return a divergent result for non-deterministic SQL - // (CURRENT_TIMESTAMP, RAND, late-arriving rows). The right - // recovery is upstream backpressure or capacity tuning, not - // silently double-billing. - logger.warn( - "Inline Arrow stash exhausted (regular + overflow); rejecting", - ); - throw ExecutionError.stashExhausted(); - } - return makeArrowMessage(inlineId, { status: result.status }); - } else { - return makeResultMessage(result?.data, { - status: result?.status, - statement_id: result?.statement_id, - }); + return this._stashAndEmitInline(result, stashUserKey, signal); } + // INLINE succeeded but the warehouse didn't return an Arrow + // attachment — degrade to a plain result message with whatever + // data the warehouse did return. + return makeResultMessage(result?.data, { + status: result?.status, + statement_id: result?.statement_id, + }); } catch (err: unknown) { // If the request was aborted, do not retry — the signal is dead and // a second statement would be billed but never read. - if (signal?.aborted) { - throw err; - } - - if (_classifyInlineRejection(err) !== "needs-json") { - throw err; - } + if (signal?.aborted) throw err; + if (_classifyInlineRejection(err) !== "needs-json") throw err; const msg = err instanceof Error ? err.message : String(err); logger.warn( @@ -470,6 +555,46 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { return makeArrowMessage(result.statement_id, { status: result.status }); } + /** + * Decode an INLINE Arrow attachment, push it onto the stash, and + * return the `arrow` SSE message that references the synthetic id. + * Pulled out of `_executeArrowStreamPath` so the cap / overflow / + * exhaustion logic lives in one place. + */ + private _stashAndEmitInline( + result: { attachment: string; status?: unknown }, + stashUserKey: string, + signal?: AbortSignal, + ): AnalyticsSseMessage { + // If the client has already disconnected, the SSE write would be + // dropped anyway — skip the decode + stash so the bytes do not + // linger in memory until TTL eviction. + if (signal?.aborted) { + throw ExecutionError.canceled(); + } + const decoded = Buffer.from(result.attachment, "base64"); + const stashBytes = new Uint8Array( + decoded.buffer, + decoded.byteOffset, + decoded.byteLength, + ); + const putResult = this.inlineArrowStash.put(stashUserKey, stashBytes); + if (putResult === null) { + // Both regular and overflow pools are at cap. Throw with a clear + // signal rather than re-running the same statement on + // EXTERNAL_LINKS: the bytes we already decoded are gone, the + // warehouse has been billed, and a second execution could return + // a divergent result for non-deterministic SQL. + this._recordStashOutcome("exhausted", stashBytes.byteLength); + logger.warn( + "Inline Arrow stash exhausted (regular + overflow); rejecting", + ); + throw ExecutionError.stashExhausted(); + } + this._recordStashOutcome(putResult.pool, stashBytes.byteLength); + return makeArrowMessage(putResult.id, { status: result.status }); + } + /** * Execute a SQL query using the current execution context. * diff --git a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts index 31a91dfd..cb52ff71 100644 --- a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts +++ b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts @@ -135,9 +135,17 @@ export class InlineArrowStash { * split across — so the largest payload we can accept is * `Math.max(maxBytes, maxOverflowBytes)`. Exceeding that throws * synchronously so the caller sees the misconfiguration loudly - * rather than burning a warehouse round-trip and then getting `null`. + * rather than burning a warehouse round-trip and then getting a + * null id. + * + * Returns `{ id, pool }` on success (`pool` ∈ {"regular", "overflow"}) + * or `null` when both pools are at cap. The pool tag lets callers + * emit accurate telemetry labels without re-introspecting the stash. */ - put(userId: string, bytes: Uint8Array): string | null { + put( + userId: string, + bytes: Uint8Array, + ): { id: string; pool: "regular" | "overflow" } | null { const largestSlot = Math.max(this.maxBytes, this.maxOverflowBytes); if (bytes.length > largestSlot) { throw new Error( @@ -168,7 +176,7 @@ export class InlineArrowStash { } else { this.totalBytes += bytes.length; } - return id; + return { id, pool: overflow ? "overflow" : "regular" }; } /** diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index fb6ecbf9..9c7d9698 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -113,7 +113,7 @@ describe("Analytics Plugin", () => { plugin.injectRoutes(router); const arrowBytes = new Uint8Array([0xff, 0xfe, 0xfd, 0xfc]); - const id = (plugin as any).inlineArrowStash.put("global", arrowBytes); + const { id } = (plugin as any).inlineArrowStash.put("global", arrowBytes); expect(id.startsWith("inline-")).toBe(true); const handler = getHandler("GET", "/arrow-result/:jobId"); @@ -175,7 +175,7 @@ describe("Analytics Plugin", () => { // (no x-forwarded-user header) — keys differ, take must return // nothing, and the entry stays put (single-user view). const bytes = new Uint8Array([1, 2, 3]); - const id = (plugin as any).inlineArrowStash.put("user-a", bytes); + const { id } = (plugin as any).inlineArrowStash.put("user-a", bytes); const handler = getHandler("GET", "/arrow-result/:jobId"); const mockReq = createMockRequest({ params: { jobId: id } }); diff --git a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts index 31d9b01d..05ae821b 100644 --- a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts @@ -5,19 +5,20 @@ function bytes(n: number): Uint8Array { return new Uint8Array(n); } -// `put()` returns `string | null` — it rejects with null when the stash is -// full. Every test below that exercises a successful put narrows via this -// helper so the non-null contract is explicit at the call site. +// `put()` returns `{ id, pool } | null` — it rejects with null when the stash +// is full. Most tests only care about the id; this helper narrows via the +// non-null contract and returns just the id. Tests that need the pool tag +// call `stash.put(...)` directly. function mustPut( stash: InlineArrowStash, userId: string, b: Uint8Array, ): string { - const id = stash.put(userId, b); - if (id === null) { + const result = stash.put(userId, b); + if (result === null) { throw new Error("test setup: stash unexpectedly rejected put"); } - return id; + return result.id; } describe("InlineArrowStash", () => { @@ -119,6 +120,17 @@ describe("InlineArrowStash", () => { expect(stash.take(b, "user-1")).toBeDefined(); }); + test("put tags the result with its pool so callers can label telemetry without re-introspecting", () => { + const stash = new InlineArrowStash({ + maxBytes: 100, + maxOverflowBytes: 100, + }); + expect(stash.put("u", bytes(80))).toMatchObject({ pool: "regular" }); + // Regular has 80/100 used; another 80 would push it to 160 > 100 so it + // spills. + expect(stash.put("u", bytes(80))).toMatchObject({ pool: "overflow" }); + }); + test("put throws when a single payload would not fit in the largest pool (caller misconfiguration)", () => { const stash = new InlineArrowStash({ maxBytes: 100, diff --git a/packages/shared/src/sse/analytics.ts b/packages/shared/src/sse/analytics.ts index 31aaf511..eba7597e 100644 --- a/packages/shared/src/sse/analytics.ts +++ b/packages/shared/src/sse/analytics.ts @@ -34,22 +34,30 @@ export const AnalyticsResultMessage = z.object({ // 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 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(), }); -// The TS-level shape narrows `data` to `Record[]` to -// match the typed `makeResultMessage` builder — the Zod schema is -// intentionally looser at runtime for performance (see comment above). -export type AnalyticsResultMessage = Omit< - z.infer, - "data" -> & { + +/** + * 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; +} /** * ARROW_STREAM result delivered via /arrow-result/:jobId. The id is either: From ed3dfeb67d856a552f03544295fc329de9500edd Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Thu, 28 May 2026 11:24:50 +0000 Subject: [PATCH 18/25] docs: regenerate api docs for clientMessage field --- docs/docs/api/appkit/Class.AppKitError.md | 39 ++++++++- .../api/appkit/Class.AuthenticationError.md | 47 ++++++++++- .../api/appkit/Class.ConfigurationError.md | 47 ++++++++++- docs/docs/api/appkit/Class.ConnectionError.md | 47 ++++++++++- docs/docs/api/appkit/Class.ExecutionError.md | 80 ++++++++++++++++++- .../api/appkit/Class.InitializationError.md | 47 ++++++++++- docs/docs/api/appkit/Class.ServerError.md | 47 ++++++++++- docs/docs/api/appkit/Class.TunnelError.md | 47 ++++++++++- docs/docs/api/appkit/Class.ValidationError.md | 47 ++++++++++- 9 files changed, 436 insertions(+), 12 deletions(-) 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 eacfdc23..c7bb2b44 100644 --- a/docs/docs/api/appkit/Class.ExecutionError.md +++ b/docs/docs/api/appkit/Class.ExecutionError.md @@ -21,6 +21,7 @@ throw new ExecutionError("Statement was canceled"); ```ts new ExecutionError(message: string, options?: { cause?: Error; + clientMessage?: string; context?: Record; errorCode?: string; }): ExecutionError; @@ -31,8 +32,9 @@ new ExecutionError(message: string, options?: { | Parameter | Type | | ------ | ------ | | `message` | `string` | -| `options?` | \{ `cause?`: `Error`; `context?`: `Record`\<`string`, `unknown`\>; `errorCode?`: `string`; \} | +| `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` | @@ -46,6 +48,28 @@ new ExecutionError(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 @@ -127,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` +(`INLINE_ARROW_STASH_EXHAUSTED`, `NOT_IMPLEMENTED`, etc.) and not on +the human string. + +##### Returns + +`string` + +#### Overrides + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + ## Methods ### toJSON() @@ -214,10 +262,33 @@ Create an execution error for closed/expired results *** +### stashExhausted() + +```ts +static stashExhausted(): ExecutionError; +``` + +Create an execution error for the inline Arrow stash being unable to +accept a payload (both regular and overflow pools at cap). + +The route deliberately surfaces this rather than silently re-running +the statement on EXTERNAL_LINKS — a second execution can be billed +and return divergent results for non-deterministic SQL. Operators +should tune stash capacity or back off load when this fires. + +#### Returns + +`ExecutionError` + +*** + ### statementFailed() ```ts -static statementFailed(errorMessage?: string, errorCode?: string): ExecutionError; +static statementFailed( + errorMessage?: string, + errorCode?: string, + clientMessage?: string): ExecutionError; ``` Create an execution error for statement failure. @@ -226,8 +297,9 @@ Create an execution error for statement failure. | Parameter | Type | Description | | ------ | ------ | ------ | -| `errorMessage?` | `string` | Human-readable error from the warehouse / SDK. | -| `errorCode?` | `string` | Structured code (e.g. "INVALID_PARAMETER_VALUE") to preserve through wrapping. Optional. | +| `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() From f1ef36b3c6effecba52872dbb2c185edfa019d7b Mon Sep 17 00:00:00 2001 From: Mario Cadenas <17888484+MarioCadenas@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:05:40 +0200 Subject: [PATCH 19/25] refactor(appkit): direct Arrow IPC streaming + hardened result delivery (#472) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(appkit): direct Arrow IPC streaming + hardened result delivery Rework of the Arrow/JSON result-delivery path on top of #329's arrow foundation (base branch is #329 merged with current main). Delivery: - ARROW_STREAM queries stream the raw Arrow IPC bytes on the /query response body — no SSE, no server-side stash, no second /arrow-result request. The inline stash was process-local (broke horizontal scaling) and buffered whole payloads; it and its route are removed. - Capability fallback centralized in result-delivery.ts (errorCode-first classification): INLINE+ARROW_STREAM (Reyden) -> EXTERNAL_LINKS -> a structured ARROW_DELIVERY_UNSUPPORTED error; the JSON path retries as ARROW_STREAM and decodes server-side when the warehouse is inline-only. - EXTERNAL_LINKS chunks pipe response.body one network read at a time (peak memory = one read buffer). Multi-chunk results resolve every chunk's links via next_chunk_index in the request's identity context (OBO-correct) — previously only chunk 0 was streamed (silent truncation). - Real column names ride an X-Appkit-Arrow-Columns header (statement-id ref + /columns endpoint for wide schemas); the client relabels the positional col_N schema. Hardening (pre-merge review): - writeChunk no longer hangs on a client disconnect mid-backpressure — it races drain vs close/error so the stream unwinds and the upstream reader is cancelled instead of leaking. - Fail-fast first-byte timeout returns 503 WAREHOUSE_UNAVAILABLE on a stuck warehouse. - /columns resolves under the user identity then the service principal (fixes OBO wide-schema fallback). - Retry backoff is abortable; arrow errors forward the actionable clientMessage. Green: typecheck, 3220 tests + 1 skipped, biome, build + publint. Multi-chunk external, OBO+external, and wide-schema /columns are unit-tested only (no warm standard warehouse to live-verify). * fix(appkit): idle-timeout backpressure, OBO warehouse identity, writeChunk close guard Addresses three defects surfaced in external review of the Arrow rework: - Idle-timeout vs backpressure (arrow-stream-processor.ts): the per-chunk idle timeout was armed before reader.read() but stayed armed across the downstream `yield`, so a slow client backpressuring writeChunk past the timeout (60s) was mistaken for an upstream stall and aborted a healthy download. The timeout now covers only the read (armed before, cleared in a finally the instant it resolves). - OBO warehouse startup ran as the service principal (analytics.ts): the direct-binary Arrow path called getWorkspaceClient() + ensureWarehouseRunning bare in the handler body, so warehouse auto-start ran as the SP even for .obo.sql requests — unlike the SSE path, which runs readiness inside the executor (user) context. Readiness now runs via the request executor (asUser for OBO) through _ensureArrowWarehouseReady, matching the SSE path. - writeChunk hang on an already-closed socket (analytics.ts): if res was already destroyed/ended when writeChunk was entered, res.write returned false while drain/close/error had already fired, so the promise never settled and wedged the for-await loop + upstream reader. It now rejects up front when res.destroyed || res.writableEnded. Regression tests added: slow-consumer-between-reads (no false abort), already-closed writeChunk, and OBO readiness routed through the user executor. Green: typecheck, 3222 tests + 1 skipped, biome, build + publint. Signed-off-by: MarioCadenas --------- Signed-off-by: MarioCadenas Co-authored-by: MarioCadenas --- docs/docs/api/appkit/Class.ExecutionError.md | 22 +- .../src/js/arrow/arrow-client.test.ts | 60 ++ .../appkit-ui/src/js/arrow/arrow-client.ts | 32 +- packages/appkit-ui/src/js/sse/connect-sse.ts | 4 +- .../__tests__/use-analytics-query.test.ts | 143 ++- packages/appkit-ui/src/react/hooks/types.ts | 8 +- .../src/react/hooks/use-analytics-query.ts | 185 +++- .../src/connectors/sql-warehouse/client.ts | 275 ++++-- .../sql-warehouse/tests/client.test.ts | 159 ++- packages/appkit/src/errors/execution.ts | 21 +- .../appkit/src/plugins/analytics/analytics.ts | 932 ++++++------------ .../plugins/analytics/inline-arrow-stash.ts | 240 ----- .../src/plugins/analytics/result-delivery.ts | 431 ++++++++ .../plugins/analytics/tests/analytics.test.ts | 574 ++++++----- .../tests/arrow-delivery.integration.test.ts | 68 ++ .../tests/inline-arrow-stash.test.ts | 184 ---- .../analytics/tests/result-delivery.test.ts | 266 +++++ .../appkit/src/plugins/analytics/types.ts | 8 + .../src/stream/arrow-stream-processor.ts | 315 +++--- packages/appkit/src/stream/defaults.ts | 11 +- packages/appkit/src/stream/stream-manager.ts | 2 +- .../tests/arrow-stream-processor.test.ts | 621 +++--------- packages/appkit/src/stream/types.ts | 2 +- packages/shared/src/sse/analytics.test.ts | 57 +- packages/shared/src/sse/analytics.ts | 63 +- tools/test-helpers.ts | 32 +- 26 files changed, 2347 insertions(+), 2368 deletions(-) create mode 100644 packages/appkit-ui/src/js/arrow/arrow-client.test.ts delete mode 100644 packages/appkit/src/plugins/analytics/inline-arrow-stash.ts create mode 100644 packages/appkit/src/plugins/analytics/result-delivery.ts create mode 100644 packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts delete mode 100644 packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts create mode 100644 packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts diff --git a/docs/docs/api/appkit/Class.ExecutionError.md b/docs/docs/api/appkit/Class.ExecutionError.md index c7bb2b44..e7c6f457 100644 --- a/docs/docs/api/appkit/Class.ExecutionError.md +++ b/docs/docs/api/appkit/Class.ExecutionError.md @@ -164,7 +164,7 @@ 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` -(`INLINE_ARROW_STASH_EXHAUSTED`, `NOT_IMPLEMENTED`, etc.) and not on +(`RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`, etc.) and not on the human string. ##### Returns @@ -262,26 +262,6 @@ Create an execution error for closed/expired results *** -### stashExhausted() - -```ts -static stashExhausted(): ExecutionError; -``` - -Create an execution error for the inline Arrow stash being unable to -accept a payload (both regular and overflow pools at cap). - -The route deliberately surfaces this rather than silently re-running -the statement on EXTERNAL_LINKS — a second execution can be billed -and return divergent results for non-deterministic SQL. Operators -should tune stash capacity or back off load when this fires. - -#### Returns - -`ExecutionError` - -*** - ### statementFailed() ```ts 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 13d9053d..010ba3bf 100644 --- a/packages/appkit-ui/src/js/sse/connect-sse.ts +++ b/packages/appkit-ui/src/js/sse/connect-sse.ts @@ -19,8 +19,8 @@ export async function connectSSE( retryDelay = 2000, maxRetries = 3, // 1 MiB — matches the server's `streamDefaults.maxEventSize`. SSE - // carries only short JSON control messages; bulk Arrow payloads flow - // over plain HTTP via `/api/analytics/arrow-result/:jobId`, so this + // 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 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 8b1acc13..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,5 +1,5 @@ import { act, renderHook, waitFor } from "@testing-library/react"; -import { beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; let lastConnectArgs: any = null; let capturedCallbacks: { @@ -52,72 +52,83 @@ describe("useAnalyticsQuery", () => { capturedCallbacks = {}; }); - test("fetches an arrow message (warehouse statement id) via /arrow-result", async () => { + 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]); - mockFetchArrow.mockResolvedValueOnce(fakeBytes); + 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 lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow", statement_id: "stmt-warehouse-1" }), - }); - await waitFor(() => { expect(result.current.data).toBe(fakeTable); }); - expect(mockFetchArrow).toHaveBeenCalledTimes(1); - expect(mockFetchArrow).toHaveBeenCalledWith( - "/api/analytics/arrow-result/stmt-warehouse-1", + // 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(mockProcessArrowBuffer).toHaveBeenCalledWith(fakeBytes); + expect(result.current.loading).toBe(false); }); - test("fetches an arrow message with synthetic inline- id through the same /arrow-result path", async () => { - // The client must treat inline and external-links responses uniformly — - // it never decodes base64 locally. The /arrow-result route on the - // server is the only place that knows which path the bytes came from. + test("relabels ARROW_STREAM columns from the X-Appkit-Arrow-Columns header", async () => { const fakeTable = { numRows: 1, schema: { fields: [] } }; - const fakeBytes = new Uint8Array([1, 2, 3, 4, 5]); - mockFetchArrow.mockResolvedValueOnce(fakeBytes); + 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 lastConnectArgs.onMessage({ - data: JSON.stringify({ - type: "arrow", - statement_id: "inline-abc-xyz", - }), - }); - await waitFor(() => { expect(result.current.data).toBe(fakeTable); }); - expect(mockFetchArrow).toHaveBeenCalledTimes(1); - expect(mockFetchArrow).toHaveBeenCalledWith( - "/api/analytics/arrow-result/inline-abc-xyz", + // 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 fetch fails", async () => { - mockFetchArrow.mockRejectedValueOnce(new Error("network")); + 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 lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow", statement_id: "stmt-1" }), - }); - await waitFor(() => { expect(result.current.error).toBe( "Unable to load data, please try again", @@ -126,27 +137,28 @@ describe("useAnalyticsQuery", () => { expect(result.current.loading).toBe(false); }); - test("rejects the retired arrow_inline message type as schema-invalid", async () => { - // arrow_inline was the prior wire shape. The discriminated union no - // longer accepts it, so it falls through to the generic error/code - // branch — but critically, it must NEVER trigger ArrowClient calls. + 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 lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow_inline", attachment: "AQID" }), - }); - await waitFor(() => { - expect( - result.current.loading || - result.current.error || - result.current.data === null, - ).toBeTruthy(); + expect(result.current.error).toBe("Result too large for JSON format"); }); - expect(mockProcessArrowBuffer).not.toHaveBeenCalled(); - expect(mockFetchArrow).not.toHaveBeenCalled(); + 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 () => { @@ -208,15 +220,14 @@ describe("useAnalyticsQuery", () => { }); 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 (e.g. INLINE_ARROW_STASH_EXHAUSTED). 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. + // 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: "ARROW_STREAM" }), + useAnalyticsQuery("q", null, { format: "JSON_ARRAY" }), ); await lastConnectArgs.onMessage({ @@ -224,7 +235,7 @@ describe("useAnalyticsQuery", () => { type: "error", error: "Server is at capacity, please retry", code: "UPSTREAM_ERROR", - errorCode: "INLINE_ARROW_STASH_EXHAUSTED", + errorCode: "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", }), }); @@ -232,7 +243,7 @@ describe("useAnalyticsQuery", () => { expect(result.current.error).toBe("Server is at capacity, please retry"); }); expect(result.current.loading).toBe(false); - expect(result.current.errorCode).toBe("INLINE_ARROW_STASH_EXHAUSTED"); + expect(result.current.errorCode).toBe("RESULT_TOO_LARGE_FOR_JSON_FALLBACK"); errorSpy.mockRestore(); }); @@ -447,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/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index 5112ac94..aa0df890 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -96,10 +96,10 @@ export interface UseAnalyticsQueryResult { error: string | null; /** * Structured upstream error code when the server attaches one to the - * SSE error payload (e.g. `INLINE_ARROW_STASH_EXHAUSTED`, - * `RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`). Prefer - * branching UI on this stable identifier rather than substring- - * matching `error`, which is a free-form sanitized message. + * 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; /** 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 a27bb6f3..20f2c99e 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -62,12 +62,21 @@ 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; @@ -124,28 +133,9 @@ async function handleAnalyticsSseMessage( return; } - // success - Arrow format. Both INLINE (server-stashed, statement_id - // prefixed with "inline-") and EXTERNAL_LINKS (warehouse statement_id) - // flow through this single branch — the /arrow-result route dispatches - // based on the id prefix so the client doesn't need to know which path - // the bytes came from. - if (msg?.type === "arrow") { - try { - const arrowData = await ArrowClient.fetchArrow( - getArrowStreamUrl(msg.statement_id), - ); - 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 = @@ -156,9 +146,9 @@ async function handleAnalyticsSseMessage( ctx.setError(errorMsg); ctx.unpublishWarehouseStatus(); // Propagate the upstream structured code so UI consumers can branch on - // a stable identifier (e.g. retry on INLINE_ARROW_STASH_EXHAUSTED, - // format-switch on RESULT_TOO_LARGE_FOR_JSON_FALLBACK) instead of - // parsing the human-readable message. + // 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); } @@ -183,8 +173,116 @@ async function handleAnalyticsSseMessage( } } +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; +} + +/** + * 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 over SSE and returns its latest result. + * 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: @@ -289,6 +387,22 @@ 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, @@ -323,7 +437,7 @@ export function useAnalyticsQuery< // forces the consumer into a clean retry path. console.warn("[useAnalyticsQuery] Malformed message received", error); setLoading(false); - setError("Unable to load data, please try again"); + setError(GENERIC_LOAD_ERROR); abortController.abort(); } }, @@ -332,26 +446,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, ]); diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 9167ef97..239de551 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -29,18 +29,42 @@ 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). * - * Bulk Arrow payloads no longer traverse SSE — the analytics route stashes - * them via `InlineArrowStash` and the client fetches over HTTP — so this - * cap is bounded by the upstream API rather than our event-size budget. - * Larger results still fall through to `disposition: "EXTERNAL_LINKS"`, - * handled by the analytics format-fallback. + * 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; @@ -167,8 +191,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, }); } @@ -287,7 +309,11 @@ export class SQLWarehouseConnector { ); break; case "SUCCEEDED": - result = this._transformDataArray(response); + result = await this._transformDataArray( + response, + workspaceClient, + signal, + ); break; case "FAILED": throw ExecutionError.statementFailed( @@ -816,7 +842,11 @@ 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, @@ -866,7 +896,11 @@ export class SQLWarehouseConnector { ); } - private _transformDataArray(response: sql.StatementResponse) { + private async _transformDataArray( + response: sql.StatementResponse, + workspaceClient: WorkspaceClient, + signal?: AbortSignal, + ) { if (response.manifest?.format === "ARROW_STREAM") { const result = response.result as | (sql.ResultData & { attachment?: string }) @@ -882,7 +916,7 @@ export class SQLWarehouseConnector { // External links: data fetched separately via statement_id. if (result?.external_links) { - return this.updateWithArrowStatus(response); + return this.updateWithArrowStatus(response, workspaceClient, signal); } // Empty result with a known schema: synthesize a zero-row Arrow IPC @@ -955,22 +989,24 @@ export class SQLWarehouseConnector { } /** - * Validate (but do not decode) a base64 Arrow IPC attachment. - * Some serverless warehouses return inline results as Arrow IPC in - * `result.attachment`. We pass the base64 string through to the client, - * which decodes it into an Arrow Table via the existing ArrowClient - * infrastructure. This keeps the wire contract for ARROW_STREAM - * consistent (client always receives an Arrow Table) and avoids - * decode/re-encode work on the server. + * 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. 64 MiB is well above the typical inline limit - // (~25 MiB hard cap on the API) but bounds memory if a server returns - // a runaway response. + // 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. @@ -986,101 +1022,146 @@ export class SQLWarehouseConnector { `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 updateWithArrowStatus(response: sql.StatementResponse): { - result: { statement_id: string; status: sql.StatementStatus }; - } { + private async updateWithArrowStatus( + response: sql.StatementResponse, + workspaceClient: WorkspaceClient, + signal?: AbortSignal, + ): Promise<{ + result: { + statement_id: string; + status: sql.StatementStatus; + columnNames?: string[]; + external_links?: sql.ExternalLink[]; + }; + }> { + 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, + ), }, }; } - 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", - }); - - logger.event()?.setContext("sql-warehouse", { - arrow_data_size_bytes: result.data.length, - arrow_job_id: jobId, - }); - - return result; - } catch (error) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: error instanceof Error ? error.message : "Unknown error", - }); - span.recordException(error as Error); + ): 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; + } - const duration = Date.now() - startTime; - this.telemetryMetrics.queryDuration.record(duration, { - operation: "arrow.getData", - status: "error", - }); + /** 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; + } - logger.error("Failed Arrow job: %s %O", jobId, 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, + ): AsyncGenerator { + return this.arrowProcessor.streamChunks(chunks, signal); + } - if (error instanceof AppKitError) { - throw error; - } - throw ExecutionError.statementFailed( - error instanceof Error ? error.message : String(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); } // create context for cancellation token diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts index c6780f3f..dc964702 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts @@ -37,6 +37,18 @@ 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]. @@ -45,7 +57,7 @@ const REAL_ARROW_ATTACHMENT = describe("SQLWarehouseConnector._transformDataArray", () => { describe("classic warehouse (JSON_ARRAY + INLINE)", () => { - test("transforms data_array rows into named objects", () => { + test("transforms data_array rows into named objects", async () => { const connector = createConnector(); // Real response shape from classic warehouse: INLINE + JSON_ARRAY const response = { @@ -78,12 +90,12 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + 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", () => { + test("parses JSON strings in STRING columns", async () => { const connector = createConnector(); const response = { statement_id: "stmt-1", @@ -102,13 +114,13 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + 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", () => { + test("returns statement_id for external links fetch", async () => { const connector = createConnector(); // Real response shape from classic warehouse: EXTERNAL_LINKS + ARROW_STREAM const response = { @@ -133,14 +145,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + 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", () => { + 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. @@ -179,14 +191,18 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + 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", () => { + test("preserves manifest and status alongside attachment", async () => { const connector = createConnector(); const response = { statement_id: "00000001-test-stmt", @@ -207,14 +223,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + 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", () => { + 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 @@ -240,7 +256,7 @@ describe("SQLWarehouseConnector._transformDataArray", () => { result: {}, } as unknown as sql.StatementResponse; - const transformed = (connector as any)._transformDataArray(response); + const transformed = await transform(connector, response); const attachment: string = transformed.result.attachment; expect(typeof attachment).toBe("string"); expect(attachment.length).toBeGreaterThan(0); @@ -255,7 +271,7 @@ describe("SQLWarehouseConnector._transformDataArray", () => { ]); }); - test("does NOT synthesize an attachment when external_links are present", () => { + test("does NOT synthesize an attachment when external_links are present", async () => { const connector = createConnector(); const response = { statement_id: "stmt-ext", @@ -271,13 +287,13 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const transformed = (connector as any)._transformDataArray(response); + 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("does NOT synthesize an attachment when schema is missing", () => { + test("does NOT synthesize an attachment when schema is missing", async () => { const connector = createConnector(); const response = { statement_id: "stmt-no-schema", @@ -286,12 +302,12 @@ describe("SQLWarehouseConnector._transformDataArray", () => { result: {}, } as unknown as sql.StatementResponse; - const transformed = (connector as any)._transformDataArray(response); + 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", () => { + 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. @@ -303,14 +319,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { result: { attachment: oversized }, } as unknown as sql.StatementResponse; - expect(() => (connector as any)._transformDataArray(response)).toThrow( + 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", () => { + test("transforms data_array like JSON_ARRAY path", async () => { const connector = createConnector(); const response = { statement_id: "stmt-1", @@ -332,7 +348,7 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); expect(result.result.data).toEqual([ { id: "1", value: "hello" }, { id: "2", value: "world" }, @@ -341,7 +357,7 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }); describe("edge cases", () => { - test("returns response unchanged when no data_array, attachment, or schema", () => { + test("returns response unchanged when no data_array, attachment, or schema", async () => { const connector = createConnector(); const response = { statement_id: "stmt-1", @@ -350,11 +366,11 @@ describe("SQLWarehouseConnector._transformDataArray", () => { result: {}, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); expect(result).toBe(response); }); - test("attachment takes priority over data_array when both present", () => { + test("attachment takes priority over data_array when both present", async () => { const connector = createConnector(); const response = { statement_id: "stmt-1", @@ -374,10 +390,105 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + 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/execution.ts b/packages/appkit/src/errors/execution.ts index 3c9df7fe..1c9acf91 100644 --- a/packages/appkit/src/errors/execution.ts +++ b/packages/appkit/src/errors/execution.ts @@ -40,7 +40,7 @@ export class ExecutionError extends AppKitError { * 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` - * (`INLINE_ARROW_STASH_EXHAUSTED`, `NOT_IMPLEMENTED`, etc.) and not on + * (`RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`, etc.) and not on * the human string. */ override get clientMessage(): string { @@ -108,23 +108,4 @@ export class ExecutionError extends AppKitError { context: { dataType }, }); } - - /** - * Create an execution error for the inline Arrow stash being unable to - * accept a payload (both regular and overflow pools at cap). - * - * The route deliberately surfaces this rather than silently re-running - * the statement on EXTERNAL_LINKS — a second execution can be billed - * and return divergent results for non-deterministic SQL. Operators - * should tune stash capacity or back off load when this fires. - */ - static stashExhausted(): ExecutionError { - return new ExecutionError( - "Inline Arrow stash exhausted; retry shortly or increase stash capacity", - { - errorCode: "INLINE_ARROW_STASH_EXHAUSTED", - clientMessage: "Server is at capacity, please retry", - }, - ); - } } diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 0c49f29c..1b5aa1d5 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1,11 +1,8 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; -import { type DataType, Type, tableFromIPC } from "apache-arrow"; import type express from "express"; import { type AgentToolDefinition, type AnalyticsSseMessage, type IAppRouter, - makeArrowMessage, makeResultMessage, type PluginExecuteConfig, type SQLTypeMarker, @@ -30,13 +27,11 @@ import { AppKitError, ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; -import type { Counter, Histogram } from "../../telemetry"; import { queryDefaults } from "./defaults"; -import { InlineArrowStash } from "./inline-arrow-stash"; import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; +import { deliverArrowBytes, deliverJsonResult } from "./result-delivery"; import { - type AnalyticsFormat, type AnalyticsQueryResponse, type AnalyticsStreamMessage, type IAnalyticsConfig, @@ -106,33 +101,6 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { private SQLClient: SQLWarehouseConnector; private queryProcessor: QueryProcessor; - /** - * Server-side stash for inline Arrow IPC payloads. - * - * INLINE ARROW_STREAM responses do not ride the SSE control channel — - * the route puts the decoded bytes here and emits an `arrow` SSE - * message with a synthetic `inline-` id, and the client fetches - * the bytes through the existing `/arrow-result/:jobId` endpoint with - * a real binary content-type. - */ - protected inlineArrowStash: InlineArrowStash = new InlineArrowStash(); - - /** - * Stash telemetry — created lazily on first stash operation because - * `this.telemetry` may not be bound at constructor time (see - * `Plugin.attachContext`). Cached via `_stashMetrics` after the first - * call; subsequent puts hit the cached counter/histogram instances. - * - * Counter labels: `result` ∈ {"regular", "overflow", "exhausted"} — - * one counter with a label is friendlier to dashboards than three - * separate metrics, and aggregates trivially in Prometheus / OTel. - */ - private _stashMetrics?: { - putCount: Counter; - putBytes: Histogram; - }; - private _stashMetricsAttempted = false; - constructor(config: IAnalyticsConfig) { super(config); this.config = config; @@ -145,19 +113,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", @@ -166,156 +121,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. - * - * Two id shapes are supported: - * - `inline-`: bytes were stashed server-side by the query route. - * Drain the stash, serve directly with the canonical Arrow content - * type. No warehouse round-trip. - * - any other id: a warehouse-issued statement id. Fetch the Arrow - * stream from the warehouse via the SDK; serve the bytes. - * - * When called via asUser(req), uses the user's Databricks credentials - * for the warehouse path. The inline path is user-scoped at the stash - * layer instead. + * 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 { - const { jobId } = req.params; - const event = logger.event(req); - event?.setComponent("analytics", "getArrowData").setContext("analytics", { - job_id: jobId, - plugin: this.name, - }); - - if (jobId.startsWith("inline-")) { - const userKey = this._stashUserKey(req); - const bytes = this.inlineArrowStash.take(jobId, userKey); - if (!bytes) { - // Already drained, expired, or never belonged to this user. 410 - // distinguishes this from "warehouse statement id not found" (404) - // so the client can surface a useful error. - logger.debug("Inline Arrow stash miss for jobId=%s", jobId); - res.status(410).json({ - error: "Inline Arrow result expired or unknown", - plugin: this.name, - }); - return; - } - logger.debug( - "Serving inline Arrow buffer: %d bytes for jobId=%s", - bytes.length, - jobId, - ); - res.setHeader("Content-Type", "application/vnd.apache.arrow.stream"); - res.setHeader("Content-Length", bytes.length.toString()); - // Inline payloads are single-use and short-lived; no public caching. + const { statementId } = req.params; + const columns = await this._resolveColumnNames(req, statementId); + if (columns && columns.length > 0) { res.setHeader("Cache-Control", "no-store"); - res.send(Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength)); + res.json({ columns }); return; } - - try { - const workspaceClient = getWorkspaceClient(); - logger.debug("Processing Arrow job request for jobId=%s", jobId); - - 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"); - - 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); - // Do not echo upstream / SDK error text to the client — it can - // include statement fragments, internal object names, and - // correlation IDs. The full detail stays in the server log above. - const errorCode = - error instanceof ExecutionError ? error.errorCode : undefined; - res.status(404).json({ - error: "Arrow result unavailable", - errorCode, - plugin: this.name, - }); - } + res.status(404).json({ + error: "Column names unavailable", + plugin: this.name, + }); } /** - * Lazy accessor for stash telemetry instruments. Returns `undefined` - * when telemetry isn't bound (factory-constructed plugins before - * `attachContext`, no-op test paths) — callers should branch on the - * return value rather than relying on it being present. - * - * Records once per call site: - * - `analytics.inline_arrow_stash.put.count{result}` — counter with - * label "regular" | "overflow" | "exhausted" - * - `analytics.inline_arrow_stash.put.bytes` — histogram of accepted - * payload sizes (label same as counter; exhausted puts record 0) + * 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 _getStashMetrics() { - if (this._stashMetricsAttempted) return this._stashMetrics; - this._stashMetricsAttempted = true; - if (!this.telemetry) return undefined; - const meter = this.telemetry.getMeter(); - this._stashMetrics = { - putCount: meter.createCounter("analytics.inline_arrow_stash.put.count", { - description: - "Count of inline-Arrow stash put attempts, labeled by pool outcome (regular | overflow | exhausted).", - unit: "1", - }), - putBytes: meter.createHistogram( - "analytics.inline_arrow_stash.put.bytes", - { - description: - "Sizes of accepted inline-Arrow stash payloads. Aggregate by `result` label for utilization analysis.", - unit: "By", - }, - ), - }; - return this._stashMetrics; - } - - private _recordStashOutcome( - result: "regular" | "overflow" | "exhausted", - bytes: number, - ): void { - const m = this._getStashMetrics(); - if (!m) return; - m.putCount.add(1, { result }); - m.putBytes.record(result === "exhausted" ? 0 : bytes, { result }); + 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; } /** - * Stash key used at put-time (in `_handleQueryRoute`) and take-time - * (in `_handleArrowRoute`). Centralized so the two sides cannot drift. - * - * Returns the user id when an `x-forwarded-user` header is present, - * otherwise `"global"` for service-principal contexts (no user header). - * Both queries from the same request resolve to the same key, so the - * subsequent /arrow-result fetch reliably hits the entry stashed - * during the SSE query. - * - * `resolveUserId` throws when no header is present — catch and degrade - * to "global" rather than letting that failure mode bubble through the - * route handler. + * 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. */ - protected _stashUserKey(req: express.Request): string { - try { - return this.resolveUserId(req) || "global"; - } catch { - return "global"; - } + async _getColumnNames(statementId: string): Promise { + return this.SQLClient.getColumnNames(getWorkspaceClient(), statementId); } /** @@ -373,31 +253,25 @@ 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, 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"; - // Stash key is always per-request user (never "global"), independent - // of the executor's cache scope. Inline Arrow payloads are single-use - // and short-lived — there is no benefit to sharing them across users, - // and per-user scoping is defense in depth on top of unguessable ids. - const stashUserKey = this._stashUserKey(req); const hashedQuery = this.queryProcessor.hashQuery(query); - // ARROW_STREAM responses reference ephemeral resources that cannot be - // safely replayed from cache: - // - EXTERNAL_LINKS pre-signed URLs expire ~15 min after issue, and - // the warehouse rotates them per execution. - // - INLINE responses point at a synthetic `inline-` job id - // backed by `InlineArrowStash`, which drains on the first - // /arrow-result fetch. A cache hit would replay an id whose bytes - // are already gone and reliably 410 the client. - // So we bypass cache for ARROW_STREAM and let every request execute - // a fresh statement. JSON_ARRAY responses still cache normally. - const cacheTtl = format === "ARROW_STREAM" ? 0 : queryDefaults.cache?.ttl; const cacheConfig = { ...queryDefaults.cache, - ttl: cacheTtl, cacheKey: [ "analytics:query", query_key, @@ -476,15 +350,15 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { try { const processedParams = await self.queryProcessor.processQueryParams(query, parameters); - // Disposition/format fallback (inline-arrow stash, JSON_ARRAY↔ - // ARROW_STREAM retries) lives in `_executeWithFormatFallback`, - // which returns the SSE message the client should receive. - return await self._executeWithFormatFallback( + // 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, - format, - stashUserKey, sig, ); } catch (err) { @@ -511,7 +385,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { throw err; } // Re-throw the original error so its structured `errorCode` (e.g. - // INLINE_ARROW_STASH_EXHAUSTED) and sanitized `clientMessage` + // 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) { @@ -531,215 +405,204 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } /** - * Execute a query with automatic disposition/format fallback. - * - * - **JSON_ARRAY** first tries `INLINE + JSON_ARRAY`. If the warehouse - * only supports `ARROW_STREAM` for `INLINE` (some serverless variants), - * retries as `INLINE + ARROW_STREAM`, decodes the Arrow IPC attachment - * server-side, and returns plain row objects — the caller's - * `JSON_ARRAY` contract is preserved. - * - **ARROW_STREAM** first tries `INLINE + ARROW_STREAM`. If the - * warehouse refuses (most classic + some serverless variants), falls - * back to `EXTERNAL_LINKS + ARROW_STREAM`. When the regular stash is - * full, the already-decoded bytes spill into the stash's overflow - * pool — never thrown away to trigger a second warehouse round-trip. - * - * INLINE Arrow attachments under the ARROW_STREAM path are decoded once - * and put on the plugin's `inlineArrowStash`; the SSE message carries the - * synthetic stash id so the client fetches the bytes out-of-band via - * `/arrow-result/`. + * 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 _executeWithFormatFallback( + private async _executeJsonArrayPath( executor: AnalyticsPlugin, query: string, processedParams: | Record | undefined, - requestedFormat: AnalyticsFormat, - stashUserKey: string, signal?: AbortSignal, ): Promise { - return requestedFormat === "JSON_ARRAY" - ? this._executeJsonArrayPath(executor, query, processedParams, signal) - : this._executeArrowStreamPath( - executor, - query, - processedParams, - stashUserKey, - signal, - ); + const result = await deliverJsonResult( + executor, + query, + processedParams, + signal, + ); + return makeResultMessage(result.data, { + status: result.status, + statement_id: result.statement_id, + }); } /** - * JSON_ARRAY path. Tries `INLINE + JSON_ARRAY` first; on a structured - * `needs-arrow` rejection (warehouse only accepts ARROW_STREAM under - * INLINE), retries as `INLINE + ARROW_STREAM` and decodes the Arrow - * IPC attachment server-side so the caller's JSON_ARRAY contract is - * preserved. + * Attach the real column names so the client can relabel the positional + * Arrow schema (Databricks encodes ARROW_STREAM columns as col_0, …). * - * Failure modes surfaced to the caller: - * - Unrelated warehouse errors (anything not classified as - * `needs-arrow`) propagate untouched. - * - The retry can throw `ExecutionError.missingData("ARROW_STREAM - * attachment")` if the warehouse returned no attachment. - * - The decode itself can throw `RESULT_TOO_LARGE_FOR_JSON_FALLBACK` - * when the attachment exceeds the row or byte cap. + * 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 async _executeJsonArrayPath( - executor: AnalyticsPlugin, - query: string, - processedParams: - | Record - | undefined, - signal?: AbortSignal, - ): Promise { - try { - const result = await executor.query( - query, - processedParams, - { disposition: "INLINE", format: "JSON_ARRAY" }, - signal, - ); - return makeResultMessage(result?.data, { - status: result?.status, - statement_id: result?.statement_id, - }); - } catch (err: unknown) { - if (signal?.aborted) throw err; - if (_classifyInlineRejection(err) !== "needs-arrow") throw err; + private _setArrowColumnsHeader( + res: express.Response, + columnsRef: { columnNames?: string[]; statementId?: string }, + ): void { + const names = columnsRef.columnNames; + if (!names || names.length === 0) return; - const msg = err instanceof Error ? err.message : String(err); + 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( - "JSON_ARRAY INLINE rejected by warehouse, retrying as ARROW_STREAM INLINE and decoding server-side: %s", - msg, + "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", ); } - - // Retry as ARROW_STREAM + INLINE so the warehouse will accept the - // request, then decode the Arrow IPC attachment to plain row - // objects so the caller still gets JSON_ARRAY-shaped data. - 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); - return makeResultMessage(rows, { - status: arrowResult.status, - statement_id: arrowResult.statement_id, - }); } /** - * ARROW_STREAM path. Tries `INLINE + ARROW_STREAM` first; on a - * structured `needs-json` rejection (warehouse refuses ARROW_STREAM - * under INLINE), falls back to `EXTERNAL_LINKS + ARROW_STREAM`. When - * INLINE succeeds with an Arrow attachment, the decoded bytes go - * through `inlineArrowStash` and the client fetches them via - * `/arrow-result/inline-` — never the SSE channel. + * 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. * - * Failure modes surfaced to the caller: - * - Unrelated warehouse errors (not classified as `needs-json`) - * propagate untouched. - * - `ExecutionError.canceled()` if the client disconnects between - * the INLINE response and the stash put. - * - `ExecutionError.stashExhausted()` if both stash pools are full — - * we never silently re-run on EXTERNAL_LINKS because that would - * double-bill the warehouse and risk a divergent result for - * non-deterministic SQL. + * 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 _executeArrowStreamPath( - executor: AnalyticsPlugin, + private async _handleArrowStreamQuery( + req: express.Request, + res: express.Response, query: string, - processedParams: - | Record - | undefined, - stashUserKey: string, - signal?: AbortSignal, - ): Promise { + isAsUser: boolean, + parameters: IAnalyticsQueryRequest["parameters"], + ): Promise { + const executor = isAsUser ? this.asUser(req) : this; + 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 { - const result = await executor.query( + // 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, + ); + + // 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( + executor, + this.SQLClient, query, processedParams, - { disposition: "INLINE", format: "ARROW_STREAM" }, + columnsRef, signal, ); - if (result?.attachment) { - return this._stashAndEmitInline(result, stashUserKey, signal); - } - // INLINE succeeded but the warehouse returned no Arrow attachment - // (rare: inline `data_array` under ARROW_STREAM). An ARROW_STREAM - // caller's client decodes the `/arrow-result` bytes into a - // TypedArrowTable, so emitting a `result` row message here would - // break that contract. Fall through to the EXTERNAL_LINKS path - // below to obtain a fetchable Arrow statement instead. - logger.warn( - "ARROW_STREAM INLINE returned no attachment; falling back to EXTERNAL_LINKS", - ); - } catch (err: unknown) { - // If the request was aborted, do not retry — the signal is dead and - // a second statement would be billed but never read. - if (signal?.aborted) throw err; - if (_classifyInlineRejection(err) !== "needs-json") throw err; + const first = await bytes.next(); + // First byte in hand — stop the fail-fast clock. + clearTimeout(failFast); - const msg = err instanceof Error ? err.message : String(err); - logger.warn( - "ARROW_STREAM INLINE rejected by warehouse, falling back to EXTERNAL_LINKS: %s", - msg, - ); - } + res.setHeader("Content-Type", "application/vnd.apache.arrow.stream"); + res.setHeader("Cache-Control", "no-store"); + this._setArrowColumnsHeader(res, columnsRef); - const result = await executor.query( - query, - processedParams, - { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, - signal, - ); - return makeArrowMessage(result.statement_id, { status: result.status }); + if (!first.done) { + await writeChunk(res, first.value); + for await (const buf of bytes) { + await writeChunk(res, buf); + } + } + res.end(); + } catch (error) { + clearTimeout(failFast); + if (res.headersSent) { + logger.error("Arrow query stream failed mid-flight: %O", error); + res.destroy(error instanceof Error ? error : new Error(String(error))); + return; + } + 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; + } + if (signal.aborted) { + res.end(); + 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); + } } /** - * Decode an INLINE Arrow attachment, push it onto the stash, and - * return the `arrow` SSE message that references the synthetic id. - * Pulled out of `_executeArrowStreamPath` so the cap / overflow / - * exhaustion logic lives in one place. + * 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. */ - private _stashAndEmitInline( - result: { attachment: string; status?: unknown }, - stashUserKey: string, - signal?: AbortSignal, - ): AnalyticsSseMessage { - // If the client has already disconnected, the SSE write would be - // dropped anyway — skip the decode + stash so the bytes do not - // linger in memory until TTL eviction. - if (signal?.aborted) { - throw ExecutionError.canceled(); - } - const decoded = Buffer.from(result.attachment, "base64"); - const stashBytes = new Uint8Array( - decoded.buffer, - decoded.byteOffset, - decoded.byteLength, - ); - const putResult = this.inlineArrowStash.put(stashUserKey, stashBytes); - if (putResult === null) { - // Both regular and overflow pools are at cap. Throw with a clear - // signal rather than re-running the same statement on - // EXTERNAL_LINKS: the bytes we already decoded are gone, the - // warehouse has been billed, and a second execution could return - // a divergent result for non-deterministic SQL. - this._recordStashOutcome("exhausted", stashBytes.byteLength); - logger.warn( - "Inline Arrow stash exhausted (regular + overflow); rejecting", - ); - throw ExecutionError.stashExhausted(); - } - this._recordStashOutcome(putResult.pool, stashBytes.byteLength); - return makeArrowMessage(putResult.id, { status: result.status }); + 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: () => {}, + }); } /** @@ -783,17 +646,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(); } @@ -859,267 +711,69 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } /** - * 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. We cap two ways — rows AND decoded bytes — - * because either dimension can blow up independently (100k×1B rows is - * fine, 10 rows × 10MB cells is not). - */ -const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; -const JSON_ARRAY_FALLBACK_MAX_BYTES = 64 * 1024 * 1024; - -/** - * Replacer used by the nested-value `JSON.stringify` path. Arrow IPC - * carries column values whose JS representations include `bigint` - * (which the spec'd `JSON.stringify` cannot serialize and throws on), - * `Uint8Array` (binary buffers — would serialize as `{"0":...,"1":...}` - * which is wrong), and `Date` (would lose timezone fidelity inside a - * nested struct). Coerce each to a string form that matches what the - * warehouse itself emits under native JSON_ARRAY. - */ -function nestedJsonReplacer(_key: string, value: unknown): unknown { - if (typeof value === "bigint") return value.toString(); - if (value instanceof Uint8Array) { - return Buffer.from(value).toString("base64"); - } - if (value instanceof Date) return value.toISOString(); - return value; -} - -/** - * Render an `apache-arrow` Decimal cell to a fixed-point string. The JS - * value is a `DecimalBigNum` whose `toString()` yields the *signed - * unscaled* integer (e.g. `12345`, `-6789`); the decimal point is placed - * `scale` digits from the right to match the warehouse's native - * JSON_ARRAY rendering (`"123.45"`, `"-67.89"`). - */ -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 the warehouse's native - * JSON_ARRAY rendering, which is ISO-8601 with millisecond precision: - * `yyyy-MM-ddTHH:mm:ss.SSSZ` for zoned TIMESTAMP / TIMESTAMP_LTZ, and the - * same without the trailing `Z` for TIMESTAMP_NTZ. This is exactly - * `Date#toISOString()` (drop the `Z` when the column carries no timezone). + * 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. * - * `arrow-js` normalizes every Timestamp unit to epoch **milliseconds** at - * `.get()` (a MICROSECOND column returns ms, not micros); the warehouse - * also truncates JSON_ARRAY timestamps to ms, so the two agree. Verified - * against a dogfood serverless warehouse: `2024-06-13T01:02:03.500Z`, - * `2024-06-13T01:02:03.000Z`, and NTZ `2024-01-02T03:04:05.000`. + * @internal exported for unit testing the backpressure/disconnect behavior. */ -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) to match the - * warehouse's native JSON_ARRAY rendering. `arrow-js` normalizes Date - * columns to epoch milliseconds at UTC midnight on `.get()`. - */ -function formatDateCell(epochMs: number): string { - return new Date(epochMs).toISOString().slice(0, 10); -} - -/** - * Parse a STRING cell that looks like JSON (`{`/`[` prefixed) into an - * object/array, mirroring `SQLWarehouseConnector._transformDataArray` so - * the JSON_ARRAY fallback produces the same nested shape the native - * JSON_ARRAY path does. Non-JSON strings (and parse failures) pass - * through unchanged. - */ -function maybeParseJsonString(value: string): unknown { - if (value && (value[0] === "{" || value[0] === "[")) { - try { - return JSON.parse(value); - } catch { - return value; - } +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"), + ); } - return value; + 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); + }); } /** - * Decode a base64 Arrow IPC attachment to plain row objects. - * - * Used by the JSON_ARRAY fallback path when a warehouse refuses - * `JSON_ARRAY + INLINE` and we have to satisfy the request via - * `ARROW_STREAM + INLINE` — the bytes come back as Arrow IPC but the - * caller's contract is JSON-shaped rows, so we convert server-side. - * - * Shape rules (chosen to match what the warehouse emits natively under - * JSON_ARRAY, so callers cannot tell which path served their query): - * - **Scalars (number / bigint / boolean)**: stringified. JSON_ARRAY - * wire format emits everything in `result.data_array` as strings; - * bigint stringification is also necessary for JSON-serializability. - * - **Strings / Date / Uint8Array**: stringified consistently — - * `Date` to ISO string, `Uint8Array` to base64. - * - **Nested (List / Struct / Map)**: `JSON.stringify` with the - * `nestedJsonReplacer` so nested bigints / Dates / binary buffers - * serialize sanely instead of throwing or producing object-as-map - * output. Apache-arrow's typed values expose `toJSON()` so the - * resulting string matches what the warehouse would have emitted. + * 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`. */ -function decodeArrowAttachmentToRows( - attachment: 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 the child vectors once. `table.getChild(name)` walks the - // schema fields on every call; without this hoist we'd pay that cost - // O(rows × cols) times. At 100k × 50 columns that's millions of - // redundant lookups on the event loop. - const columns = table.schema.fields.map( - (f) => [f.name, 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; - } - // Branch on the Arrow field type first: several Databricks types - // decode to JS values whose default stringification does NOT match - // the warehouse's native JSON_ARRAY rendering. Without these the - // fallback silently diverges from the direct JSON_ARRAY path - // (decimals lose scale, temporals come back as raw epoch numbers). - switch (type.typeId) { - case Type.Decimal: - row[name] = formatDecimalCell( - v as { toString(): string }, - (type as DataType & { scale: number }).scale, - ); - continue; - case Type.Timestamp: - row[name] = formatTimestampCell( - Number(v), - (type as DataType & { timezone?: string | null }).timezone != null, - ); - continue; - case Type.Date: - row[name] = formatDateCell(Number(v)); - continue; - } - if ( - typeof v === "number" || - typeof v === "bigint" || - typeof v === "boolean" - ) { - row[name] = String(v); - } else if (typeof v === "string") { - // Match `_transformDataArray`: STRING columns holding JSON are - // parsed into objects/arrays so both paths yield the same shape. - row[name] = maybeParseJsonString(v); - } else if (v instanceof Date) { - row[name] = v.toISOString(); - } else if (v instanceof Uint8Array) { - row[name] = Buffer.from(v).toString("base64"); - } else { - // Nested types (List, Struct, Map). Use the replacer so nested - // bigint / Date / Uint8Array values serialize predictably — - // bare `JSON.stringify` throws on bigint and emits broken - // shapes for binary buffers. - row[name] = JSON.stringify(v, nestedJsonReplacer); - } - } - rows.push(row); - } - return rows; -} +const DEFAULT_ARROW_FIRST_BYTE_TIMEOUT_MS = 120_000; /** - * Classify a warehouse rejection of an INLINE statement. - * - * Two distinct rejection modes are observed in the wild: - * - * - **needs-arrow**: warehouse refuses `JSON_ARRAY + INLINE`, only accepts - * `ARROW_STREAM + INLINE`. Example message: - * `"Inline disposition only supports ARROW_STREAM format."` - * Action: retry as `ARROW_STREAM + INLINE` and decode server-side. - * - * - **needs-json**: warehouse refuses `ARROW_STREAM + INLINE`, only accepts - * `JSON_ARRAY + INLINE`. Examples: - * `"The format field must be JSON_ARRAY when the disposition field is INLINE."` - * `"ARROW_STREAM not supported with INLINE disposition"` - * `"ExternalLinks disposition is not yet implemented."` (same family — - * the warehouse rejected the disposition/format combo we sent). - * Action: retry as `ARROW_STREAM + EXTERNAL_LINKS`. - * - * The structured `errorCode` (INVALID_PARAMETER_VALUE / NOT_IMPLEMENTED) - * gates the classification so unrelated SQL errors don't trigger a retry. - * Message matching is case-insensitive — warehouses are inconsistent about - * casing of "Inline"/"INLINE". + * 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. */ -type InlineRejection = "needs-arrow" | "needs-json" | null; - -function _classifyInlineRejection(err: unknown): InlineRejection { - const msg = err instanceof Error ? err.message : String(err); - const lower = msg.toLowerCase(); - - const structuredCode = - err instanceof ExecutionError ? err.errorCode : undefined; - const hasCode = - structuredCode === "INVALID_PARAMETER_VALUE" || - structuredCode === "NOT_IMPLEMENTED" || - lower.includes("invalid_parameter_value") || - lower.includes("not_implemented"); - if (!hasCode) return null; - - // Must mention the inline disposition to count as a disposition-rejection. - if (!lower.includes("inline")) return null; - - // "needs-arrow": warehouse only supports ARROW_STREAM for INLINE. - if ( - /only supports\s+arrow_stream/i.test(msg) || - /must be\s+arrow_stream/i.test(msg) - ) { - return "needs-arrow"; - } - - // "needs-json": warehouse only supports JSON_ARRAY for 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"; - } - - return null; -} +const MAX_ARROW_COLUMNS_HEADER_BYTES = 6000; /** * @internal diff --git a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts deleted file mode 100644 index cb52ff71..00000000 --- a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts +++ /dev/null @@ -1,240 +0,0 @@ -import { randomUUID } from "node:crypto"; - -/** - * Server-side stash for inline Arrow IPC payloads. - * - * When a warehouse returns ARROW_STREAM + INLINE results, the bytes are - * stashed here and a synthetic "inline-" job id is emitted on the - * SSE control channel. The client then fetches the bytes out-of-band via - * `/arrow-result/`, which drains the stash and serves the payload as - * `application/vnd.apache.arrow.stream`. - * - * Keeps multi-MiB Arrow blobs off SSE, lets the existing /arrow-result - * pipeline handle both inline and EXTERNAL_LINKS results uniformly, and - * delivers the bytes with a real binary content-type instead of base64 - * inside JSON inside SSE framing. - * - * Properties: - * - **Drain-on-read**: a successful `take()` removes the entry. There is - * no replay path — a lost client connection means the bytes are gone. - * - **TTL bounded**: entries past their expiry are evicted on every - * `put()` and `take()`. No background timer. - * - **Per-user keyed**: `take()` only returns bytes if the requesting - * user matches the user that originally put them. Defense in depth on - * top of unguessable ids. - * - **Memory bounded with overflow slot**: total stashed bytes are capped - * at `maxBytes`. When `put()` cannot fit a payload, it spills into a - * separate overflow pool capped at `maxOverflowBytes` (default - * `maxBytes / 4` — kept small because overflow only needs to bridge - * the immediate `/arrow-result` fetch). Overflow entries behave like - * regular ones except (a) they do not count against the regular cap - * and (b) they expire on a much shorter TTL (`overflowTtlMs`, - * default 30s) — they exist to absorb already-decoded bytes for a - * single request, not to hold them long-term. Only when both pools - * are full does `put()` return `null` and the caller has to fall - * back to a different delivery path. Memory is bounded above by - * `maxBytes + maxOverflowBytes`. - * - * Caveat (multi-replica deployments): this stash is process-local. A - * subsequent `GET /arrow-result/inline-*` that lands on a different - * replica than the one that stashed the bytes will 410. Deployments - * that run more than one replica need sticky sessions (route both - * requests in the same logical session to the same replica) or a - * shared external store, neither of which is in scope here. - */ -interface InlineArrowStashOptions { - /** Regular-pool entries older than this are dropped on the next gc tick. */ - ttlMs?: number; - /** - * Overflow-pool entries older than this are dropped on the next gc - * tick. Defaults to 30s — overflow exists solely to bridge the - * immediate `/arrow-result` fetch that follows the SSE `arrow` - * message, so it should drain on the order of seconds, not minutes. - * A short TTL bounds the cross-user memory pressure that overflow - * can sustain in a multi-tenant deployment. - */ - overflowTtlMs?: number; - /** - * Hard cap on total bytes held in the regular pool. `put()` spills to - * the overflow pool when this cap would be exceeded; entries already - * in the stash are not evicted to fit new ones. - */ - maxBytes?: number; - /** - * Hard cap on total bytes held in the overflow pool. Overflow holds - * bytes that have already been decoded for an in-flight request — its - * purpose is to avoid throwing those bytes away and double-billing - * the warehouse. Defaults to `maxBytes / 4` (kept small because the - * pool only needs to absorb transient spillover, not hold long-term - * state). `put()` returns `null` only when both regular and overflow - * pools are at cap. - */ - maxOverflowBytes?: number; - /** - * Minimum interval between gc passes. `gc()` is O(N) in entry count, - * so on hot paths we skip when the previous pass was recent enough. - * Defaults to 5s. Set to 0 to disable throttling (test seam). - */ - gcMinIntervalMs?: number; - /** Test seam: override the synthetic-id generator. */ - idGenerator?: () => string; - /** Test seam: override the clock. */ - now?: () => number; -} - -interface StashEntry { - userId: string; - bytes: Uint8Array; - expiresAt: number; - insertedAt: number; - /** True for entries in the overflow pool (do not count against maxBytes). */ - overflow: boolean; -} - -export class InlineArrowStash { - private entries = new Map(); - private totalBytes = 0; - private overflowBytes = 0; - private lastGcAt = 0; - private readonly ttlMs: number; - private readonly overflowTtlMs: number; - private readonly maxBytes: number; - private readonly maxOverflowBytes: number; - private readonly gcMinIntervalMs: number; - private readonly idGenerator: () => string; - private readonly now: () => number; - - constructor(opts: InlineArrowStashOptions = {}) { - this.ttlMs = opts.ttlMs ?? 10 * 60 * 1000; - this.overflowTtlMs = opts.overflowTtlMs ?? 30 * 1000; - this.maxBytes = opts.maxBytes ?? 256 * 1024 * 1024; - this.maxOverflowBytes = - opts.maxOverflowBytes ?? Math.floor(this.maxBytes / 4); - this.gcMinIntervalMs = opts.gcMinIntervalMs ?? 5000; - this.idGenerator = opts.idGenerator ?? randomUUID; - this.now = opts.now ?? Date.now; - } - - /** - * Stash a payload and return its synthetic job id. - * - * Tries the regular pool first; if it would overflow, spills into the - * overflow pool (sized at `maxOverflowBytes`). Returns `null` only when - * both pools are at cap — the caller then has no choice but to fall - * back to a different delivery path (e.g. EXTERNAL_LINKS). - * - * The overflow pool exists because the caller has *already decoded* - * these bytes for an in-flight request: throwing them away would force - * a second warehouse round-trip (extra latency, double billing, and - * potentially divergent results for non-deterministic SQL). Holding - * them transiently in a bounded overflow region — they are drained - * single-use on the next `/arrow-result` fetch — is strictly safer - * than re-execution. - * - * A single payload can only land in one pool — the pools are not - * split across — so the largest payload we can accept is - * `Math.max(maxBytes, maxOverflowBytes)`. Exceeding that throws - * synchronously so the caller sees the misconfiguration loudly - * rather than burning a warehouse round-trip and then getting a - * null id. - * - * Returns `{ id, pool }` on success (`pool` ∈ {"regular", "overflow"}) - * or `null` when both pools are at cap. The pool tag lets callers - * emit accurate telemetry labels without re-introspecting the stash. - */ - put( - userId: string, - bytes: Uint8Array, - ): { id: string; pool: "regular" | "overflow" } | null { - const largestSlot = Math.max(this.maxBytes, this.maxOverflowBytes); - if (bytes.length > largestSlot) { - throw new Error( - `Inline Arrow payload (${bytes.length} bytes) exceeds largest stash slot (${largestSlot}); cannot fit in either pool`, - ); - } - this.gc(); - const fitsRegular = this.totalBytes + bytes.length <= this.maxBytes; - const fitsOverflow = - !fitsRegular && - this.overflowBytes + bytes.length <= this.maxOverflowBytes; - if (!fitsRegular && !fitsOverflow) { - // Both pools are full — refuse rather than evicting any issued id. - return null; - } - const id = `inline-${this.idGenerator()}`; - const now = this.now(); - const overflow = !fitsRegular; - this.entries.set(id, { - userId, - bytes, - expiresAt: now + (overflow ? this.overflowTtlMs : this.ttlMs), - insertedAt: now, - overflow, - }); - if (overflow) { - this.overflowBytes += bytes.length; - } else { - this.totalBytes += bytes.length; - } - return { id, pool: overflow ? "overflow" : "regular" }; - } - - /** - * Drain a payload from the stash. Returns `undefined` if the id is - * unknown, expired, or belongs to a different user. - */ - take(id: string, userId: string): Uint8Array | undefined { - this.gc(); - const entry = this.entries.get(id); - if (!entry) return undefined; - if (entry.userId !== userId) return undefined; - this.entries.delete(id); - if (entry.overflow) { - this.overflowBytes -= entry.bytes.length; - } else { - this.totalBytes -= entry.bytes.length; - } - return entry.bytes; - } - - /** Inspection helpers (primarily for tests). */ - size(): number { - return this.totalBytes; - } - /** Bytes currently held in the overflow pool. */ - overflowSize(): number { - return this.overflowBytes; - } - count(): number { - return this.entries.size; - } - - /** Drop all entries (used in plugin shutdown). */ - clear(): void { - this.entries.clear(); - this.totalBytes = 0; - this.overflowBytes = 0; - } - - private gc(): void { - const now = this.now(); - if (now - this.lastGcAt < this.gcMinIntervalMs) { - // Skip the O(N) sweep — recent enough to assume nothing - // significant has expired since the last pass. Worst case an - // entry lingers an extra `gcMinIntervalMs` past its TTL, which - // is negligible relative to either pool's intended lifetime. - return; - } - this.lastGcAt = now; - for (const [id, entry] of this.entries) { - if (entry.expiresAt <= now) { - this.entries.delete(id); - if (entry.overflow) { - this.overflowBytes -= entry.bytes.length; - } else { - this.totalBytes -= entry.bytes.length; - } - } - } - } -} 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..60fa957a --- /dev/null +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -0,0 +1,431 @@ +import type { sql } from "@databricks/sdk-experimental"; +import { type DataType, Type, tableFromIPC } from "apache-arrow"; +import type { SQLTypeMarker } from "shared"; +import { ExecutionError } from "../../errors"; +import { createLogger } from "../../logging/logger"; + +/** + * 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; + } + | undefined + >; +} + +/** Streams already-resolved EXTERNAL_LINKS chunks; the connector provides it. */ +export interface ArrowChunkStreamer { + streamExternalLinks( + chunks: sql.ExternalLink[], + signal?: AbortSignal, + ): AsyncGenerator; +} + +/** + * 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; + +export function classifyDispositionRejection( + err: unknown, +): DispositionRejection { + const msg = err instanceof Error ? err.message : String(err); + const lower = msg.toLowerCase(); + + const structuredCode = + err instanceof ExecutionError ? err.errorCode : undefined; + const hasCapabilityCode = + structuredCode === "INVALID_PARAMETER_VALUE" || + structuredCode === "NOT_IMPLEMENTED" || + lower.includes("invalid_parameter_value") || + lower.includes("not_implemented"); + // Auth / permission / SQL errors carry other codes → never fall back. + if (!hasCapabilityCode) return null; + + // 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, +): AsyncGenerator { + try { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "ARROW_STREAM" }, + signal, + ); + if (result?.attachment) { + 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; + if (classifyDispositionRejection(err) !== "needs-json-inline") 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"); + } + 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. + yield* streamer.streamExternalLinks(ext.external_links, signal); +} + +/** 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; + +/** + * Replacer for nested `JSON.stringify`. Arrow values include `bigint` (which + * spec'd `JSON.stringify` throws on), `Uint8Array` (would serialize as an + * index map), and `Date` (would lose timezone in a nested struct). Coerce + * each to the string form the warehouse itself emits under native JSON_ARRAY. + */ +function nestedJsonReplacer(_key: string, value: unknown): unknown { + if (typeof value === "bigint") return value.toString(); + if (value instanceof Uint8Array) return Buffer.from(value).toString("base64"); + if (value instanceof Date) return value.toISOString(); + return value; +} + +/** 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); +} + +/** 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: + row[name] = formatDecimalCell( + v as { toString(): string }, + (type as DataType & { scale: number }).scale, + ); + continue; + case Type.Timestamp: + row[name] = formatTimestampCell( + Number(v), + (type as DataType & { timezone?: string | null }).timezone != null, + ); + continue; + case Type.Date: + row[name] = formatDateCell(Number(v)); + 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 Date) { + row[name] = v.toISOString(); + } else if (v instanceof Uint8Array) { + row[name] = Buffer.from(v).toString("base64"); + } else { + row[name] = JSON.stringify(v, nestedJsonReplacer); + } + } + 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 d66fa364..86154c66 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -19,8 +19,7 @@ import { 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 { InlineArrowStash } from "../inline-arrow-stash"; +import { AnalyticsPlugin, analytics, writeChunk } from "../analytics"; import type { IAnalyticsConfig } from "../types"; // Mock CacheManager singleton with actual caching behavior @@ -105,100 +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("/arrow-result/inline-* drains the stash and serves bytes as application/vnd.apache.arrow.stream", async () => { - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - plugin.injectRoutes(router); - - const arrowBytes = new Uint8Array([0xff, 0xfe, 0xfd, 0xfc]); - const { id } = (plugin as any).inlineArrowStash.put("global", arrowBytes); - expect(id.startsWith("inline-")).toBe(true); - - const handler = getHandler("GET", "/arrow-result/:jobId"); - const mockReq = createMockRequest({ params: { jobId: id } }); - const mockRes = createMockResponse(); - - await handler(mockReq, mockRes); - - expect(mockRes.setHeader).toHaveBeenCalledWith( - "Content-Type", - "application/vnd.apache.arrow.stream", - ); - expect(mockRes.setHeader).toHaveBeenCalledWith( - "Content-Length", - String(arrowBytes.length), - ); - expect(mockRes.setHeader).toHaveBeenCalledWith( - "Cache-Control", - "no-store", - ); - expect(mockRes.send).toHaveBeenCalledTimes(1); - const sentBuf = (mockRes.send as any).mock.calls[0][0] as Buffer; - expect(Buffer.isBuffer(sentBuf)).toBe(true); - expect(Array.from(sentBuf)).toEqual(Array.from(arrowBytes)); - - // Drain-on-read: a second fetch must return 410, not the bytes again. - const secondRes = createMockResponse(); - await handler(mockReq, secondRes); - expect(secondRes.status).toHaveBeenCalledWith(410); - }); - - test("/arrow-result/inline-* returns 410 when the stash entry never existed", async () => { - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - plugin.injectRoutes(router); - - const handler = getHandler("GET", "/arrow-result/:jobId"); - const mockReq = createMockRequest({ - params: { jobId: "inline-does-not-exist" }, - }); - const mockRes = createMockResponse(); - - await handler(mockReq, mockRes); - - expect(mockRes.status).toHaveBeenCalledWith(410); - expect(mockRes.json).toHaveBeenCalledWith( - expect.objectContaining({ - error: expect.stringMatching(/expired or unknown/), - }), - ); - }); - - test("/arrow-result/inline-* returns 410 when the stash entry belongs to a different user", async () => { - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - plugin.injectRoutes(router); - - // Stash entry keyed to user-a, but the request resolves to "global" - // (no x-forwarded-user header) — keys differ, take must return - // nothing, and the entry stays put (single-user view). - const bytes = new Uint8Array([1, 2, 3]); - const { id } = (plugin as any).inlineArrowStash.put("user-a", bytes); - - const handler = getHandler("GET", "/arrow-result/:jobId"); - const mockReq = createMockRequest({ params: { jobId: id } }); - const mockRes = createMockResponse(); - - await handler(mockReq, mockRes); - - expect(mockRes.status).toHaveBeenCalledWith(410); - // The entry must still be there for the real owner. - expect((plugin as any).inlineArrowStash.take(id, "user-a")).toBeDefined(); - }); - test("/query/:query_key should return 400 when query_key is missing", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -834,7 +739,7 @@ describe("Analytics Plugin", () => { }); }); - test("/query/:query_key should fall back ARROW_STREAM from INLINE to EXTERNAL_LINKS when warehouse rejects INLINE", async () => { + 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(); @@ -843,6 +748,11 @@ describe("Analytics Plugin", () => { 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( @@ -851,10 +761,21 @@ describe("Analytics Plugin", () => { ), ) .mockResolvedValueOnce({ - result: { statement_id: "stmt-1", status: { state: "SUCCEEDED" } }, + 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"); @@ -866,16 +787,227 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); - // First call: INLINE (rejected) + // First call INLINE (rejected), second EXTERNAL_LINKS (fallback). expect(executeMock.mock.calls[0][1]).toMatchObject({ disposition: "INLINE", format: "ARROW_STREAM", }); - // Second call: EXTERNAL_LINKS (fallback) 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", + ); + // 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]); + }); + + 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" }, + }); + 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", + ); + }); + + 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("ARROW_STREAM: a schema too wide for the header advertises a columns-ref instead", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (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: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + 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"] }); + }); + + 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("/query/:query_key falls back on a structured ExecutionError.errorCode without scanning the message", async () => { @@ -1039,7 +1171,7 @@ describe("Analytics Plugin", () => { } }); - test("/query/:query_key stashes ARROW_STREAM INLINE bytes and emits an arrow message with a synthetic inline- id", async () => { + 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(); @@ -1049,90 +1181,19 @@ describe("Analytics Plugin", () => { }); // 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 }, - }); - (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 route should not fall back to EXTERNAL_LINKS — INLINE succeeded. - expect(executeMock).toHaveBeenCalledTimes(1); - expect(executeMock.mock.calls[0][1]).toMatchObject({ - disposition: "INLINE", - format: "ARROW_STREAM", - }); - // SSE payload: unified `arrow` message with an inline- prefixed id. - // The base64 attachment must NOT appear on the SSE channel. - 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":"arrow"'); - expect(payload).toMatch(/"statement_id":"inline-[^"]+"/); - expect(payload).not.toContain("arrow_inline"); - expect(payload).not.toContain(fakeAttachment); - - // The decoded bytes should be in the stash, keyed by the same - // synthetic id; a subsequent /arrow-result fetch will drain them. - const idMatch = payload?.match(/"statement_id":"(inline-[^"]+)"/); - expect(idMatch).not.toBeNull(); - const inlineId = idMatch?.[1]; - const stashed = (plugin as any).inlineArrowStash.take(inlineId, "global"); - expect(stashed).toBeDefined(); - expect(Array.from(stashed)).toEqual(Array.from(arrowBytes)); - }); - - test("/query/:query_key spills the already-decoded bytes into the stash overflow pool when the regular pool is full — single execution, no double-billing", async () => { - // The regular stash put may refuse new entries when at cap. We must - // NOT respond by re-executing the same statement with EXTERNAL_LINKS: - // the warehouse has already been billed, the bytes are already - // decoded server-side, and a second execution can return a divergent - // result for non-deterministic SQL (CURRENT_TIMESTAMP, RAND, late - // rows). The stash's overflow pool absorbs the bytes so the client - // still gets an inline- id pointing at the original result. - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - - (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ - query: "SELECT * FROM test", - isAsUser: false, - }); - - const fakeAttachment = Buffer.from(new Uint8Array([1, 2, 3])).toString( - "base64", - ); - const executeMock = vi.fn().mockResolvedValueOnce({ - result: { attachment: fakeAttachment, row_count: 1 }, + result: { + attachment: fakeAttachment, + row_count: 1, + columnNames: ["name", "totalSpend"], + }, }); (plugin as any).SQLClient.executeStatement = executeMock; - // Force the regular pool to be at cap so the put spills into overflow. - // We construct a tiny stash and pre-fill the regular pool. - const tinyStash = new InlineArrowStash({ - maxBytes: 4, - maxOverflowBytes: 64, - }); - tinyStash.put("filler", new Uint8Array([9, 9, 9, 9])); - (plugin as any).inlineArrowStash = tinyStash; - plugin.injectRoutes(router); const handler = getHandler("POST", "/query/:query_key"); @@ -1144,83 +1205,34 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); - // Single execution: no EXTERNAL_LINKS retry. + // 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", }); - // SSE payload carries an inline- id pointing at the overflow entry. - 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"), + // 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", ); - expect(payload).toBeDefined(); - expect(payload).toContain('"type":"arrow"'); - expect(payload).toMatch(/"statement_id":"inline-[^"]+"/); - - // Bytes landed in the overflow pool, regular pool size unchanged. - expect(tinyStash.overflowSize()).toBe(3); - expect(tinyStash.size()).toBe(4); - }); - - test("/query/:query_key surfaces a stable error when both regular and overflow pools are exhausted — never silently double-bills", async () => { - // When even the overflow pool cannot fit the payload, the route - // surfaces INLINE_ARROW_STASH_EXHAUSTED instead of re-running the - // statement on EXTERNAL_LINKS. The previous behavior (silent retry) - // double-billed the warehouse and could return divergent results. - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - - (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ - query: "SELECT * FROM test", - isAsUser: false, - }); - - const fakeAttachment = Buffer.from(new Uint8Array([1, 2, 3])).toString( - "base64", + // 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"])), ); - const executeMock = vi.fn().mockResolvedValueOnce({ - result: { attachment: fakeAttachment, row_count: 1 }, - }); - (plugin as any).SQLClient.executeStatement = executeMock; - - // Force both put paths to refuse: spy returns null unconditionally. - vi.spyOn((plugin as any).inlineArrowStash, "put").mockReturnValue(null); - - 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); - - // Single execution: never re-runs on EXTERNAL_LINKS. - expect(executeMock).toHaveBeenCalledTimes(1); + expect(mockRes.json).not.toHaveBeenCalled(); + expect(mockRes.end).toHaveBeenCalled(); - // SSE error payload carries the stable errorCode for UI branching. - // The writer emits separate lines (`id:`, `event: error`, `data: ...`) - // — join them so we can match across the framing. - const writeCalls = (mockRes.write as any).mock.calls.map( - (c: any[]) => c[0] as string, + const written = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as Buffer, ); - const errorFrame = writeCalls - .filter((s: string) => s.startsWith("data: ")) - .join("\n"); - expect(errorFrame).toContain("INLINE_ARROW_STASH_EXHAUSTED"); - // Client-safe message reaches the wire; raw upstream text does not. - expect(errorFrame).toContain("Server is at capacity"); - expect(errorFrame).not.toContain("Inline Arrow stash exhausted"); + 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 () => { @@ -1687,4 +1699,74 @@ describe("Analytics Plugin", () => { 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/inline-arrow-stash.test.ts b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts deleted file mode 100644 index 05ae821b..00000000 --- a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts +++ /dev/null @@ -1,184 +0,0 @@ -import { describe, expect, test } from "vitest"; -import { InlineArrowStash } from "../inline-arrow-stash"; - -function bytes(n: number): Uint8Array { - return new Uint8Array(n); -} - -// `put()` returns `{ id, pool } | null` — it rejects with null when the stash -// is full. Most tests only care about the id; this helper narrows via the -// non-null contract and returns just the id. Tests that need the pool tag -// call `stash.put(...)` directly. -function mustPut( - stash: InlineArrowStash, - userId: string, - b: Uint8Array, -): string { - const result = stash.put(userId, b); - if (result === null) { - throw new Error("test setup: stash unexpectedly rejected put"); - } - return result.id; -} - -describe("InlineArrowStash", () => { - test("put returns an inline-prefixed synthetic id", () => { - const stash = new InlineArrowStash({ idGenerator: () => "abc" }); - const id = mustPut(stash, "user-1", bytes(100)); - expect(id).toBe("inline-abc"); - }); - - test("take drains the entry", () => { - const stash = new InlineArrowStash(); - const id = mustPut(stash, "user-1", bytes(100)); - expect(stash.count()).toBe(1); - expect(stash.size()).toBe(100); - - const got = stash.take(id, "user-1"); - expect(got).toBeDefined(); - expect(got?.length).toBe(100); - expect(stash.count()).toBe(0); - expect(stash.size()).toBe(0); - // Drain-on-read: second take returns undefined. - expect(stash.take(id, "user-1")).toBeUndefined(); - }); - - test("take returns undefined for unknown id", () => { - const stash = new InlineArrowStash(); - expect(stash.take("inline-nope", "user-1")).toBeUndefined(); - }); - - test("take returns undefined when userId does not match", () => { - const stash = new InlineArrowStash(); - const id = mustPut(stash, "user-1", bytes(100)); - expect(stash.take(id, "user-2")).toBeUndefined(); - // Entry is still there for the right user. - expect(stash.take(id, "user-1")).toBeDefined(); - }); - - test("entries past TTL are evicted on next gc tick", () => { - let clock = 0; - // gcMinIntervalMs: 0 disables gc throttling so the test can drive - // the clock past TTL on a sub-throttle scale without the gc pass - // being skipped. - const stash = new InlineArrowStash({ - ttlMs: 1000, - gcMinIntervalMs: 0, - now: () => clock, - }); - const id = mustPut(stash, "user-1", bytes(50)); - clock = 999; - expect(stash.take(id, "user-1")).toBeDefined(); - - const id2 = mustPut(stash, "user-1", bytes(50)); - clock = 2000; - // Bump the clock past TTL and trigger gc via another put. - mustPut(stash, "user-2", bytes(10)); - expect(stash.take(id2, "user-1")).toBeUndefined(); - }); - - test("put spills into the overflow pool when the regular pool is at cap — every issued id remains valid", () => { - let seq = 0; - const stash = new InlineArrowStash({ - maxBytes: 200, - maxOverflowBytes: 200, - idGenerator: () => String(seq++), - }); - const a = mustPut(stash, "user-1", bytes(80)); - const b = mustPut(stash, "user-1", bytes(80)); - expect(stash.size()).toBe(160); - expect(stash.overflowSize()).toBe(0); - - // The third 80-byte entry would push the regular pool to 240 (>200), - // so it spills into overflow. Both prior entries must survive. - const c = mustPut(stash, "user-1", bytes(80)); - expect(stash.size()).toBe(160); - expect(stash.overflowSize()).toBe(80); - expect(stash.take(a, "user-1")).toBeDefined(); - expect(stash.take(b, "user-1")).toBeDefined(); - expect(stash.take(c, "user-1")).toBeDefined(); - // After draining the overflow entry, the counter reflects it. - expect(stash.overflowSize()).toBe(0); - }); - - test("put returns null only when both regular and overflow pools are full", () => { - let seq = 0; - const stash = new InlineArrowStash({ - maxBytes: 100, - maxOverflowBytes: 100, - idGenerator: () => String(seq++), - }); - const a = mustPut(stash, "user-1", bytes(100)); // fills regular - const b = mustPut(stash, "user-1", bytes(100)); // fills overflow - expect(stash.size()).toBe(100); - expect(stash.overflowSize()).toBe(100); - - // Both pools at cap — refuse rather than evict. - const c = stash.put("user-1", bytes(50)); - expect(c).toBeNull(); - expect(stash.take(a, "user-1")).toBeDefined(); - expect(stash.take(b, "user-1")).toBeDefined(); - }); - - test("put tags the result with its pool so callers can label telemetry without re-introspecting", () => { - const stash = new InlineArrowStash({ - maxBytes: 100, - maxOverflowBytes: 100, - }); - expect(stash.put("u", bytes(80))).toMatchObject({ pool: "regular" }); - // Regular has 80/100 used; another 80 would push it to 160 > 100 so it - // spills. - expect(stash.put("u", bytes(80))).toMatchObject({ pool: "overflow" }); - }); - - test("put throws when a single payload would not fit in the largest pool (caller misconfiguration)", () => { - const stash = new InlineArrowStash({ - maxBytes: 100, - maxOverflowBytes: 100, - }); - // 300 > max(maxBytes, maxOverflowBytes) (100); pools don't split, - // so no individual put can ever succeed. - expect(() => stash.put("user-1", bytes(300))).toThrow( - /exceeds largest stash slot/, - ); - }); - - test("overflow entries expire on a shorter TTL than regular entries", () => { - let clock = 0; - const stash = new InlineArrowStash({ - maxBytes: 80, - maxOverflowBytes: 80, - ttlMs: 600_000, - overflowTtlMs: 30_000, - gcMinIntervalMs: 0, - now: () => clock, - }); - const reg = mustPut(stash, "user-1", bytes(80)); // fills regular - const ovf = mustPut(stash, "user-1", bytes(80)); // spills to overflow - expect(stash.size()).toBe(80); - expect(stash.overflowSize()).toBe(80); - - // 45 s in: overflow expired, regular still alive. - clock = 45_000; - expect(stash.take(ovf, "user-1")).toBeUndefined(); - expect(stash.take(reg, "user-1")).toBeDefined(); - }); - - test("synthetic ids are unique across puts", () => { - const stash = new InlineArrowStash(); - const a = mustPut(stash, "user-1", bytes(10)); - const b = mustPut(stash, "user-1", bytes(10)); - expect(a).not.toBe(b); - expect(a.startsWith("inline-")).toBe(true); - expect(b.startsWith("inline-")).toBe(true); - }); - - test("clear drops every entry", () => { - const stash = new InlineArrowStash(); - mustPut(stash, "user-1", bytes(10)); - mustPut(stash, "user-2", bytes(20)); - stash.clear(); - expect(stash.count()).toBe(0); - expect(stash.size()).toBe(0); - }); -}); 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..758e9011 --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts @@ -0,0 +1,266 @@ +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. + expect(streamer.streamExternalLinks).toHaveBeenCalledWith(links, 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("arrowDeliveryUnsupported", () => { + test("is a structured, client-actionable error", () => { + const err = arrowDeliveryUnsupported(); + expect(err.errorCode).toBe("ARROW_DELIVERY_UNSUPPORTED"); + expect(err.clientMessage).toMatch(/JSON_ARRAY/); + }); +}); 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..34752018 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -4,243 +4,190 @@ import { createLogger } from "../logging/logger"; const logger = createLogger("stream:arrow"); -type ResultManifest = sql.ResultManifest; type ExternalLink = sql.ExternalLink; 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. - * - * 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. + * Stream Arrow chunks in array order, piping each chunk's response body. * - * @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 { + ): 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); + } } /** - * 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; + ): AsyncGenerator { + const 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); } - } 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]; - } - - 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; + 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) + }`, + ); } - 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 706af342..c8863b4f 100644 --- a/packages/appkit/src/stream/defaults.ts +++ b/packages/appkit/src/stream/defaults.ts @@ -1,12 +1,9 @@ export const streamDefaults = { bufferSize: 100, - // 1 MiB. SSE is used only for short JSON control messages — JSON_ARRAY - // result rows (already row-size-bounded by the warehouse) and the small - // `arrow` envelope (statement id + status) for ARROW_STREAM. Bulk Arrow - // payloads do not traverse SSE; they are fetched over HTTP via - // `/api/analytics/arrow-result/:jobId`, which dispatches to the warehouse - // (EXTERNAL_LINKS) or the server-side `InlineArrowStash` (INLINE) based - // on the id prefix. + // 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 diff --git a/packages/appkit/src/stream/stream-manager.ts b/packages/appkit/src/stream/stream-manager.ts index e83fa5ea..cf28ead5 100644 --- a/packages/appkit/src/stream/stream-manager.ts +++ b/packages/appkit/src/stream/stream-manager.ts @@ -309,7 +309,7 @@ export class StreamManager { error instanceof AppKitError ? error.clientMessage : "Internal server error"; - // Upstream structured code (e.g. INLINE_ARROW_STASH_EXHAUSTED, + // 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; 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..68b10e3f 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,146 @@ 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(); - }); - }); - - 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("throws when no chunks are provided", async () => { + await expect(drain(processor.streamChunks([]))).rejects.toThrow(); }); - 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); - }); + test("streams chunks in array order", async () => { + const bytes = await drain(processor.streamChunks(mockChunks(3))); + expect(bytes).toEqual([100, 101, 102]); }); - 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 { - 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 { + 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 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"); - }); - - 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); + const p = new ArrowStreamProcessor({ timeout: 50, retries: 1 }); + const gen = p.streamChunks(mockChunks(1)); - 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("should include HTTP status in error message for failed responses", async () => { - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: false, - status: 403, - statusText: "Forbidden", - }); + 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; - const singleRetryProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 1, - timeout: 5000, - retries: 1, - }); + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const drained = drain(p.streamChunks(mockChunks(1), controller.signal)); - const chunks = createMockChunks(1); - - await expect( - singleRetryProcessor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/403 Forbidden/); - }); - }); -}); - -describe("Semaphore (via ArrowStreamProcessor)", () => { - let originalFetch: typeof globalThis.fetch; - - beforeEach(() => { - originalFetch = globalThis.fetch; - vi.clearAllMocks(); - }); + // Let the first attempt fail and enter the (~1s) backoff, then abort. + await new Promise((r) => setTimeout(r, 50)); + controller.abort(); - afterEach(() => { - globalThis.fetch = originalFetch; + await expect(drained).rejects.toThrow(); + // Aborting mid-backoff cancels immediately — no second fetch attempt. + expect(fetchMock).toHaveBeenCalledTimes(1); }); - 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 { - ok: true, - arrayBuffer: async () => new Uint8Array([index + 1]).buffer, - }; - }); - - 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" }, - ]; - - await processor.processChunks(chunks as any, { columns: [] }); - - // With concurrency of 1, they should complete in order - expect(order).toHaveLength(3); - expect(globalThis.fetch).toHaveBeenCalledTimes(3); + 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(); }); }); diff --git a/packages/appkit/src/stream/types.ts b/packages/appkit/src/stream/types.ts index 154d388d..78b14e35 100644 --- a/packages/appkit/src/stream/types.ts +++ b/packages/appkit/src/stream/types.ts @@ -26,7 +26,7 @@ export interface SSEError { error: string; code: SSEErrorCode; /** - * Upstream-domain structured code (e.g. `INLINE_ARROW_STASH_EXHAUSTED`, + * 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. */ diff --git a/packages/shared/src/sse/analytics.test.ts b/packages/shared/src/sse/analytics.test.ts index f66437c3..cdc8a70e 100644 --- a/packages/shared/src/sse/analytics.test.ts +++ b/packages/shared/src/sse/analytics.test.ts @@ -1,9 +1,5 @@ import { describe, expect, test } from "vitest"; -import { - AnalyticsSseMessage, - makeArrowMessage, - makeResultMessage, -} from "./analytics"; +import { AnalyticsSseMessage, makeResultMessage } from "./analytics"; describe("AnalyticsSseMessage schema", () => { test("accepts a result message with rows", () => { @@ -18,42 +14,11 @@ describe("AnalyticsSseMessage schema", () => { expect(() => AnalyticsSseMessage.parse({ type: "result" })).not.toThrow(); }); - test("accepts an arrow message with warehouse statement_id", () => { - const parsed = AnalyticsSseMessage.parse({ - type: "arrow", - statement_id: "stmt-1", - }); - expect(parsed.type).toBe("arrow"); - }); - - test("accepts an arrow message with synthetic inline- id", () => { - // Inline Arrow payloads are stashed server-side and surfaced through the - // same `arrow` message variant — the `inline-` prefix tells the - // /arrow-result handler to drain the stash instead of hitting the - // warehouse. The schema must accept both id shapes transparently. - const parsed = AnalyticsSseMessage.parse({ - type: "arrow", - statement_id: "inline-abc-123", - }); - expect(parsed.statement_id).toBe("inline-abc-123"); - }); - - test("rejects an arrow message with empty statement_id", () => { - expect(() => - AnalyticsSseMessage.parse({ type: "arrow", statement_id: "" }), - ).toThrow(); - }); - - test("rejects an arrow message with no statement_id", () => { - expect(() => AnalyticsSseMessage.parse({ type: "arrow" })).toThrow(); - }); - - test("rejects the retired arrow_inline message type", () => { - // arrow_inline was the prior wire shape (base64 payload on the SSE - // channel). The current protocol routes all Arrow payloads through - // /arrow-result; the type must no longer parse. + 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_inline", attachment: "AQID" }), + AnalyticsSseMessage.parse({ type: "arrow", statement_id: "stmt-1" }), ).toThrow(); }); @@ -69,19 +34,9 @@ describe("AnalyticsSseMessage schema", () => { }); }); -describe("typed builders", () => { +describe("typed builder", () => { test("makeResultMessage roundtrips through the schema", () => { const msg = makeResultMessage([{ id: 1 }], { statement_id: "s-1" }); expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); }); - - test("makeArrowMessage roundtrips through the schema", () => { - const msg = makeArrowMessage("stmt-2"); - expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); - }); - - test("makeArrowMessage accepts synthetic inline- ids", () => { - const msg = makeArrowMessage("inline-some-uuid"); - expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); - }); }); diff --git a/packages/shared/src/sse/analytics.ts b/packages/shared/src/sse/analytics.ts index eba7597e..41022672 100644 --- a/packages/shared/src/sse/analytics.ts +++ b/packages/shared/src/sse/analytics.ts @@ -3,25 +3,21 @@ import { z } from "zod"; /** * Wire protocol for analytics SSE messages emitted by `/api/analytics/query`. * - * These schemas are the single source of truth for the contract between the - * server (`AnalyticsPlugin._handleQueryRoute`) and the client - * (`useAnalyticsQuery`). Both sides validate with the same schema: - * - * - Server uses the typed builders (`makeResultMessage`, `makeArrowMessage`) - * to construct messages with compile-time guarantees that all required - * fields are present. - * - Client calls `AnalyticsSseMessage.parse(JSON.parse(event.data))` to fail - * loudly on a malformed payload instead of silently treating an undefined - * field as data. + * 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. * - * Arrow payloads — inline or external-links — never traverse the SSE control - * channel; both flow through `/api/analytics/arrow-result/:jobId` and are - * differentiated by an `inline-` prefix on the job id (see - * `InlineArrowStash`). The wire shape from the client's perspective is - * therefore uniform: an `arrow` message carries an id, the client fetches. + * 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: * - * Adding a new message variant requires a schema update here, which keeps - * server and client in lockstep. + * - 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). */ @@ -60,30 +56,16 @@ export interface AnalyticsResultMessage { } /** - * ARROW_STREAM result delivered via /arrow-result/:jobId. The id is either: - * - the warehouse-issued `statement_id` for EXTERNAL_LINKS responses, or - * - a synthetic `inline-` id pointing at the server-side - * `InlineArrowStash` for INLINE responses. - * - * Both shapes are fetched the same way; the prefix tells the route handler - * which path to take. + * 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 AnalyticsArrowMessage = z.object({ - type: z.literal("arrow"), - statement_id: z.string().min(1), - status: z.unknown().optional(), -}); -export type AnalyticsArrowMessage = z.infer; - -/** Discriminated union of every message the analytics SSE stream may emit. */ -export const AnalyticsSseMessage = z.discriminatedUnion("type", [ - AnalyticsResultMessage, - AnalyticsArrowMessage, -]); +export const AnalyticsSseMessage = AnalyticsResultMessage; export type AnalyticsSseMessage = z.infer; // ──────────────────────────────────────────────────────────────────────────── -// Typed builders — call from the server route handler. The compiler enforces +// 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. // ──────────────────────────────────────────────────────────────────────────── @@ -94,10 +76,3 @@ export function makeResultMessage( ): AnalyticsResultMessage { return { type: "result", data, ...extras }; } - -export function makeArrowMessage( - statement_id: string, - extras: { status?: unknown } = {}, -): AnalyticsArrowMessage { - return { type: "arrow", statement_id, ...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; From 525a9bd8339551ea22f5fc1e73ff954168fc85b3 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 7 Jul 2026 18:49:21 +0200 Subject: [PATCH 20/25] test(playground): serve Arrow IPC from the e2e mock for ARROW_STREAM charts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The direct-streaming rework made `format="arrow"` charts fetch raw Arrow IPC bytes over HTTP (fetchArrowDirect), but the Playwright mock (setupMockAPI) only ever returned an SSE/JSON response. Arrow-format charts therefore decoded SSE text as Arrow, failed, and rendered empty — dropping the `bar-chart-apps_list` count below 3 and failing arrow-analytics.spec.ts "charts render with mock data (no empty states)". setupMockAPI now detects `format: "ARROW_STREAM"` in the request body and responds with real Arrow IPC bytes (tableFromJSON + tableToIPC, Content-Type application/vnd.apache.arrow.stream) built from the same mock data; array/object cells are stringified so Arrow can infer primitive column types (chart columns like name/totalSpend are unaffected). The JSON/SSE path is unchanged. Verified with the CI env (APPKIT_E2E_TEST, DATABRICKS_*=e2e-mock): arrow-analytics 4/4 and the full integration suite 24/24 pass. Signed-off-by: MarioCadenas --- apps/dev-playground/tests/utils/test-utils.ts | 98 ++++++++++--------- 1 file changed, 51 insertions(+), 47 deletions(-) diff --git a/apps/dev-playground/tests/utils/test-utils.ts b/apps/dev-playground/tests/utils/test-utils.ts index a4449023..b45e814a 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,68 @@ 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). + * Array/object cells are stringified so Arrow can infer a primitive column type, + * mirroring how the warehouse serializes complex types; chart-relevant columns + * (names, numbers) are unaffected. + */ +function toArrowIPC(rows: readonly unknown[]): Buffer { + const flat = rows.map((row) => + Object.fromEntries( + Object.entries(row as Record).map(([k, v]) => [ + k, + v !== null && typeof v === "object" ? JSON.stringify(v) : v, + ]), + ), + ); + return Buffer.from(tableToIPC(tableFromJSON(flat), "stream")); +} + 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), - }); - } - 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), - }); + // 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("sql_helpers_test")) { + + if (requestFormat === "ARROW_STREAM" && data.length > 0) { return route.fulfill({ status: 200, - headers: getSSEHeaders(), - body: createSSEResponse(mockAnalyticsData.sqlHelpersTest), + headers: { + "Content-Type": "application/vnd.apache.arrow.stream", + "Cache-Control": "no-store", + }, + body: toArrowIPC(data), }); } - // Default empty response for unknown queries return route.fulfill({ status: 200, headers: getSSEHeaders(), - body: createSSEResponse([]), + body: createSSEResponse(data), }); }); From f3b6c371e45beea7d586068a7765beb8b5ece8c4 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 7 Jul 2026 20:31:13 +0200 Subject: [PATCH 21/25] fix(appkit-ui): keep zod out of the browser bundle in useAnalyticsQuery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit useAnalyticsQuery validated SSE messages at runtime with `AnalyticsSseMessage.safeParse` (a zod schema value), pulling the entire zod library (~60 KB gz / ~303 KB min) into the `./react` initial bundle — a regression this branch introduced (main validates with a type-only import and no zod). The `shared` sideEffects hint removed zod from `./js`, but it can't help `./react` because the schema value is genuinely called. The SSE wire schema is intentionally loose (`data` is an optional array of unknown), so a structural check suffices — no schema validator needed to read our own same-origin server's messages. Replaced the `.safeParse` with a plain `parsed.type === "result"` check (missing/non-array data normalizes to []), matching main. Verified via esbuild metafile: zod is now 0 in the initial bundle of both ./react and ./js; apache-arrow stays lazy (197 KB). Typecheck, 16 hook tests, and biome pass. Signed-off-by: MarioCadenas --- .../src/react/hooks/use-analytics-query.ts | 39 +++++++------------ 1 file changed, 14 insertions(+), 25 deletions(-) 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 20f2c99e..1c63ffe1 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -6,7 +6,6 @@ import { useRef, useState, } from "react"; -import { AnalyticsSseMessage } from "shared"; import { ArrowClient, connectSSE } from "@/js"; import type { AnalyticsFormat, @@ -115,20 +114,15 @@ async function handleAnalyticsSseMessage( return; } - // The error/code branch below predates the SSE wire schema and can fire - // for messages that don't match any AnalyticsSseMessage variant (e.g. - // server-side error events from executeStream). Validate the known - // result/arrow variants first; fall through to error handling otherwise. - const validated = AnalyticsSseMessage.safeParse(parsed); - const msg = validated.success ? validated.data : null; - - // success - JSON format. The wire schema makes `data` optional (e.g. an - // empty result set may omit it), so normalize the missing case to an - // explicit empty array rather than letting `undefined` bleed into the - // hook's `T | null` state. - if (msg?.type === "result") { + // 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((msg.data ?? []) as ResultType); + ctx.setData((Array.isArray(parsed.data) ? parsed.data : []) as ResultType); ctx.unpublishWarehouseStatus(); return; } @@ -160,17 +154,12 @@ async function handleAnalyticsSseMessage( return; } - // The payload matched neither AnalyticsSseMessage nor an error event — - // surface a generic error rather than silently dropping it. - if (!validated.success) { - console.error( - "[useAnalyticsQuery] Malformed SSE payload", - validated.error.flatten(), - ); - ctx.setLoading(false); - ctx.setError(GENERIC_LOAD_ERROR); - ctx.unpublishWarehouseStatus(); - } + // 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 { From 35415242696e706b7c23a79c1256a82f952b80a3 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 8 Jul 2026 15:29:36 +0200 Subject: [PATCH 22/25] =?UTF-8?q?feat(appkit):=20harden=20arrow=20delivery?= =?UTF-8?q?=20=E2=80=94=20caching,=20capability=20memo,=20link=20refresh?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Robustness pass on the arrow result-delivery path from a full-feature review. - Cache INLINE arrow like the JSON path: the INLINE+ARROW_STREAM attempt now runs through a caching executor (per-user TTL cache + retry), so an arrow chart isn't a fresh warehouse execution on every render. Only bounded (<=25 MiB) inline attachments are cached; EXTERNAL_LINKS carry expiring pre-signed URLs and are never cached. Uses cache.getOrExecute directly so the thrown errorCode the capability fallback classifies on survives. - Memoize warehouse arrow capability (inline vs external) per warehouse id, in process: a standard warehouse rejects INLINE arrow on every query, so after the first we skip that doomed probe and go straight to EXTERNAL_LINKS. - Re-mint expired chunk links mid-download: EXTERNAL_LINKS pre-signed URLs expire in <=15 min, so a large/slow result can outlast its earliest links. The streamer now re-resolves a chunk's link via getStatementResultChunkN (bound to the caller's identity context) before retrying a failed fetch, instead of retrying the same expired URL. - Loosen fallback classification: any capability-CODED rejection (INVALID_PARAMETER_VALUE / NOT_IMPLEMENTED) of INLINE+ARROW now falls back to EXTERNAL_LINKS, rather than requiring an exact message phrase — a Databricks message reword no longer 500s every standard warehouse. - Zero-row external_links ([]): treat an empty array as a zero-row result and synthesize an empty Arrow table, instead of routing to streamChunks([]) which rejected with a 500. - Quiet routine client disconnects: an aborted signal now tears down without an ERROR log whether it fires before or after headers (checked before the headers-sent branch), so unmounts/param-changes don't spam alerting. Fail-fast 503 ordering preserved. - e2e mock now mirrors a real warehouse: positional col_N schema + real names in the X-Appkit-Arrow-Columns header, so the client relabel path (the live bug) is exercised end-to-end rather than only unit-tested. Green: pnpm -r typecheck, full unit suite (3421 pass / 1 skip), biome, build + publint, playground e2e (arrow-analytics 4/4, full integration 24/24). Multi- chunk external, OBO+external, and wide-schema paths remain unit-tested only (no warm standard warehouse to live-verify). Signed-off-by: MarioCadenas --- apps/dev-playground/tests/utils/test-utils.ts | 46 ++++-- .../src/connectors/sql-warehouse/client.ts | 45 +++++- .../sql-warehouse/tests/client.test.ts | 26 ++++ .../appkit/src/plugins/analytics/analytics.ts | 130 ++++++++++++++-- .../src/plugins/analytics/result-delivery.ts | 145 +++++++++++++----- .../plugins/analytics/tests/analytics.test.ts | 64 +++++++- .../analytics/tests/result-delivery.test.ts | 102 +++++++++++- .../src/stream/arrow-stream-processor.ts | 37 ++++- .../tests/arrow-stream-processor.test.ts | 65 ++++++++ 9 files changed, 583 insertions(+), 77 deletions(-) diff --git a/apps/dev-playground/tests/utils/test-utils.ts b/apps/dev-playground/tests/utils/test-utils.ts index b45e814a..f608a0e0 100644 --- a/apps/dev-playground/tests/utils/test-utils.ts +++ b/apps/dev-playground/tests/utils/test-utils.ts @@ -46,20 +46,35 @@ function resolveAnalyticsMock(url: string): readonly unknown[] { /** * 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). - * Array/object cells are stringified so Arrow can infer a primitive column type, - * mirroring how the warehouse serializes complex types; chart-relevant columns - * (names, numbers) are unaffected. + * + * 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[]): Buffer { - const flat = rows.map((row) => - Object.fromEntries( - Object.entries(row as Record).map(([k, v]) => [ - k, +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 Buffer.from(tableToIPC(tableFromJSON(flat), "stream")); + ); + }); + return { + body: Buffer.from(tableToIPC(tableFromJSON(positional), "stream")), + columnNames, + }; } export async function setupMockAPI(page: Page) { @@ -76,13 +91,20 @@ export async function setupMockAPI(page: Page) { } if (requestFormat === "ARROW_STREAM" && data.length > 0) { + const { body, columnNames } = toArrowIPC(data); return route.fulfill({ status: 200, 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: toArrowIPC(data), + body, }); } diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 239de551..f7dab68a 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, @@ -914,8 +917,12 @@ export class SQLWarehouseConnector { return this._validateArrowAttachment(response, result.attachment); } - // External links: data fetched separately via statement_id. - if (result?.external_links) { + // 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); } @@ -1054,6 +1061,7 @@ export class SQLWarehouseConnector { status: sql.StatementStatus; columnNames?: string[]; external_links?: sql.ExternalLink[]; + refreshChunkLink?: RefreshChunkLink; }; }> { const statementId = response.statement_id as string; @@ -1076,6 +1084,13 @@ export class SQLWarehouseConnector { 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, + ), }, }; } @@ -1132,6 +1147,27 @@ export class SQLWarehouseConnector { return undefined; } + /** + * 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); + }; + } + /** * 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 @@ -1142,8 +1178,9 @@ export class SQLWarehouseConnector { streamExternalLinks( chunks: sql.ExternalLink[], signal?: AbortSignal, + refresh?: RefreshChunkLink, ): AsyncGenerator { - return this.arrowProcessor.streamChunks(chunks, signal); + return this.arrowProcessor.streamChunks(chunks, signal, refresh); } /** diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts index dc964702..f0b03c07 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts @@ -293,6 +293,32 @@ describe("SQLWarehouseConnector._transformDataArray", () => { 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 = { diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 1b5aa1d5..4ad05cc3 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -30,7 +30,12 @@ import type { PluginManifest } from "../../registry"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; -import { deliverArrowBytes, deliverJsonResult } from "./result-delivery"; +import { + type ArrowCapability, + deliverArrowBytes, + deliverJsonResult, + type QueryExecutor, +} from "./result-delivery"; import { type AnalyticsQueryResponse, type AnalyticsStreamMessage, @@ -101,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; @@ -260,7 +275,14 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { // 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, isAsUser, parameters); + await this._handleArrowStreamQuery( + req, + res, + query_key, + query, + isAsUser, + parameters, + ); return; } @@ -474,11 +496,13 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { 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); @@ -510,16 +534,35 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { 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( - executor, + // 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. @@ -538,11 +581,10 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { res.end(); } catch (error) { clearTimeout(failFast); - if (res.headersSent) { - logger.error("Arrow query stream failed mid-flight: %O", error); - res.destroy(error instanceof Error ? error : new Error(String(error))); - return; - } + // 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", @@ -556,8 +598,18 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { }); 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) { - res.end(); + 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); @@ -605,6 +657,66 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { }); } + /** + * 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. * diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts index 60fa957a..f1270a62 100644 --- a/packages/appkit/src/plugins/analytics/result-delivery.ts +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -3,6 +3,7 @@ import { type DataType, 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. @@ -41,6 +42,7 @@ export interface QueryExecutor { columnNames?: string[]; statement_id?: string; status?: unknown; + refreshChunkLink?: RefreshChunkLink; } | undefined >; @@ -51,9 +53,29 @@ 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. * @@ -69,21 +91,40 @@ type DispositionRejection = | "external-links-unsupported" // EXTERNAL_LINKS not implemented | null; -export function classifyDispositionRejection( - err: unknown, -): DispositionRejection { - const msg = err instanceof Error ? err.message : String(err); - const lower = msg.toLowerCase(); - +/** + * 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; - const hasCapabilityCode = + if ( structuredCode === "INVALID_PARAMETER_VALUE" || - structuredCode === "NOT_IMPLEMENTED" || + 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"); + lower.includes("not_implemented") + ); +} + +export function classifyDispositionRejection( + err: unknown, +): DispositionRejection { // Auth / permission / SQL errors carry other codes → never fall back. - if (!hasCapabilityCode) return null; + 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 ( @@ -150,38 +191,54 @@ export async function* deliverArrowBytes( processedParams: Record | undefined, out: { columnNames?: string[]; statementId?: string }, signal?: AbortSignal, + opts?: ArrowDeliveryOptions, ): AsyncGenerator { - try { - const result = await executor.query( - query, - processedParams, - { disposition: "INLINE", format: "ARROW_STREAM" }, - signal, - ); - if (result?.attachment) { - 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, + // 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), ); - 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; - if (classifyDispositionRejection(err) !== "needs-json-inline") throw err; - logger.warn( - "ARROW_STREAM INLINE rejected by warehouse, falling back to EXTERNAL_LINKS: %s", - errMessage(err), - ); } let ext: Awaited>; @@ -205,13 +262,19 @@ export async function* deliverArrowBytes( 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. - yield* streamer.streamExternalLinks(ext.external_links, signal); + // 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. */ diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 86154c66..35677792 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -861,6 +861,12 @@ describe("Analytics Plugin", () => { 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(); @@ -1095,7 +1101,7 @@ describe("Analytics Plugin", () => { }); }); - test("/query/:query_key does NOT fall back when only one of INLINE/ARROW_STREAM appears in the error", async () => { + 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(); @@ -1104,13 +1110,14 @@ describe("Analytics Plugin", () => { isAsUser: false, }); - // Realistic non-format error that mentions just one of the keywords — - // e.g. an unrelated INVALID_PARAMETER_VALUE about a different param. + // 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":"INVALID_PARAMETER_VALUE","message":"INLINE is not a valid value for parameter `mode`"}', + '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; @@ -1126,9 +1133,9 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); - // The retry interceptor may attempt the query multiple times, but the - // analytics plugin must never escalate to EXTERNAL_LINKS for an error - // that doesn't actually indicate a format/disposition rejection. + // 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", @@ -1137,6 +1144,49 @@ describe("Analytics Plugin", () => { } }); + 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(); diff --git a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts index 758e9011..19a222da 100644 --- a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts @@ -163,8 +163,13 @@ describe("deliverArrowBytes — normal warehouse (fallback to EXTERNAL_LINKS)", { disposition: "INLINE", format: "ARROW_STREAM" }, { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, ]); - // Streams the links resolved in-context (no re-fetch): passed straight through. - expect(streamer.streamExternalLinks).toHaveBeenCalledWith(links, undefined); + // 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]); }); @@ -257,6 +262,99 @@ describe("decodeArrowAttachmentToRows", () => { }); }); +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(); diff --git a/packages/appkit/src/stream/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index 34752018..94eb8257 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -6,6 +6,19 @@ const logger = createLogger("stream:arrow"); type ExternalLink = sql.ExternalLink; +/** + * Re-mint a chunk's pre-signed URL. DBSQL external links expire in <= 15 min, + * so a large result whose tail chunks are reached after the earlier chunks + * 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 { /** Idle timeout (ms) for a single chunk: no progress → abort the download. */ timeout: number; @@ -46,13 +59,14 @@ export class ArrowStreamProcessor { async *streamChunks( chunks: ExternalLink[], signal?: AbortSignal, + refresh?: RefreshChunkLink, ): AsyncGenerator { if (chunks.length === 0) { throw ValidationError.missingField("chunks"); } for (const chunk of chunks) { - yield* this.streamChunkBody(chunk, signal); + yield* this.streamChunkBody(chunk, signal, refresh); } } @@ -67,8 +81,9 @@ export class ArrowStreamProcessor { private async *streamChunkBody( chunk: ExternalLink, signal?: AbortSignal, + refresh?: RefreshChunkLink, ): AsyncGenerator { - const externalLink = chunk.external_link; + let externalLink = chunk.external_link; if (!externalLink) { // A missing link cannot be fixed by retrying — fail loudly. throw ExecutionError.statementFailed( @@ -115,6 +130,24 @@ export class ArrowStreamProcessor { if (signal?.aborted) throw ExecutionError.canceled(); if (attempt < this.options.retries - 1) { 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, + ); + } + } } } } 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 68b10e3f..eb367929 100644 --- a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts +++ b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts @@ -197,4 +197,69 @@ describe("ArrowStreamProcessor.streamChunks", () => { ).rejects.toThrow(); expect(fetchMock).not.toHaveBeenCalled(); }); + + 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, + 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 refresh = vi.fn(async () => { + throw new Error("re-resolve failed"); + }); + + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const bytes = await drain( + p.streamChunks(mockChunks(1), undefined, refresh), + ); + + expect(refresh).toHaveBeenCalled(); + // Refresh failure is swallowed; the retry still runs against the original. + expect(bytes).toEqual([7]); + }, 10000); }); From 45885794680efb8a6fc2160bbecbbe0d56837636 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 8 Jul 2026 15:49:44 +0200 Subject: [PATCH 23/25] refactor(appkit): extract shared statement-error rethrow helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The executeStatement and _pollForStatementResult catch blocks carried a verbatim-identical error-rethrow tail (AbortError passthrough, AppKitError passthrough, and SDK-errorCode-preserving ExecutionError wrap). Extract it into a private `_rethrowStatementError(error): never` helper. The `never` return type keeps control flow identical — both copies were the last statement in their catch. Behavior-preserving. Signed-off-by: MarioCadenas --- .../src/connectors/sql-warehouse/client.ts | 59 +++++++++---------- 1 file changed, 27 insertions(+), 32 deletions(-) diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index f7dab68a..9e924635 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -369,24 +369,7 @@ export class SQLWarehouseConnector { ); } - if (error instanceof Error && error.name === "AbortError") { - throw error; - } - if (error instanceof AppKitError) { - throw error; - } - // Preserve the SDK's structured ApiError.errorCode (e.g. - // "INVALID_PARAMETER_VALUE", "BAD_REQUEST") through the wrap so - // callers can branch on a stable identifier rather than - // substring-matching the message. - 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, - ); + this._rethrowStatementError(error); } finally { // remove abort handler if (abortHandler && signal) { @@ -877,20 +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; - } - 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, - ); + this._rethrowStatementError(error); } finally { span.end(); } @@ -1201,6 +1171,31 @@ export class SQLWarehouseConnector { return arrowColumnNames(response); } + /** + * 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, + ); + } + // create context for cancellation token private _createContext(signal?: AbortSignal) { return new Context({ From 6f91df301a7ed92f0c761ba40c3a492018d06adc Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 8 Jul 2026 16:43:16 +0200 Subject: [PATCH 24/25] chore(appkit): bump bundle-size baseline for inline-arrow feature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The +7.2% appkit tarball growth is the PR's own new source (arrow-schema.ts, result-delivery.ts, expanded analytics/client). apache-arrow is external and not in the tarball, so the increase is justified own-code — acknowledge it by updating the committed baseline. Signed-off-by: MarioCadenas --- bundle-size-baseline.json | 98 +++++++++++++++++++-------------------- 1 file changed, 49 insertions(+), 49 deletions(-) diff --git a/bundle-size-baseline.json b/bundle-size-baseline.json index 1c3f942c..d7a90b7b 100644 --- a/bundle-size-baseline.json +++ b/bundle-size-baseline.json @@ -3,25 +3,25 @@ { "name": "@databricks/appkit", "tarball": { - "packed": 678473, - "unpacked": 2398820 + "packed": 727721, + "unpacked": 2550712 }, "dist": { "total": { - "raw": 2394404, - "gzip": 802685 + "raw": 2546267, + "gzip": 855702 }, "js": { - "raw": 705697, - "gzip": 246586 + "raw": 750162, + "gzip": 262624 }, "types": { - "raw": 274553, - "gzip": 93038 + "raw": 287585, + "gzip": 98642 }, "maps": { - "raw": 1403359, - "gzip": 459243 + "raw": 1497725, + "gzip": 490618 }, "css": { "raw": 0, @@ -31,22 +31,22 @@ "raw": 10795, "gzip": 3818 }, - "fileCount": 522 + "fileCount": 533 }, "entries": [ { "id": ".", - "gzip": 78284, + "gzip": 83090, "composition": { - "initialGzip": 75710, + "initialGzip": 80516, "lazyGzip": 2574, - "totalGzip": 78284, - "own": 249483, + "totalGzip": 83090, + "own": 263971, "nodeModules": null, "chunks": [ { "label": "index.js", - "gzip": 71612, + "gzip": 76418, "kind": "initial" }, { @@ -64,17 +64,17 @@ }, { "id": "./beta", - "gzip": 39881, + "gzip": 40062, "composition": { - "initialGzip": 39650, + "initialGzip": 39831, "lazyGzip": 231, - "totalGzip": 39881, - "own": 119365, + "totalGzip": 40062, + "own": 119920, "nodeModules": null, "chunks": [ { "label": "beta.js", - "gzip": 30501, + "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": 19056, + "gzip": 19068, "composition": { - "initialGzip": 19056, + "initialGzip": 19068, "lazyGzip": 0, - "totalGzip": 19056, - "own": 54837, + "totalGzip": 19068, + "own": 55003, "nodeModules": null, "chunks": [ { "label": "index.js", - "gzip": 19056, + "gzip": 19068, "kind": "initial" } ] @@ -128,25 +128,25 @@ { "name": "@databricks/appkit-ui", "tarball": { - "packed": 303755, - "unpacked": 1288650 + "packed": 313659, + "unpacked": 1314754 }, "dist": { "total": { - "raw": 1284681, - "gzip": 424469 + "raw": 1310746, + "gzip": 434328 }, "js": { - "raw": 371553, - "gzip": 122929 + "raw": 377356, + "gzip": 125186 }, "types": { - "raw": 208062, - "gzip": 74449 + "raw": 210262, + "gzip": 75512 }, "maps": { - "raw": 688206, - "gzip": 223745 + "raw": 706268, + "gzip": 230284 }, "css": { "raw": 16860, @@ -156,22 +156,22 @@ "raw": 0, "gzip": 0 }, - "fileCount": 475 + "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": 47055, + "gzip": 47663, "composition": { - "initialGzip": 605334, + "initialGzip": 439553, "lazyGzip": 49772, - "totalGzip": 655106, - "own": 170537, - "nodeModules": 1909652, + "totalGzip": 489325, + "own": 171715, + "nodeModules": 1402952, "chunks": [ { "label": "index.js", - "gzip": 603184, + "gzip": 437403, "kind": "initial" }, { From c9409bb10dfd63e3a535971977e365baad37f565 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 9 Jul 2026 11:48:32 +0200 Subject: [PATCH 25/25] fix(appkit): format nested Arrow leaves by type in JSON_ARRAY fallback decodeArrowAttachmentToRows correctly formatted top-level Decimal/Timestamp/ Date scalars, but nested values inside STRUCT/ARRAY/MAP fell through to a type-blind JSON.stringify replacer that only handled bigint/Uint8Array/Date. Nested decimals serialized as double-quoted unscaled integers (123.45 -> "12345") and nested timestamps as raw epoch-ms, diverging from the shape a warehouse emits natively under JSON_ARRAY. Add a type-aware renderNestedValue that recurses List/Struct/Map applying the same decimal/timestamp/date formatters by child field type, and route the top-level switch through a shared formatScalarCell. Output now matches the native data_array wire shape exactly. Verified against real bytes captured from the Reyden serverless warehouse (refuses INLINE+JSON_ARRAY): two round-trip fixtures assert all decoded columns equal the native JSON_ARRAY data_array for the same query. Signed-off-by: MarioCadenas --- .../src/plugins/analytics/result-delivery.ts | 119 +++++++++++++----- .../analytics/tests/result-delivery.test.ts | 61 +++++++++ 2 files changed, 152 insertions(+), 28 deletions(-) diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts index f1270a62..2a323104 100644 --- a/packages/appkit/src/plugins/analytics/result-delivery.ts +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -1,5 +1,5 @@ import type { sql } from "@databricks/sdk-experimental"; -import { type DataType, Type, tableFromIPC } from "apache-arrow"; +import { type DataType, type Field, Type, tableFromIPC } from "apache-arrow"; import type { SQLTypeMarker } from "shared"; import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; @@ -354,19 +354,6 @@ export async function deliverJsonResult( const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; const JSON_ARRAY_FALLBACK_MAX_BYTES = 64 * 1024 * 1024; -/** - * Replacer for nested `JSON.stringify`. Arrow values include `bigint` (which - * spec'd `JSON.stringify` throws on), `Uint8Array` (would serialize as an - * index map), and `Date` (would lose timezone in a nested struct). Coerce - * each to the string form the warehouse itself emits under native JSON_ARRAY. - */ -function nestedJsonReplacer(_key: string, value: unknown): unknown { - if (typeof value === "bigint") return value.toString(); - if (value instanceof Uint8Array) return Buffer.from(value).toString("base64"); - if (value instanceof Date) return value.toISOString(); - return value; -} - /** Render an apache-arrow Decimal cell to a fixed-point string. */ function formatDecimalCell( value: { toString(): string }, @@ -396,6 +383,84 @@ 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] === "[")) { @@ -457,19 +522,19 @@ export function decodeArrowAttachmentToRows( } switch (type.typeId) { case Type.Decimal: - row[name] = formatDecimalCell( - v as { toString(): string }, - (type as DataType & { scale: number }).scale, - ); - continue; case Type.Timestamp: - row[name] = formatTimestampCell( - Number(v), - (type as DataType & { timezone?: string | null }).timezone != null, - ); - continue; case Type.Date: - row[name] = formatDateCell(Number(v)); + 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 ( @@ -480,12 +545,10 @@ export function decodeArrowAttachmentToRows( row[name] = String(v); } else if (typeof v === "string") { row[name] = maybeParseJsonString(v); - } else if (v instanceof Date) { - row[name] = v.toISOString(); } else if (v instanceof Uint8Array) { row[name] = Buffer.from(v).toString("base64"); } else { - row[name] = JSON.stringify(v, nestedJsonReplacer); + row[name] = formatScalarCell(v, type); } } rows.push(row); diff --git a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts index 19a222da..65ffd228 100644 --- a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts @@ -362,3 +362,64 @@ describe("arrowDeliveryUnsupported", () => { 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]); + }); +});