From 840a67b058d6cd53d519509a8d511d14bed8bf9b Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Wed, 29 Apr 2026 15:10:30 +0000 Subject: [PATCH 01/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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 8d0c6c57fb34c81d0c91fa6b831c13cf4cf0d680 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Thu, 21 May 2026 22:15:17 +0000 Subject: [PATCH 18/20] feat(appkit): allowlist DBSQL error codes whose messages are safe to forward MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the error-sanitization work on #329. The clientMessage default of "Query execution failed" was applied uniformly to every ExecutionError.statementFailed call, which hid genuinely user-facing errors — the user's own SQL "Table 'foo' not found", "Syntax error near ',' on line 3", etc. — and forced users to debug in the dark. The fix is *not* to remove sanitization (CWE-209 is real — control plane messages carry correlation IDs, internal RPC paths, stack traces), but to distinguish the two sources by `ServiceErrorCode`: 1. **DBR (data plane)** — populates `errorDisplayMessage` for user-authored SQL errors and surfaces them via the gateway untouched. Source: `sqlgateway/scheduler/src/main/scala/DriverRequests.scala:181-208`, `ErrorInfo.withMessage(command.errorDisplayMessage).withErrorReturnedFromDataPlane(true)`. Codes: BAD_REQUEST, NOT_FOUND, ALREADY_EXISTS, DEADLINE_EXCEEDED, CANCELLED, UNAUTHENTICATED. 2. **Control plane** — wraps internal RPC failures, scheduler exceptions, capacity rejections. Source: `sqlgateway/scheduler/src/main/scala/utils/CommandUtils.scala`, `exceptions.scala`. Codes: INTERNAL_ERROR, IO_ERROR, UNKNOWN, RESOURCE_EXHAUSTED, SERVICE_UNDER_MAINTENANCE, TEMPORARILY_UNAVAILABLE, WORKSPACE_TEMPORARILY_UNAVAILABLE, ABORTED. `dbsql-error-allowlist.ts` enumerates source-1 codes. The Set is default-deny: any code not on the list — including future SDK additions the allowlist hasn't yet been reviewed against — collapses to the generic clientMessage. An explicit `clientMessage` argument to `statementFailed()` always wins. Server-side `.message` always preserves the raw upstream text for `logger.error` capture; sanitization is purely at the wire boundary. Validation: 2,129 appkit tests pass (+31 new); tsc clean; biome clean. Co-authored-by: Isaac Signed-off-by: James Broadhead --- .../src/errors/dbsql-error-allowlist.ts | 97 ++++++++++++ packages/appkit/src/errors/execution.ts | 42 +++-- .../tests/dbsql-error-allowlist.test.ts | 54 +++++++ .../appkit/src/errors/tests/errors.test.ts | 148 ++++++++++++++++++ 4 files changed, 330 insertions(+), 11 deletions(-) create mode 100644 packages/appkit/src/errors/dbsql-error-allowlist.ts create mode 100644 packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts diff --git a/packages/appkit/src/errors/dbsql-error-allowlist.ts b/packages/appkit/src/errors/dbsql-error-allowlist.ts new file mode 100644 index 00000000..080cc781 --- /dev/null +++ b/packages/appkit/src/errors/dbsql-error-allowlist.ts @@ -0,0 +1,97 @@ +/** + * Allowlist of Databricks SQL Statement Execution API error codes whose + * `error.message` is safe to forward to clients verbatim. + * + * Background — the warehouse populates `error.message` from two distinct + * sources, distinguishable only by `error.error_code`: + * + * 1. **DBR (data plane)** sets `errorDisplayMessage` for user-authored + * SQL errors (parse, semantic, type, missing-table) and surfaces them + * via the gateway untouched. These messages are inherently user-facing + * — they reference the user's own SQL — and were filtered through + * DBR's Thrift layer with the express purpose of being shown back to + * the user. They are *safe to passthrough*. + * + * Source of truth: DBR -> SQL Gateway scheduler + * `sqlgateway/scheduler/src/main/scala/DriverRequests.scala:181-208` + * populates `ErrorInfo.withMessage(command.errorDisplayMessage)` and + * sets `withErrorReturnedFromDataPlane(true)` for these. + * + * 2. **Control plane (proxy / scheduler / metaservice)** sets + * `error.message` from internal RPC failures, scheduler exceptions, + * capacity rejections, etc. These messages can include correlation + * IDs, internal pod / service names, stack traces, and storage paths + * — they are *not safe to passthrough* (CWE-209). + * + * Source: `sqlgateway/scheduler/src/main/scala/utils/CommandUtils.scala` + * and `exceptions.scala` wrap internal failures with + * `CommandExecutionFailed`. + * + * The mapping below classifies each `ServiceErrorCode` value from + * `@databricks/sdk-experimental` (`apis/sql/model.d.ts`) by which source + * dominates in practice. When the SDK adds a new code, default it to + * `false` here so it stays out of the passthrough until reviewed. + * + * SDK type reference: + * ``` + * type ServiceErrorCode = + * | "ABORTED" | "ALREADY_EXISTS" | "BAD_REQUEST" | "CANCELLED" + * | "DEADLINE_EXCEEDED" | "INTERNAL_ERROR" | "IO_ERROR" | "NOT_FOUND" + * | "RESOURCE_EXHAUSTED" | "SERVICE_UNDER_MAINTENANCE" + * | "TEMPORARILY_UNAVAILABLE" | "UNAUTHENTICATED" | "UNKNOWN" + * | "WORKSPACE_TEMPORARILY_UNAVAILABLE"; + * ``` + */ + +/** ServiceErrorCode values for which `error.message` may be forwarded verbatim. */ +const SAFE_PASSTHROUGH_CODES: ReadonlySet = new Set([ + // User-authored SQL errors flow through DBR's errorDisplayMessage. + // Examples: "Table or view 'foo' not found", "Syntax error at or near + // ',' on line 3", "Cannot resolve column 'col_x' given input columns". + "BAD_REQUEST", + // Catalog lookup misses — "Table 'foo' not found", "Schema 'bar' not + // found". DBR sets these via the catalog layer; messages reference + // the user's own SQL. + "NOT_FOUND", + // Catalog conflicts — "Table 'foo' already exists". Same provenance. + "ALREADY_EXISTS", + // Query timeout. Message is typically "Query timed out after Ns" — + // user-actionable, no internal state. + "DEADLINE_EXCEEDED", + // Client / admin cancellation. Message is "Statement was canceled". + "CANCELLED", + // Authentication failures don't carry internal state; the SDK surfaces + // a generic "Permission denied" or "Authentication required" string. + "UNAUTHENTICATED", +]); + +/** + * Codes that explicitly carry internal control-plane wording and must + * NEVER pass through to clients. Listed for documentation; classification + * is by absence from `SAFE_PASSTHROUGH_CODES`. + * + * - `INTERNAL_ERROR` — control-plane stack traces, RPC failures. + * - `IO_ERROR` — storage / network paths (may leak bucket names, pod IPs). + * - `UNKNOWN` — by definition unclassified; default-deny. + * - `RESOURCE_EXHAUSTED` — mixed; control plane sometimes leaks + * internal load data ("scheduler-pod-3 at 92% CPU"). Default-deny. + * - `SERVICE_UNDER_MAINTENANCE` / `TEMPORARILY_UNAVAILABLE` / + * `WORKSPACE_TEMPORARILY_UNAVAILABLE` — typically generic, but the + * gateway has been observed to include workspace identifiers and + * scheduler pod names. Default-deny. + * - `ABORTED` — server-side abort; can carry transaction / lock context + * identifying internal locking subsystems. Default-deny. + */ + +/** + * Returns true when the given upstream `error_code` belongs to the + * passthrough allowlist — i.e. the corresponding `error.message` was + * authored by DBR for the user's own SQL and is safe to forward. + * + * Returns false (default-deny) for any unrecognized code, including + * codes added to the SDK after this allowlist was last reviewed. + */ +export function isSqlErrorPassthrough(errorCode: string | undefined): boolean { + if (!errorCode) return false; + return SAFE_PASSTHROUGH_CODES.has(errorCode); +} diff --git a/packages/appkit/src/errors/execution.ts b/packages/appkit/src/errors/execution.ts index 3c9df7fe..3dc2944f 100644 --- a/packages/appkit/src/errors/execution.ts +++ b/packages/appkit/src/errors/execution.ts @@ -1,4 +1,5 @@ import { AppKitError } from "./base"; +import { isSqlErrorPassthrough } from "./dbsql-error-allowlist"; /** * Error thrown when an operation execution fails. @@ -50,16 +51,22 @@ export class ExecutionError extends AppKitError { /** * Create an execution error for statement failure. * @param errorMessage Human-readable error from the warehouse / SDK. - * Goes into `.message` for server logs only — *never* echoed to the - * client. Pass `clientMessage` explicitly if a sanitized text should - * reach the UI. - * @param errorCode Structured code (e.g. "INVALID_PARAMETER_VALUE") to - * preserve through wrapping. Optional. Forwarded on SSE error - * payloads so UI can branch on it instead of substring-matching - * `error`. - * @param clientMessage Optional client-safe replacement for `.message`. - * Defaults to "Query execution failed" via the `clientMessage` - * getter. Set this only when the upstream text is known-safe. + * When `errorCode` is on the DBSQL safe-passthrough allowlist (see + * `dbsql-error-allowlist.ts`), this text is forwarded to clients + * verbatim — it's authored by DBR for the user's own SQL and is + * inherently user-facing. For codes NOT on the allowlist (control + * plane errors carrying correlation IDs, internal paths, stack + * traces), `.message` goes into server logs only and the client + * sees a generic "Query execution failed". An explicit + * `clientMessage` argument always wins over both paths. + * @param errorCode Structured code (e.g. "BAD_REQUEST", + * "INVALID_PARAMETER_VALUE") to preserve through wrapping. Optional. + * Forwarded on SSE error payloads so UI can branch on it instead of + * substring-matching `error`. Also drives the passthrough decision. + * @param clientMessage Explicit client-safe message that always wins — + * bypasses the allowlist check. Use when the upstream text is + * known-safe regardless of code (e.g. constructed in-process from a + * trusted template). */ static statementFailed( errorMessage?: string, @@ -69,7 +76,20 @@ export class ExecutionError extends AppKitError { const message = errorMessage ? `Statement failed: ${errorMessage}` : "Statement failed: Unknown error"; - return new ExecutionError(message, { errorCode, clientMessage }); + + // Allowlist-driven passthrough: BAD_REQUEST / NOT_FOUND / + // ALREADY_EXISTS / etc. carry DBR-authored messages that *are* the + // user's own SQL error ("Table 'foo' not found", "Syntax error + // near ',' on line 3"). Hiding these forces users to debug in the + // dark; surfacing them is the entire reason an error message + // exists. Any code not on the allowlist defaults to generic. + const inferredClient = + clientMessage ?? (isSqlErrorPassthrough(errorCode) ? message : undefined); + + return new ExecutionError(message, { + errorCode, + clientMessage: inferredClient, + }); } /** diff --git a/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts b/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts new file mode 100644 index 00000000..f65856dd --- /dev/null +++ b/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, test } from "vitest"; +import { isSqlErrorPassthrough } from "../dbsql-error-allowlist"; + +describe("isSqlErrorPassthrough", () => { + // Allowlist members — DBR-authored, user-facing messages. + test.each([ + ["BAD_REQUEST"], + ["NOT_FOUND"], + ["ALREADY_EXISTS"], + ["DEADLINE_EXCEEDED"], + ["CANCELLED"], + ["UNAUTHENTICATED"], + ])("%s is allowlisted (DBR-authored user-facing message)", (code) => { + expect(isSqlErrorPassthrough(code)).toBe(true); + }); + + // Explicit denylist — control-plane sourced, internal state risk. + test.each([ + ["INTERNAL_ERROR"], + ["IO_ERROR"], + ["UNKNOWN"], + ["RESOURCE_EXHAUSTED"], + ["SERVICE_UNDER_MAINTENANCE"], + ["TEMPORARILY_UNAVAILABLE"], + ["WORKSPACE_TEMPORARILY_UNAVAILABLE"], + ["ABORTED"], + ])( + "%s is denied (control-plane sourced, may carry internal state)", + (code) => { + expect(isSqlErrorPassthrough(code)).toBe(false); + }, + ); + + test("undefined and empty string are denied (default-deny on missing source)", () => { + expect(isSqlErrorPassthrough(undefined)).toBe(false); + expect(isSqlErrorPassthrough("")).toBe(false); + }); + + test("any unrecognized code (e.g. new SDK additions) is denied", () => { + // When the SDK ships a new ServiceErrorCode variant, the allowlist + // must default-deny it until a human reviews the upstream message + // source. This test pins the default-deny behavior. + expect(isSqlErrorPassthrough("BRAND_NEW_CODE")).toBe(false); + expect(isSqlErrorPassthrough("PERMISSION_DENIED")).toBe(false); + expect(isSqlErrorPassthrough("FOOBAR")).toBe(false); + }); + + test("case-sensitive match — codes from the SDK are SHOUTY_CASE, lowercased input is rejected", () => { + // Defensive: if an upstream layer ever lowercases the code, we + // want the allowlist to fail closed rather than silently match. + expect(isSqlErrorPassthrough("bad_request")).toBe(false); + expect(isSqlErrorPassthrough("Bad_Request")).toBe(false); + }); +}); diff --git a/packages/appkit/src/errors/tests/errors.test.ts b/packages/appkit/src/errors/tests/errors.test.ts index 347ce1c0..03cc8001 100644 --- a/packages/appkit/src/errors/tests/errors.test.ts +++ b/packages/appkit/src/errors/tests/errors.test.ts @@ -303,6 +303,154 @@ describe("ExecutionError", () => { expect(error.message).toBe("No chunks or schema found in response"); expect(error.context?.dataType).toBe("chunks or schema"); }); + + describe("statementFailed clientMessage passthrough", () => { + test("BAD_REQUEST passes through the warehouse-authored message — user's own SQL error", () => { + // DBR sets errorDisplayMessage from the SQL parser / semantic + // analyzer; "Table not found" is the user's own SQL artifact + // and the entire point of an error message is to show it back. + const error = ExecutionError.statementFailed( + "Table or view 'users' not found", + "BAD_REQUEST", + ); + expect(error.clientMessage).toBe( + "Statement failed: Table or view 'users' not found", + ); + expect(error.errorCode).toBe("BAD_REQUEST"); + }); + + test("NOT_FOUND passes through (catalog miss is user-authored)", () => { + const error = ExecutionError.statementFailed( + "Schema 'analytics' not found", + "NOT_FOUND", + ); + expect(error.clientMessage).toBe( + "Statement failed: Schema 'analytics' not found", + ); + }); + + test("ALREADY_EXISTS passes through (catalog conflict is user-authored)", () => { + const error = ExecutionError.statementFailed( + "Table 'foo' already exists", + "ALREADY_EXISTS", + ); + expect(error.clientMessage).toBe( + "Statement failed: Table 'foo' already exists", + ); + }); + + test("DEADLINE_EXCEEDED passes through (timeout text is user-actionable)", () => { + const error = ExecutionError.statementFailed( + "Query timed out after 30s", + "DEADLINE_EXCEEDED", + ); + expect(error.clientMessage).toBe( + "Statement failed: Query timed out after 30s", + ); + }); + + test("CANCELLED passes through (cancellation is initiated by user/admin)", () => { + const error = ExecutionError.statementFailed( + "Statement was canceled", + "CANCELLED", + ); + expect(error.clientMessage).toBe( + "Statement failed: Statement was canceled", + ); + }); + + test("UNAUTHENTICATED passes through (SDK auth messages are generic)", () => { + const error = ExecutionError.statementFailed( + "Permission denied", + "UNAUTHENTICATED", + ); + expect(error.clientMessage).toBe("Statement failed: Permission denied"); + }); + + test("INTERNAL_ERROR collapses to generic — control-plane text may carry stack traces, correlation IDs, internal paths (CWE-209)", () => { + const error = ExecutionError.statementFailed( + "RPC to warehouse-prod-us-east-1.svc failed: corrId=abc-123, stack: at scheduler.process()", + "INTERNAL_ERROR", + ); + // .message keeps the raw text for server logs. + expect(error.message).toContain("corrId=abc-123"); + // .clientMessage collapses to the generic default. + expect(error.clientMessage).toBe("Query execution failed"); + }); + + test("IO_ERROR collapses to generic (may leak storage paths / bucket names)", () => { + const error = ExecutionError.statementFailed( + "Failed to read s3://internal-staging-bucket-3/path/to/shard", + "IO_ERROR", + ); + expect(error.clientMessage).toBe("Query execution failed"); + expect(error.message).toContain("internal-staging-bucket-3"); + }); + + test("UNKNOWN collapses to generic (unclassified by definition)", () => { + const error = ExecutionError.statementFailed( + "Something went wrong: scheduler-pod-3 returned 500", + "UNKNOWN", + ); + expect(error.clientMessage).toBe("Query execution failed"); + }); + + test("RESOURCE_EXHAUSTED collapses to generic (mixed: may leak load data)", () => { + const error = ExecutionError.statementFailed( + "Cluster at 92% capacity, scheduler-pod-3", + "RESOURCE_EXHAUSTED", + ); + expect(error.clientMessage).toBe("Query execution failed"); + }); + + test("absent errorCode falls back to generic (default-deny on unknown source)", () => { + const error = ExecutionError.statementFailed("Some upstream text"); + expect(error.clientMessage).toBe("Query execution failed"); + }); + + test("unrecognized errorCode (post-SDK-update) defaults to generic", () => { + // When the SDK adds a new code we haven't reviewed yet, we must + // NOT passthrough by default — security relies on the allowlist + // being explicit. + const error = ExecutionError.statementFailed( + "scheduler internal: stack overflow at sun.misc.Unsafe.park()", + "SOME_NEW_CODE_NOT_IN_ALLOWLIST", + ); + expect(error.clientMessage).toBe("Query execution failed"); + }); + + test("explicit clientMessage argument wins over the allowlist (caller knows best)", () => { + // Even on an allowlisted code, an explicit clientMessage takes + // precedence — for cases where the caller has constructed a + // better-tuned message in-process. + const error = ExecutionError.statementFailed( + "Table or view 'users' not found", + "BAD_REQUEST", + "We could not find the requested table.", + ); + expect(error.clientMessage).toBe( + "We could not find the requested table.", + ); + }); + + test("server-side .message always preserves the raw upstream text for log debugging", () => { + // Regardless of passthrough decision, the raw text stays on + // `.message` for `logger.error(err)` calls to capture full detail. + const safe = ExecutionError.statementFailed( + "Table 'users' not found", + "BAD_REQUEST", + ); + expect(safe.message).toBe("Statement failed: Table 'users' not found"); + + const unsafe = ExecutionError.statementFailed( + "Internal RPC failed: corrId=abc", + "INTERNAL_ERROR", + ); + expect(unsafe.message).toBe( + "Statement failed: Internal RPC failed: corrId=abc", + ); + }); + }); }); describe("InitializationError", () => { From a5c8e4a4377ac581b5a93ddc4702e8498c6260f3 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Fri, 22 May 2026 10:58:05 +0000 Subject: [PATCH 19/20] fix(appkit): expand allowlist to 5xx-class user-facing service errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reviewer pushback on the original allowlist: TEMPORARILY_UNAVAILABLE, RESOURCE_EXHAUSTED, SERVICE_UNDER_MAINTENANCE, WORKSPACE_TEMPORARILY_- UNAVAILABLE, and ABORTED carry messages that ARE the entire point of having an error string at the wire — telling users "Query execution failed" when the actual cause is a quota / maintenance / shard-moving condition forces them to guess. Allowlisting these. Verified against universe sources before reclassifying: - TEMPORARILY_UNAVAILABLE: stable template "DBSQL temporarily unavailable. Please try again in a few minutes." constructed at `sqlgateway/scheduler/api/src/main/scala/client/SchedulerRpcClient.scala` and `client/linkstore/SqlGatewayClerk.scala`. - RESOURCE_EXHAUSTED: stable user-actionable templates from `SchedulerRpcValidatorHook.scala`: "The maximum number of warehouses has been reached. Please contact Databricks support.", "You've hit the limit for warehouses for free usage. Stop or delete existing warehouses to free up capacity." - WORKSPACE_TEMPORARILY_UNAVAILABLE: workspace-level shard-routing messages at `sqlgateway/proxy/src/main/scala/thrift/Exceptions.scala` — names a shard ID, which is internal topology metadata but low-sensitivity and operationally useful for retry. - ABORTED: short reason strings like "version mismatch" at `ReydenWarehouseMonitor`. - SERVICE_UNDER_MAINTENANCE: maintenance-window template messages. INTERNAL_ERROR, IO_ERROR, UNKNOWN stay denied — these wrap raw `ex.getMessage` and interpolate orgIds, warehouse IDs, stack traces. `EndpointModel.scala` emits "SQL ${conf.warehouse} cannot be created due to repeated collisions on unique IDs." — exactly the kind of internal-state leak the sanitization exists to prevent. Reframed the module docs from "DBR vs control plane" to "designed-for-user vs designed-for-debugging" — the same code can be constructed at user-facing or debug-facing sites, but the contested 5xx-class codes are overwhelmingly the former in practice. Co-authored-by: Isaac Signed-off-by: James Broadhead --- .../src/errors/dbsql-error-allowlist.ts | 147 ++++++++++++------ .../tests/dbsql-error-allowlist.test.ts | 35 +++-- .../appkit/src/errors/tests/errors.test.ts | 52 ++++++- 3 files changed, 166 insertions(+), 68 deletions(-) diff --git a/packages/appkit/src/errors/dbsql-error-allowlist.ts b/packages/appkit/src/errors/dbsql-error-allowlist.ts index 080cc781..34cbefd5 100644 --- a/packages/appkit/src/errors/dbsql-error-allowlist.ts +++ b/packages/appkit/src/errors/dbsql-error-allowlist.ts @@ -2,37 +2,66 @@ * Allowlist of Databricks SQL Statement Execution API error codes whose * `error.message` is safe to forward to clients verbatim. * - * Background — the warehouse populates `error.message` from two distinct - * sources, distinguishable only by `error.error_code`: + * Framing: each `ServiceErrorCode` value falls into one of two camps, + * based on what wording is interpolated into `error.message` at the + * gateway construction site: * - * 1. **DBR (data plane)** sets `errorDisplayMessage` for user-authored - * SQL errors (parse, semantic, type, missing-table) and surfaces them - * via the gateway untouched. These messages are inherently user-facing - * — they reference the user's own SQL — and were filtered through - * DBR's Thrift layer with the express purpose of being shown back to - * the user. They are *safe to passthrough*. + * 1. **Designed-for-user** — the gateway interpolates a stable template + * string authored for end users (e.g. "DBSQL temporarily unavailable. + * Please try again in a few minutes.") or a DBR-stamped + * `errorDisplayMessage` covering the user's own SQL. These messages + * are *the entire point* of having an error string — hiding them + * forces users to debug in the dark. *Safe to passthrough.* * - * Source of truth: DBR -> SQL Gateway scheduler - * `sqlgateway/scheduler/src/main/scala/DriverRequests.scala:181-208` - * populates `ErrorInfo.withMessage(command.errorDisplayMessage)` and - * sets `withErrorReturnedFromDataPlane(true)` for these. + * 2. **Designed-for-debugging** — the gateway passes through a wrapped + * `ex.getMessage` or interpolates internal identifiers (orgId, unique + * warehouse IDs, scheduler pod names, internal column structure). + * These are intended for operator log inspection, not for clients. + * *Not safe to passthrough* (CWE-209). * - * 2. **Control plane (proxy / scheduler / metaservice)** sets - * `error.message` from internal RPC failures, scheduler exceptions, - * capacity rejections, etc. These messages can include correlation - * IDs, internal pod / service names, stack traces, and storage paths - * — they are *not safe to passthrough* (CWE-209). + * Source-of-truth construction sites for the classifications below: * - * Source: `sqlgateway/scheduler/src/main/scala/utils/CommandUtils.scala` - * and `exceptions.scala` wrap internal failures with - * `CommandExecutionFailed`. + * - DBR data-plane (`BAD_REQUEST`, `NOT_FOUND`, `ALREADY_EXISTS`): + * `sqlgateway/scheduler/src/main/scala/DriverRequests.scala:181-208` + * populates `ErrorInfo.withMessage(command.errorDisplayMessage)` and + * sets `withErrorReturnedFromDataPlane(true)`. * - * The mapping below classifies each `ServiceErrorCode` value from - * `@databricks/sdk-experimental` (`apis/sql/model.d.ts`) by which source - * dominates in practice. When the SDK adds a new code, default it to - * `false` here so it stays out of the passthrough until reviewed. + * - User-facing service messages (`TEMPORARILY_UNAVAILABLE`, + * `WORKSPACE_TEMPORARILY_UNAVAILABLE`, `SERVICE_UNDER_MAINTENANCE`): + * - `sqlgateway/scheduler/api/src/main/scala/client/SchedulerRpcClient.scala` + * constructs `SchedulerTemporarilyUnavailableException` with the + * literal "DBSQL temporarily unavailable. Please try again in a + * few minutes." + * - `sqlgateway/scheduler/api/src/main/scala/client/linkstore/SqlGatewayClerk.scala` + * constructs with `internalMessage = "DBSQL temporarily unavailable..."` + * - `sqlgateway/proxy/src/main/scala/thrift/Exceptions.scala` + * constructs `WORKSPACE_TEMPORARILY_UNAVAILABLE` from + * `NotServedByShardException.getMessage` — names a shard but no + * credentials / PII / stack traces. * - * SDK type reference: + * - User-actionable quota messages (`RESOURCE_EXHAUSTED`): + * `sqlgateway/scheduler/src/main/scala/SchedulerRpcValidatorHook.scala` + * emits stable templates: "The maximum number of warehouses has been + * reached. Please contact Databricks support.", "You've hit the limit + * for warehouses for free usage. Stop or delete existing warehouses + * to free up capacity.", etc. + * + * - Concurrency conflict (`ABORTED`): short reason strings like + * "version mismatch" (Reyden warehouse monitor). No internal context. + * + * - Genuinely-internal-only (`INTERNAL_ERROR`, `IO_ERROR`, `UNKNOWN`): + * `sqlgateway/database/src/main/scala/endpoints/EndpointModel.scala` + * interpolates `s"SQL ${conf.warehouse} cannot be created due to + * repeated collisions on unique IDs."` — exposes internal database + * state. `SchedulerRpcContextValidatorHook.workspaceNotOwned(orgId)` + * interpolates orgId. Test fixtures show stack traces interpolated + * into `internalMessage`. *Default-deny.* + * + * - `UNAUTHENTICATED` / `CANCELLED` / `DEADLINE_EXCEEDED`: SDK-level + * codes with generic message templates. + * + * SDK type reference (`@databricks/sdk-experimental`, + * `apis/sql/model.d.ts`): * ``` * type ServiceErrorCode = * | "ABORTED" | "ALREADY_EXISTS" | "BAD_REQUEST" | "CANCELLED" @@ -41,46 +70,62 @@ * | "TEMPORARILY_UNAVAILABLE" | "UNAUTHENTICATED" | "UNKNOWN" * | "WORKSPACE_TEMPORARILY_UNAVAILABLE"; * ``` + * + * **Default-deny on unknown codes** — when the SDK ships a new variant + * the allowlist hasn't been reviewed against, the new code collapses to + * the generic clientMessage until someone audits the construction sites + * and updates the Set below. */ /** ServiceErrorCode values for which `error.message` may be forwarded verbatim. */ const SAFE_PASSTHROUGH_CODES: ReadonlySet = new Set([ - // User-authored SQL errors flow through DBR's errorDisplayMessage. - // Examples: "Table or view 'foo' not found", "Syntax error at or near - // ',' on line 3", "Cannot resolve column 'col_x' given input columns". + // DBR-authored: user's own SQL errors. Parse, semantic, type, missing + // table/column. Examples: "Table or view 'foo' not found", "Syntax + // error at or near ',' on line 3". "BAD_REQUEST", - // Catalog lookup misses — "Table 'foo' not found", "Schema 'bar' not - // found". DBR sets these via the catalog layer; messages reference - // the user's own SQL. + // DBR catalog miss — "Table 'foo' not found", "Schema 'bar' not found". "NOT_FOUND", - // Catalog conflicts — "Table 'foo' already exists". Same provenance. + // DBR catalog conflict — "Table 'foo' already exists". "ALREADY_EXISTS", - // Query timeout. Message is typically "Query timed out after Ns" — - // user-actionable, no internal state. + // Stable template: "Query timed out after Ns". User-actionable. "DEADLINE_EXCEEDED", - // Client / admin cancellation. Message is "Statement was canceled". + // "Statement was canceled". User/admin-initiated; no internal state. "CANCELLED", - // Authentication failures don't carry internal state; the SDK surfaces - // a generic "Permission denied" or "Authentication required" string. + // SDK generic — "Permission denied", "Authentication required". "UNAUTHENTICATED", + // Stable user-facing template: "DBSQL temporarily unavailable. Please + // try again in a few minutes." Hiding this would force users to + // distinguish "warehouse issue" from "user error" by guesswork. + "TEMPORARILY_UNAVAILABLE", + // Workspace-level service messages from shard-routing failures. The + // message may name a shard ID — that's internal topology metadata, + // low sensitivity, and operationally useful for the user to know + // the workspace is in motion. + "WORKSPACE_TEMPORARILY_UNAVAILABLE", + // Maintenance windows — generic template messages. + "SERVICE_UNDER_MAINTENANCE", + // Stable user-actionable templates from the warehouse quota path: + // "The maximum number of warehouses has been reached. Please contact + // Databricks support.", "Stop or delete existing warehouses to free + // up capacity.", etc. Telling users "Query execution failed" when + // they've actually hit a quota is unhelpful. + "RESOURCE_EXHAUSTED", + // Concurrency conflict — short reason strings ("version mismatch", + // "concurrent modification"). User-relevant for retry decisions. + "ABORTED", ]); /** - * Codes that explicitly carry internal control-plane wording and must - * NEVER pass through to clients. Listed for documentation; classification - * is by absence from `SAFE_PASSTHROUGH_CODES`. + * Explicitly NOT on the allowlist (documented for future-reviewer + * context; classification is by absence from `SAFE_PASSTHROUGH_CODES`): * - * - `INTERNAL_ERROR` — control-plane stack traces, RPC failures. - * - `IO_ERROR` — storage / network paths (may leak bucket names, pod IPs). - * - `UNKNOWN` — by definition unclassified; default-deny. - * - `RESOURCE_EXHAUSTED` — mixed; control plane sometimes leaks - * internal load data ("scheduler-pod-3 at 92% CPU"). Default-deny. - * - `SERVICE_UNDER_MAINTENANCE` / `TEMPORARILY_UNAVAILABLE` / - * `WORKSPACE_TEMPORARILY_UNAVAILABLE` — typically generic, but the - * gateway has been observed to include workspace identifiers and - * scheduler pod names. Default-deny. - * - `ABORTED` — server-side abort; can carry transaction / lock context - * identifying internal locking subsystems. Default-deny. + * - `INTERNAL_ERROR` — construction sites interpolate unique warehouse + * IDs, orgIds, and wrap `ex.getMessage` from arbitrary exceptions. + * Example: `s"SQL ${conf.warehouse} cannot be created due to repeated + * collisions on unique IDs."` (`EndpointModel.scala`). High leak risk. + * - `IO_ERROR` — storage / network failures; messages typically include + * bucket names, paths, internal hostnames. + * - `UNKNOWN` — unclassified by definition; can be anything. */ /** diff --git a/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts b/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts index f65856dd..115f3e25 100644 --- a/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts +++ b/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts @@ -2,30 +2,37 @@ import { describe, expect, test } from "vitest"; import { isSqlErrorPassthrough } from "../dbsql-error-allowlist"; describe("isSqlErrorPassthrough", () => { - // Allowlist members — DBR-authored, user-facing messages. + // Allowlist members — designed-for-user messages. test.each([ + // DBR data plane — user's own SQL errors. ["BAD_REQUEST"], ["NOT_FOUND"], ["ALREADY_EXISTS"], + // SDK-level codes with generic templates. ["DEADLINE_EXCEEDED"], ["CANCELLED"], ["UNAUTHENTICATED"], - ])("%s is allowlisted (DBR-authored user-facing message)", (code) => { - expect(isSqlErrorPassthrough(code)).toBe(true); - }); - - // Explicit denylist — control-plane sourced, internal state risk. - test.each([ - ["INTERNAL_ERROR"], - ["IO_ERROR"], - ["UNKNOWN"], - ["RESOURCE_EXHAUSTED"], - ["SERVICE_UNDER_MAINTENANCE"], + // 5xx-class service messages — stable user-facing templates + // ("DBSQL temporarily unavailable. Please try again in a few + // minutes."). Hiding these forces users to guess whether a failure + // is on their side or ours. ["TEMPORARILY_UNAVAILABLE"], ["WORKSPACE_TEMPORARILY_UNAVAILABLE"], + ["SERVICE_UNDER_MAINTENANCE"], + // Quota — stable user-actionable templates ("Stop or delete + // existing warehouses to free up capacity."). + ["RESOURCE_EXHAUSTED"], + // Concurrency conflict — short reason strings, user-relevant for + // retry decisions. ["ABORTED"], - ])( - "%s is denied (control-plane sourced, may carry internal state)", + ])("%s is allowlisted (designed-for-user message)", (code) => { + expect(isSqlErrorPassthrough(code)).toBe(true); + }); + + // Explicit denylist — designed-for-debugging, interpolates internal + // identifiers, stack traces, or storage paths. + test.each([["INTERNAL_ERROR"], ["IO_ERROR"], ["UNKNOWN"]])( + "%s is denied (designed-for-debugging, interpolates internal state)", (code) => { expect(isSqlErrorPassthrough(code)).toBe(false); }, diff --git a/packages/appkit/src/errors/tests/errors.test.ts b/packages/appkit/src/errors/tests/errors.test.ts index 03cc8001..296e0ebc 100644 --- a/packages/appkit/src/errors/tests/errors.test.ts +++ b/packages/appkit/src/errors/tests/errors.test.ts @@ -395,12 +395,58 @@ describe("ExecutionError", () => { expect(error.clientMessage).toBe("Query execution failed"); }); - test("RESOURCE_EXHAUSTED collapses to generic (mixed: may leak load data)", () => { + test("RESOURCE_EXHAUSTED passes through (quota templates are user-actionable)", () => { + // SchedulerRpcValidatorHook.scala emits stable templates like + // "The maximum number of warehouses has been reached. Please + // contact Databricks support." Telling users "Query execution + // failed" when they've actually hit a quota is unhelpful. const error = ExecutionError.statementFailed( - "Cluster at 92% capacity, scheduler-pod-3", + "The maximum number of warehouses has been reached. Please contact Databricks support.", "RESOURCE_EXHAUSTED", ); - expect(error.clientMessage).toBe("Query execution failed"); + expect(error.clientMessage).toBe( + "Statement failed: The maximum number of warehouses has been reached. Please contact Databricks support.", + ); + }); + + test("TEMPORARILY_UNAVAILABLE passes through (stable service-down template)", () => { + // Source: SchedulerRpcClient.scala — fixed string authored for + // end users. + const error = ExecutionError.statementFailed( + "DBSQL temporarily unavailable. Please try again in a few minutes.", + "TEMPORARILY_UNAVAILABLE", + ); + expect(error.clientMessage).toBe( + "Statement failed: DBSQL temporarily unavailable. Please try again in a few minutes.", + ); + }); + + test("SERVICE_UNDER_MAINTENANCE passes through (maintenance messages are user-facing by design)", () => { + const error = ExecutionError.statementFailed( + "DBSQL is undergoing scheduled maintenance.", + "SERVICE_UNDER_MAINTENANCE", + ); + expect(error.clientMessage).toBe( + "Statement failed: DBSQL is undergoing scheduled maintenance.", + ); + }); + + test("WORKSPACE_TEMPORARILY_UNAVAILABLE passes through (shard topology is low-sensitivity, retry-relevant)", () => { + const error = ExecutionError.statementFailed( + "Workspace temporarily unavailable: shard moving", + "WORKSPACE_TEMPORARILY_UNAVAILABLE", + ); + expect(error.clientMessage).toBe( + "Statement failed: Workspace temporarily unavailable: shard moving", + ); + }); + + test("ABORTED passes through (conflict reason strings are retry-relevant)", () => { + const error = ExecutionError.statementFailed( + "version mismatch", + "ABORTED", + ); + expect(error.clientMessage).toBe("Statement failed: version mismatch"); }); test("absent errorCode falls back to generic (default-deny on unknown source)", () => { From b6ba29c697d677fcfe32814dde20a27beb416353 Mon Sep 17 00:00:00 2001 From: James Broadhead Date: Fri, 22 May 2026 11:20:51 +0000 Subject: [PATCH 20/20] fix(appkit): allowlist INTERNAL_ERROR + IO_ERROR; plumb requestId for triage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Threat model walkthrough — the recipient of these messages is an authenticated workspace user, not an anonymous internet visitor. They already know: - Their workspace URL and orgId (visible in every SDK call) - The warehouses they have grants on (listable via the SDK) - Their catalog contents Interpolated identifiers in INTERNAL_ERROR / IO_ERROR messages (warehouse names, orgIds, internal class names) are not new information to this audience — the classic CWE-209 leak-to-anonymous-attacker scenario does not apply. The cost of denial (every internal error becomes "Query execution failed" with zero signal) is certain; the benefit (defense-in-depth against unbounded wrapped `ex.getMessage` content) is speculative. Allowlisting both, with a requestId on the SSE error frame as the safety net: - `SSEError` now carries `requestId` (the SSE event id, same value the server logs in its `Stream execution failed` line — quoting it in a support ticket lets staff grep logs directly). - `SSEWriter.writeError` populates `requestId` from the event id. - `stream-manager.ts` logs `requestId=...` alongside `rawMsg`, `errorCode`, `upstreamCode` so a user-supplied requestId pinpoints the exact log line. - Analytics `/arrow-result` 410/404 paths generate their own requestId and emit it in both the JSON response and the server log line. - `useAnalyticsQuery` exposes `requestId: string | null` on `UseAnalyticsQueryResult` so UI code can render "Error ref: ". Only `UNKNOWN` remains denied — unclassified by definition, so the contents are unbounded in both shape and source. Validation: 2,133 appkit + 292 appkit-ui tests pass; tsc clean. Co-authored-by: Isaac Signed-off-by: James Broadhead --- .../__tests__/use-analytics-query.test.ts | 16 ++-- packages/appkit-ui/src/react/hooks/types.ts | 9 ++ .../src/react/hooks/use-analytics-query.ts | 12 ++- .../src/errors/dbsql-error-allowlist.ts | 93 +++++++++++++------ .../tests/dbsql-error-allowlist.test.ts | 30 +++--- .../appkit/src/errors/tests/errors.test.ts | 42 +++++---- .../appkit/src/plugins/analytics/analytics.ts | 20 +++- packages/appkit/src/stream/sse-writer.ts | 8 ++ packages/appkit/src/stream/stream-manager.ts | 7 +- .../appkit/src/stream/tests/stream.test.ts | 11 ++- packages/appkit/src/stream/types.ts | 8 ++ 11 files changed, 182 insertions(+), 74 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 ba4c962b..5417ac39 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,12 +184,12 @@ describe("useAnalyticsQuery", () => { expect(result.current.data).toBeNull(); }); - test("a server error event carrying a structured errorCode exposes it on the hook return value", async () => { - // The SSE error broadcaster forwards an `errorCode` field for - // UI branching (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. + test("a server error event carrying a structured errorCode + requestId exposes both on the hook return value", async () => { + // The SSE error broadcaster forwards `errorCode` (for UI branching) + // and `requestId` (for support triage — same id appears in the + // server-side logger.error line). The hook surfaces both so + // consumers can render "Error ref: " alongside the + // human message and branch behavior on the stable errorCode. const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); const { result } = renderHook(() => @@ -202,6 +202,7 @@ describe("useAnalyticsQuery", () => { error: "Server is at capacity, please retry", code: "UPSTREAM_ERROR", errorCode: "INLINE_ARROW_STASH_EXHAUSTED", + requestId: "11111111-2222-3333-4444-555555555555", }), }); @@ -210,6 +211,9 @@ describe("useAnalyticsQuery", () => { }); expect(result.current.loading).toBe(false); expect(result.current.errorCode).toBe("INLINE_ARROW_STASH_EXHAUSTED"); + expect(result.current.requestId).toBe( + "11111111-2222-3333-4444-555555555555", + ); errorSpy.mockRestore(); }); diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index a603e10a..feb1d8bf 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -73,6 +73,15 @@ export interface UseAnalyticsQueryResult { * matching `error`, which is a free-form sanitized message. */ errorCode: string | null; + /** + * Server-side correlation id for support triage. When non-null, the + * server emitted this id on the error frame and the same id appears + * in the backend's `logger.error` line. Surface it in error UIs + * ("Error ref: abc-123 — quote when filing a support ticket") so + * staff can grep server logs directly instead of asking the user to + * reproduce. + */ + requestId: 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 a3d2cf30..ff2b18ba 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -122,6 +122,7 @@ export function useAnalyticsQuery< const [loading, setLoading] = useState(false); const [error, setError] = useState(null); const [errorCode, setErrorCode] = useState(null); + const [requestId, setRequestId] = useState(null); const abortControllerRef = useRef(null); if (!queryKey || queryKey.trim().length === 0) { @@ -170,6 +171,7 @@ export function useAnalyticsQuery< setLoading(true); setError(null); setErrorCode(null); + setRequestId(null); setData(null); const abortController = new AbortController(); @@ -247,10 +249,16 @@ export function useAnalyticsQuery< if (typeof parsed.errorCode === "string") { setErrorCode(parsed.errorCode); } + // requestId is the SSE event id of the error frame — same + // value the server logs in its `Stream execution failed` + // line, so users can quote it in support tickets. + if (typeof parsed.requestId === "string") { + setRequestId(parsed.requestId); + } if (parsed.code) { console.error( - `[useAnalyticsQuery] Code: ${parsed.code}, Message: ${errorMsg}`, + `[useAnalyticsQuery] Code: ${parsed.code}, Message: ${errorMsg}, requestId: ${parsed.requestId ?? "n/a"}`, ); } return; @@ -321,5 +329,5 @@ export function useAnalyticsQuery< // Enable HMR for query updates in dev mode useQueryHMR(queryKey, start); - return { data, loading, error, errorCode }; + return { data, loading, error, errorCode, requestId }; } diff --git a/packages/appkit/src/errors/dbsql-error-allowlist.ts b/packages/appkit/src/errors/dbsql-error-allowlist.ts index 34cbefd5..780797ff 100644 --- a/packages/appkit/src/errors/dbsql-error-allowlist.ts +++ b/packages/appkit/src/errors/dbsql-error-allowlist.ts @@ -2,22 +2,35 @@ * Allowlist of Databricks SQL Statement Execution API error codes whose * `error.message` is safe to forward to clients verbatim. * - * Framing: each `ServiceErrorCode` value falls into one of two camps, - * based on what wording is interpolated into `error.message` at the - * gateway construction site: - * - * 1. **Designed-for-user** — the gateway interpolates a stable template - * string authored for end users (e.g. "DBSQL temporarily unavailable. - * Please try again in a few minutes.") or a DBR-stamped - * `errorDisplayMessage` covering the user's own SQL. These messages - * are *the entire point* of having an error string — hiding them - * forces users to debug in the dark. *Safe to passthrough.* - * - * 2. **Designed-for-debugging** — the gateway passes through a wrapped - * `ex.getMessage` or interpolates internal identifiers (orgId, unique - * warehouse IDs, scheduler pod names, internal column structure). - * These are intended for operator log inspection, not for clients. - * *Not safe to passthrough* (CWE-209). + * Threat model — the recipient of these messages is an *authenticated + * workspace user*, not an anonymous internet visitor. They already know + * the workspace URL, their orgId, their session token, the warehouses + * they have grants on, and the catalog contents they can see. Names of + * resources within the workspace, internal correlation IDs, and class + * names from Databricks-internal packages are not new information to + * this audience — the classic CWE-209 "stack-trace-on-500-page-leaks- + * to-attacker" scenario does not apply. + * + * Given that, the bar for passthrough is: does the message help the + * user (or their support contact) understand and act on the failure? + * Defaults skew toward yes because: + * + * - The cost of denial is *certain* (every error becomes "Query + * execution failed" with no signal; users open tickets, support has + * to reproduce or grep server logs to find the actual error). + * - The cost of allowing is *speculative* — depends on whatever the + * most recent wrapped exception happened to put in its `.toString()`. + * + * Server-side, the route handlers always `logger.error(rawMsg)` the + * full upstream text before sanitization — operators retain complete + * visibility regardless of what reaches the wire. SSE error frames + * also carry a `requestId` (the SSE event id) so users can quote it + * in support tickets and staff can grep logs against it directly. + * + * **Currently denied**: `UNKNOWN` only — by definition unclassified, + * so we can't reason about its contents and a default-deny is the + * minimum safe stance. Every other `ServiceErrorCode` variant is on + * the allowlist below with a documented rationale. * * Source-of-truth construction sites for the classifications below: * @@ -49,13 +62,25 @@ * - Concurrency conflict (`ABORTED`): short reason strings like * "version mismatch" (Reyden warehouse monitor). No internal context. * - * - Genuinely-internal-only (`INTERNAL_ERROR`, `IO_ERROR`, `UNKNOWN`): + * - `INTERNAL_ERROR` / `IO_ERROR`: * `sqlgateway/database/src/main/scala/endpoints/EndpointModel.scala` * interpolates `s"SQL ${conf.warehouse} cannot be created due to - * repeated collisions on unique IDs."` — exposes internal database - * state. `SchedulerRpcContextValidatorHook.workspaceNotOwned(orgId)` - * interpolates orgId. Test fixtures show stack traces interpolated - * into `internalMessage`. *Default-deny.* + * repeated collisions on unique IDs."`; + * `SchedulerRpcContextValidatorHook.workspaceNotOwned(orgId)` + * interpolates orgId. These identifiers are not sensitive to a + * workspace user (they own / can list both). The wrapped + * `ex.getMessage` content is unbounded but for the typical case is + * either operational (RPC timeout, lock contention) or shape-of-error + * information the user needs to act. Allowlisted with the requestId + * plumbing as the safety net for triage. *Not* allowlisted if a + * future review finds construction sites that interpolate user + * credentials, OAuth tokens, or other secrets — open an issue + flip + * the entry. + * + * - `UNKNOWN`: unclassified by definition. We have no way to reason + * about its contents — could be anything from "RPC connect failed" + * to a panic dump from an upstream library. Default-deny is the + * only defensible stance until the SDK gives us a finer code. * * - `UNAUTHENTICATED` / `CANCELLED` / `DEADLINE_EXCEEDED`: SDK-level * codes with generic message templates. @@ -113,19 +138,31 @@ const SAFE_PASSTHROUGH_CODES: ReadonlySet = new Set([ // Concurrency conflict — short reason strings ("version mismatch", // "concurrent modification"). User-relevant for retry decisions. "ABORTED", + // Wrapped internal exceptions. Construction sites interpolate + // workspace-internal identifiers (warehouse names, orgIds) and + // `ex.getMessage` from arbitrary causes. To an authenticated + // workspace user these identifiers are not sensitive — they can + // already enumerate them. The wrapped `ex.getMessage` content is + // unbounded but in practice is operational ("RPC timed out", + // "deadlock detected") or shape-of-failure detail the user needs to + // act on. The `requestId` we emit on the SSE error frame is the + // safety net for unhappy cases — users can quote it in tickets and + // staff can grep logs against it for the full unsanitized message. + "INTERNAL_ERROR", + // Storage / network failures. Same logic as INTERNAL_ERROR: path / + // bucket names visible to a workspace user with catalog access are + // not new information, and the wrapped error text usually tells + // them whether to retry, switch warehouses, or escalate. + "IO_ERROR", ]); /** * Explicitly NOT on the allowlist (documented for future-reviewer * context; classification is by absence from `SAFE_PASSTHROUGH_CODES`): * - * - `INTERNAL_ERROR` — construction sites interpolate unique warehouse - * IDs, orgIds, and wrap `ex.getMessage` from arbitrary exceptions. - * Example: `s"SQL ${conf.warehouse} cannot be created due to repeated - * collisions on unique IDs."` (`EndpointModel.scala`). High leak risk. - * - `IO_ERROR` — storage / network failures; messages typically include - * bucket names, paths, internal hostnames. - * - `UNKNOWN` — unclassified by definition; can be anything. + * - `UNKNOWN` — unclassified by definition; could be anything from a + * library panic to an unhandled `case` in the gateway. Default-deny + * is the only defensible stance until the SDK gives us a finer code. */ /** diff --git a/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts b/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts index 115f3e25..63827144 100644 --- a/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts +++ b/packages/appkit/src/errors/tests/dbsql-error-allowlist.test.ts @@ -2,7 +2,10 @@ import { describe, expect, test } from "vitest"; import { isSqlErrorPassthrough } from "../dbsql-error-allowlist"; describe("isSqlErrorPassthrough", () => { - // Allowlist members — designed-for-user messages. + // Allowlist members — passthrough is net-positive given the + // workspace-user threat model (recipient already knows the workspace's + // resources). Server-side full-detail logging + the requestId on the + // SSE error frame backstop the unhappy cases. test.each([ // DBR data plane — user's own SQL errors. ["BAD_REQUEST"], @@ -22,21 +25,24 @@ describe("isSqlErrorPassthrough", () => { // Quota — stable user-actionable templates ("Stop or delete // existing warehouses to free up capacity."). ["RESOURCE_EXHAUSTED"], - // Concurrency conflict — short reason strings, user-relevant for - // retry decisions. + // Concurrency conflict — short reason strings, user-relevant. ["ABORTED"], - ])("%s is allowlisted (designed-for-user message)", (code) => { + // Wrapped internal exceptions — interpolated identifiers (orgId, + // warehouse name) are non-sensitive to an authenticated workspace + // user; the wrapped `ex.getMessage` content is operationally + // useful for triage. RequestId on the SSE error frame is the + // safety net for unhappy cases. + ["INTERNAL_ERROR"], + ["IO_ERROR"], + ])("%s is allowlisted", (code) => { expect(isSqlErrorPassthrough(code)).toBe(true); }); - // Explicit denylist — designed-for-debugging, interpolates internal - // identifiers, stack traces, or storage paths. - test.each([["INTERNAL_ERROR"], ["IO_ERROR"], ["UNKNOWN"]])( - "%s is denied (designed-for-debugging, interpolates internal state)", - (code) => { - expect(isSqlErrorPassthrough(code)).toBe(false); - }, - ); + // Only UNKNOWN stays denied — unclassified by definition, so the + // wrapped content is unbounded in shape AND in source. + test("UNKNOWN is denied (unclassified — cannot reason about contents)", () => { + expect(isSqlErrorPassthrough("UNKNOWN")).toBe(false); + }); test("undefined and empty string are denied (default-deny on missing source)", () => { expect(isSqlErrorPassthrough(undefined)).toBe(false); diff --git a/packages/appkit/src/errors/tests/errors.test.ts b/packages/appkit/src/errors/tests/errors.test.ts index 296e0ebc..5a13fb69 100644 --- a/packages/appkit/src/errors/tests/errors.test.ts +++ b/packages/appkit/src/errors/tests/errors.test.ts @@ -367,27 +367,32 @@ describe("ExecutionError", () => { expect(error.clientMessage).toBe("Statement failed: Permission denied"); }); - test("INTERNAL_ERROR collapses to generic — control-plane text may carry stack traces, correlation IDs, internal paths (CWE-209)", () => { + test("INTERNAL_ERROR passes through (workspace identifiers are non-sensitive to authenticated user; requestId backstops triage)", () => { + // Interpolated warehouse names and orgIds are not new information + // to a workspace user — they can enumerate both via the SDK. + // Wrapped `ex.getMessage` is operationally useful for the user + // to act ("RPC timed out", "deadlock detected") and the + // requestId on the SSE error frame keys server logs for triage. const error = ExecutionError.statementFailed( - "RPC to warehouse-prod-us-east-1.svc failed: corrId=abc-123, stack: at scheduler.process()", + "SQL my-warehouse cannot be created due to repeated collisions on unique IDs.", "INTERNAL_ERROR", ); - // .message keeps the raw text for server logs. - expect(error.message).toContain("corrId=abc-123"); - // .clientMessage collapses to the generic default. - expect(error.clientMessage).toBe("Query execution failed"); + expect(error.clientMessage).toBe( + "Statement failed: SQL my-warehouse cannot be created due to repeated collisions on unique IDs.", + ); }); - test("IO_ERROR collapses to generic (may leak storage paths / bucket names)", () => { + test("IO_ERROR passes through (storage path / network failure detail is actionable, not credential-bearing)", () => { const error = ExecutionError.statementFailed( - "Failed to read s3://internal-staging-bucket-3/path/to/shard", + "Failed to read manifest after 3 retries", "IO_ERROR", ); - expect(error.clientMessage).toBe("Query execution failed"); - expect(error.message).toContain("internal-staging-bucket-3"); + expect(error.clientMessage).toBe( + "Statement failed: Failed to read manifest after 3 retries", + ); }); - test("UNKNOWN collapses to generic (unclassified by definition)", () => { + test("UNKNOWN collapses to generic (unclassified by definition — cannot reason about contents)", () => { const error = ExecutionError.statementFailed( "Something went wrong: scheduler-pod-3 returned 500", "UNKNOWN", @@ -481,20 +486,23 @@ describe("ExecutionError", () => { test("server-side .message always preserves the raw upstream text for log debugging", () => { // Regardless of passthrough decision, the raw text stays on - // `.message` for `logger.error(err)` calls to capture full detail. + // `.message` for `logger.error(err)` calls to capture full + // detail. For UNKNOWN (the only denied code), clientMessage + // collapses but .message still has the full content. const safe = ExecutionError.statementFailed( "Table 'users' not found", "BAD_REQUEST", ); expect(safe.message).toBe("Statement failed: Table 'users' not found"); - const unsafe = ExecutionError.statementFailed( - "Internal RPC failed: corrId=abc", - "INTERNAL_ERROR", + const denied = ExecutionError.statementFailed( + "Unclassified failure with corrId=abc", + "UNKNOWN", ); - expect(unsafe.message).toBe( - "Statement failed: Internal RPC failed: corrId=abc", + expect(denied.message).toBe( + "Statement failed: Unclassified failure with corrId=abc", ); + expect(denied.clientMessage).toBe("Query execution failed"); }); }); }); diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 5ba121d1..dca3ddd8 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1,3 +1,4 @@ +import { randomUUID } from "node:crypto"; import type { WorkspaceClient } from "@databricks/sdk-experimental"; import { tableFromIPC } from "apache-arrow"; import type express from "express"; @@ -146,9 +147,15 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { // 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); + const requestId = randomUUID(); + logger.debug( + "Inline Arrow stash miss for jobId=%s requestId=%s", + jobId, + requestId, + ); res.status(410).json({ error: "Inline Arrow result expired or unknown", + requestId, plugin: this.name, }); return; @@ -183,15 +190,18 @@ 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 requestId = randomUUID(); + logger.error("Arrow job error (requestId=%s): %O", requestId, error); + // Do not echo upstream / SDK error text to the client — the + // detail stays in the server log above, keyed by requestId so + // a user quoting it in a support ticket lets staff grep + // directly to this line. const errorCode = error instanceof ExecutionError ? error.errorCode : undefined; res.status(404).json({ error: "Arrow result unavailable", errorCode, + requestId, plugin: this.name, }); } diff --git a/packages/appkit/src/stream/sse-writer.ts b/packages/appkit/src/stream/sse-writer.ts index 2b49d1de..bf365993 100644 --- a/packages/appkit/src/stream/sse-writer.ts +++ b/packages/appkit/src/stream/sse-writer.ts @@ -43,9 +43,17 @@ export class SSEWriter { ): void { if (res.writableEnded) return; + // The SSE `id:` line and the in-payload `requestId` are + // intentionally the same value: the wire-level id is invisible to + // most React consumers (the event-source `data` is what reaches + // app code), so we duplicate it inside the JSON so the hook can + // surface it. Server-side logs use this same id in the + // `Stream execution failed` line — a user quoting it lets staff + // grep logs directly. const errorData: SSEError = { error, code, + requestId: eventId, ...(errorCode ? { errorCode } : {}), }; diff --git a/packages/appkit/src/stream/stream-manager.ts b/packages/appkit/src/stream/stream-manager.ts index 2cacbcb9..2ac0ea96 100644 --- a/packages/appkit/src/stream/stream-manager.ts +++ b/packages/appkit/src/stream/stream-manager.ts @@ -309,17 +309,22 @@ export class StreamManager { if (errorCode === SSEErrorCode.STREAM_ABORTED) { logger.info("Stream aborted by client (code=%s)", errorCode); } else { + // Log line includes the requestId we'll emit to the client — + // a user quoting their requestId in a support ticket lets + // staff grep this exact line for the full raw message. logger.error( - "Stream execution failed: %s (code=%s upstreamCode=%s)", + "Stream execution failed: %s (code=%s upstreamCode=%s requestId=%s)", rawMsg, errorCode, upstreamCode ?? "n/a", + errorEventId, ); } const payload: Record = { error: clientMsg, code: errorCode, + requestId: errorEventId, }; if (upstreamCode) payload.errorCode = upstreamCode; diff --git a/packages/appkit/src/stream/tests/stream.test.ts b/packages/appkit/src/stream/tests/stream.test.ts index a9924203..243178eb 100644 --- a/packages/appkit/src/stream/tests/stream.test.ts +++ b/packages/appkit/src/stream/tests/stream.test.ts @@ -271,9 +271,14 @@ describe("StreamManager", () => { await streamManager.stream(mockRes as any, generator); expect(events).toContain("event: error\n"); - expect(events).toContain( - 'data: {"error":"Internal server error","code":"INTERNAL_ERROR"}\n\n', - ); + // Match the error payload shape rather than the exact frame string — + // the frame now also carries a per-error requestId (randomUUID) for + // support triage, so we can't pin the full data line literally. + const dataLines = events.filter((e: string) => e.startsWith("data: ")); + const errorLine = dataLines.find((e: string) => e.includes('"error"')); + expect(errorLine).toContain('"error":"Internal server error"'); + expect(errorLine).toContain('"code":"INTERNAL_ERROR"'); + expect(errorLine).toMatch(/"requestId":"[0-9a-f-]+"/); }); test("should not crash if client disconnects during error", async () => { diff --git a/packages/appkit/src/stream/types.ts b/packages/appkit/src/stream/types.ts index ba0d4915..89a573ba 100644 --- a/packages/appkit/src/stream/types.ts +++ b/packages/appkit/src/stream/types.ts @@ -31,6 +31,14 @@ export interface SSEError { * the human-readable `error` string. */ errorCode?: string; + /** + * Correlation id for support triage. Set to the SSE event id of the + * error frame — the same value that appears in the server-side + * `logger.error("Stream execution failed: ...")` line, so a user + * quoting this id lets staff grep logs directly. Stable per-error, + * unique across the process lifetime. + */ + requestId?: string; } export interface BufferedEvent {