diff --git a/docs/docs/api/appkit/Class.ExecutionError.md b/docs/docs/api/appkit/Class.ExecutionError.md index c7bb2b44..e7c6f457 100644 --- a/docs/docs/api/appkit/Class.ExecutionError.md +++ b/docs/docs/api/appkit/Class.ExecutionError.md @@ -164,7 +164,7 @@ get clientMessage(): string; Execution errors default to a generic message — the raw warehouse / SDK text in `.message` often includes statement fragments, internal paths, and correlation IDs. UI code should branch on `errorCode` -(`INLINE_ARROW_STASH_EXHAUSTED`, `NOT_IMPLEMENTED`, etc.) and not on +(`RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`, etc.) and not on the human string. ##### Returns @@ -262,26 +262,6 @@ Create an execution error for closed/expired results *** -### stashExhausted() - -```ts -static stashExhausted(): ExecutionError; -``` - -Create an execution error for the inline Arrow stash being unable to -accept a payload (both regular and overflow pools at cap). - -The route deliberately surfaces this rather than silently re-running -the statement on EXTERNAL_LINKS — a second execution can be billed -and return divergent results for non-deterministic SQL. Operators -should tune stash capacity or back off load when this fires. - -#### Returns - -`ExecutionError` - -*** - ### statementFailed() ```ts diff --git a/packages/appkit-ui/src/js/arrow/arrow-client.test.ts b/packages/appkit-ui/src/js/arrow/arrow-client.test.ts new file mode 100644 index 00000000..8f9771f0 --- /dev/null +++ b/packages/appkit-ui/src/js/arrow/arrow-client.test.ts @@ -0,0 +1,60 @@ +import { Table, tableToIPC, Utf8, vectorFromArray } from "apache-arrow"; +import { describe, expect, test } from "vitest"; +import { ArrowClient } from "./arrow-client"; + +/** Build an Arrow IPC stream (2 cols, 2 rows) with the given field names. */ +function ipc(names: [string, string]): Uint8Array { + const table = new Table({ + [names[0]]: vectorFromArray(["a", "b"], new Utf8()), + [names[1]]: vectorFromArray(["1", "2"], new Utf8()), + }); + return tableToIPC(table, "stream"); +} + +describe("ArrowClient.processArrowBuffer column relabeling", () => { + test("relabels positional col_N to the provided manifest names", async () => { + const table = await ArrowClient.processArrowBuffer( + ipc(["col_0", "col_1"]), + ["name", "totalSpend"], + ); + expect(table.schema.fields.map((f) => f.name)).toEqual([ + "name", + "totalSpend", + ]); + // Data is preserved under the new names. + expect(table.getChild("name")?.get(0)).toBe("a"); + expect(table.getChild("totalSpend")?.get(1)).toBe("2"); + }); + + test("no-op when no names are provided", async () => { + const table = await ArrowClient.processArrowBuffer(ipc(["col_0", "col_1"])); + expect(table.schema.fields.map((f) => f.name)).toEqual(["col_0", "col_1"]); + }); + + test("no-op when the name count does not match", async () => { + const table = await ArrowClient.processArrowBuffer( + ipc(["col_0", "col_1"]), + ["only_one"], + ); + expect(table.schema.fields.map((f) => f.name)).toEqual(["col_0", "col_1"]); + }); + + test("ignores non-unique names rather than dropping a column", async () => { + const table = await ArrowClient.processArrowBuffer( + ipc(["col_0", "col_1"]), + ["dup", "dup"], + ); + expect(table.schema.fields.map((f) => f.name)).toEqual(["col_0", "col_1"]); + }); + + test("already-correct names are left as-is", async () => { + const table = await ArrowClient.processArrowBuffer( + ipc(["name", "totalSpend"]), + ["name", "totalSpend"], + ); + expect(table.schema.fields.map((f) => f.name)).toEqual([ + "name", + "totalSpend", + ]); + }); +}); diff --git a/packages/appkit-ui/src/js/arrow/arrow-client.ts b/packages/appkit-ui/src/js/arrow/arrow-client.ts index 3ee17363..4554243f 100644 --- a/packages/appkit-ui/src/js/arrow/arrow-client.ts +++ b/packages/appkit-ui/src/js/arrow/arrow-client.ts @@ -1,4 +1,4 @@ -import type { Field, Table } from "apache-arrow"; +import type { Field, Table, Vector } from "apache-arrow"; import { DATE_FIELD_PATTERNS, METADATA_DATE_PATTERNS, @@ -23,14 +23,40 @@ export class ArrowClient { * Lazily loads the Apache Arrow library on first use. * * @param buffer - The Arrow IPC format buffer + * @param columnNames - Optional real column names to relabel the schema + * with. Databricks encodes ARROW_STREAM schemas positionally (`col_0`, + * …); the analytics server sends the manifest names in a response header + * so charts can look columns up by their real name. Relabels in-place by + * rebuilding the Table around the existing vectors (no re-decode). A + * no-op when the names already match, the count differs, or the names + * aren't unique. * @returns Promise resolving to an Arrow Table */ - static async processArrowBuffer(buffer: Uint8Array): Promise { + static async processArrowBuffer( + buffer: Uint8Array, + columnNames?: string[], + ): Promise
{ try { const arrow = await getArrowModule(); // Initialize type ID sets now that Arrow is loaded await initializeTypeIdSets(); - return arrow.tableFromIPC(buffer); + const table = arrow.tableFromIPC(buffer); + + if (!columnNames || columnNames.length !== table.numCols) return table; + const fields = table.schema.fields; + const targetNames = columnNames.map((n, i) => + n && n.length > 0 ? n : fields[i].name, + ); + if (targetNames.every((n, i) => n === fields[i].name)) return table; + if (new Set(targetNames).size !== targetNames.length) return table; + + const children: Record = {}; + for (let i = 0; i < fields.length; i++) { + const child = table.getChildAt(i); + if (!child) return table; + children[targetNames[i]] = child; + } + return new arrow.Table(children); } catch (error) { throw new Error( `Failed to process Arrow buffer: ${ diff --git a/packages/appkit-ui/src/js/sse/connect-sse.ts b/packages/appkit-ui/src/js/sse/connect-sse.ts index 13d9053d..010ba3bf 100644 --- a/packages/appkit-ui/src/js/sse/connect-sse.ts +++ b/packages/appkit-ui/src/js/sse/connect-sse.ts @@ -19,8 +19,8 @@ export async function connectSSE( retryDelay = 2000, maxRetries = 3, // 1 MiB — matches the server's `streamDefaults.maxEventSize`. SSE - // carries only short JSON control messages; bulk Arrow payloads flow - // over plain HTTP via `/api/analytics/arrow-result/:jobId`, so this + // carries only short JSON control messages; bulk Arrow payloads stream + // back as the raw HTTP response body of the query request, so this // buffer never needs to hold multi-MiB attachments. maxBufferSize = 1 * 1024 * 1024, timeout = 300000, // 5 minutes diff --git a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts index 8b1acc13..4c5f1dd5 100644 --- a/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-analytics-query.test.ts @@ -1,5 +1,5 @@ import { act, renderHook, waitFor } from "@testing-library/react"; -import { beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; let lastConnectArgs: any = null; let capturedCallbacks: { @@ -52,72 +52,83 @@ describe("useAnalyticsQuery", () => { capturedCallbacks = {}; }); - test("fetches an arrow message (warehouse statement id) via /arrow-result", async () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + test("fetches ARROW_STREAM results as raw Arrow bytes directly from the query endpoint (no SSE)", async () => { const fakeTable = { numRows: 1, schema: { fields: [] } }; const fakeBytes = new Uint8Array([1, 2, 3]); - mockFetchArrow.mockResolvedValueOnce(fakeBytes); + const fetchMock = vi.fn().mockResolvedValue({ + ok: true, + arrayBuffer: async () => fakeBytes.buffer, + headers: { get: () => null }, + }); + vi.stubGlobal("fetch", fetchMock); mockProcessArrowBuffer.mockResolvedValueOnce(fakeTable); const { result } = renderHook(() => useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), ); - await lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow", statement_id: "stmt-warehouse-1" }), - }); - await waitFor(() => { expect(result.current.data).toBe(fakeTable); }); - expect(mockFetchArrow).toHaveBeenCalledTimes(1); - expect(mockFetchArrow).toHaveBeenCalledWith( - "/api/analytics/arrow-result/stmt-warehouse-1", + // ARROW_STREAM never opens an SSE stream — the bytes come straight back + // on a direct POST to the query endpoint. + expect(mockConnectSSE).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(1); + const [url, init] = fetchMock.mock.calls[0]; + expect(String(url)).toContain("/api/analytics/query/q"); + expect(init.method).toBe("POST"); + // No column-names header → decode with the raw Arrow schema names. + expect(mockProcessArrowBuffer).toHaveBeenCalledWith( + expect.any(Uint8Array), + undefined, ); - expect(mockProcessArrowBuffer).toHaveBeenCalledWith(fakeBytes); + expect(result.current.loading).toBe(false); }); - test("fetches an arrow message with synthetic inline- id through the same /arrow-result path", async () => { - // The client must treat inline and external-links responses uniformly — - // it never decodes base64 locally. The /arrow-result route on the - // server is the only place that knows which path the bytes came from. + test("relabels ARROW_STREAM columns from the X-Appkit-Arrow-Columns header", async () => { const fakeTable = { numRows: 1, schema: { fields: [] } }; - const fakeBytes = new Uint8Array([1, 2, 3, 4, 5]); - mockFetchArrow.mockResolvedValueOnce(fakeBytes); + const names = ["name", "totalSpend"]; + const fetchMock = vi.fn().mockResolvedValue({ + ok: true, + arrayBuffer: async () => new Uint8Array([1, 2, 3]).buffer, + headers: { + get: (h: string) => + h === "X-Appkit-Arrow-Columns" + ? encodeURIComponent(JSON.stringify(names)) + : null, + }, + }); + vi.stubGlobal("fetch", fetchMock); mockProcessArrowBuffer.mockResolvedValueOnce(fakeTable); const { result } = renderHook(() => useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), ); - await lastConnectArgs.onMessage({ - data: JSON.stringify({ - type: "arrow", - statement_id: "inline-abc-xyz", - }), - }); - await waitFor(() => { expect(result.current.data).toBe(fakeTable); }); - expect(mockFetchArrow).toHaveBeenCalledTimes(1); - expect(mockFetchArrow).toHaveBeenCalledWith( - "/api/analytics/arrow-result/inline-abc-xyz", + // The parsed manifest names are handed to the decoder for relabeling. + expect(mockProcessArrowBuffer).toHaveBeenCalledWith( + expect.any(Uint8Array), + names, ); }); - test("surfaces an error when the arrow fetch fails", async () => { - mockFetchArrow.mockRejectedValueOnce(new Error("network")); + test("surfaces an error when the ARROW_STREAM fetch fails", async () => { + const fetchMock = vi.fn().mockRejectedValue(new Error("boom")); + vi.stubGlobal("fetch", fetchMock); const { result } = renderHook(() => useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), ); - await lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow", statement_id: "stmt-1" }), - }); - await waitFor(() => { expect(result.current.error).toBe( "Unable to load data, please try again", @@ -126,27 +137,28 @@ describe("useAnalyticsQuery", () => { expect(result.current.loading).toBe(false); }); - test("rejects the retired arrow_inline message type as schema-invalid", async () => { - // arrow_inline was the prior wire shape. The discriminated union no - // longer accepts it, so it falls through to the generic error/code - // branch — but critically, it must NEVER trigger ArrowClient calls. + test("surfaces a structured errorCode from an ARROW_STREAM JSON error body", async () => { + // On a pre-first-byte failure the server responds with a JSON + // `{ error, errorCode }` body; the hook exposes both so consumers can + // branch on the stable code (e.g. RESULT_TOO_LARGE_FOR_JSON_FALLBACK). + const fetchMock = vi.fn().mockResolvedValue({ + ok: false, + json: async () => ({ + error: "Result too large for JSON format", + errorCode: "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", + }), + }); + vi.stubGlobal("fetch", fetchMock); + const { result } = renderHook(() => useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), ); - await lastConnectArgs.onMessage({ - data: JSON.stringify({ type: "arrow_inline", attachment: "AQID" }), - }); - await waitFor(() => { - expect( - result.current.loading || - result.current.error || - result.current.data === null, - ).toBeTruthy(); + expect(result.current.error).toBe("Result too large for JSON format"); }); - expect(mockProcessArrowBuffer).not.toHaveBeenCalled(); - expect(mockFetchArrow).not.toHaveBeenCalled(); + expect(result.current.errorCode).toBe("RESULT_TOO_LARGE_FOR_JSON_FALLBACK"); + expect(result.current.loading).toBe(false); }); test("normalizes an empty result message (no data field) to []", async () => { @@ -208,15 +220,14 @@ describe("useAnalyticsQuery", () => { }); test("a server error event carrying a structured errorCode exposes it on the hook return value", async () => { - // The SSE error broadcaster forwards an `errorCode` field for - // UI branching (e.g. INLINE_ARROW_STASH_EXHAUSTED). The hook - // surfaces both the human `error` text AND the structured - // `errorCode` so consumers can branch on the stable identifier - // instead of substring-matching the sanitized human message. + // The SSE error broadcaster forwards an `errorCode` field for UI + // branching. The hook surfaces both the human `error` text AND the + // structured `errorCode` so consumers can branch on the stable + // identifier instead of substring-matching the sanitized human message. const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); const { result } = renderHook(() => - useAnalyticsQuery("q", null, { format: "ARROW_STREAM" }), + useAnalyticsQuery("q", null, { format: "JSON_ARRAY" }), ); await lastConnectArgs.onMessage({ @@ -224,7 +235,7 @@ describe("useAnalyticsQuery", () => { type: "error", error: "Server is at capacity, please retry", code: "UPSTREAM_ERROR", - errorCode: "INLINE_ARROW_STASH_EXHAUSTED", + errorCode: "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", }), }); @@ -232,7 +243,7 @@ describe("useAnalyticsQuery", () => { expect(result.current.error).toBe("Server is at capacity, please retry"); }); expect(result.current.loading).toBe(false); - expect(result.current.errorCode).toBe("INLINE_ARROW_STASH_EXHAUSTED"); + expect(result.current.errorCode).toBe("RESULT_TOO_LARGE_FOR_JSON_FALLBACK"); errorSpy.mockRestore(); }); @@ -447,25 +458,5 @@ describe("useAnalyticsQuery", () => { expect(result.current.data).toBeNull(); }); - - test("ignores late arrow envelope after the controller was aborted", async () => { - const { result } = renderHook(() => - useAnalyticsQuery("test_query", null, { format: "ARROW_STREAM" }), - ); - - await waitFor(() => expect(capturedCallbacks.signal).toBeDefined()); - - markAborted(); - - await act(async () => { - await capturedCallbacks.onMessage?.({ - data: JSON.stringify({ type: "arrow", statement_id: "stmt-123" }), - }); - }); - - expect(mockFetchArrow).not.toHaveBeenCalled(); - expect(result.current.data).toBeNull(); - expect(result.current.error).toBeNull(); - }); }); }); diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index 5112ac94..aa0df890 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -96,10 +96,10 @@ export interface UseAnalyticsQueryResult { error: string | null; /** * Structured upstream error code when the server attaches one to the - * SSE error payload (e.g. `INLINE_ARROW_STASH_EXHAUSTED`, - * `RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`). Prefer - * branching UI on this stable identifier rather than substring- - * matching `error`, which is a free-form sanitized message. + * error payload (e.g. `RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, + * `ARROW_DELIVERY_UNSUPPORTED`, `NOT_IMPLEMENTED`). Prefer branching UI + * on this stable identifier rather than substring-matching `error`, + * which is a free-form sanitized message. */ errorCode: string | null; /** diff --git a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts index a27bb6f3..20f2c99e 100644 --- a/packages/appkit-ui/src/react/hooks/use-analytics-query.ts +++ b/packages/appkit-ui/src/react/hooks/use-analytics-query.ts @@ -62,12 +62,21 @@ function getDevMode(): string { return dev ? `?dev=${dev}` : ""; } -function getArrowStreamUrl(id: string): string { - return `/api/analytics/arrow-result/${id}`; -} - const GENERIC_LOAD_ERROR = "Unable to load data, please try again"; +/** Map a fetch/SSE transport error to a user-facing message. */ +function userFacingFetchError(error: unknown): string { + if (error instanceof Error) { + if (error.name === "AbortError") { + return "Request timed out, please try again"; + } + if (error.message.includes("Failed to fetch")) { + return "Network error. Please check your connection."; + } + } + return GENERIC_LOAD_ERROR; +} + interface AnalyticsQuerySseContext { setLoading: (loading: boolean) => void; setError: (error: string | null) => void; @@ -124,28 +133,9 @@ async function handleAnalyticsSseMessage( return; } - // success - Arrow format. Both INLINE (server-stashed, statement_id - // prefixed with "inline-") and EXTERNAL_LINKS (warehouse statement_id) - // flow through this single branch — the /arrow-result route dispatches - // based on the id prefix so the client doesn't need to know which path - // the bytes came from. - if (msg?.type === "arrow") { - try { - const arrowData = await ArrowClient.fetchArrow( - getArrowStreamUrl(msg.statement_id), - ); - const table = await ArrowClient.processArrowBuffer(arrowData); - ctx.setLoading(false); - ctx.setData(table as ResultType); - ctx.unpublishWarehouseStatus(); - } catch (error) { - console.error("[useAnalyticsQuery] Failed to fetch Arrow data", error); - ctx.setLoading(false); - ctx.setError(GENERIC_LOAD_ERROR); - ctx.unpublishWarehouseStatus(); - } - return; - } + // NOTE: ARROW_STREAM no longer flows over SSE — the server streams the + // raw Arrow IPC bytes back as the query response body, handled by + // `fetchArrowDirect` instead of this SSE handler. if (parsed.type === "error" || parsed.error || parsed.code) { const errorMsg = @@ -156,9 +146,9 @@ async function handleAnalyticsSseMessage( ctx.setError(errorMsg); ctx.unpublishWarehouseStatus(); // Propagate the upstream structured code so UI consumers can branch on - // a stable identifier (e.g. retry on INLINE_ARROW_STASH_EXHAUSTED, - // format-switch on RESULT_TOO_LARGE_FOR_JSON_FALLBACK) instead of - // parsing the human-readable message. + // a stable identifier (e.g. format-switch on + // RESULT_TOO_LARGE_FOR_JSON_FALLBACK or ARROW_DELIVERY_UNSUPPORTED) + // instead of parsing the human-readable message. if (typeof parsed.errorCode === "string") { ctx.setErrorCode(parsed.errorCode); } @@ -183,8 +173,116 @@ async function handleAnalyticsSseMessage( } } +interface ArrowDirectContext { + url: string; + payload: string; + signal: AbortSignal; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; + setErrorCode: (code: string | null) => void; + setData: (data: unknown) => void; + unpublishWarehouseStatus: () => void; +} + +/** + * Fetch the real column names for a statement from the fallback endpoint, + * used when a very wide schema's names didn't fit in the response header. + * Returns undefined on any failure so decoding falls back to the raw Arrow + * schema names. + */ +async function fetchArrowColumns( + statementId: string, + signal: AbortSignal, +): Promise { + try { + const res = await fetch( + `/api/analytics/columns/${encodeURIComponent(statementId)}`, + { signal }, + ); + if (!res.ok) return undefined; + const body = (await res.json()) as { columns?: unknown }; + return Array.isArray(body.columns) ? (body.columns as string[]) : undefined; + } catch { + return undefined; + } +} + +/** + * Fetch an ARROW_STREAM query result as raw Arrow IPC bytes directly from + * the query endpoint (no SSE, no second /arrow-result request) and decode + * it into a Table. The server streams the bytes back as the POST response + * body; errors before the first byte arrive as a JSON `{ error, errorCode }`. + */ +async function fetchArrowDirect(ctx: ArrowDirectContext): Promise { + try { + const response = await fetch(ctx.url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: ctx.payload, + signal: ctx.signal, + }); + if (ctx.signal.aborted) return; + + if (!response.ok) { + let message = GENERIC_LOAD_ERROR; + let code: string | null = null; + try { + const body = (await response.json()) as { + error?: string; + errorCode?: string; + }; + if (body.error) message = body.error; + if (typeof body.errorCode === "string") code = body.errorCode; + } catch { + // Non-JSON error body — keep the generic message. + } + ctx.setLoading(false); + ctx.setError(message); + if (code) ctx.setErrorCode(code); + ctx.unpublishWarehouseStatus(); + return; + } + + const buffer = await response.arrayBuffer(); + if (ctx.signal.aborted) return; + // Databricks encodes ARROW_STREAM columns positionally (col_0, …); the + // server sends the real manifest names so we can relabel the decoded + // Table (charts look columns up by name). Normally inline in the + // `X-Appkit-Arrow-Columns` header; for very wide schemas the header + // carries only a statement-id reference and we fetch the names. + let columnNames: string[] | undefined; + const header = response.headers.get("X-Appkit-Arrow-Columns"); + if (header) { + try { + columnNames = JSON.parse(decodeURIComponent(header)); + } catch { + // Malformed header — fall back to the raw Arrow schema names. + } + } else { + const ref = response.headers.get("X-Appkit-Arrow-Columns-Ref"); + if (ref) { + columnNames = await fetchArrowColumns(ref, ctx.signal); + } + } + const table = await ArrowClient.processArrowBuffer( + new Uint8Array(buffer), + columnNames, + ); + ctx.setData(table); + ctx.setLoading(false); + ctx.unpublishWarehouseStatus(); + } catch (error) { + if (ctx.signal.aborted) return; + ctx.setLoading(false); + ctx.unpublishWarehouseStatus(); + ctx.setError(userFacingFetchError(error)); + } +} + /** - * Subscribe to an analytics query over SSE and returns its latest result. + * Subscribe to an analytics query and return its latest result. JSON_ARRAY + * results stream over SSE (with warehouse-readiness progress); ARROW_STREAM + * results are fetched as raw Arrow bytes directly from the query endpoint. * Integration hook between client and analytics plugin. * * The return type is automatically inferred based on the format: @@ -289,6 +387,22 @@ export function useAnalyticsQuery< const abortController = new AbortController(); abortControllerRef.current = abortController; + // ARROW_STREAM: the server streams raw Arrow IPC bytes back on the query + // response body (no SSE). Fetch and decode directly. + if (format === "ARROW_STREAM") { + void fetchArrowDirect({ + url: urlSuffix, + payload, + signal: abortController.signal, + setLoading, + setError, + setErrorCode, + setData: (table) => setData(table as ResultType), + unpublishWarehouseStatus, + }); + return; + } + const sseContext: AnalyticsQuerySseContext = { setLoading, setError, @@ -323,7 +437,7 @@ export function useAnalyticsQuery< // forces the consumer into a clean retry path. console.warn("[useAnalyticsQuery] Malformed message received", error); setLoading(false); - setError("Unable to load data, please try again"); + setError(GENERIC_LOAD_ERROR); abortController.abort(); } }, @@ -332,26 +446,21 @@ export function useAnalyticsQuery< setLoading(false); unpublishWarehouseStatus(); - let userMessage = GENERIC_LOAD_ERROR; if (error instanceof Error) { - if (error.name === "AbortError") { - userMessage = "Request timed out, please try again"; - } else if (error.message.includes("Failed to fetch")) { - userMessage = "Network error. Please check your connection."; - } console.error("[useAnalyticsQuery] Error", { queryKey, error: error.message, stack: error.stack, }); } - setError(userMessage); + setError(userFacingFetchError(error)); }, }); }, [ queryKey, payload, urlSuffix, + format, publishWarehouseStatus, unpublishWarehouseStatus, ]); diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 9167ef97..239de551 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -29,18 +29,42 @@ import { WarehouseStatusEmitter } from "./warehouse-status-emitter"; const logger = createLogger("connectors:sql-warehouse"); +/** + * Real column names from a statement's result manifest (positional + * `column_N` fallback for any blank ones). Databricks encodes ARROW_STREAM + * schemas positionally (`col_0`, …); these names let consumers relabel an + * Arrow result to match the JSON path. Returns `undefined` when the manifest + * carries no columns. + */ +function arrowColumnNames( + response: sql.StatementResponse, +): string[] | undefined { + const cols = response.manifest?.schema?.columns; + if (!cols || cols.length === 0) return undefined; + return cols.map((c, i) => + c.name && c.name.length > 0 ? c.name : `column_${i}`, + ); +} + /** * Maximum size for inline Arrow IPC attachments (25 MiB decoded — the * Databricks Statement Execution API hard cap on INLINE responses). * - * Bulk Arrow payloads no longer traverse SSE — the analytics route stashes - * them via `InlineArrowStash` and the client fetches over HTTP — so this - * cap is bounded by the upstream API rather than our event-size budget. - * Larger results still fall through to `disposition: "EXTERNAL_LINKS"`, - * handled by the analytics format-fallback. + * The analytics route streams the decoded attachment straight back on the + * query response body; this cap bounds the single in-memory copy. Larger + * results fall through to `disposition: "EXTERNAL_LINKS"`, which streams + * chunk-by-chunk and is not bound by this limit. */ const MAX_INLINE_ATTACHMENT_BYTES = 25 * 1024 * 1024; +/** + * Safety cap on how many additional EXTERNAL_LINKS chunks + * {@link SQLWarehouseConnector._resolveAllExternalLinks} will follow when the + * manifest omits `total_chunk_count`. High enough to cover any real result; + * only bounds a misbehaving warehouse with a cyclic `next_chunk_index`. + */ +const MAX_EXTERNAL_CHUNK_FOLLOWS = 10_000; + interface SQLWarehouseConfig { timeout?: number; telemetry?: TelemetryOptions; @@ -167,8 +191,6 @@ export class SQLWarehouseConnector { if (!this._arrowProcessor) { this._arrowProcessor = new ArrowStreamProcessor({ timeout: this.config.timeout || executeStatementDefaults.timeout, - maxConcurrentDownloads: - ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS, retries: ArrowStreamProcessor.DEFAULT_RETRIES, }); } @@ -287,7 +309,11 @@ export class SQLWarehouseConnector { ); break; case "SUCCEEDED": - result = this._transformDataArray(response); + result = await this._transformDataArray( + response, + workspaceClient, + signal, + ); break; case "FAILED": throw ExecutionError.statementFailed( @@ -816,7 +842,11 @@ export class SQLWarehouseConnector { "poll.duration_ms": elapsedTime, }); span.setStatus({ code: SpanStatusCode.OK }); - return this._transformDataArray(response); + return this._transformDataArray( + response, + workspaceClient, + signal, + ); case "FAILED": throw ExecutionError.statementFailed( status.error?.message, @@ -866,7 +896,11 @@ export class SQLWarehouseConnector { ); } - private _transformDataArray(response: sql.StatementResponse) { + private async _transformDataArray( + response: sql.StatementResponse, + workspaceClient: WorkspaceClient, + signal?: AbortSignal, + ) { if (response.manifest?.format === "ARROW_STREAM") { const result = response.result as | (sql.ResultData & { attachment?: string }) @@ -882,7 +916,7 @@ export class SQLWarehouseConnector { // External links: data fetched separately via statement_id. if (result?.external_links) { - return this.updateWithArrowStatus(response); + return this.updateWithArrowStatus(response, workspaceClient, signal); } // Empty result with a known schema: synthesize a zero-row Arrow IPC @@ -955,22 +989,24 @@ export class SQLWarehouseConnector { } /** - * Validate (but do not decode) a base64 Arrow IPC attachment. - * Some serverless warehouses return inline results as Arrow IPC in - * `result.attachment`. We pass the base64 string through to the client, - * which decodes it into an Arrow Table via the existing ArrowClient - * infrastructure. This keeps the wire contract for ARROW_STREAM - * consistent (client always receives an Arrow Table) and avoids - * decode/re-encode work on the server. + * Validate a base64 Arrow IPC attachment (INLINE ARROW_STREAM) and attach + * the real column names from the result manifest. + * + * Databricks warehouses encode the Arrow schema positionally (`col_0`, …); + * the real, aliased names live only in `manifest.schema.columns`. Rather + * than decode/re-encode the payload server-side, we carry the names on the + * result (`columnNames`) so the route can hand them to the client in a + * response header and the client relabels the decoded Table — the one + * mechanism used for both INLINE and EXTERNAL_LINKS. */ private _validateArrowAttachment( response: sql.StatementResponse, attachment: string, ) { // Cap the size to protect against unbounded inline payloads from - // misbehaving warehouses. 64 MiB is well above the typical inline limit - // (~25 MiB hard cap on the API) but bounds memory if a server returns - // a runaway response. + // misbehaving warehouses. `MAX_INLINE_ATTACHMENT_BYTES` tracks the API's + // ~25 MiB inline hard cap and bounds memory if a server returns a runaway + // response. // // Strip whitespace (rare but legal in base64) and account for trailing // `=` padding so the byte count is exact rather than an upper bound. @@ -986,101 +1022,146 @@ export class SQLWarehouseConnector { `Inline Arrow attachment exceeds maximum size (${decodedSize} > ${MAX_INLINE_ATTACHMENT_BYTES} bytes)`, ); } + + const columnNames = arrowColumnNames(response); + if (columnNames) { + return { + ...response, + result: { + ...(response.result as sql.ResultData & { + attachment?: string; + columnNames?: string[]; + }), + // `statement_id` is a top-level field, not on `ResultData` — carry it + // onto the result (as the EXTERNAL_LINKS path does) so the route can + // advertise it in `X-Appkit-Arrow-Columns-Ref` for wide inline schemas. + statement_id: response.statement_id, + columnNames, + }, + }; + } + return response; } - private updateWithArrowStatus(response: sql.StatementResponse): { - result: { statement_id: string; status: sql.StatementStatus }; - } { + private async updateWithArrowStatus( + response: sql.StatementResponse, + workspaceClient: WorkspaceClient, + signal?: AbortSignal, + ): Promise<{ + result: { + statement_id: string; + status: sql.StatementStatus; + columnNames?: string[]; + external_links?: sql.ExternalLink[]; + }; + }> { + const statementId = response.statement_id as string; return { result: { - statement_id: response.statement_id as string, + statement_id: statementId, status: { state: response.status?.state, error: response.status?.error, } as sql.StatementStatus, + columnNames: arrowColumnNames(response), + // Resolve the pre-signed links for EVERY chunk in the caller's own + // execution context. Streaming these directly (see + // {@link streamExternalLinks}) avoids a second `getStatement` under + // the ambient service-principal context — which would fail for + // `.obo.sql` (user-owned) statements. + external_links: await this._resolveAllExternalLinks( + workspaceClient, + statementId, + response, + signal, + ), }, }; } - async getArrowData( + /** + * Resolve pre-signed links for EVERY chunk of an EXTERNAL_LINKS result. + * + * The execute/getStatement response carries only the first chunk's links + * (each link, except the last, exposes `next_chunk_index`); the remaining + * chunks are fetched with `getStatementResultChunkN`. Runs in the caller's + * identity context (user creds for `.obo.sql`), so there is no cross-identity + * fetch. Only the tiny link metadata is resolved eagerly — the bytes still + * stream one chunk at a time downstream. Without this a multi-chunk result + * would be silently truncated to its first chunk. + */ + private async _resolveAllExternalLinks( workspaceClient: WorkspaceClient, - jobId: string, + statementId: string, + response: sql.StatementResponse, signal?: AbortSignal, - ): Promise> { - const startTime = Date.now(); - - return this.telemetry.startActiveSpan( - "arrow.getData", - { - kind: SpanKind.CLIENT, - attributes: { - "db.system": "databricks", - "arrow.job_id": jobId, - }, - }, - async (span: Span) => { - try { - const response = - await workspaceClient.statementExecution.getStatement( - { statement_id: jobId }, - this._createContext(signal), - ); - - const chunks = response.result?.external_links; - const schema = response.manifest?.schema; - - if (!chunks || !schema) { - throw ExecutionError.missingData("chunks or schema"); - } - - span.setAttribute("arrow.chunk_count", chunks.length); - - const result = await this.arrowProcessor.processChunks( - chunks, - schema, - signal, - ); - - span.setAttribute("arrow.data_size_bytes", result.data.length); - span.setStatus({ code: SpanStatusCode.OK }); - - const duration = Date.now() - startTime; - this.telemetryMetrics.queryDuration.record(duration, { - operation: "arrow.getData", - status: "success", - }); - - logger.event()?.setContext("sql-warehouse", { - arrow_data_size_bytes: result.data.length, - arrow_job_id: jobId, - }); - - return result; - } catch (error) { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: error instanceof Error ? error.message : "Unknown error", - }); - span.recordException(error as Error); + ): Promise { + const first = response.result?.external_links; + if (!first || first.length === 0) return first; + + const links: sql.ExternalLink[] = [...first]; + // Bound the follow loop so a warehouse returning a cyclic/never-ending + // `next_chunk_index` can't spin forever. The manifest's chunk count is the + // natural bound; fall back to a generous safety cap if it's absent (real + // results still terminate earlier when `next_chunk_index` becomes null) so + // a missing count doesn't silently truncate a genuine multi-chunk result. + const maxFetches = + response.manifest?.total_chunk_count ?? MAX_EXTERNAL_CHUNK_FOLLOWS; + let next = this._nextChunkIndex(first); + for (let fetches = 0; next != null && fetches < maxFetches; fetches++) { + if (signal?.aborted) throw ExecutionError.canceled(); + const chunk = + await workspaceClient.statementExecution.getStatementResultChunkN( + { statement_id: statementId, chunk_index: next }, + this._createContext(signal), + ); + const chunkLinks = chunk.external_links ?? []; + if (chunkLinks.length === 0) break; + links.push(...chunkLinks); + next = this._nextChunkIndex(chunkLinks); + } + return links; + } - const duration = Date.now() - startTime; - this.telemetryMetrics.queryDuration.record(duration, { - operation: "arrow.getData", - status: "error", - }); + /** The `next_chunk_index` advertised by a chunk's links, if any. */ + private _nextChunkIndex(links: sql.ExternalLink[]): number | undefined { + for (const link of links) { + if (link.next_chunk_index != null) return link.next_chunk_index; + } + return undefined; + } - logger.error("Failed Arrow job: %s %O", jobId, error); + /** + * Stream already-resolved EXTERNAL_LINKS chunks as raw Arrow IPC bytes, one + * chunk in memory at a time. Takes the `external_links` the caller already + * has from `executeStatement` (see {@link updateWithArrowStatus}), so there + * is no extra `getStatement` round-trip and no execution-context mismatch — + * the pre-signed URLs need no auth to download. + */ + streamExternalLinks( + chunks: sql.ExternalLink[], + signal?: AbortSignal, + ): AsyncGenerator { + return this.arrowProcessor.streamChunks(chunks, signal); + } - if (error instanceof AppKitError) { - throw error; - } - throw ExecutionError.statementFailed( - error instanceof Error ? error.message : String(error), - ); - } - }, + /** + * Fetch just the real column names for a completed statement from its result + * manifest — no result data. Backs the `/columns/:statementId` fallback the + * client uses when a very wide schema's names exceed the response-header + * size limit. Stateless: re-derives from the warehouse, no server cache. + */ + async getColumnNames( + workspaceClient: WorkspaceClient, + jobId: string, + signal?: AbortSignal, + ): Promise { + const response = await workspaceClient.statementExecution.getStatement( + { statement_id: jobId }, + this._createContext(signal), ); + return arrowColumnNames(response); } // create context for cancellation token diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts index c6780f3f..dc964702 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts @@ -37,6 +37,18 @@ function createConnector() { return new SQLWarehouseConnector({ timeout: 30000 }); } +// `_transformDataArray` is async — it paginates multi-chunk EXTERNAL_LINKS +// results. The workspace client is only touched when following +// `next_chunk_index`, so a bare stub suffices for the inline / JSON / +// single-chunk cases; the multi-chunk tests pass a real mock. +function transform( + connector: SQLWarehouseConnector, + response: sql.StatementResponse, + workspaceClient: unknown = {}, +) { + return (connector as any)._transformDataArray(response, workspaceClient); +} + // Real base64 Arrow IPC from a serverless warehouse returning // `SELECT 1 AS test_col, 2 AS test_col2` with INLINE + ARROW_STREAM. // Contains schema (two INT columns) + one record batch with values [1, 2]. @@ -45,7 +57,7 @@ const REAL_ARROW_ATTACHMENT = describe("SQLWarehouseConnector._transformDataArray", () => { describe("classic warehouse (JSON_ARRAY + INLINE)", () => { - test("transforms data_array rows into named objects", () => { + test("transforms data_array rows into named objects", async () => { const connector = createConnector(); // Real response shape from classic warehouse: INLINE + JSON_ARRAY const response = { @@ -78,12 +90,12 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); expect(result.result.data).toEqual([{ test_col: "1", test_col2: "2" }]); expect(result.result.data_array).toBeUndefined(); }); - test("parses JSON strings in STRING columns", () => { + test("parses JSON strings in STRING columns", async () => { const connector = createConnector(); const response = { statement_id: "stmt-1", @@ -102,13 +114,13 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); expect(result.result.data[0].metadata).toEqual({ key: "value" }); }); }); describe("classic warehouse (EXTERNAL_LINKS + ARROW_STREAM)", () => { - test("returns statement_id for external links fetch", () => { + test("returns statement_id for external links fetch", async () => { const connector = createConnector(); // Real response shape from classic warehouse: EXTERNAL_LINKS + ARROW_STREAM const response = { @@ -133,14 +145,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); expect(result.result.statement_id).toBe("stmt-1"); expect(result.result.data).toBeUndefined(); }); }); describe("serverless warehouse (INLINE + ARROW_STREAM with attachment)", () => { - test("passes attachment through unchanged for client-side decoding", () => { + test("passes attachment through unchanged for client-side decoding", async () => { const connector = createConnector(); // Real response shape from serverless warehouse: INLINE + ARROW_STREAM // Data arrives in result.attachment as base64-encoded Arrow IPC, not data_array. @@ -179,14 +191,18 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); expect(result.result.data).toBeUndefined(); // Preserves other result fields expect(result.result.row_count).toBe(1); + // `statement_id` is projected onto the result so the route can advertise + // an `X-Appkit-Arrow-Columns-Ref` for wide inline schemas (it lives on + // the top-level response, not on `ResultData`). + expect(result.result.statement_id).toBe("00000001-test-stmt"); }); - test("preserves manifest and status alongside attachment", () => { + test("preserves manifest and status alongside attachment", async () => { const connector = createConnector(); const response = { statement_id: "00000001-test-stmt", @@ -207,14 +223,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); // Manifest, statement_id, and attachment are all preserved expect(result.manifest.format).toBe("ARROW_STREAM"); expect(result.statement_id).toBe("00000001-test-stmt"); expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); }); - test("synthesizes an empty Arrow IPC attachment for empty results so the client always gets a Table", () => { + test("synthesizes an empty Arrow IPC attachment for empty results so the client always gets a Table", async () => { const connector = createConnector(); // Empty result: no attachment, no data_array, no external_links — but // the manifest still describes the schema. The connector should fill in @@ -240,7 +256,7 @@ describe("SQLWarehouseConnector._transformDataArray", () => { result: {}, } as unknown as sql.StatementResponse; - const transformed = (connector as any)._transformDataArray(response); + const transformed = await transform(connector, response); const attachment: string = transformed.result.attachment; expect(typeof attachment).toBe("string"); expect(attachment.length).toBeGreaterThan(0); @@ -255,7 +271,7 @@ describe("SQLWarehouseConnector._transformDataArray", () => { ]); }); - test("does NOT synthesize an attachment when external_links are present", () => { + test("does NOT synthesize an attachment when external_links are present", async () => { const connector = createConnector(); const response = { statement_id: "stmt-ext", @@ -271,13 +287,13 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const transformed = (connector as any)._transformDataArray(response); + const transformed = await transform(connector, response); // External-links path returns the statement_id projection — no attachment. expect(transformed.result.attachment).toBeUndefined(); expect(transformed.result.statement_id).toBe("stmt-ext"); }); - test("does NOT synthesize an attachment when schema is missing", () => { + test("does NOT synthesize an attachment when schema is missing", async () => { const connector = createConnector(); const response = { statement_id: "stmt-no-schema", @@ -286,12 +302,12 @@ describe("SQLWarehouseConnector._transformDataArray", () => { result: {}, } as unknown as sql.StatementResponse; - const transformed = (connector as any)._transformDataArray(response); + const transformed = await transform(connector, response); // Without a schema we cannot build a Table — pass through unchanged. expect(transformed.result?.attachment).toBeUndefined(); }); - test("rejects oversized attachments to bound memory", () => { + test("rejects oversized attachments to bound memory", async () => { const connector = createConnector(); // 25 MiB decoded cap (Databricks API hard cap on INLINE) → 36 MiB of // base64 chars decodes to ~27 MiB, comfortably above the limit. @@ -303,14 +319,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { result: { attachment: oversized }, } as unknown as sql.StatementResponse; - expect(() => (connector as any)._transformDataArray(response)).toThrow( + await expect(transform(connector, response)).rejects.toThrow( /exceeds maximum size/, ); }); }); describe("ARROW_STREAM with data_array (hypothetical inline variant)", () => { - test("transforms data_array like JSON_ARRAY path", () => { + test("transforms data_array like JSON_ARRAY path", async () => { const connector = createConnector(); const response = { statement_id: "stmt-1", @@ -332,7 +348,7 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); expect(result.result.data).toEqual([ { id: "1", value: "hello" }, { id: "2", value: "world" }, @@ -341,7 +357,7 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }); describe("edge cases", () => { - test("returns response unchanged when no data_array, attachment, or schema", () => { + test("returns response unchanged when no data_array, attachment, or schema", async () => { const connector = createConnector(); const response = { statement_id: "stmt-1", @@ -350,11 +366,11 @@ describe("SQLWarehouseConnector._transformDataArray", () => { result: {}, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); expect(result).toBe(response); }); - test("attachment takes priority over data_array when both present", () => { + test("attachment takes priority over data_array when both present", async () => { const connector = createConnector(); const response = { statement_id: "stmt-1", @@ -374,10 +390,105 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }, } as unknown as sql.StatementResponse; - const result = (connector as any)._transformDataArray(response); + const result = await transform(connector, response); // Should pass attachment through (client decodes), not transform data_array expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); expect(result.result.data).toBeUndefined(); }); }); + + describe("multi-chunk EXTERNAL_LINKS pagination", () => { + function multiChunkResponse(totalChunks: number): sql.StatementResponse { + return { + statement_id: "stmt-multi", + status: { state: "SUCCEEDED" }, + manifest: { + format: "ARROW_STREAM", + total_chunk_count: totalChunks, + schema: { columns: [{ name: "x", type_name: "INT" }] }, + }, + result: { + external_links: [ + { + chunk_index: 0, + external_link: "https://example.com/chunk0", + next_chunk_index: 1, + }, + ], + }, + } as unknown as sql.StatementResponse; + } + + test("follows next_chunk_index to resolve every chunk's links", async () => { + const connector = createConnector(); + const getStatementResultChunkN = vi + .fn() + .mockResolvedValueOnce({ + external_links: [ + { + chunk_index: 1, + external_link: "https://example.com/chunk1", + next_chunk_index: 2, + }, + ], + }) + .mockResolvedValueOnce({ + external_links: [ + { chunk_index: 2, external_link: "https://example.com/chunk2" }, + ], + }); + const workspaceClient = { + statementExecution: { getStatementResultChunkN }, + }; + + const result = await transform( + connector, + multiChunkResponse(3), + workspaceClient, + ); + + expect(getStatementResultChunkN).toHaveBeenCalledTimes(2); + expect(getStatementResultChunkN).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ statement_id: "stmt-multi", chunk_index: 1 }), + expect.anything(), + ); + expect( + result.result.external_links.map( + (l: sql.ExternalLink) => l.external_link, + ), + ).toEqual([ + "https://example.com/chunk0", + "https://example.com/chunk1", + "https://example.com/chunk2", + ]); + }); + + test("is bounded by total_chunk_count when next_chunk_index never terminates", async () => { + const connector = createConnector(); + // Misbehaving warehouse: always advertises another chunk. + const getStatementResultChunkN = vi.fn().mockResolvedValue({ + external_links: [ + { + chunk_index: 1, + external_link: "https://example.com/loop", + next_chunk_index: 99, + }, + ], + }); + const workspaceClient = { + statementExecution: { getStatementResultChunkN }, + }; + + const result = await transform( + connector, + multiChunkResponse(2), + workspaceClient, + ); + + // Terminates (no hang) — capped at total_chunk_count fetches. + expect(getStatementResultChunkN).toHaveBeenCalledTimes(2); + expect(result.result.external_links.length).toBeGreaterThan(0); + }); + }); }); diff --git a/packages/appkit/src/errors/execution.ts b/packages/appkit/src/errors/execution.ts index 3c9df7fe..1c9acf91 100644 --- a/packages/appkit/src/errors/execution.ts +++ b/packages/appkit/src/errors/execution.ts @@ -40,7 +40,7 @@ export class ExecutionError extends AppKitError { * Execution errors default to a generic message — the raw warehouse / * SDK text in `.message` often includes statement fragments, internal * paths, and correlation IDs. UI code should branch on `errorCode` - * (`INLINE_ARROW_STASH_EXHAUSTED`, `NOT_IMPLEMENTED`, etc.) and not on + * (`RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, `NOT_IMPLEMENTED`, etc.) and not on * the human string. */ override get clientMessage(): string { @@ -108,23 +108,4 @@ export class ExecutionError extends AppKitError { context: { dataType }, }); } - - /** - * Create an execution error for the inline Arrow stash being unable to - * accept a payload (both regular and overflow pools at cap). - * - * The route deliberately surfaces this rather than silently re-running - * the statement on EXTERNAL_LINKS — a second execution can be billed - * and return divergent results for non-deterministic SQL. Operators - * should tune stash capacity or back off load when this fires. - */ - static stashExhausted(): ExecutionError { - return new ExecutionError( - "Inline Arrow stash exhausted; retry shortly or increase stash capacity", - { - errorCode: "INLINE_ARROW_STASH_EXHAUSTED", - clientMessage: "Server is at capacity, please retry", - }, - ); - } } diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 0c49f29c..1b5aa1d5 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1,11 +1,8 @@ -import type { WorkspaceClient } from "@databricks/sdk-experimental"; -import { type DataType, Type, tableFromIPC } from "apache-arrow"; import type express from "express"; import { type AgentToolDefinition, type AnalyticsSseMessage, type IAppRouter, - makeArrowMessage, makeResultMessage, type PluginExecuteConfig, type SQLTypeMarker, @@ -30,13 +27,11 @@ import { AppKitError, ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; -import type { Counter, Histogram } from "../../telemetry"; import { queryDefaults } from "./defaults"; -import { InlineArrowStash } from "./inline-arrow-stash"; import manifest from "./manifest.json"; import { QueryProcessor } from "./query"; +import { deliverArrowBytes, deliverJsonResult } from "./result-delivery"; import { - type AnalyticsFormat, type AnalyticsQueryResponse, type AnalyticsStreamMessage, type IAnalyticsConfig, @@ -106,33 +101,6 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { private SQLClient: SQLWarehouseConnector; private queryProcessor: QueryProcessor; - /** - * Server-side stash for inline Arrow IPC payloads. - * - * INLINE ARROW_STREAM responses do not ride the SSE control channel — - * the route puts the decoded bytes here and emits an `arrow` SSE - * message with a synthetic `inline-` id, and the client fetches - * the bytes through the existing `/arrow-result/:jobId` endpoint with - * a real binary content-type. - */ - protected inlineArrowStash: InlineArrowStash = new InlineArrowStash(); - - /** - * Stash telemetry — created lazily on first stash operation because - * `this.telemetry` may not be bound at constructor time (see - * `Plugin.attachContext`). Cached via `_stashMetrics` after the first - * call; subsequent puts hit the cached counter/histogram instances. - * - * Counter labels: `result` ∈ {"regular", "overflow", "exhausted"} — - * one counter with a label is friendlier to dashboards than three - * separate metrics, and aggregates trivially in Prometheus / OTel. - */ - private _stashMetrics?: { - putCount: Counter; - putBytes: Histogram; - }; - private _stashMetricsAttempted = false; - constructor(config: IAnalyticsConfig) { super(config); this.config = config; @@ -145,19 +113,6 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } injectRoutes(router: IAppRouter) { - // Arrow data downloads always run as service principal and bypass the - // interceptor chain (execute/executeStream). The original query execution - // handles OBO via executeStream(); this endpoint fetches pre-computed - // results by job ID. - this.route(router, { - name: "arrow", - method: "get", - path: "/arrow-result/:jobId", - handler: async (req: express.Request, res: express.Response) => { - await this._handleArrowRoute(req, res); - }, - }); - this.route(router, { name: "query", method: "post", @@ -166,156 +121,81 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { await this._handleQueryRoute(req, res); }, }); + + // Column-names fallback for very wide Arrow schemas whose names don't fit + // in the `X-Appkit-Arrow-Columns` response header (see + // `_setArrowColumnsHeader`). The client hits this with the statement id + // from `X-Appkit-Arrow-Columns-Ref`. + this.route(router, { + name: "arrow-columns", + method: "get", + path: "/columns/:statementId", + handler: async (req: express.Request, res: express.Response) => { + await this._handleColumnsRoute(req, res); + }, + }); } /** - * Handle Arrow data download requests. - * - * Two id shapes are supported: - * - `inline-`: bytes were stashed server-side by the query route. - * Drain the stash, serve directly with the canonical Arrow content - * type. No warehouse round-trip. - * - any other id: a warehouse-issued statement id. Fetch the Arrow - * stream from the warehouse via the SDK; serve the bytes. - * - * When called via asUser(req), uses the user's Databricks credentials - * for the warehouse path. The inline path is user-scoped at the stash - * layer instead. + * Column-names fallback endpoint. Re-derives the real column names from the + * statement's result manifest (stateless — no server cache), for the client + * to relabel a positional Arrow schema when the names were too large for the + * response header. */ - async _handleArrowRoute( + async _handleColumnsRoute( req: express.Request, res: express.Response, ): Promise { - const { jobId } = req.params; - const event = logger.event(req); - event?.setComponent("analytics", "getArrowData").setContext("analytics", { - job_id: jobId, - plugin: this.name, - }); - - if (jobId.startsWith("inline-")) { - const userKey = this._stashUserKey(req); - const bytes = this.inlineArrowStash.take(jobId, userKey); - if (!bytes) { - // Already drained, expired, or never belonged to this user. 410 - // distinguishes this from "warehouse statement id not found" (404) - // so the client can surface a useful error. - logger.debug("Inline Arrow stash miss for jobId=%s", jobId); - res.status(410).json({ - error: "Inline Arrow result expired or unknown", - plugin: this.name, - }); - return; - } - logger.debug( - "Serving inline Arrow buffer: %d bytes for jobId=%s", - bytes.length, - jobId, - ); - res.setHeader("Content-Type", "application/vnd.apache.arrow.stream"); - res.setHeader("Content-Length", bytes.length.toString()); - // Inline payloads are single-use and short-lived; no public caching. + const { statementId } = req.params; + const columns = await this._resolveColumnNames(req, statementId); + if (columns && columns.length > 0) { res.setHeader("Cache-Control", "no-store"); - res.send(Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength)); + res.json({ columns }); return; } - - try { - const workspaceClient = getWorkspaceClient(); - logger.debug("Processing Arrow job request for jobId=%s", jobId); - - const result = await this.getArrowData(workspaceClient, jobId); - - res.setHeader("Content-Type", "application/octet-stream"); - res.setHeader("Content-Length", result.data.length.toString()); - res.setHeader("Cache-Control", "public, max-age=3600"); - - logger.debug( - "Sending Arrow buffer: %d bytes for job %s", - result.data.length, - jobId, - ); - res.send(Buffer.from(result.data)); - } catch (error) { - logger.error("Arrow job error: %O", error); - // Do not echo upstream / SDK error text to the client — it can - // include statement fragments, internal object names, and - // correlation IDs. The full detail stays in the server log above. - const errorCode = - error instanceof ExecutionError ? error.errorCode : undefined; - res.status(404).json({ - error: "Arrow result unavailable", - errorCode, - plugin: this.name, - }); - } + res.status(404).json({ + error: "Column names unavailable", + plugin: this.name, + }); } /** - * Lazy accessor for stash telemetry instruments. Returns `undefined` - * when telemetry isn't bound (factory-constructed plugins before - * `attachContext`, no-op test paths) — callers should branch on the - * return value rather than relying on it being present. - * - * Records once per call site: - * - `analytics.inline_arrow_stash.put.count{result}` — counter with - * label "regular" | "overflow" | "exhausted" - * - `analytics.inline_arrow_stash.put.bytes` — histogram of accepted - * payload sizes (label same as counter; exhausted puts record 0) + * Resolve a statement's real column names, trying the user's identity first + * (required for `.obo.sql` statements, which the service principal cannot + * `getStatement`) then falling back to the service principal (for + * SP-executed statements). Returns undefined if neither identity can read it, + * so the client falls back to the raw positional Arrow schema names. */ - private _getStashMetrics() { - if (this._stashMetricsAttempted) return this._stashMetrics; - this._stashMetricsAttempted = true; - if (!this.telemetry) return undefined; - const meter = this.telemetry.getMeter(); - this._stashMetrics = { - putCount: meter.createCounter("analytics.inline_arrow_stash.put.count", { - description: - "Count of inline-Arrow stash put attempts, labeled by pool outcome (regular | overflow | exhausted).", - unit: "1", - }), - putBytes: meter.createHistogram( - "analytics.inline_arrow_stash.put.bytes", - { - description: - "Sizes of accepted inline-Arrow stash payloads. Aggregate by `result` label for utilization analysis.", - unit: "By", - }, - ), - }; - return this._stashMetrics; - } - - private _recordStashOutcome( - result: "regular" | "overflow" | "exhausted", - bytes: number, - ): void { - const m = this._getStashMetrics(); - if (!m) return; - m.putCount.add(1, { result }); - m.putBytes.record(result === "exhausted" ? 0 : bytes, { result }); + private async _resolveColumnNames( + req: express.Request, + statementId: string, + ): Promise { + const attempts: Array<() => Promise> = [ + () => this.asUser(req)._getColumnNames(statementId), + () => this._getColumnNames(statementId), + ]; + for (const attempt of attempts) { + try { + const columns = await attempt(); + if (columns && columns.length > 0) return columns; + } catch (error) { + logger.debug( + "Arrow column-names lookup attempt failed for %s: %O", + statementId, + error, + ); + } + } + return undefined; } /** - * Stash key used at put-time (in `_handleQueryRoute`) and take-time - * (in `_handleArrowRoute`). Centralized so the two sides cannot drift. - * - * Returns the user id when an `x-forwarded-user` header is present, - * otherwise `"global"` for service-principal contexts (no user header). - * Both queries from the same request resolve to the same key, so the - * subsequent /arrow-result fetch reliably hits the entry stashed - * during the SSE query. - * - * `resolveUserId` throws when no header is present — catch and degrade - * to "global" rather than letting that failure mode bubble through the - * route handler. + * Fetch column names in the current execution context. Proxied by `asUser`, + * so `getWorkspaceClient()` resolves to the user's client when invoked via + * `this.asUser(req)` and the service principal's otherwise. */ - protected _stashUserKey(req: express.Request): string { - try { - return this.resolveUserId(req) || "global"; - } catch { - return "global"; - } + async _getColumnNames(statementId: string): Promise { + return this.SQLClient.getColumnNames(getWorkspaceClient(), statementId); } /** @@ -373,31 +253,25 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { const { query, isAsUser } = queryResult; + // ARROW_STREAM streams the raw Arrow IPC bytes back as the HTTP response + // body — no SSE, no server-side stash, no second /arrow-result request. + // INLINE attachments are piped straight through (the bytes are already in + // hand from executeStatement); a warehouse that refuses INLINE falls back + // to EXTERNAL_LINKS and streams those chunks. JSON keeps the SSE path + // below (it carries warehouse-readiness progress + cached rows). + if (format === "ARROW_STREAM") { + await this._handleArrowStreamQuery(req, res, query, isAsUser, parameters); + return; + } + // get execution context - user-scoped if .obo.sql, otherwise service principal const executor = isAsUser ? this.asUser(req) : this; const executorKey = isAsUser ? this.resolveUserId(req) : "global"; - // Stash key is always per-request user (never "global"), independent - // of the executor's cache scope. Inline Arrow payloads are single-use - // and short-lived — there is no benefit to sharing them across users, - // and per-user scoping is defense in depth on top of unguessable ids. - const stashUserKey = this._stashUserKey(req); const hashedQuery = this.queryProcessor.hashQuery(query); - // ARROW_STREAM responses reference ephemeral resources that cannot be - // safely replayed from cache: - // - EXTERNAL_LINKS pre-signed URLs expire ~15 min after issue, and - // the warehouse rotates them per execution. - // - INLINE responses point at a synthetic `inline-` job id - // backed by `InlineArrowStash`, which drains on the first - // /arrow-result fetch. A cache hit would replay an id whose bytes - // are already gone and reliably 410 the client. - // So we bypass cache for ARROW_STREAM and let every request execute - // a fresh statement. JSON_ARRAY responses still cache normally. - const cacheTtl = format === "ARROW_STREAM" ? 0 : queryDefaults.cache?.ttl; const cacheConfig = { ...queryDefaults.cache, - ttl: cacheTtl, cacheKey: [ "analytics:query", query_key, @@ -476,15 +350,15 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { try { const processedParams = await self.queryProcessor.processQueryParams(query, parameters); - // Disposition/format fallback (inline-arrow stash, JSON_ARRAY↔ - // ARROW_STREAM retries) lives in `_executeWithFormatFallback`, - // which returns the SSE message the client should receive. - return await self._executeWithFormatFallback( + // JSON_ARRAY path: tries INLINE + JSON_ARRAY and, if the + // warehouse only accepts ARROW_STREAM for INLINE, retries as + // ARROW_STREAM and decodes server-side — returning the SSE + // `result` message with plain rows. (ARROW_STREAM requests are + // handled earlier via `_handleArrowStreamQuery`.) + return await self._executeJsonArrayPath( executor, query, processedParams, - format, - stashUserKey, sig, ); } catch (err) { @@ -511,7 +385,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { throw err; } // Re-throw the original error so its structured `errorCode` (e.g. - // INLINE_ARROW_STASH_EXHAUSTED) and sanitized `clientMessage` + // RESULT_TOO_LARGE_FOR_JSON_FALLBACK) and sanitized `clientMessage` // survive to the SSE error payload. Fall back to a generic // statement failure only if the original wasn't an AppKitError. if (originalError instanceof AppKitError) { @@ -531,215 +405,204 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } /** - * Execute a query with automatic disposition/format fallback. - * - * - **JSON_ARRAY** first tries `INLINE + JSON_ARRAY`. If the warehouse - * only supports `ARROW_STREAM` for `INLINE` (some serverless variants), - * retries as `INLINE + ARROW_STREAM`, decodes the Arrow IPC attachment - * server-side, and returns plain row objects — the caller's - * `JSON_ARRAY` contract is preserved. - * - **ARROW_STREAM** first tries `INLINE + ARROW_STREAM`. If the - * warehouse refuses (most classic + some serverless variants), falls - * back to `EXTERNAL_LINKS + ARROW_STREAM`. When the regular stash is - * full, the already-decoded bytes spill into the stash's overflow - * pool — never thrown away to trigger a second warehouse round-trip. - * - * INLINE Arrow attachments under the ARROW_STREAM path are decoded once - * and put on the plugin's `inlineArrowStash`; the SSE message carries the - * synthetic stash id so the client fetches the bytes out-of-band via - * `/arrow-result/`. + * JSON_ARRAY SSE path. Delegates the disposition/format fallback to + * {@link deliverJsonResult} (INLINE JSON_ARRAY → on `needs-arrow-inline`, + * INLINE ARROW_STREAM decoded to rows) and wraps the rows in a `result` + * message. External links are never used for the JSON fallback. */ - private async _executeWithFormatFallback( + private async _executeJsonArrayPath( executor: AnalyticsPlugin, query: string, processedParams: | Record | undefined, - requestedFormat: AnalyticsFormat, - stashUserKey: string, signal?: AbortSignal, ): Promise { - return requestedFormat === "JSON_ARRAY" - ? this._executeJsonArrayPath(executor, query, processedParams, signal) - : this._executeArrowStreamPath( - executor, - query, - processedParams, - stashUserKey, - signal, - ); + const result = await deliverJsonResult( + executor, + query, + processedParams, + signal, + ); + return makeResultMessage(result.data, { + status: result.status, + statement_id: result.statement_id, + }); } /** - * JSON_ARRAY path. Tries `INLINE + JSON_ARRAY` first; on a structured - * `needs-arrow` rejection (warehouse only accepts ARROW_STREAM under - * INLINE), retries as `INLINE + ARROW_STREAM` and decodes the Arrow - * IPC attachment server-side so the caller's JSON_ARRAY contract is - * preserved. + * Attach the real column names so the client can relabel the positional + * Arrow schema (Databricks encodes ARROW_STREAM columns as col_0, …). * - * Failure modes surfaced to the caller: - * - Unrelated warehouse errors (anything not classified as - * `needs-arrow`) propagate untouched. - * - The retry can throw `ExecutionError.missingData("ARROW_STREAM - * attachment")` if the warehouse returned no attachment. - * - The decode itself can throw `RESULT_TOO_LARGE_FOR_JSON_FALLBACK` - * when the attachment exceeds the row or byte cap. + * Small schemas ride `X-Appkit-Arrow-Columns` directly. A very wide schema + * whose URL-encoded names would blow the HTTP header size limit instead + * advertises the statement id in `X-Appkit-Arrow-Columns-Ref`, and the + * client fetches the names from `GET /columns/:statementId`. */ - private async _executeJsonArrayPath( - executor: AnalyticsPlugin, - query: string, - processedParams: - | Record - | undefined, - signal?: AbortSignal, - ): Promise { - try { - const result = await executor.query( - query, - processedParams, - { disposition: "INLINE", format: "JSON_ARRAY" }, - signal, - ); - return makeResultMessage(result?.data, { - status: result?.status, - statement_id: result?.statement_id, - }); - } catch (err: unknown) { - if (signal?.aborted) throw err; - if (_classifyInlineRejection(err) !== "needs-arrow") throw err; + private _setArrowColumnsHeader( + res: express.Response, + columnsRef: { columnNames?: string[]; statementId?: string }, + ): void { + const names = columnsRef.columnNames; + if (!names || names.length === 0) return; - const msg = err instanceof Error ? err.message : String(err); + const encoded = encodeURIComponent(JSON.stringify(names)); + if (encoded.length <= MAX_ARROW_COLUMNS_HEADER_BYTES) { + res.setHeader("X-Appkit-Arrow-Columns", encoded); + return; + } + if (columnsRef.statementId) { + res.setHeader("X-Appkit-Arrow-Columns-Ref", columnsRef.statementId); + } else { logger.warn( - "JSON_ARRAY INLINE rejected by warehouse, retrying as ARROW_STREAM INLINE and decoding server-side: %s", - msg, + "Arrow column names exceed the header limit and no statement id is available for the fallback endpoint; client will fall back to the raw schema names", ); } - - // Retry as ARROW_STREAM + INLINE so the warehouse will accept the - // request, then decode the Arrow IPC attachment to plain row - // objects so the caller still gets JSON_ARRAY-shaped data. - const arrowResult = await executor.query( - query, - processedParams, - { disposition: "INLINE", format: "ARROW_STREAM" }, - signal, - ); - if (!arrowResult?.attachment) { - throw ExecutionError.missingData("ARROW_STREAM attachment"); - } - const rows = decodeArrowAttachmentToRows(arrowResult.attachment); - return makeResultMessage(rows, { - status: arrowResult.status, - statement_id: arrowResult.statement_id, - }); } /** - * ARROW_STREAM path. Tries `INLINE + ARROW_STREAM` first; on a - * structured `needs-json` rejection (warehouse refuses ARROW_STREAM - * under INLINE), falls back to `EXTERNAL_LINKS + ARROW_STREAM`. When - * INLINE succeeds with an Arrow attachment, the decoded bytes go - * through `inlineArrowStash` and the client fetches them via - * `/arrow-result/inline-` — never the SSE channel. + * ARROW_STREAM query handler: stream the raw Arrow IPC bytes back as the + * HTTP response body — no SSE, no server-side stash, no second + * `/arrow-result` request. * - * Failure modes surfaced to the caller: - * - Unrelated warehouse errors (not classified as `needs-json`) - * propagate untouched. - * - `ExecutionError.canceled()` if the client disconnects between - * the INLINE response and the stash put. - * - `ExecutionError.stashExhausted()` if both stash pools are full — - * we never silently re-run on EXTERNAL_LINKS because that would - * double-bill the warehouse and risk a divergent result for - * non-deterministic SQL. + * The first chunk is pulled before headers are sent so a failure still + * yields a clean JSON error; once bytes are in flight a mid-stream failure + * can only abort the socket. Warehouse readiness is awaited (no SSE + * progress on this path) — a no-op for a warm warehouse, a blocking wait + * on a cold start. Runs under the user's context for `.obo.sql` queries. */ - private async _executeArrowStreamPath( - executor: AnalyticsPlugin, + private async _handleArrowStreamQuery( + req: express.Request, + res: express.Response, query: string, - processedParams: - | Record - | undefined, - stashUserKey: string, - signal?: AbortSignal, - ): Promise { + isAsUser: boolean, + parameters: IAnalyticsQueryRequest["parameters"], + ): Promise { + const executor = isAsUser ? this.asUser(req) : this; + const abortController = new AbortController(); + const onClose = () => abortController.abort(); + res.on("close", onClose); + const signal = abortController.signal; + + // Fail-fast: bound the wait for the first byte (warehouse readiness + + // execute + first chunk) so a stuck/overloaded warehouse returns a clear + // 503 instead of hanging until the client gives up. Cleared once the + // first chunk arrives — a legitimately long stream is never interrupted. + const firstByteTimeoutMs = + this.config.arrowFirstByteTimeoutMs ?? + DEFAULT_ARROW_FIRST_BYTE_TIMEOUT_MS; + let timedOut = false; + const failFast = setTimeout(() => { + timedOut = true; + abortController.abort(); + }, firstByteTimeoutMs); + try { - const result = await executor.query( + // Run warehouse readiness in the SAME identity context as the query: + // for `.obo.sql`, `executor` is the `asUser(req)` proxy, so + // `getWorkspaceClient()` inside resolves to the user's client (matching + // the SSE path). Calling it bare here would auto-start the warehouse as + // the service principal even for OBO requests. + await executor._ensureArrowWarehouseReady(signal); + + const processedParams = await this.queryProcessor.processQueryParams( + query, + parameters, + ); + + // Populated by `deliverArrowBytes` from the result manifest before the + // first chunk is yielded, so the header below carries the real names. + const columnsRef: { columnNames?: string[]; statementId?: string } = {}; + const bytes = deliverArrowBytes( + executor, + this.SQLClient, query, processedParams, - { disposition: "INLINE", format: "ARROW_STREAM" }, + columnsRef, signal, ); - if (result?.attachment) { - return this._stashAndEmitInline(result, stashUserKey, signal); - } - // INLINE succeeded but the warehouse returned no Arrow attachment - // (rare: inline `data_array` under ARROW_STREAM). An ARROW_STREAM - // caller's client decodes the `/arrow-result` bytes into a - // TypedArrowTable, so emitting a `result` row message here would - // break that contract. Fall through to the EXTERNAL_LINKS path - // below to obtain a fetchable Arrow statement instead. - logger.warn( - "ARROW_STREAM INLINE returned no attachment; falling back to EXTERNAL_LINKS", - ); - } catch (err: unknown) { - // If the request was aborted, do not retry — the signal is dead and - // a second statement would be billed but never read. - if (signal?.aborted) throw err; - if (_classifyInlineRejection(err) !== "needs-json") throw err; + const first = await bytes.next(); + // First byte in hand — stop the fail-fast clock. + clearTimeout(failFast); - const msg = err instanceof Error ? err.message : String(err); - logger.warn( - "ARROW_STREAM INLINE rejected by warehouse, falling back to EXTERNAL_LINKS: %s", - msg, - ); - } + res.setHeader("Content-Type", "application/vnd.apache.arrow.stream"); + res.setHeader("Cache-Control", "no-store"); + this._setArrowColumnsHeader(res, columnsRef); - const result = await executor.query( - query, - processedParams, - { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, - signal, - ); - return makeArrowMessage(result.statement_id, { status: result.status }); + if (!first.done) { + await writeChunk(res, first.value); + for await (const buf of bytes) { + await writeChunk(res, buf); + } + } + res.end(); + } catch (error) { + clearTimeout(failFast); + if (res.headersSent) { + logger.error("Arrow query stream failed mid-flight: %O", error); + res.destroy(error instanceof Error ? error : new Error(String(error))); + return; + } + if (timedOut) { + logger.warn( + "Arrow query timed out before first byte after %dms", + firstByteTimeoutMs, + ); + res.status(503).json({ + error: + "The SQL warehouse is starting or overloaded and did not respond in time. Please retry.", + errorCode: "WAREHOUSE_UNAVAILABLE", + plugin: this.name, + }); + return; + } + if (signal.aborted) { + res.end(); + return; + } + logger.error("Arrow query error: %O", error); + // Do not echo upstream / SDK error text — it can include statement + // fragments and correlation ids. Keep the structured code so the + // client can branch (e.g. RESULT_TOO_LARGE_FOR_JSON_FALLBACK, + // ARROW_DELIVERY_UNSUPPORTED). + const errorCode = + error instanceof ExecutionError ? error.errorCode : undefined; + res.status(500).json({ + // `clientMessage` is the sanitized, actionable text (e.g. "Re-run with + // JSON_ARRAY" for ARROW_DELIVERY_UNSUPPORTED); it never carries raw + // warehouse/SDK strings. Fall back to a generic message otherwise. + error: + error instanceof AppKitError + ? error.clientMessage + : "Unable to execute query", + errorCode, + plugin: this.name, + }); + } finally { + res.off("close", onClose); + } } /** - * Decode an INLINE Arrow attachment, push it onto the stash, and - * return the `arrow` SSE message that references the synthetic id. - * Pulled out of `_executeArrowStreamPath` so the cap / overflow / - * exhaustion logic lives in one place. + * Await SQL warehouse readiness for the direct-binary Arrow path. There is no + * SSE progress channel here — readiness is simply awaited, bounded by the + * caller's abort signal / fail-fast timeout. Invoked via the request executor + * (`asUser(req)` for `.obo.sql`) so `getWorkspaceClient()` resolves in the + * correct identity context rather than defaulting to the service principal. */ - private _stashAndEmitInline( - result: { attachment: string; status?: unknown }, - stashUserKey: string, - signal?: AbortSignal, - ): AnalyticsSseMessage { - // If the client has already disconnected, the SSE write would be - // dropped anyway — skip the decode + stash so the bytes do not - // linger in memory until TTL eviction. - if (signal?.aborted) { - throw ExecutionError.canceled(); - } - const decoded = Buffer.from(result.attachment, "base64"); - const stashBytes = new Uint8Array( - decoded.buffer, - decoded.byteOffset, - decoded.byteLength, - ); - const putResult = this.inlineArrowStash.put(stashUserKey, stashBytes); - if (putResult === null) { - // Both regular and overflow pools are at cap. Throw with a clear - // signal rather than re-running the same statement on - // EXTERNAL_LINKS: the bytes we already decoded are gone, the - // warehouse has been billed, and a second execution could return - // a divergent result for non-deterministic SQL. - this._recordStashOutcome("exhausted", stashBytes.byteLength); - logger.warn( - "Inline Arrow stash exhausted (regular + overflow); rejecting", - ); - throw ExecutionError.stashExhausted(); - } - this._recordStashOutcome(putResult.pool, stashBytes.byteLength); - return makeArrowMessage(putResult.id, { status: result.status }); + async _ensureArrowWarehouseReady(signal: AbortSignal): Promise { + const workspaceClient = getWorkspaceClient(); + const warehouseId = await getWarehouseId(); + const startupTimeoutMs = + this.config.warehouseStartupTimeoutMs ?? + DEFAULT_WAREHOUSE_STARTUP_TIMEOUT_MS; + const autoStart = this.config.autoStartWarehouse ?? true; + await this.SQLClient.ensureWarehouseRunning(workspaceClient, warehouseId, { + signal, + timeoutMs: startupTimeoutMs, + autoStart, + onStatus: () => {}, + }); } /** @@ -783,17 +646,6 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { return response.result; } - /** - * Get Arrow-formatted data for a completed query job. - */ - protected async getArrowData( - workspaceClient: WorkspaceClient, - jobId: string, - signal?: AbortSignal, - ): Promise> { - return await this.SQLClient.getArrowData(workspaceClient, jobId, signal); - } - async shutdown(): Promise { this.streamManager.abortAll(); } @@ -859,267 +711,69 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } /** - * Hard caps on the server-side JSON_ARRAY fallback. The materializer - * builds every row as a plain JS object on the Node main thread - * (O(rows × cols) allocations), so a runaway result blocks the event - * loop and pressures GC. We cap two ways — rows AND decoded bytes — - * because either dimension can blow up independently (100k×1B rows is - * fine, 10 rows × 10MB cells is not). - */ -const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; -const JSON_ARRAY_FALLBACK_MAX_BYTES = 64 * 1024 * 1024; - -/** - * Replacer used by the nested-value `JSON.stringify` path. Arrow IPC - * carries column values whose JS representations include `bigint` - * (which the spec'd `JSON.stringify` cannot serialize and throws on), - * `Uint8Array` (binary buffers — would serialize as `{"0":...,"1":...}` - * which is wrong), and `Date` (would lose timezone fidelity inside a - * nested struct). Coerce each to a string form that matches what the - * warehouse itself emits under native JSON_ARRAY. - */ -function nestedJsonReplacer(_key: string, value: unknown): unknown { - if (typeof value === "bigint") return value.toString(); - if (value instanceof Uint8Array) { - return Buffer.from(value).toString("base64"); - } - if (value instanceof Date) return value.toISOString(); - return value; -} - -/** - * Render an `apache-arrow` Decimal cell to a fixed-point string. The JS - * value is a `DecimalBigNum` whose `toString()` yields the *signed - * unscaled* integer (e.g. `12345`, `-6789`); the decimal point is placed - * `scale` digits from the right to match the warehouse's native - * JSON_ARRAY rendering (`"123.45"`, `"-67.89"`). - */ -function formatDecimalCell( - value: { toString(): string }, - scale: number, -): string { - const unscaled = value.toString(); - if (scale <= 0) return unscaled; - const negative = unscaled.startsWith("-"); - let digits = negative ? unscaled.slice(1) : unscaled; - if (digits.length <= scale) digits = digits.padStart(scale + 1, "0"); - const point = digits.length - scale; - const out = `${digits.slice(0, point)}.${digits.slice(point)}`; - return negative ? `-${out}` : out; -} - -/** - * Render an `apache-arrow` Timestamp cell to the warehouse's native - * JSON_ARRAY rendering, which is ISO-8601 with millisecond precision: - * `yyyy-MM-ddTHH:mm:ss.SSSZ` for zoned TIMESTAMP / TIMESTAMP_LTZ, and the - * same without the trailing `Z` for TIMESTAMP_NTZ. This is exactly - * `Date#toISOString()` (drop the `Z` when the column carries no timezone). + * Write one chunk to the response honoring backpressure: if the socket + * buffer is full (`res.write` returns false), wait for `drain` before + * resolving so a slow client can't balloon Node's internal write queue and + * defeat the constant-memory goal of streaming. * - * `arrow-js` normalizes every Timestamp unit to epoch **milliseconds** at - * `.get()` (a MICROSECOND column returns ms, not micros); the warehouse - * also truncates JSON_ARRAY timestamps to ms, so the two agree. Verified - * against a dogfood serverless warehouse: `2024-06-13T01:02:03.500Z`, - * `2024-06-13T01:02:03.000Z`, and NTZ `2024-01-02T03:04:05.000`. + * @internal exported for unit testing the backpressure/disconnect behavior. */ -function formatTimestampCell(epochMs: number, hasTimezone: boolean): string { - const iso = new Date(epochMs).toISOString(); - return hasTimezone ? iso : iso.slice(0, -1); -} - -/** - * Render an `apache-arrow` Date cell to `yyyy-MM-dd` (UTC) to match the - * warehouse's native JSON_ARRAY rendering. `arrow-js` normalizes Date - * columns to epoch milliseconds at UTC midnight on `.get()`. - */ -function formatDateCell(epochMs: number): string { - return new Date(epochMs).toISOString().slice(0, 10); -} - -/** - * Parse a STRING cell that looks like JSON (`{`/`[` prefixed) into an - * object/array, mirroring `SQLWarehouseConnector._transformDataArray` so - * the JSON_ARRAY fallback produces the same nested shape the native - * JSON_ARRAY path does. Non-JSON strings (and parse failures) pass - * through unchanged. - */ -function maybeParseJsonString(value: string): unknown { - if (value && (value[0] === "{" || value[0] === "[")) { - try { - return JSON.parse(value); - } catch { - return value; - } +export function writeChunk( + res: express.Response, + bytes: Uint8Array, +): Promise { + const buf = Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength); + // If the socket is already gone, `res.write` won't return true and the + // `drain`/`close`/`error` events have already fired — so the promise below + // would never settle, wedging the for-await loop and the upstream reader. + // Reject up front instead. + if (res.destroyed || res.writableEnded) { + return Promise.reject( + new DOMException("The response stream closed", "AbortError"), + ); } - return value; + if (res.write(buf)) return Promise.resolve(); + // Backpressured: resolve on `drain`, but also settle on `close`/`error`. A + // client that disconnects mid-backpressure never emits `drain` on the + // destroyed socket, so waiting on `drain` alone would wedge this promise — + // and with it the awaiting for-await loop and the upstream Arrow reader — + // forever. Rejecting instead unwinds the stream so its `finally` can cancel + // the reader. + return new Promise((resolve, reject) => { + const cleanup = () => { + res.off("drain", onDrain); + res.off("close", onClose); + res.off("error", onClose); + }; + const onDrain = () => { + cleanup(); + resolve(); + }; + const onClose = () => { + cleanup(); + reject(new DOMException("The response stream closed", "AbortError")); + }; + res.once("drain", onDrain); + res.once("close", onClose); + res.once("error", onClose); + }); } /** - * Decode a base64 Arrow IPC attachment to plain row objects. - * - * Used by the JSON_ARRAY fallback path when a warehouse refuses - * `JSON_ARRAY + INLINE` and we have to satisfy the request via - * `ARROW_STREAM + INLINE` — the bytes come back as Arrow IPC but the - * caller's contract is JSON-shaped rows, so we convert server-side. - * - * Shape rules (chosen to match what the warehouse emits natively under - * JSON_ARRAY, so callers cannot tell which path served their query): - * - **Scalars (number / bigint / boolean)**: stringified. JSON_ARRAY - * wire format emits everything in `result.data_array` as strings; - * bigint stringification is also necessary for JSON-serializability. - * - **Strings / Date / Uint8Array**: stringified consistently — - * `Date` to ISO string, `Uint8Array` to base64. - * - **Nested (List / Struct / Map)**: `JSON.stringify` with the - * `nestedJsonReplacer` so nested bigints / Dates / binary buffers - * serialize sanely instead of throwing or producing object-as-map - * output. Apache-arrow's typed values expose `toJSON()` so the - * resulting string matches what the warehouse would have emitted. + * Fail-fast ceiling on the wait for the first Arrow byte (warehouse + * readiness + execute + first chunk). Past this a stuck/overloaded warehouse + * yields a clear 503 instead of hanging. Override per plugin via + * `arrowFirstByteTimeoutMs`. */ -function decodeArrowAttachmentToRows( - attachment: string, -): Record[] { - const decoded = Buffer.from(attachment, "base64"); - if (decoded.byteLength > JSON_ARRAY_FALLBACK_MAX_BYTES) { - throw ExecutionError.statementFailed( - `Result attachment is ${decoded.byteLength} bytes; JSON_ARRAY fallback materializer caps at ${JSON_ARRAY_FALLBACK_MAX_BYTES} bytes. Re-issue the query with format="ARROW_STREAM" to stream the full result.`, - "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", - "Result too large for JSON format. Re-run with ARROW_STREAM format.", - ); - } - const table = tableFromIPC( - new Uint8Array(decoded.buffer, decoded.byteOffset, decoded.byteLength), - ); - if (table.numRows > JSON_ARRAY_FALLBACK_MAX_ROWS) { - throw ExecutionError.statementFailed( - `Result has ${table.numRows} rows; JSON_ARRAY fallback materializer caps at ${JSON_ARRAY_FALLBACK_MAX_ROWS}. Re-issue the query with format="ARROW_STREAM" to stream the full result.`, - "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", - `Result too large for JSON format (over ${JSON_ARRAY_FALLBACK_MAX_ROWS} rows). Re-run with ARROW_STREAM format.`, - ); - } - // Resolve the child vectors once. `table.getChild(name)` walks the - // schema fields on every call; without this hoist we'd pay that cost - // O(rows × cols) times. At 100k × 50 columns that's millions of - // redundant lookups on the event loop. - const columns = table.schema.fields.map( - (f) => [f.name, table.getChild(f.name), f.type] as const, - ); - const rows: Record[] = []; - for (let i = 0; i < table.numRows; i++) { - const row: Record = {}; - for (const [name, col, type] of columns) { - const v = col?.get(i); - if (v == null) { - row[name] = null; - continue; - } - // Branch on the Arrow field type first: several Databricks types - // decode to JS values whose default stringification does NOT match - // the warehouse's native JSON_ARRAY rendering. Without these the - // fallback silently diverges from the direct JSON_ARRAY path - // (decimals lose scale, temporals come back as raw epoch numbers). - switch (type.typeId) { - case Type.Decimal: - row[name] = formatDecimalCell( - v as { toString(): string }, - (type as DataType & { scale: number }).scale, - ); - continue; - case Type.Timestamp: - row[name] = formatTimestampCell( - Number(v), - (type as DataType & { timezone?: string | null }).timezone != null, - ); - continue; - case Type.Date: - row[name] = formatDateCell(Number(v)); - continue; - } - if ( - typeof v === "number" || - typeof v === "bigint" || - typeof v === "boolean" - ) { - row[name] = String(v); - } else if (typeof v === "string") { - // Match `_transformDataArray`: STRING columns holding JSON are - // parsed into objects/arrays so both paths yield the same shape. - row[name] = maybeParseJsonString(v); - } else if (v instanceof Date) { - row[name] = v.toISOString(); - } else if (v instanceof Uint8Array) { - row[name] = Buffer.from(v).toString("base64"); - } else { - // Nested types (List, Struct, Map). Use the replacer so nested - // bigint / Date / Uint8Array values serialize predictably — - // bare `JSON.stringify` throws on bigint and emits broken - // shapes for binary buffers. - row[name] = JSON.stringify(v, nestedJsonReplacer); - } - } - rows.push(row); - } - return rows; -} +const DEFAULT_ARROW_FIRST_BYTE_TIMEOUT_MS = 120_000; /** - * Classify a warehouse rejection of an INLINE statement. - * - * Two distinct rejection modes are observed in the wild: - * - * - **needs-arrow**: warehouse refuses `JSON_ARRAY + INLINE`, only accepts - * `ARROW_STREAM + INLINE`. Example message: - * `"Inline disposition only supports ARROW_STREAM format."` - * Action: retry as `ARROW_STREAM + INLINE` and decode server-side. - * - * - **needs-json**: warehouse refuses `ARROW_STREAM + INLINE`, only accepts - * `JSON_ARRAY + INLINE`. Examples: - * `"The format field must be JSON_ARRAY when the disposition field is INLINE."` - * `"ARROW_STREAM not supported with INLINE disposition"` - * `"ExternalLinks disposition is not yet implemented."` (same family — - * the warehouse rejected the disposition/format combo we sent). - * Action: retry as `ARROW_STREAM + EXTERNAL_LINKS`. - * - * The structured `errorCode` (INVALID_PARAMETER_VALUE / NOT_IMPLEMENTED) - * gates the classification so unrelated SQL errors don't trigger a retry. - * Message matching is case-insensitive — warehouses are inconsistent about - * casing of "Inline"/"INLINE". + * Byte ceiling for the `X-Appkit-Arrow-Columns` header value. Beyond this (a + * very wide schema) the names are served via the `/columns/:statementId` + * fallback endpoint instead of the header. Kept well under the common ~8 KiB + * per-header limit. */ -type InlineRejection = "needs-arrow" | "needs-json" | null; - -function _classifyInlineRejection(err: unknown): InlineRejection { - const msg = err instanceof Error ? err.message : String(err); - const lower = msg.toLowerCase(); - - const structuredCode = - err instanceof ExecutionError ? err.errorCode : undefined; - const hasCode = - structuredCode === "INVALID_PARAMETER_VALUE" || - structuredCode === "NOT_IMPLEMENTED" || - lower.includes("invalid_parameter_value") || - lower.includes("not_implemented"); - if (!hasCode) return null; - - // Must mention the inline disposition to count as a disposition-rejection. - if (!lower.includes("inline")) return null; - - // "needs-arrow": warehouse only supports ARROW_STREAM for INLINE. - if ( - /only supports\s+arrow_stream/i.test(msg) || - /must be\s+arrow_stream/i.test(msg) - ) { - return "needs-arrow"; - } - - // "needs-json": warehouse only supports JSON_ARRAY for INLINE. - if ( - /only supports\s+json_array/i.test(msg) || - /must be\s+json_array/i.test(msg) || - /arrow_stream\s+(is\s+|was\s+)?not\s+supported/i.test(msg) - ) { - return "needs-json"; - } - - return null; -} +const MAX_ARROW_COLUMNS_HEADER_BYTES = 6000; /** * @internal diff --git a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts b/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts deleted file mode 100644 index cb52ff71..00000000 --- a/packages/appkit/src/plugins/analytics/inline-arrow-stash.ts +++ /dev/null @@ -1,240 +0,0 @@ -import { randomUUID } from "node:crypto"; - -/** - * Server-side stash for inline Arrow IPC payloads. - * - * When a warehouse returns ARROW_STREAM + INLINE results, the bytes are - * stashed here and a synthetic "inline-" job id is emitted on the - * SSE control channel. The client then fetches the bytes out-of-band via - * `/arrow-result/`, which drains the stash and serves the payload as - * `application/vnd.apache.arrow.stream`. - * - * Keeps multi-MiB Arrow blobs off SSE, lets the existing /arrow-result - * pipeline handle both inline and EXTERNAL_LINKS results uniformly, and - * delivers the bytes with a real binary content-type instead of base64 - * inside JSON inside SSE framing. - * - * Properties: - * - **Drain-on-read**: a successful `take()` removes the entry. There is - * no replay path — a lost client connection means the bytes are gone. - * - **TTL bounded**: entries past their expiry are evicted on every - * `put()` and `take()`. No background timer. - * - **Per-user keyed**: `take()` only returns bytes if the requesting - * user matches the user that originally put them. Defense in depth on - * top of unguessable ids. - * - **Memory bounded with overflow slot**: total stashed bytes are capped - * at `maxBytes`. When `put()` cannot fit a payload, it spills into a - * separate overflow pool capped at `maxOverflowBytes` (default - * `maxBytes / 4` — kept small because overflow only needs to bridge - * the immediate `/arrow-result` fetch). Overflow entries behave like - * regular ones except (a) they do not count against the regular cap - * and (b) they expire on a much shorter TTL (`overflowTtlMs`, - * default 30s) — they exist to absorb already-decoded bytes for a - * single request, not to hold them long-term. Only when both pools - * are full does `put()` return `null` and the caller has to fall - * back to a different delivery path. Memory is bounded above by - * `maxBytes + maxOverflowBytes`. - * - * Caveat (multi-replica deployments): this stash is process-local. A - * subsequent `GET /arrow-result/inline-*` that lands on a different - * replica than the one that stashed the bytes will 410. Deployments - * that run more than one replica need sticky sessions (route both - * requests in the same logical session to the same replica) or a - * shared external store, neither of which is in scope here. - */ -interface InlineArrowStashOptions { - /** Regular-pool entries older than this are dropped on the next gc tick. */ - ttlMs?: number; - /** - * Overflow-pool entries older than this are dropped on the next gc - * tick. Defaults to 30s — overflow exists solely to bridge the - * immediate `/arrow-result` fetch that follows the SSE `arrow` - * message, so it should drain on the order of seconds, not minutes. - * A short TTL bounds the cross-user memory pressure that overflow - * can sustain in a multi-tenant deployment. - */ - overflowTtlMs?: number; - /** - * Hard cap on total bytes held in the regular pool. `put()` spills to - * the overflow pool when this cap would be exceeded; entries already - * in the stash are not evicted to fit new ones. - */ - maxBytes?: number; - /** - * Hard cap on total bytes held in the overflow pool. Overflow holds - * bytes that have already been decoded for an in-flight request — its - * purpose is to avoid throwing those bytes away and double-billing - * the warehouse. Defaults to `maxBytes / 4` (kept small because the - * pool only needs to absorb transient spillover, not hold long-term - * state). `put()` returns `null` only when both regular and overflow - * pools are at cap. - */ - maxOverflowBytes?: number; - /** - * Minimum interval between gc passes. `gc()` is O(N) in entry count, - * so on hot paths we skip when the previous pass was recent enough. - * Defaults to 5s. Set to 0 to disable throttling (test seam). - */ - gcMinIntervalMs?: number; - /** Test seam: override the synthetic-id generator. */ - idGenerator?: () => string; - /** Test seam: override the clock. */ - now?: () => number; -} - -interface StashEntry { - userId: string; - bytes: Uint8Array; - expiresAt: number; - insertedAt: number; - /** True for entries in the overflow pool (do not count against maxBytes). */ - overflow: boolean; -} - -export class InlineArrowStash { - private entries = new Map(); - private totalBytes = 0; - private overflowBytes = 0; - private lastGcAt = 0; - private readonly ttlMs: number; - private readonly overflowTtlMs: number; - private readonly maxBytes: number; - private readonly maxOverflowBytes: number; - private readonly gcMinIntervalMs: number; - private readonly idGenerator: () => string; - private readonly now: () => number; - - constructor(opts: InlineArrowStashOptions = {}) { - this.ttlMs = opts.ttlMs ?? 10 * 60 * 1000; - this.overflowTtlMs = opts.overflowTtlMs ?? 30 * 1000; - this.maxBytes = opts.maxBytes ?? 256 * 1024 * 1024; - this.maxOverflowBytes = - opts.maxOverflowBytes ?? Math.floor(this.maxBytes / 4); - this.gcMinIntervalMs = opts.gcMinIntervalMs ?? 5000; - this.idGenerator = opts.idGenerator ?? randomUUID; - this.now = opts.now ?? Date.now; - } - - /** - * Stash a payload and return its synthetic job id. - * - * Tries the regular pool first; if it would overflow, spills into the - * overflow pool (sized at `maxOverflowBytes`). Returns `null` only when - * both pools are at cap — the caller then has no choice but to fall - * back to a different delivery path (e.g. EXTERNAL_LINKS). - * - * The overflow pool exists because the caller has *already decoded* - * these bytes for an in-flight request: throwing them away would force - * a second warehouse round-trip (extra latency, double billing, and - * potentially divergent results for non-deterministic SQL). Holding - * them transiently in a bounded overflow region — they are drained - * single-use on the next `/arrow-result` fetch — is strictly safer - * than re-execution. - * - * A single payload can only land in one pool — the pools are not - * split across — so the largest payload we can accept is - * `Math.max(maxBytes, maxOverflowBytes)`. Exceeding that throws - * synchronously so the caller sees the misconfiguration loudly - * rather than burning a warehouse round-trip and then getting a - * null id. - * - * Returns `{ id, pool }` on success (`pool` ∈ {"regular", "overflow"}) - * or `null` when both pools are at cap. The pool tag lets callers - * emit accurate telemetry labels without re-introspecting the stash. - */ - put( - userId: string, - bytes: Uint8Array, - ): { id: string; pool: "regular" | "overflow" } | null { - const largestSlot = Math.max(this.maxBytes, this.maxOverflowBytes); - if (bytes.length > largestSlot) { - throw new Error( - `Inline Arrow payload (${bytes.length} bytes) exceeds largest stash slot (${largestSlot}); cannot fit in either pool`, - ); - } - this.gc(); - const fitsRegular = this.totalBytes + bytes.length <= this.maxBytes; - const fitsOverflow = - !fitsRegular && - this.overflowBytes + bytes.length <= this.maxOverflowBytes; - if (!fitsRegular && !fitsOverflow) { - // Both pools are full — refuse rather than evicting any issued id. - return null; - } - const id = `inline-${this.idGenerator()}`; - const now = this.now(); - const overflow = !fitsRegular; - this.entries.set(id, { - userId, - bytes, - expiresAt: now + (overflow ? this.overflowTtlMs : this.ttlMs), - insertedAt: now, - overflow, - }); - if (overflow) { - this.overflowBytes += bytes.length; - } else { - this.totalBytes += bytes.length; - } - return { id, pool: overflow ? "overflow" : "regular" }; - } - - /** - * Drain a payload from the stash. Returns `undefined` if the id is - * unknown, expired, or belongs to a different user. - */ - take(id: string, userId: string): Uint8Array | undefined { - this.gc(); - const entry = this.entries.get(id); - if (!entry) return undefined; - if (entry.userId !== userId) return undefined; - this.entries.delete(id); - if (entry.overflow) { - this.overflowBytes -= entry.bytes.length; - } else { - this.totalBytes -= entry.bytes.length; - } - return entry.bytes; - } - - /** Inspection helpers (primarily for tests). */ - size(): number { - return this.totalBytes; - } - /** Bytes currently held in the overflow pool. */ - overflowSize(): number { - return this.overflowBytes; - } - count(): number { - return this.entries.size; - } - - /** Drop all entries (used in plugin shutdown). */ - clear(): void { - this.entries.clear(); - this.totalBytes = 0; - this.overflowBytes = 0; - } - - private gc(): void { - const now = this.now(); - if (now - this.lastGcAt < this.gcMinIntervalMs) { - // Skip the O(N) sweep — recent enough to assume nothing - // significant has expired since the last pass. Worst case an - // entry lingers an extra `gcMinIntervalMs` past its TTL, which - // is negligible relative to either pool's intended lifetime. - return; - } - this.lastGcAt = now; - for (const [id, entry] of this.entries) { - if (entry.expiresAt <= now) { - this.entries.delete(id); - if (entry.overflow) { - this.overflowBytes -= entry.bytes.length; - } else { - this.totalBytes -= entry.bytes.length; - } - } - } - } -} diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts new file mode 100644 index 00000000..60fa957a --- /dev/null +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -0,0 +1,431 @@ +import type { sql } from "@databricks/sdk-experimental"; +import { type DataType, Type, tableFromIPC } from "apache-arrow"; +import type { SQLTypeMarker } from "shared"; +import { ExecutionError } from "../../errors"; +import { createLogger } from "../../logging/logger"; + +/** + * Centralized disposition/format fallback for analytics result delivery. + * + * Two warehouse capability profiles must both work, and their supported + * combinations are mutually exclusive per warehouse: + * + * - **Reyden**: `ARROW_STREAM + INLINE` (and `JSON_ARRAY + INLINE`); does NOT + * support `EXTERNAL_LINKS`. + * - **Normal SQL warehouse**: `JSON_ARRAY + INLINE` and + * `ARROW_STREAM + EXTERNAL_LINKS`; rejects `ARROW_STREAM + INLINE`. + * + * The fallback logic lives here (rather than inline in the plugin route) so it + * can be unit-tested against a fake executor without HTTP, and so the brittle + * "which rejection is this" decision is expressed once. + */ + +const logger = createLogger("analytics:delivery"); + +/** + * Minimal query surface the fallback needs. The analytics plugin (and its + * `asUser(req)` proxy) satisfy it, so the same delivery logic runs unchanged + * under service-principal and on-behalf-of identities. + */ +export interface QueryExecutor { + query( + query: string, + parameters: Record | undefined, + formatParameters: { disposition: string; format: string }, + signal?: AbortSignal, + ): Promise< + | { + attachment?: string; + data?: Record[]; + external_links?: sql.ExternalLink[]; + columnNames?: string[]; + statement_id?: string; + status?: unknown; + } + | undefined + >; +} + +/** Streams already-resolved EXTERNAL_LINKS chunks; the connector provides it. */ +export interface ArrowChunkStreamer { + streamExternalLinks( + chunks: sql.ExternalLink[], + signal?: AbortSignal, + ): AsyncGenerator; +} + +/** + * How a warehouse rejected a disposition/format combination. + * + * Only the two structured error codes Databricks emits for capability + * mismatches (`INVALID_PARAMETER_VALUE`, `NOT_IMPLEMENTED`) gate a fallback — + * auth, permission, and SQL errors never match, so they propagate untouched. + * Message matching only *disambiguates* which capability is missing once the + * code has confirmed it's a capability error. + */ +type DispositionRejection = + | "needs-arrow-inline" // INLINE requires ARROW_STREAM + | "needs-json-inline" // INLINE requires JSON_ARRAY + | "external-links-unsupported" // EXTERNAL_LINKS not implemented + | null; + +export function classifyDispositionRejection( + err: unknown, +): DispositionRejection { + const msg = err instanceof Error ? err.message : String(err); + const lower = msg.toLowerCase(); + + const structuredCode = + err instanceof ExecutionError ? err.errorCode : undefined; + const hasCapabilityCode = + structuredCode === "INVALID_PARAMETER_VALUE" || + structuredCode === "NOT_IMPLEMENTED" || + lower.includes("invalid_parameter_value") || + lower.includes("not_implemented"); + // Auth / permission / SQL errors carry other codes → never fall back. + if (!hasCapabilityCode) return null; + + // EXTERNAL_LINKS disposition not implemented (Reyden). + if ( + /external[_\s]?links/.test(lower) && + (lower.includes("not yet implemented") || + lower.includes("not implemented") || + lower.includes("not supported")) + ) { + return "external-links-unsupported"; + } + + // Remaining cases are INLINE disposition/format mismatches. + if (!lower.includes("inline")) return null; + + if ( + /only supports\s+arrow_stream/i.test(msg) || + /must be\s+arrow_stream/i.test(msg) + ) { + return "needs-arrow-inline"; + } + if ( + /only supports\s+json_array/i.test(msg) || + /must be\s+json_array/i.test(msg) || + /arrow_stream\s+(is\s+|was\s+)?not\s+supported/i.test(msg) + ) { + return "needs-json-inline"; + } + + return null; +} + +/** Structured error: the warehouse supports no Arrow delivery mode at all. */ +export function arrowDeliveryUnsupported(): ExecutionError { + return ExecutionError.statementFailed( + "Warehouse supports neither ARROW_STREAM+INLINE nor ARROW_STREAM+EXTERNAL_LINKS", + "ARROW_DELIVERY_UNSUPPORTED", + 'This warehouse cannot return Arrow results. Re-run the query with format="JSON_ARRAY".', + ); +} + +const errMessage = (err: unknown): string => + err instanceof Error ? err.message : String(err); + +/** + * Produce the raw Arrow IPC bytes for an ARROW_STREAM query, applying the + * capability fallback: + * + * 1. `INLINE + ARROW_STREAM` — on success the base64 attachment is decoded + * once and yielded as a `Uint8Array` view (no Arrow parsing, no stash). + * 2. On a `needs-json-inline` rejection, `EXTERNAL_LINKS + ARROW_STREAM` — + * the pre-signed chunks resolved in the executor's own context are + * streamed one at a time. + * 3. If EXTERNAL_LINKS is also unsupported, a structured + * `ARROW_DELIVERY_UNSUPPORTED` error is thrown. + * + * `out` is populated with the real column names and statement id before the + * first chunk is yielded, so the caller can emit them (e.g. as a header) + * ahead of the body. + */ +export async function* deliverArrowBytes( + executor: QueryExecutor, + streamer: ArrowChunkStreamer, + query: string, + processedParams: Record | undefined, + out: { columnNames?: string[]; statementId?: string }, + signal?: AbortSignal, +): AsyncGenerator { + try { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "ARROW_STREAM" }, + signal, + ); + if (result?.attachment) { + out.columnNames = result.columnNames; + out.statementId = result.statement_id; + // Decode base64 once; yield a view over the Buffer (no copy, no parse). + const decoded = Buffer.from(result.attachment, "base64"); + yield new Uint8Array( + decoded.buffer, + decoded.byteOffset, + decoded.byteLength, + ); + return; + } + // INLINE succeeded but returned no attachment (rare: inline data_array + // under ARROW_STREAM). Fall through to EXTERNAL_LINKS. + logger.warn( + "ARROW_STREAM INLINE returned no attachment; falling back to EXTERNAL_LINKS", + ); + } catch (err: unknown) { + if (signal?.aborted) throw err; + if (classifyDispositionRejection(err) !== "needs-json-inline") throw err; + logger.warn( + "ARROW_STREAM INLINE rejected by warehouse, falling back to EXTERNAL_LINKS: %s", + errMessage(err), + ); + } + + let ext: Awaited>; + try { + ext = await executor.query( + query, + processedParams, + { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, + signal, + ); + } catch (err: unknown) { + if (signal?.aborted) throw err; + // Neither INLINE nor EXTERNAL_LINKS Arrow is supported — surface a clear, + // actionable error instead of a raw warehouse rejection. + if (classifyDispositionRejection(err) === "external-links-unsupported") { + throw arrowDeliveryUnsupported(); + } + throw err; + } + + if (!ext?.external_links) { + throw ExecutionError.missingData("external_links"); + } + out.columnNames = ext.columnNames; + out.statementId = ext.statement_id; + // Stream the pre-signed links resolved in THIS request's execution context + // (user creds for `.obo.sql`, service principal otherwise). Pre-signed URLs + // need no auth to download, so there is no second `getStatement` under a + // mismatched identity. + yield* streamer.streamExternalLinks(ext.external_links, signal); +} + +/** JSON row result plus the metadata the SSE `result` message forwards. */ +interface JsonResult { + data: Record[] | undefined; + status: unknown; + statement_id?: string; +} + +/** + * Produce JSON rows for a JSON_ARRAY query, applying the capability fallback: + * + * 1. `INLINE + JSON_ARRAY` — the native, common path. + * 2. On a `needs-arrow-inline` rejection (warehouse only accepts ARROW_STREAM + * for INLINE, e.g. Reyden), retry `INLINE + ARROW_STREAM` and decode the + * attachment to plain rows server-side so the caller's JSON contract holds. + * + * EXTERNAL_LINKS is never used for the JSON fallback. + */ +export async function deliverJsonResult( + executor: QueryExecutor, + query: string, + processedParams: Record | undefined, + signal?: AbortSignal, +): Promise { + try { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "JSON_ARRAY" }, + signal, + ); + return { + data: result?.data, + status: result?.status, + statement_id: result?.statement_id, + }; + } catch (err: unknown) { + if (signal?.aborted) throw err; + if (classifyDispositionRejection(err) !== "needs-arrow-inline") throw err; + logger.warn( + "JSON_ARRAY INLINE rejected by warehouse, retrying as ARROW_STREAM INLINE and decoding server-side: %s", + errMessage(err), + ); + } + + const arrowResult = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "ARROW_STREAM" }, + signal, + ); + if (!arrowResult?.attachment) { + throw ExecutionError.missingData("ARROW_STREAM attachment"); + } + const rows = decodeArrowAttachmentToRows( + arrowResult.attachment, + arrowResult.columnNames, + ); + return { + data: rows, + status: arrowResult.status, + statement_id: arrowResult.statement_id, + }; +} + +// ──────────────────────────────────────────────────────────────────────────── +// Server-side Arrow → JSON row materializer (JSON fallback only) +// ──────────────────────────────────────────────────────────────────────────── + +/** + * Hard caps on the server-side JSON_ARRAY fallback. The materializer builds + * every row as a plain JS object on the Node main thread (O(rows × cols) + * allocations), so a runaway result blocks the event loop and pressures GC. + * Cap on rows AND decoded bytes — either dimension can blow up independently. + */ +const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; +const JSON_ARRAY_FALLBACK_MAX_BYTES = 64 * 1024 * 1024; + +/** + * Replacer for nested `JSON.stringify`. Arrow values include `bigint` (which + * spec'd `JSON.stringify` throws on), `Uint8Array` (would serialize as an + * index map), and `Date` (would lose timezone in a nested struct). Coerce + * each to the string form the warehouse itself emits under native JSON_ARRAY. + */ +function nestedJsonReplacer(_key: string, value: unknown): unknown { + if (typeof value === "bigint") return value.toString(); + if (value instanceof Uint8Array) return Buffer.from(value).toString("base64"); + if (value instanceof Date) return value.toISOString(); + return value; +} + +/** Render an apache-arrow Decimal cell to a fixed-point string. */ +function formatDecimalCell( + value: { toString(): string }, + scale: number, +): string { + const unscaled = value.toString(); + if (scale <= 0) return unscaled; + const negative = unscaled.startsWith("-"); + let digits = negative ? unscaled.slice(1) : unscaled; + if (digits.length <= scale) digits = digits.padStart(scale + 1, "0"); + const point = digits.length - scale; + const out = `${digits.slice(0, point)}.${digits.slice(point)}`; + return negative ? `-${out}` : out; +} + +/** + * Render an apache-arrow Timestamp cell to ISO-8601 ms precision — `Z` for + * zoned columns, no `Z` for TIMESTAMP_NTZ — matching native JSON_ARRAY. + */ +function formatTimestampCell(epochMs: number, hasTimezone: boolean): string { + const iso = new Date(epochMs).toISOString(); + return hasTimezone ? iso : iso.slice(0, -1); +} + +/** Render an apache-arrow Date cell to `yyyy-MM-dd` (UTC). */ +function formatDateCell(epochMs: number): string { + return new Date(epochMs).toISOString().slice(0, 10); +} + +/** Parse a STRING cell that looks like JSON into an object/array. */ +function maybeParseJsonString(value: string): unknown { + if (value && (value[0] === "{" || value[0] === "[")) { + try { + return JSON.parse(value); + } catch { + return value; + } + } + return value; +} + +/** + * Decode a base64 Arrow IPC attachment to plain row objects, matching what the + * warehouse emits natively under JSON_ARRAY so callers cannot tell which path + * served the query. + * + * The output key uses the manifest name (`columnNames[i]`) when available — + * Databricks encodes the Arrow schema positionally (col_0, …) — while the + * vector is still fetched by the Arrow field's own name. + */ +export function decodeArrowAttachmentToRows( + attachment: string, + columnNames?: string[], +): Record[] { + const decoded = Buffer.from(attachment, "base64"); + if (decoded.byteLength > JSON_ARRAY_FALLBACK_MAX_BYTES) { + throw ExecutionError.statementFailed( + `Result attachment is ${decoded.byteLength} bytes; JSON_ARRAY fallback materializer caps at ${JSON_ARRAY_FALLBACK_MAX_BYTES} bytes. Re-issue the query with format="ARROW_STREAM" to stream the full result.`, + "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", + "Result too large for JSON format. Re-run with ARROW_STREAM format.", + ); + } + const table = tableFromIPC( + new Uint8Array(decoded.buffer, decoded.byteOffset, decoded.byteLength), + ); + if (table.numRows > JSON_ARRAY_FALLBACK_MAX_ROWS) { + throw ExecutionError.statementFailed( + `Result has ${table.numRows} rows; JSON_ARRAY fallback materializer caps at ${JSON_ARRAY_FALLBACK_MAX_ROWS}. Re-issue the query with format="ARROW_STREAM" to stream the full result.`, + "RESULT_TOO_LARGE_FOR_JSON_FALLBACK", + `Result too large for JSON format (over ${JSON_ARRAY_FALLBACK_MAX_ROWS} rows). Re-run with ARROW_STREAM format.`, + ); + } + // Resolve child vectors once (getChild walks fields on every call). The + // output key uses the manifest name; the vector is fetched by field name. + const columns = table.schema.fields.map((f, i) => { + const outName = + columnNames?.[i] && columnNames[i].length > 0 ? columnNames[i] : f.name; + return [outName, table.getChild(f.name), f.type] as const; + }); + const rows: Record[] = []; + for (let i = 0; i < table.numRows; i++) { + const row: Record = {}; + for (const [name, col, type] of columns) { + const v = col?.get(i); + if (v == null) { + row[name] = null; + continue; + } + switch (type.typeId) { + case Type.Decimal: + row[name] = formatDecimalCell( + v as { toString(): string }, + (type as DataType & { scale: number }).scale, + ); + continue; + case Type.Timestamp: + row[name] = formatTimestampCell( + Number(v), + (type as DataType & { timezone?: string | null }).timezone != null, + ); + continue; + case Type.Date: + row[name] = formatDateCell(Number(v)); + continue; + } + if ( + typeof v === "number" || + typeof v === "bigint" || + typeof v === "boolean" + ) { + row[name] = String(v); + } else if (typeof v === "string") { + row[name] = maybeParseJsonString(v); + } else if (v instanceof Date) { + row[name] = v.toISOString(); + } else if (v instanceof Uint8Array) { + row[name] = Buffer.from(v).toString("base64"); + } else { + row[name] = JSON.stringify(v, nestedJsonReplacer); + } + } + rows.push(row); + } + return rows; +} diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index d66fa364..86154c66 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -19,8 +19,7 @@ import { import { sql } from "shared"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { ServiceContext } from "../../../context/service-context"; -import { AnalyticsPlugin, analytics } from "../analytics"; -import { InlineArrowStash } from "../inline-arrow-stash"; +import { AnalyticsPlugin, analytics, writeChunk } from "../analytics"; import type { IAnalyticsConfig } from "../types"; // Mock CacheManager singleton with actual caching behavior @@ -105,100 +104,6 @@ describe("Analytics Plugin", () => { ); }); - test("should register GET route for arrow results", () => { - const plugin = new AnalyticsPlugin(config); - const { router } = createMockRouter(); - - plugin.injectRoutes(router); - - expect(router.get).toHaveBeenCalledTimes(1); - expect(router.get).toHaveBeenCalledWith( - "/arrow-result/:jobId", - expect.any(Function), - ); - }); - - test("/arrow-result/inline-* drains the stash and serves bytes as application/vnd.apache.arrow.stream", async () => { - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - plugin.injectRoutes(router); - - const arrowBytes = new Uint8Array([0xff, 0xfe, 0xfd, 0xfc]); - const { id } = (plugin as any).inlineArrowStash.put("global", arrowBytes); - expect(id.startsWith("inline-")).toBe(true); - - const handler = getHandler("GET", "/arrow-result/:jobId"); - const mockReq = createMockRequest({ params: { jobId: id } }); - const mockRes = createMockResponse(); - - await handler(mockReq, mockRes); - - expect(mockRes.setHeader).toHaveBeenCalledWith( - "Content-Type", - "application/vnd.apache.arrow.stream", - ); - expect(mockRes.setHeader).toHaveBeenCalledWith( - "Content-Length", - String(arrowBytes.length), - ); - expect(mockRes.setHeader).toHaveBeenCalledWith( - "Cache-Control", - "no-store", - ); - expect(mockRes.send).toHaveBeenCalledTimes(1); - const sentBuf = (mockRes.send as any).mock.calls[0][0] as Buffer; - expect(Buffer.isBuffer(sentBuf)).toBe(true); - expect(Array.from(sentBuf)).toEqual(Array.from(arrowBytes)); - - // Drain-on-read: a second fetch must return 410, not the bytes again. - const secondRes = createMockResponse(); - await handler(mockReq, secondRes); - expect(secondRes.status).toHaveBeenCalledWith(410); - }); - - test("/arrow-result/inline-* returns 410 when the stash entry never existed", async () => { - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - plugin.injectRoutes(router); - - const handler = getHandler("GET", "/arrow-result/:jobId"); - const mockReq = createMockRequest({ - params: { jobId: "inline-does-not-exist" }, - }); - const mockRes = createMockResponse(); - - await handler(mockReq, mockRes); - - expect(mockRes.status).toHaveBeenCalledWith(410); - expect(mockRes.json).toHaveBeenCalledWith( - expect.objectContaining({ - error: expect.stringMatching(/expired or unknown/), - }), - ); - }); - - test("/arrow-result/inline-* returns 410 when the stash entry belongs to a different user", async () => { - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - plugin.injectRoutes(router); - - // Stash entry keyed to user-a, but the request resolves to "global" - // (no x-forwarded-user header) — keys differ, take must return - // nothing, and the entry stays put (single-user view). - const bytes = new Uint8Array([1, 2, 3]); - const { id } = (plugin as any).inlineArrowStash.put("user-a", bytes); - - const handler = getHandler("GET", "/arrow-result/:jobId"); - const mockReq = createMockRequest({ params: { jobId: id } }); - const mockRes = createMockResponse(); - - await handler(mockReq, mockRes); - - expect(mockRes.status).toHaveBeenCalledWith(410); - // The entry must still be there for the real owner. - expect((plugin as any).inlineArrowStash.take(id, "user-a")).toBeDefined(); - }); - test("/query/:query_key should return 400 when query_key is missing", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -834,7 +739,7 @@ describe("Analytics Plugin", () => { }); }); - test("/query/:query_key should fall back ARROW_STREAM from INLINE to EXTERNAL_LINKS when warehouse rejects INLINE", async () => { + test("/query/:query_key falls back ARROW_STREAM INLINE→EXTERNAL_LINKS and streams the preserved links in-context", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -843,6 +748,11 @@ describe("Analytics Plugin", () => { isAsUser: false, }); + // INLINE rejected → EXTERNAL_LINKS. The result carries the pre-signed + // `external_links` resolved in this request's context, so the route + // streams them directly — no second getStatement under the ambient + // service-principal context. + const links = [{ external_link: "https://example.com/chunk-0" }]; const executeMock = vi .fn() .mockRejectedValueOnce( @@ -851,10 +761,21 @@ describe("Analytics Plugin", () => { ), ) .mockResolvedValueOnce({ - result: { statement_id: "stmt-1", status: { state: "SUCCEEDED" } }, + result: { + statement_id: "stmt-1", + status: { state: "SUCCEEDED" }, + columnNames: ["group_key", "cost_usd"], + external_links: links, + }, }); (plugin as any).SQLClient.executeStatement = executeMock; + const extBytes = new Uint8Array([9, 8, 7]); + const streamExternalLinksMock = vi.fn(function* (_chunks: unknown) { + yield extBytes; + }); + (plugin as any).SQLClient.streamExternalLinks = streamExternalLinksMock; + plugin.injectRoutes(router); const handler = getHandler("POST", "/query/:query_key"); @@ -866,16 +787,227 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); - // First call: INLINE (rejected) + // First call INLINE (rejected), second EXTERNAL_LINKS (fallback). expect(executeMock.mock.calls[0][1]).toMatchObject({ disposition: "INLINE", format: "ARROW_STREAM", }); - // Second call: EXTERNAL_LINKS (fallback) expect(executeMock.mock.calls[1][1]).toMatchObject({ disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM", }); + // Streams the pre-signed links from the execute response (no re-fetch). + expect(streamExternalLinksMock).toHaveBeenCalledTimes(1); + expect(streamExternalLinksMock.mock.calls[0][0]).toBe(links); + // Bytes stream on the response body with the Arrow content type — no + // JSON error, no missingData throw. + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Content-Type", + "application/vnd.apache.arrow.stream", + ); + // Manifest names (from EXTERNAL_LINKS result) ride the header too. + expect(mockRes.setHeader).toHaveBeenCalledWith( + "X-Appkit-Arrow-Columns", + encodeURIComponent(JSON.stringify(["group_key", "cost_usd"])), + ); + expect(mockRes.json).not.toHaveBeenCalled(); + const written = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as Buffer, + ); + expect(written).toHaveLength(1); + expect(Array.from(written[0])).toEqual([9, 8, 7]); + }); + + test("OBO: .obo.sql ARROW_STREAM external-links streams under the user context", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + // `.obo.sql` → isAsUser true → the route must run through asUser(req). + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: true, + }); + + const links = [{ external_link: "https://example.com/obo-chunk-0" }]; + // Fake user-context executor: INLINE rejected, EXTERNAL_LINKS resolves + // with the pre-signed links (resolved with the user's identity). + const userExecutorQuery = vi.fn(async (_q, _p, fp) => { + if (fp.disposition === "INLINE") { + throw new Error( + "INVALID_PARAMETER_VALUE: The format field must be JSON_ARRAY when the disposition field is INLINE.", + ); + } + return { + external_links: links, + columnNames: ["a"], + statement_id: "obo-stmt", + }; + }); + // Warehouse readiness must run through the user-context executor too, so + // `getWorkspaceClient()` resolves to the user (not the SP) for `.obo.sql`. + const ensureReadyMock = vi.fn().mockResolvedValue(undefined); + const asUserSpy = vi.spyOn(plugin as any, "asUser").mockReturnValue({ + query: userExecutorQuery, + _ensureArrowWarehouseReady: ensureReadyMock, + }); + + const streamExternalLinksMock = vi.fn(function* (_chunks: unknown) { + yield new Uint8Array([1, 2, 3]); + }); + (plugin as any).SQLClient.streamExternalLinks = streamExternalLinksMock; + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // The user context (asUser) executor ran the queries AND the warehouse + // readiness — not the SP `this`. + expect(asUserSpy).toHaveBeenCalledWith(mockReq); + expect(userExecutorQuery).toHaveBeenCalled(); + expect(ensureReadyMock).toHaveBeenCalled(); + // The links the user-context executor resolved are streamed directly — + // no re-fetch under a different identity. + expect(streamExternalLinksMock).toHaveBeenCalledTimes(1); + expect(streamExternalLinksMock.mock.calls[0][0]).toBe(links); + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Content-Type", + "application/vnd.apache.arrow.stream", + ); + }); + + test("ARROW_STREAM: a stuck warehouse fails fast with a 503 WAREHOUSE_UNAVAILABLE", async () => { + const plugin = new AnalyticsPlugin({ + ...config, + arrowFirstByteTimeoutMs: 20, + }); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + (plugin as any).SQLClient.ensureWarehouseRunning = vi + .fn() + .mockResolvedValue(undefined); + // Warehouse never produces a first byte — only settles when aborted. + (plugin as any).SQLClient.executeStatement = vi.fn( + (_c: unknown, _i: unknown, signal?: AbortSignal) => + new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => + reject( + new DOMException("The operation was aborted.", "AbortError"), + ), + ); + }), + ); + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(503); + expect(mockRes.json).toHaveBeenCalledWith( + expect.objectContaining({ errorCode: "WAREHOUSE_UNAVAILABLE" }), + ); + }); + + test("ARROW_STREAM: a schema too wide for the header advertises a columns-ref instead", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + // 2000 columns → the encoded header value exceeds the size cap. + const wideNames = Array.from({ length: 2000 }, (_, i) => `column_${i}`); + (plugin as any).SQLClient.executeStatement = vi.fn().mockResolvedValue({ + result: { + attachment: Buffer.from([1, 2, 3]).toString("base64"), + columnNames: wideNames, + statement_id: "stmt-wide", + }, + }); + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + const setHeaderCalls = (mockRes.setHeader as any).mock.calls.map( + (c: any[]) => c[0], + ); + expect(setHeaderCalls).toContain("X-Appkit-Arrow-Columns-Ref"); + expect(setHeaderCalls).not.toContain("X-Appkit-Arrow-Columns"); + expect(mockRes.setHeader).toHaveBeenCalledWith( + "X-Appkit-Arrow-Columns-Ref", + "stmt-wide", + ); + }); + + test("GET /columns/:statementId returns the manifest column names", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + (plugin as any).SQLClient.getColumnNames = vi + .fn() + .mockResolvedValue(["name", "spend"]); + + plugin.injectRoutes(router); + const handler = getHandler("GET", "/columns/:statementId"); + const mockReq = createMockRequest({ + params: { statementId: "stmt-wide" }, + }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + expect(mockRes.json).toHaveBeenCalledWith({ columns: ["name", "spend"] }); + }); + + test("GET /columns/:statementId falls back to the service principal when the user identity can't read it (OBO)", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + // User identity 404s (e.g. the statement was executed by the SP, which + // the user can't `getStatement`); the SP `this` resolves it. + const spGetColumns = vi.fn().mockResolvedValue(["a", "b"]); + (plugin as any).SQLClient.getColumnNames = spGetColumns; + const userGetColumnNames = vi + .fn() + .mockRejectedValue(new Error("RESOURCE_DOES_NOT_EXIST")); + const asUserSpy = vi + .spyOn(plugin as any, "asUser") + .mockReturnValue({ _getColumnNames: userGetColumnNames }); + + plugin.injectRoutes(router); + const handler = getHandler("GET", "/columns/:statementId"); + const mockReq = createMockRequest({ params: { statementId: "stmt-x" } }); + const mockRes = createMockResponse(); + + await handler(mockReq, mockRes); + + // Tried the user identity first, then fell back to the service principal. + expect(asUserSpy).toHaveBeenCalledWith(mockReq); + expect(userGetColumnNames).toHaveBeenCalledWith("stmt-x"); + expect(spGetColumns).toHaveBeenCalled(); + expect(mockRes.json).toHaveBeenCalledWith({ columns: ["a", "b"] }); }); test("/query/:query_key falls back on a structured ExecutionError.errorCode without scanning the message", async () => { @@ -1039,7 +1171,7 @@ describe("Analytics Plugin", () => { } }); - test("/query/:query_key stashes ARROW_STREAM INLINE bytes and emits an arrow message with a synthetic inline- id", async () => { + test("/query/:query_key streams ARROW_STREAM INLINE bytes directly on the response body (no SSE, no stash)", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -1049,90 +1181,19 @@ describe("Analytics Plugin", () => { }); // Real base64 so the route can decode it via Buffer.from(..., "base64"). + // `columnNames` is what the connector attaches from the manifest (the + // Arrow schema itself is positional col_N). const arrowBytes = new Uint8Array([1, 2, 3, 4, 5]); const fakeAttachment = Buffer.from(arrowBytes).toString("base64"); const executeMock = vi.fn().mockResolvedValue({ - result: { attachment: fakeAttachment, row_count: 1 }, - }); - (plugin as any).SQLClient.executeStatement = executeMock; - - plugin.injectRoutes(router); - - const handler = getHandler("POST", "/query/:query_key"); - const mockReq = createMockRequest({ - params: { query_key: "test_query" }, - body: { parameters: {}, format: "ARROW_STREAM" }, - }); - const mockRes = createMockResponse(); - - await handler(mockReq, mockRes); - - // The route should not fall back to EXTERNAL_LINKS — INLINE succeeded. - expect(executeMock).toHaveBeenCalledTimes(1); - expect(executeMock.mock.calls[0][1]).toMatchObject({ - disposition: "INLINE", - format: "ARROW_STREAM", - }); - // SSE payload: unified `arrow` message with an inline- prefixed id. - // The base64 attachment must NOT appear on the SSE channel. - const writeCalls = (mockRes.write as any).mock.calls.map( - (c: any[]) => c[0] as string, - ); - // Skip the leading warehouse_status event the route always emits; - // the terminal result/arrow/error payload is the one under test. - const payload = writeCalls.find( - (s: string) => - s.startsWith("data: ") && !s.includes("warehouse_status"), - ); - expect(payload).toBeDefined(); - expect(payload).toContain('"type":"arrow"'); - expect(payload).toMatch(/"statement_id":"inline-[^"]+"/); - expect(payload).not.toContain("arrow_inline"); - expect(payload).not.toContain(fakeAttachment); - - // The decoded bytes should be in the stash, keyed by the same - // synthetic id; a subsequent /arrow-result fetch will drain them. - const idMatch = payload?.match(/"statement_id":"(inline-[^"]+)"/); - expect(idMatch).not.toBeNull(); - const inlineId = idMatch?.[1]; - const stashed = (plugin as any).inlineArrowStash.take(inlineId, "global"); - expect(stashed).toBeDefined(); - expect(Array.from(stashed)).toEqual(Array.from(arrowBytes)); - }); - - test("/query/:query_key spills the already-decoded bytes into the stash overflow pool when the regular pool is full — single execution, no double-billing", async () => { - // The regular stash put may refuse new entries when at cap. We must - // NOT respond by re-executing the same statement with EXTERNAL_LINKS: - // the warehouse has already been billed, the bytes are already - // decoded server-side, and a second execution can return a divergent - // result for non-deterministic SQL (CURRENT_TIMESTAMP, RAND, late - // rows). The stash's overflow pool absorbs the bytes so the client - // still gets an inline- id pointing at the original result. - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - - (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ - query: "SELECT * FROM test", - isAsUser: false, - }); - - const fakeAttachment = Buffer.from(new Uint8Array([1, 2, 3])).toString( - "base64", - ); - const executeMock = vi.fn().mockResolvedValueOnce({ - result: { attachment: fakeAttachment, row_count: 1 }, + result: { + attachment: fakeAttachment, + row_count: 1, + columnNames: ["name", "totalSpend"], + }, }); (plugin as any).SQLClient.executeStatement = executeMock; - // Force the regular pool to be at cap so the put spills into overflow. - // We construct a tiny stash and pre-fill the regular pool. - const tinyStash = new InlineArrowStash({ - maxBytes: 4, - maxOverflowBytes: 64, - }); - tinyStash.put("filler", new Uint8Array([9, 9, 9, 9])); - (plugin as any).inlineArrowStash = tinyStash; - plugin.injectRoutes(router); const handler = getHandler("POST", "/query/:query_key"); @@ -1144,83 +1205,34 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); - // Single execution: no EXTERNAL_LINKS retry. + // INLINE succeeded — a single execution, no EXTERNAL_LINKS fallback. expect(executeMock).toHaveBeenCalledTimes(1); expect(executeMock.mock.calls[0][1]).toMatchObject({ disposition: "INLINE", format: "ARROW_STREAM", }); - // SSE payload carries an inline- id pointing at the overflow entry. - const writeCalls = (mockRes.write as any).mock.calls.map( - (c: any[]) => c[0] as string, - ); - // Skip the leading warehouse_status event the route always emits; - // the terminal result/arrow/error payload is the one under test. - const payload = writeCalls.find( - (s: string) => - s.startsWith("data: ") && !s.includes("warehouse_status"), + // Bytes stream straight back on the response body with the Arrow + // content type — no SSE frames, no stash, no JSON error. + expect(mockRes.setHeader).toHaveBeenCalledWith( + "Content-Type", + "application/vnd.apache.arrow.stream", ); - expect(payload).toBeDefined(); - expect(payload).toContain('"type":"arrow"'); - expect(payload).toMatch(/"statement_id":"inline-[^"]+"/); - - // Bytes landed in the overflow pool, regular pool size unchanged. - expect(tinyStash.overflowSize()).toBe(3); - expect(tinyStash.size()).toBe(4); - }); - - test("/query/:query_key surfaces a stable error when both regular and overflow pools are exhausted — never silently double-bills", async () => { - // When even the overflow pool cannot fit the payload, the route - // surfaces INLINE_ARROW_STASH_EXHAUSTED instead of re-running the - // statement on EXTERNAL_LINKS. The previous behavior (silent retry) - // double-billed the warehouse and could return divergent results. - const plugin = new AnalyticsPlugin(config); - const { router, getHandler } = createMockRouter(); - - (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ - query: "SELECT * FROM test", - isAsUser: false, - }); - - const fakeAttachment = Buffer.from(new Uint8Array([1, 2, 3])).toString( - "base64", + // Manifest column names ride a response header so the client can relabel + // the positional Arrow schema. + expect(mockRes.setHeader).toHaveBeenCalledWith( + "X-Appkit-Arrow-Columns", + encodeURIComponent(JSON.stringify(["name", "totalSpend"])), ); - const executeMock = vi.fn().mockResolvedValueOnce({ - result: { attachment: fakeAttachment, row_count: 1 }, - }); - (plugin as any).SQLClient.executeStatement = executeMock; - - // Force both put paths to refuse: spy returns null unconditionally. - vi.spyOn((plugin as any).inlineArrowStash, "put").mockReturnValue(null); - - plugin.injectRoutes(router); - - const handler = getHandler("POST", "/query/:query_key"); - const mockReq = createMockRequest({ - params: { query_key: "test_query" }, - body: { parameters: {}, format: "ARROW_STREAM" }, - }); - const mockRes = createMockResponse(); - - await handler(mockReq, mockRes); - - // Single execution: never re-runs on EXTERNAL_LINKS. - expect(executeMock).toHaveBeenCalledTimes(1); + expect(mockRes.json).not.toHaveBeenCalled(); + expect(mockRes.end).toHaveBeenCalled(); - // SSE error payload carries the stable errorCode for UI branching. - // The writer emits separate lines (`id:`, `event: error`, `data: ...`) - // — join them so we can match across the framing. - const writeCalls = (mockRes.write as any).mock.calls.map( - (c: any[]) => c[0] as string, + const written = (mockRes.write as any).mock.calls.map( + (c: any[]) => c[0] as Buffer, ); - const errorFrame = writeCalls - .filter((s: string) => s.startsWith("data: ")) - .join("\n"); - expect(errorFrame).toContain("INLINE_ARROW_STASH_EXHAUSTED"); - // Client-safe message reaches the wire; raw upstream text does not. - expect(errorFrame).toContain("Server is at capacity"); - expect(errorFrame).not.toContain("Inline Arrow stash exhausted"); + expect(written).toHaveLength(1); + expect(Buffer.isBuffer(written[0])).toBe(true); + expect(Array.from(written[0])).toEqual(Array.from(arrowBytes)); }); test("/query/:query_key falls back JSON_ARRAY to ARROW_STREAM INLINE when warehouse refuses JSON_ARRAY for INLINE", async () => { @@ -1687,4 +1699,74 @@ describe("Analytics Plugin", () => { expect(Object.keys(entries)).toEqual(["query"]); }); }); + + describe("writeChunk backpressure", () => { + // A response stub that is always backpressured (`write` → false) and lets + // the test fire lifecycle events. The shared `createMockResponse` returns a + // truthy `write`, so it never exercises the backpressure branch. + function backpressuredRes() { + const listeners: Record void>> = {}; + return { + write: vi.fn(() => false), + once: vi.fn(function (this: unknown, ev: string, h: () => void) { + listeners[ev] ??= []; + listeners[ev].push(h); + return this; + }), + off: vi.fn(function (this: unknown, ev: string, h: () => void) { + if (listeners[ev]) { + listeners[ev] = listeners[ev].filter((x) => x !== h); + } + return this; + }), + emit(ev: string) { + for (const h of [...(listeners[ev] ?? [])]) h(); + }, + listenerCount(ev: string) { + return (listeners[ev] ?? []).length; + }, + }; + } + + test("resolves once the socket drains and detaches its listeners", async () => { + const res = backpressuredRes(); + const pending = writeChunk(res as never, new Uint8Array([1, 2, 3])); + res.emit("drain"); + await expect(pending).resolves.toBeUndefined(); + expect(res.listenerCount("drain")).toBe(0); + expect(res.listenerCount("close")).toBe(0); + expect(res.listenerCount("error")).toBe(0); + }); + + test("rejects (does not hang) when the client disconnects mid-backpressure", async () => { + const res = backpressuredRes(); + const pending = writeChunk(res as never, new Uint8Array([1, 2, 3])); + // The socket closes while backpressured — `drain` will never fire. If + // writeChunk only awaited `drain`, this promise (and the stream feeding + // it) would hang forever. + res.emit("close"); + await expect(pending).rejects.toThrow(); + expect(res.listenerCount("drain")).toBe(0); + expect(res.listenerCount("close")).toBe(0); + expect(res.listenerCount("error")).toBe(0); + }); + + test("rejects up front if the socket is already closed when called", async () => { + // Socket already destroyed before writeChunk runs: `close`/`error` have + // already fired, so attaching listeners would never settle. The guard + // must reject without writing or waiting. + const res = { + destroyed: true, + writableEnded: false, + write: vi.fn(() => false), + once: vi.fn(), + off: vi.fn(), + }; + await expect( + writeChunk(res as never, new Uint8Array([1, 2, 3])), + ).rejects.toThrow(); + expect(res.write).not.toHaveBeenCalled(); + expect(res.once).not.toHaveBeenCalled(); + }); + }); }); diff --git a/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts b/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts new file mode 100644 index 00000000..83e37b4a --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts @@ -0,0 +1,68 @@ +import { WorkspaceClient } from "@databricks/sdk-experimental"; +import { tableFromIPC } from "apache-arrow"; +import { describe, expect, test } from "vitest"; +import { SQLWarehouseConnector } from "../../../connectors"; +import { deliverArrowBytes, type QueryExecutor } from "../result-delivery"; + +/** + * Gated live integration test for the Arrow delivery fallback against a REAL + * warehouse. Skipped unless `APPKIT_INTEGRATION_WAREHOUSE_ID` is set (with + * `DATABRICKS_HOST`/`DATABRICKS_CONFIG_PROFILE` auth in the environment), so it + * never runs in CI. Run locally, e.g.: + * + * DATABRICKS_CONFIG_PROFILE=dogfood \ + * APPKIT_INTEGRATION_WAREHOUSE_ID=86e2ab2c5d30b12a \ + * pnpm exec vitest run arrow-delivery.integration + * + * Verifies the full capability fallback end-to-end: on a normal warehouse the + * INLINE attempt is rejected and EXTERNAL_LINKS chunks stream through; on + * Reyden the inline attachment streams. Column names are resolved from the + * manifest either way. + * + * NOTE: this drives the service-principal identity. Full `.obo.sql` (user) + * validation requires a user token / `x-forwarded-user` request and is + * verified manually; the wiring is covered by the mocked OBO test in + * `analytics.test.ts`. + */ +const warehouseId = process.env.APPKIT_INTEGRATION_WAREHOUSE_ID; + +describe.runIf(!!warehouseId)("arrow delivery (live warehouse)", () => { + test("ARROW_STREAM delivers valid Arrow with real column names", async () => { + const client = new WorkspaceClient({}); + const connector = new SQLWarehouseConnector({ timeout: 120_000 }); + + const executor: QueryExecutor = { + query: async (statement, _params, fp, signal) => { + const res = (await connector.executeStatement( + client, + { + statement, + warehouse_id: warehouseId as string, + wait_timeout: "50s", + on_wait_timeout: "CONTINUE", + disposition: fp.disposition as never, + format: fp.format as never, + }, + signal, + )) as { result?: unknown }; + return res.result as never; + }, + }; + + const out: { columnNames?: string[]; statementId?: string } = {}; + const chunks: Buffer[] = []; + for await (const bytes of deliverArrowBytes( + executor, + connector, + "SELECT id AS my_id, cast(id AS string) AS my_name FROM range(1000)", + undefined, + out, + )) { + chunks.push(Buffer.from(bytes)); + } + + const table = tableFromIPC(new Uint8Array(Buffer.concat(chunks))); + expect(table.numRows).toBe(1000); + expect(out.columnNames).toEqual(["my_id", "my_name"]); + }, 180_000); +}); diff --git a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts b/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts deleted file mode 100644 index 05ae821b..00000000 --- a/packages/appkit/src/plugins/analytics/tests/inline-arrow-stash.test.ts +++ /dev/null @@ -1,184 +0,0 @@ -import { describe, expect, test } from "vitest"; -import { InlineArrowStash } from "../inline-arrow-stash"; - -function bytes(n: number): Uint8Array { - return new Uint8Array(n); -} - -// `put()` returns `{ id, pool } | null` — it rejects with null when the stash -// is full. Most tests only care about the id; this helper narrows via the -// non-null contract and returns just the id. Tests that need the pool tag -// call `stash.put(...)` directly. -function mustPut( - stash: InlineArrowStash, - userId: string, - b: Uint8Array, -): string { - const result = stash.put(userId, b); - if (result === null) { - throw new Error("test setup: stash unexpectedly rejected put"); - } - return result.id; -} - -describe("InlineArrowStash", () => { - test("put returns an inline-prefixed synthetic id", () => { - const stash = new InlineArrowStash({ idGenerator: () => "abc" }); - const id = mustPut(stash, "user-1", bytes(100)); - expect(id).toBe("inline-abc"); - }); - - test("take drains the entry", () => { - const stash = new InlineArrowStash(); - const id = mustPut(stash, "user-1", bytes(100)); - expect(stash.count()).toBe(1); - expect(stash.size()).toBe(100); - - const got = stash.take(id, "user-1"); - expect(got).toBeDefined(); - expect(got?.length).toBe(100); - expect(stash.count()).toBe(0); - expect(stash.size()).toBe(0); - // Drain-on-read: second take returns undefined. - expect(stash.take(id, "user-1")).toBeUndefined(); - }); - - test("take returns undefined for unknown id", () => { - const stash = new InlineArrowStash(); - expect(stash.take("inline-nope", "user-1")).toBeUndefined(); - }); - - test("take returns undefined when userId does not match", () => { - const stash = new InlineArrowStash(); - const id = mustPut(stash, "user-1", bytes(100)); - expect(stash.take(id, "user-2")).toBeUndefined(); - // Entry is still there for the right user. - expect(stash.take(id, "user-1")).toBeDefined(); - }); - - test("entries past TTL are evicted on next gc tick", () => { - let clock = 0; - // gcMinIntervalMs: 0 disables gc throttling so the test can drive - // the clock past TTL on a sub-throttle scale without the gc pass - // being skipped. - const stash = new InlineArrowStash({ - ttlMs: 1000, - gcMinIntervalMs: 0, - now: () => clock, - }); - const id = mustPut(stash, "user-1", bytes(50)); - clock = 999; - expect(stash.take(id, "user-1")).toBeDefined(); - - const id2 = mustPut(stash, "user-1", bytes(50)); - clock = 2000; - // Bump the clock past TTL and trigger gc via another put. - mustPut(stash, "user-2", bytes(10)); - expect(stash.take(id2, "user-1")).toBeUndefined(); - }); - - test("put spills into the overflow pool when the regular pool is at cap — every issued id remains valid", () => { - let seq = 0; - const stash = new InlineArrowStash({ - maxBytes: 200, - maxOverflowBytes: 200, - idGenerator: () => String(seq++), - }); - const a = mustPut(stash, "user-1", bytes(80)); - const b = mustPut(stash, "user-1", bytes(80)); - expect(stash.size()).toBe(160); - expect(stash.overflowSize()).toBe(0); - - // The third 80-byte entry would push the regular pool to 240 (>200), - // so it spills into overflow. Both prior entries must survive. - const c = mustPut(stash, "user-1", bytes(80)); - expect(stash.size()).toBe(160); - expect(stash.overflowSize()).toBe(80); - expect(stash.take(a, "user-1")).toBeDefined(); - expect(stash.take(b, "user-1")).toBeDefined(); - expect(stash.take(c, "user-1")).toBeDefined(); - // After draining the overflow entry, the counter reflects it. - expect(stash.overflowSize()).toBe(0); - }); - - test("put returns null only when both regular and overflow pools are full", () => { - let seq = 0; - const stash = new InlineArrowStash({ - maxBytes: 100, - maxOverflowBytes: 100, - idGenerator: () => String(seq++), - }); - const a = mustPut(stash, "user-1", bytes(100)); // fills regular - const b = mustPut(stash, "user-1", bytes(100)); // fills overflow - expect(stash.size()).toBe(100); - expect(stash.overflowSize()).toBe(100); - - // Both pools at cap — refuse rather than evict. - const c = stash.put("user-1", bytes(50)); - expect(c).toBeNull(); - expect(stash.take(a, "user-1")).toBeDefined(); - expect(stash.take(b, "user-1")).toBeDefined(); - }); - - test("put tags the result with its pool so callers can label telemetry without re-introspecting", () => { - const stash = new InlineArrowStash({ - maxBytes: 100, - maxOverflowBytes: 100, - }); - expect(stash.put("u", bytes(80))).toMatchObject({ pool: "regular" }); - // Regular has 80/100 used; another 80 would push it to 160 > 100 so it - // spills. - expect(stash.put("u", bytes(80))).toMatchObject({ pool: "overflow" }); - }); - - test("put throws when a single payload would not fit in the largest pool (caller misconfiguration)", () => { - const stash = new InlineArrowStash({ - maxBytes: 100, - maxOverflowBytes: 100, - }); - // 300 > max(maxBytes, maxOverflowBytes) (100); pools don't split, - // so no individual put can ever succeed. - expect(() => stash.put("user-1", bytes(300))).toThrow( - /exceeds largest stash slot/, - ); - }); - - test("overflow entries expire on a shorter TTL than regular entries", () => { - let clock = 0; - const stash = new InlineArrowStash({ - maxBytes: 80, - maxOverflowBytes: 80, - ttlMs: 600_000, - overflowTtlMs: 30_000, - gcMinIntervalMs: 0, - now: () => clock, - }); - const reg = mustPut(stash, "user-1", bytes(80)); // fills regular - const ovf = mustPut(stash, "user-1", bytes(80)); // spills to overflow - expect(stash.size()).toBe(80); - expect(stash.overflowSize()).toBe(80); - - // 45 s in: overflow expired, regular still alive. - clock = 45_000; - expect(stash.take(ovf, "user-1")).toBeUndefined(); - expect(stash.take(reg, "user-1")).toBeDefined(); - }); - - test("synthetic ids are unique across puts", () => { - const stash = new InlineArrowStash(); - const a = mustPut(stash, "user-1", bytes(10)); - const b = mustPut(stash, "user-1", bytes(10)); - expect(a).not.toBe(b); - expect(a.startsWith("inline-")).toBe(true); - expect(b.startsWith("inline-")).toBe(true); - }); - - test("clear drops every entry", () => { - const stash = new InlineArrowStash(); - mustPut(stash, "user-1", bytes(10)); - mustPut(stash, "user-2", bytes(20)); - stash.clear(); - expect(stash.count()).toBe(0); - expect(stash.size()).toBe(0); - }); -}); diff --git a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts new file mode 100644 index 00000000..758e9011 --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts @@ -0,0 +1,266 @@ +import { Table, tableToIPC, Utf8, vectorFromArray } from "apache-arrow"; +import { describe, expect, test, vi } from "vitest"; +import { ExecutionError } from "../../../errors"; +import { + type ArrowChunkStreamer, + arrowDeliveryUnsupported, + classifyDispositionRejection, + decodeArrowAttachmentToRows, + deliverArrowBytes, + deliverJsonResult, + type QueryExecutor, +} from "../result-delivery"; + +/** Build a base64 Arrow IPC stream (2 cols, given names + rows). */ +function arrowBase64(names: [string, string]): string { + const table = new Table({ + [names[0]]: vectorFromArray(["r0", "r1"], new Utf8()), + [names[1]]: vectorFromArray(["v0", "v1"], new Utf8()), + }); + return Buffer.from(tableToIPC(table, "stream")).toString("base64"); +} + +/** Fake executor whose `query` is driven by the requested disposition/format. */ +function executorFrom( + handler: (fp: { disposition: string; format: string }) => Promise, +): { + executor: QueryExecutor; + calls: { disposition: string; format: string }[]; +} { + const calls: { disposition: string; format: string }[] = []; + const executor: QueryExecutor = { + query: (_q, _p, fp) => { + calls.push(fp); + return handler(fp) as ReturnType; + }, + }; + return { executor, calls }; +} + +async function collect(gen: AsyncGenerator): Promise { + const out: Uint8Array[] = []; + for await (const b of gen) out.push(b); + return out; +} + +const reject = (code: string, message: string) => + ExecutionError.statementFailed(message, code); + +describe("classifyDispositionRejection", () => { + test("normal warehouse rejecting ARROW+INLINE → needs-json-inline", () => { + expect( + classifyDispositionRejection( + reject( + "INVALID_PARAMETER_VALUE", + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + ), + ), + ).toBe("needs-json-inline"); + }); + + test("warehouse requiring ARROW for INLINE → needs-arrow-inline", () => { + expect( + classifyDispositionRejection( + reject( + "INVALID_PARAMETER_VALUE", + "Inline disposition only supports ARROW_STREAM format.", + ), + ), + ).toBe("needs-arrow-inline"); + }); + + test("EXTERNAL_LINKS not implemented (Reyden) → external-links-unsupported", () => { + expect( + classifyDispositionRejection( + reject( + "NOT_IMPLEMENTED", + "ExternalLinks disposition is not yet implemented.", + ), + ), + ).toBe("external-links-unsupported"); + }); + + test("permission error does NOT fall back", () => { + expect( + classifyDispositionRejection( + reject("PERMISSION_DENIED", "User does not have SELECT on table t"), + ), + ).toBeNull(); + }); + + test("plain SQL error does NOT fall back", () => { + expect( + classifyDispositionRejection(new Error("Table or view not found: foo")), + ).toBeNull(); + }); + + test("capability code without an inline/external signal does NOT fall back", () => { + expect( + classifyDispositionRejection( + reject("INVALID_PARAMETER_VALUE", "Some unrelated parameter problem"), + ), + ).toBeNull(); + }); +}); + +describe("deliverArrowBytes — Reyden (ARROW + INLINE)", () => { + test("streams the inline attachment and never touches external links", async () => { + const b64 = arrowBase64(["col_0", "col_1"]); + const { executor, calls } = executorFrom(async () => ({ + attachment: b64, + columnNames: ["name", "spend"], + statement_id: "stmt-inline", + })); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () {}), + }; + const out: { columnNames?: string[]; statementId?: string } = {}; + + const chunks = await collect( + deliverArrowBytes(executor, streamer, "SELECT 1", undefined, out), + ); + + expect(calls).toEqual([{ disposition: "INLINE", format: "ARROW_STREAM" }]); + expect(streamer.streamExternalLinks).not.toHaveBeenCalled(); + expect(out.columnNames).toEqual(["name", "spend"]); + expect(out.statementId).toBe("stmt-inline"); + // Bytes are the decoded attachment. + expect( + Buffer.concat(chunks.map((c) => Buffer.from(c))).toString("base64"), + ).toBe(b64); + }); +}); + +describe("deliverArrowBytes — normal warehouse (fallback to EXTERNAL_LINKS)", () => { + test("INLINE rejected → streams the external links from the execute response", async () => { + const links = [{ external_link: "https://example.com/chunk-0" }]; + const { executor, calls } = executorFrom(async (fp) => { + if (fp.disposition === "INLINE") { + throw reject( + "INVALID_PARAMETER_VALUE", + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + ); + } + return { + external_links: links, + columnNames: ["a", "b"], + statement_id: "stmt-ext", + }; + }); + const extBytes = new Uint8Array([9, 8, 7]); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () { + yield extBytes; + }), + }; + const out: { columnNames?: string[]; statementId?: string } = {}; + + const chunks = await collect( + deliverArrowBytes(executor, streamer, "SELECT 1", undefined, out), + ); + + expect(calls).toEqual([ + { disposition: "INLINE", format: "ARROW_STREAM" }, + { disposition: "EXTERNAL_LINKS", format: "ARROW_STREAM" }, + ]); + // Streams the links resolved in-context (no re-fetch): passed straight through. + expect(streamer.streamExternalLinks).toHaveBeenCalledWith(links, undefined); + expect(out.columnNames).toEqual(["a", "b"]); + expect(chunks).toEqual([extBytes]); + }); + + test("both INLINE and EXTERNAL_LINKS unsupported → ARROW_DELIVERY_UNSUPPORTED", async () => { + const { executor } = executorFrom(async (fp) => { + if (fp.disposition === "INLINE") { + throw reject( + "INVALID_PARAMETER_VALUE", + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + ); + } + throw reject( + "NOT_IMPLEMENTED", + "ExternalLinks disposition is not yet implemented.", + ); + }); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () {}), + }; + + await expect( + collect(deliverArrowBytes(executor, streamer, "SELECT 1", undefined, {})), + ).rejects.toMatchObject({ errorCode: "ARROW_DELIVERY_UNSUPPORTED" }); + }); + + test("auth/permission error on INLINE propagates without a fallback", async () => { + const authErr = reject("PERMISSION_DENIED", "no access"); + const { executor, calls } = executorFrom(async () => { + throw authErr; + }); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: vi.fn(async function* () {}), + }; + + await expect( + collect(deliverArrowBytes(executor, streamer, "SELECT 1", undefined, {})), + ).rejects.toBe(authErr); + // Only the INLINE attempt — no EXTERNAL_LINKS fallback. + expect(calls).toEqual([{ disposition: "INLINE", format: "ARROW_STREAM" }]); + }); +}); + +describe("deliverJsonResult", () => { + test("INLINE JSON_ARRAY success returns rows directly", async () => { + const { executor, calls } = executorFrom(async () => ({ + data: [{ id: 1 }], + status: { state: "SUCCEEDED" }, + statement_id: "s1", + })); + const result = await deliverJsonResult(executor, "SELECT 1", undefined); + expect(result.data).toEqual([{ id: 1 }]); + expect(calls).toEqual([{ disposition: "INLINE", format: "JSON_ARRAY" }]); + }); + + test("needs-arrow-inline → retries as ARROW and decodes rows with manifest names", async () => { + const b64 = arrowBase64(["col_0", "col_1"]); + const { executor, calls } = executorFrom(async (fp) => { + if (fp.format === "JSON_ARRAY") { + throw reject( + "INVALID_PARAMETER_VALUE", + "Inline disposition only supports ARROW_STREAM format.", + ); + } + return { attachment: b64, columnNames: ["name", "spend"] }; + }); + const result = await deliverJsonResult(executor, "SELECT 1", undefined); + expect(calls).toEqual([ + { disposition: "INLINE", format: "JSON_ARRAY" }, + { disposition: "INLINE", format: "ARROW_STREAM" }, + ]); + // Rows carry the manifest names, not the positional col_N names. + expect(result.data).toEqual([ + { name: "r0", spend: "v0" }, + { name: "r1", spend: "v1" }, + ]); + }); +}); + +describe("decodeArrowAttachmentToRows", () => { + test("relabels positional columns from manifest names", () => { + const rows = decodeArrowAttachmentToRows(arrowBase64(["col_0", "col_1"]), [ + "name", + "spend", + ]); + expect(rows).toEqual([ + { name: "r0", spend: "v0" }, + { name: "r1", spend: "v1" }, + ]); + }); +}); + +describe("arrowDeliveryUnsupported", () => { + test("is a structured, client-actionable error", () => { + const err = arrowDeliveryUnsupported(); + expect(err.errorCode).toBe("ARROW_DELIVERY_UNSUPPORTED"); + expect(err.clientMessage).toMatch(/JSON_ARRAY/); + }); +}); diff --git a/packages/appkit/src/plugins/analytics/types.ts b/packages/appkit/src/plugins/analytics/types.ts index 9e3dd1d2..95c987d1 100644 --- a/packages/appkit/src/plugins/analytics/types.ts +++ b/packages/appkit/src/plugins/analytics/types.ts @@ -15,6 +15,14 @@ export interface IAnalyticsConfig extends BasePluginConfig { * `ConfigurationError`. */ autoStartWarehouse?: boolean; + /** + * Fail-fast ceiling (ms) for an `ARROW_STREAM` query to produce its first + * byte (warehouse readiness + execute + first chunk). Past this, a stuck or + * overloaded warehouse returns a `503` (`WAREHOUSE_UNAVAILABLE`) instead of + * hanging until the client disconnects. Defaults to 2 min. Once the first + * byte arrives the stream is not time-bounded. + */ + arrowFirstByteTimeoutMs?: number; } /** diff --git a/packages/appkit/src/stream/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index d3118e7e..34752018 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -4,243 +4,190 @@ import { createLogger } from "../logging/logger"; const logger = createLogger("stream:arrow"); -type ResultManifest = sql.ResultManifest; type ExternalLink = sql.ExternalLink; interface ArrowStreamOptions { - maxConcurrentDownloads: number; + /** Idle timeout (ms) for a single chunk: no progress → abort the download. */ timeout: number; + /** Attempts to establish a chunk's response before any bytes are yielded. */ retries: number; } -/** - * Result from zero-copy Arrow chunk processing. - * Contains raw IPC bytes without server-side parsing. - */ -interface ArrowRawResult { - /** Concatenated raw Arrow IPC bytes */ - data: Uint8Array; - /** Schema from Databricks manifest (not parsed from Arrow) */ - schema: ResultManifest["schema"]; -} - const BACKOFF_MULTIPLIER = 1000; +/** + * Streams Arrow IPC bytes for a completed statement's EXTERNAL_LINKS chunks + * without ever buffering a whole chunk — each chunk's response body is piped + * through as it arrives, so peak memory is a single network read rather than + * a full chunk (let alone the full result). No Arrow parsing on the server; + * the client parses the concatenated IPC bytes. + */ export class ArrowStreamProcessor { - static readonly DEFAULT_MAX_CONCURRENT_DOWNLOADS = 5; static readonly DEFAULT_TIMEOUT = 30000; static readonly DEFAULT_RETRIES = 3; - constructor( - private options: ArrowStreamOptions = { - maxConcurrentDownloads: - ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS, - timeout: ArrowStreamProcessor.DEFAULT_TIMEOUT, - retries: ArrowStreamProcessor.DEFAULT_RETRIES, - }, - ) { + private options: ArrowStreamOptions; + + constructor(options?: Partial) { this.options = { - maxConcurrentDownloads: - options.maxConcurrentDownloads ?? - ArrowStreamProcessor.DEFAULT_MAX_CONCURRENT_DOWNLOADS, - timeout: options.timeout ?? ArrowStreamProcessor.DEFAULT_TIMEOUT, - retries: options.retries ?? ArrowStreamProcessor.DEFAULT_RETRIES, + timeout: options?.timeout ?? ArrowStreamProcessor.DEFAULT_TIMEOUT, + retries: options?.retries ?? ArrowStreamProcessor.DEFAULT_RETRIES, }; } /** - * Process Arrow chunks using zero-copy proxy pattern. - * - * Downloads raw IPC bytes from external links and concatenates them - * without parsing into Arrow Tables on the server. This reduces: - * - Memory usage by ~50% (no parsed Table representation) - * - CPU usage (no tableFromIPC/tableToIPC calls) - * - * The client is responsible for parsing the IPC bytes. + * Stream Arrow chunks in array order, piping each chunk's response body. * - * @param chunks - External links to Arrow IPC data - * @param schema - Schema from Databricks manifest - * @param signal - Optional abort signal - * @returns Raw concatenated IPC bytes with schema + * Yields network-sized pieces as they arrive (not whole chunks), so peak + * memory is one read buffer. Chunks are fetched sequentially; the + * concatenation of everything yielded is byte-identical to the raw Arrow + * result, so the client parses it exactly as a buffered response. */ - async processChunks( + async *streamChunks( chunks: ExternalLink[], - schema: ResultManifest["schema"], signal?: AbortSignal, - ): Promise { + ): AsyncGenerator { if (chunks.length === 0) { throw ValidationError.missingField("chunks"); } - const buffers = await this.downloadChunksRaw(chunks, signal); - const data = this.concatenateBuffers(buffers); - - return { data, schema }; - } - - /** - * Download all chunks as raw bytes with concurrency control. - */ - private async downloadChunksRaw( - chunks: ExternalLink[], - signal?: AbortSignal, - ): Promise { - const semaphore = new Semaphore(this.options.maxConcurrentDownloads); - - const downloadPromises = chunks.map(async (chunk) => { - await semaphore.acquire(); - try { - return await this.downloadChunkRaw(chunk, signal); - } finally { - semaphore.release(); - } - }); - - return Promise.all(downloadPromises); + for (const chunk of chunks) { + yield* this.streamChunkBody(chunk, signal); + } } /** - * Download a single chunk as raw bytes with retry logic. + * Pipe one chunk's response body. + * + * Retry applies only while *establishing* the response (connection failure + * or non-2xx) — once bytes have been yielded downstream we cannot re-fetch, + * so a mid-body failure propagates and the caller aborts the response. An + * idle timeout, reset before every read, aborts a stalled download. */ - private async downloadChunkRaw( + private async *streamChunkBody( chunk: ExternalLink, signal?: AbortSignal, - ): Promise { - let lastError: Error | null = null; + ): AsyncGenerator { + const externalLink = chunk.external_link; + if (!externalLink) { + // A missing link cannot be fixed by retrying — fail loudly. + throw ExecutionError.statementFailed( + `External link missing for chunk ${chunk.chunk_index}`, + ); + } + + let response: Response | undefined; + let controller: AbortController | undefined; + let onOuterAbort: (() => void) | undefined; + let lastError: unknown; for (let attempt = 0; attempt < this.options.retries; attempt++) { - const timeoutController = new AbortController(); - const timeoutId = setTimeout(() => { - timeoutController.abort(); - }, this.options.timeout); + if (signal?.aborted) throw ExecutionError.canceled(); - const combinedSignal = signal - ? this.combineAbortSignals(signal, timeoutController.signal) - : timeoutController.signal; + const attemptController = new AbortController(); + const listener = () => attemptController.abort(); + signal?.addEventListener("abort", listener, { once: true }); + const timer = setTimeout( + () => attemptController.abort(), + this.options.timeout, + ); try { - const externalLink = chunk.external_link; - if (!externalLink) { - logger.error("External link is required for chunk: %O", chunk); - continue; - } - - const response = await fetch(externalLink, { - signal: combinedSignal, + const r = await fetch(externalLink, { + signal: attemptController.signal, }); - - if (!response.ok) { + clearTimeout(timer); + if (!r.ok) { throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunk_index}: ${response.status} ${response.statusText}`, + `Failed to download chunk ${chunk.chunk_index}: ${r.status} ${r.statusText}`, ); } - - const arrayBuffer = await response.arrayBuffer(); - return new Uint8Array(arrayBuffer); + // Keep this attempt's controller alive to drive the body read + idle + // timeout below. + response = r; + controller = attemptController; + onOuterAbort = listener; + break; } catch (error) { - lastError = error as Error; - - if (timeoutController.signal.aborted) { - lastError = new Error( - `Chunk ${chunk.chunk_index} download timed out after ${this.options.timeout}ms`, - ); - } - - if (signal?.aborted) { - throw ExecutionError.canceled(); - } - + clearTimeout(timer); + signal?.removeEventListener("abort", listener); + lastError = error; + if (signal?.aborted) throw ExecutionError.canceled(); if (attempt < this.options.retries - 1) { - await this.delay(2 ** attempt * BACKOFF_MULTIPLIER); + await this.delay(2 ** attempt * BACKOFF_MULTIPLIER, signal); } - } finally { - clearTimeout(timeoutId); } } - throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunk_index} after ${this.options.retries} attempts: ${lastError?.message}`, - ); - } - - /** - * Concatenate multiple Uint8Array buffers into a single buffer. - * Pre-allocates the result array for efficiency. - */ - private concatenateBuffers(buffers: Uint8Array[]): Uint8Array { - if (buffers.length === 0) { - throw ValidationError.missingField("buffers"); - } - - if (buffers.length === 1) { - return buffers[0]; - } - - const totalLength = buffers.reduce((sum, buf) => sum + buf.length, 0); - const result = new Uint8Array(totalLength); - - let offset = 0; - for (const buffer of buffers) { - result.set(buffer, offset); - offset += buffer.length; + if (!response || !controller) { + throw ExecutionError.statementFailed( + `Failed to download chunk ${chunk.chunk_index} after ${this.options.retries} attempts: ${ + lastError instanceof Error ? lastError.message : String(lastError) + }`, + ); } - return result; - } - - /** - * Combines multiple AbortSignals into one. - * The combined signal aborts when any of the input signals abort. - */ - private combineAbortSignals(...signals: AbortSignal[]): AbortSignal { - const controller = new AbortController(); - - for (const signal of signals) { - if (signal.aborted) { - controller.abort(); - return controller.signal; + const body = response.body; + let idleTimer: ReturnType | undefined; + try { + if (!body) return; // empty chunk (0 rows) + const reader = body.getReader(); + try { + while (true) { + // The idle timeout must cover ONLY the upstream read, never the + // downstream `yield`. Arm it right before `reader.read()` and clear + // it the instant the read resolves — if it stayed armed across the + // yield, a slow client backpressuring `writeChunk` would look like an + // upstream stall and abort a perfectly healthy download. + idleTimer = setTimeout( + () => controller.abort(), + this.options.timeout, + ); + let done: boolean; + let value: Uint8Array | undefined; + try { + ({ done, value } = await reader.read()); + } finally { + clearTimeout(idleTimer); + } + if (done) break; + if (value && value.byteLength > 0) yield value; + } + } finally { + // Release the stream on early exit (downstream abort / error). + reader.cancel().catch(() => {}); } - signal.addEventListener("abort", () => controller.abort(), { - once: true, - }); + } catch (error) { + if (signal?.aborted) throw ExecutionError.canceled(); + logger.error( + "Failed streaming chunk %s body: %O", + chunk.chunk_index, + error, + ); + throw error instanceof ExecutionError + ? error + : ExecutionError.statementFailed( + `Failed streaming chunk ${chunk.chunk_index}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } finally { + clearTimeout(idleTimer); + if (onOuterAbort) signal?.removeEventListener("abort", onOuterAbort); } - - return controller.signal; } - private delay(ms: number): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); - } -} - -class Semaphore { - private permits: number; - private waiting: (() => void)[] = []; - - constructor(permits: number) { - this.permits = permits; - } - - async acquire(): Promise { - if (this.permits > 0) { - this.permits--; - return; - } - - return new Promise((resolve) => { - this.waiting.push(resolve); + private delay(ms: number, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + signal?.removeEventListener("abort", onAbort); + resolve(); + }, ms); + const onAbort = () => { + clearTimeout(timer); + reject(ExecutionError.canceled()); + }; + signal?.addEventListener("abort", onAbort, { once: true }); }); } - - release(): void { - if (this.waiting.length > 0) { - const next = this.waiting.shift(); - - if (next) { - next(); - } - } else { - this.permits++; - } - } } diff --git a/packages/appkit/src/stream/defaults.ts b/packages/appkit/src/stream/defaults.ts index 706af342..c8863b4f 100644 --- a/packages/appkit/src/stream/defaults.ts +++ b/packages/appkit/src/stream/defaults.ts @@ -1,12 +1,9 @@ export const streamDefaults = { bufferSize: 100, - // 1 MiB. SSE is used only for short JSON control messages — JSON_ARRAY - // result rows (already row-size-bounded by the warehouse) and the small - // `arrow` envelope (statement id + status) for ARROW_STREAM. Bulk Arrow - // payloads do not traverse SSE; they are fetched over HTTP via - // `/api/analytics/arrow-result/:jobId`, which dispatches to the warehouse - // (EXTERNAL_LINKS) or the server-side `InlineArrowStash` (INLINE) based - // on the id prefix. + // 1 MiB. SSE carries only short JSON control messages — JSON_ARRAY result + // rows (already row-size-bounded by the warehouse) plus warehouse-readiness + // and error events. ARROW_STREAM never uses SSE: the raw Arrow IPC bytes + // stream back on the query response body (`_handleArrowStreamQuery`). maxEventSize: 1 * 1024 * 1024, bufferTTL: 10 * 60 * 1000, // 10 minutes cleanupInterval: 5 * 60 * 1000, // 5 minutes diff --git a/packages/appkit/src/stream/stream-manager.ts b/packages/appkit/src/stream/stream-manager.ts index e83fa5ea..cf28ead5 100644 --- a/packages/appkit/src/stream/stream-manager.ts +++ b/packages/appkit/src/stream/stream-manager.ts @@ -309,7 +309,7 @@ export class StreamManager { error instanceof AppKitError ? error.clientMessage : "Internal server error"; - // Upstream structured code (e.g. INLINE_ARROW_STASH_EXHAUSTED, + // Upstream structured code (e.g. RESULT_TOO_LARGE_FOR_JSON_FALLBACK, // NOT_IMPLEMENTED). UI should branch on this, not on `error`. const upstreamCode = error instanceof ExecutionError ? error.errorCode : undefined; diff --git a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts index 8a95fbdf..68b10e3f 100644 --- a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts +++ b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts @@ -1,11 +1,17 @@ -import type { sql } from "@databricks/sdk-experimental"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { ArrowStreamProcessor } from "../arrow-stream-processor"; -type ResultSchema = sql.ResultManifest["schema"]; +/** A ReadableStream that emits the given pieces in order, then closes. */ +function streamOf(...pieces: Uint8Array[]): ReadableStream { + return new ReadableStream({ + start(controller) { + for (const piece of pieces) controller.enqueue(piece); + controller.close(); + }, + }); +} -// Helper to create mock chunks -function createMockChunks(count: number) { +function mockChunks(count: number) { return Array.from({ length: count }, (_, i) => ({ chunk_index: i, external_link: `https://example.com/chunk-${i}`, @@ -14,46 +20,33 @@ function createMockChunks(count: number) { })); } -// Helper to create mock schema -function createMockSchema(): ResultSchema { - return { - columns: [ - { name: "id", type_name: "INT" }, - { name: "name", type_name: "STRING" }, - ], - } as ResultSchema; +/** Drain a byte generator to a flat number[]. */ +async function drain( + gen: AsyncGenerator, +): Promise { + const out: number[] = []; + for await (const piece of gen) out.push(...piece); + return out; } -describe("ArrowStreamProcessor", () => { +describe("ArrowStreamProcessor.streamChunks", () => { let processor: ArrowStreamProcessor; let originalFetch: typeof globalThis.fetch; beforeEach(() => { - processor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 3, - timeout: 5000, - retries: 3, - }); - + processor = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); originalFetch = globalThis.fetch; - - // Default mock: successful fetch returning raw bytes - globalThis.fetch = vi.fn().mockImplementation(async (url: string) => { - // Extract chunk index from URL to create unique data - const match = url.match(/chunk-(\d+)/); - const chunkIndex = match ? parseInt(match[1], 10) : 0; - - // Return raw bytes (simulating Arrow IPC data) + // Default: each chunk's body is a single piece [100 + index]. + globalThis.fetch = vi.fn(async (input: string | URL | Request) => { + const match = String(input).match(/chunk-(\d+)/); + const index = match ? Number.parseInt(match[1], 10) : 0; return { ok: true, status: 200, statusText: "OK", - arrayBuffer: async () => - new Uint8Array([chunkIndex + 1, chunkIndex + 2, chunkIndex + 3]) - .buffer, - }; + body: streamOf(new Uint8Array([100 + index])), + } as unknown as Response; }); - vi.clearAllMocks(); }); @@ -62,484 +55,146 @@ describe("ArrowStreamProcessor", () => { vi.restoreAllMocks(); }); - describe("constructor", () => { - test("should use default options when not provided", () => { - const defaultProcessor = new ArrowStreamProcessor(); - expect(defaultProcessor).toBeDefined(); - }); - - test("should accept custom options", () => { - const customProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 10, - timeout: 60000, - retries: 5, - }); - expect(customProcessor).toBeDefined(); - }); - - test("should use defaults for missing option properties", () => { - const partialProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 2, - } as any); - expect(partialProcessor).toBeDefined(); - }); - }); - - describe("processChunks", () => { - test("should throw error when no chunks provided", async () => { - await expect( - processor.processChunks([], createMockSchema()), - ).rejects.toThrow("Missing required field: chunks"); - }); - - test("should process single chunk successfully", async () => { - const chunks = createMockChunks(1); - const schema = createMockSchema(); - - const result = await processor.processChunks(chunks, schema); - - expect(result).toHaveProperty("data"); - expect(result).toHaveProperty("schema", schema); - expect(result.data).toBeInstanceOf(Uint8Array); - expect(globalThis.fetch).toHaveBeenCalledTimes(1); - }); - - test("should process multiple chunks successfully", async () => { - const chunks = createMockChunks(5); - const schema = createMockSchema(); - - const result = await processor.processChunks(chunks, schema); - - expect(result.data).toBeInstanceOf(Uint8Array); - expect(result.schema).toBe(schema); - expect(globalThis.fetch).toHaveBeenCalledTimes(5); - }); - - test("should concatenate raw bytes from multiple chunks", async () => { - const chunks = createMockChunks(3); - const schema = createMockSchema(); - - const result = await processor.processChunks(chunks, schema); - - // Each chunk returns 3 bytes: [chunkIndex+1, chunkIndex+2, chunkIndex+3] - // Chunk 0: [1, 2, 3], Chunk 1: [2, 3, 4], Chunk 2: [3, 4, 5] - expect(result.data).toEqual(new Uint8Array([1, 2, 3, 2, 3, 4, 3, 4, 5])); - }); - - test("should return single chunk data without modification", async () => { - const chunks = createMockChunks(1); - const schema = createMockSchema(); - - const result = await processor.processChunks(chunks, schema); - - // Single chunk returns [1, 2, 3] - expect(result.data).toEqual(new Uint8Array([1, 2, 3])); - }); - - test("should pass abort signal to fetch", async () => { - const chunks = createMockChunks(1); - const schema = createMockSchema(); - const abortController = new AbortController(); - - await processor.processChunks(chunks, schema, abortController.signal); - - expect(globalThis.fetch).toHaveBeenCalledWith( - expect.any(String), - expect.objectContaining({ - signal: expect.any(AbortSignal), - }), - ); - }); + test("throws when no chunks are provided", async () => { + await expect(drain(processor.streamChunks([]))).rejects.toThrow(); }); - describe("concurrent downloads", () => { - test("should limit concurrent downloads with semaphore", async () => { - const maxConcurrent = 2; - const limitedProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: maxConcurrent, - timeout: 5000, - retries: 1, - }); - - let currentConcurrent = 0; - let maxObservedConcurrent = 0; - - globalThis.fetch = vi.fn().mockImplementation(async () => { - currentConcurrent++; - maxObservedConcurrent = Math.max( - maxObservedConcurrent, - currentConcurrent, - ); - - // Simulate network delay - await new Promise((resolve) => setTimeout(resolve, 10)); - - currentConcurrent--; - - return { - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; - }); - - const chunks = createMockChunks(10); - await limitedProcessor.processChunks(chunks, createMockSchema()); - - expect(maxObservedConcurrent).toBeLessThanOrEqual(maxConcurrent); - expect(globalThis.fetch).toHaveBeenCalledTimes(10); - }); + test("streams chunks in array order", async () => { + const bytes = await drain(processor.streamChunks(mockChunks(3))); + expect(bytes).toEqual([100, 101, 102]); }); - describe("retry logic", () => { - test("should retry on fetch failure", async () => { - let attempts = 0; - globalThis.fetch = vi.fn().mockImplementation(async () => { - attempts++; - if (attempts < 3) { - throw new Error("Network error"); - } - return { - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; - }); - - const chunks = createMockChunks(1); - const result = await processor.processChunks(chunks, createMockSchema()); - - expect(attempts).toBe(3); - expect(result.data).toBeDefined(); - }); - - test("should throw after exhausting retries", async () => { - globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network error")); - - const chunks = createMockChunks(1); - - await expect( - processor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/Failed to download chunk 0 after 3 attempts/); - - expect(globalThis.fetch).toHaveBeenCalledTimes(3); - }); - - test("should retry on non-ok response", async () => { - let attempts = 0; - globalThis.fetch = vi.fn().mockImplementation(async () => { - attempts++; - if (attempts < 2) { - return { - ok: false, - status: 500, - statusText: "Internal Server Error", - }; - } - return { - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; - }); - - const chunks = createMockChunks(1); - const result = await processor.processChunks(chunks, createMockSchema()); - - expect(attempts).toBe(2); - expect(result.data).toBeDefined(); - }); - - test("should use exponential backoff between retries", async () => { - vi.useFakeTimers(); - - let attempts = 0; - globalThis.fetch = vi.fn().mockImplementation(async () => { - attempts++; - if (attempts < 3) { - throw new Error("Network error"); - } - return { + test("pipes a chunk body piece-by-piece (no whole-chunk buffering)", async () => { + globalThis.fetch = vi.fn( + async () => + ({ ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; - }); - - const chunks = createMockChunks(1); - const promise = processor.processChunks(chunks, createMockSchema()); - - // First attempt - immediate - await vi.advanceTimersByTimeAsync(0); - expect(attempts).toBe(1); - - // First retry after 1000ms (2^0 * 1000) - await vi.advanceTimersByTimeAsync(1000); - expect(attempts).toBe(2); - - // Second retry after 2000ms (2^1 * 1000) - await vi.advanceTimersByTimeAsync(2000); - expect(attempts).toBe(3); - - await promise; - vi.useRealTimers(); - }); - }); - - describe("timeout handling", () => { - test("should timeout slow requests", async () => { - const shortTimeoutProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 1, - timeout: 50, - retries: 1, - }); - - globalThis.fetch = vi.fn().mockImplementation( - (_url: string, options?: { signal?: AbortSignal }) => - new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - resolve({ - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }); - }, 5000); // Much longer than timeout - - // Listen for abort (from timeout) - options?.signal?.addEventListener("abort", () => { - clearTimeout(timeout); - reject(new DOMException("Aborted", "AbortError")); - }); - }), - ); - - const chunks = createMockChunks(1); - - // The processor should timeout and reject - await expect( - shortTimeoutProcessor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/timed out|Failed to download/); - }); - - test("should clear timeout after successful fetch", async () => { - const clearTimeoutSpy = vi.spyOn(globalThis, "clearTimeout"); - - const chunks = createMockChunks(1); - await processor.processChunks(chunks, createMockSchema()); - - expect(clearTimeoutSpy).toHaveBeenCalled(); - clearTimeoutSpy.mockRestore(); - }); + status: 200, + body: streamOf(new Uint8Array([1, 2]), new Uint8Array([3, 4, 5])), + }) as unknown as Response, + ); + // Collect per-yield pieces (not flattened) to prove the body streams in + // multiple reads rather than arriving as one buffer. + const pieces: number[][] = []; + for await (const piece of processor.streamChunks(mockChunks(1))) { + pieces.push([...piece]); + } + expect(pieces).toEqual([ + [1, 2], + [3, 4, 5], + ]); }); - describe("abort signal handling", () => { - test("should abort immediately if signal already aborted", async () => { - const abortController = new AbortController(); - abortController.abort(); - - // Mock fetch to check if it receives an aborted signal - globalThis.fetch = vi - .fn() - .mockImplementation( - async (_url: string, options?: { signal?: AbortSignal }) => { - // If signal is already aborted, throw - if (options?.signal?.aborted) { - throw new DOMException("Aborted", "AbortError"); - } - return { - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }; + test("does not abort a healthy download when the consumer stalls between reads (backpressure)", async () => { + // Body emits [1] immediately, then nothing until the test enqueues more. + // It errors if the attempt's abort signal fires — mirroring a real fetch + // being aborted. If the idle timer stayed armed across the downstream + // `yield`, a slow consumer would trip it and kill this healthy stream. + let bodyController!: ReadableStreamDefaultController; + globalThis.fetch = vi.fn( + async (_input: string | URL | Request, init?: RequestInit) => { + const body = new ReadableStream({ + start(c) { + bodyController = c; + c.enqueue(new Uint8Array([1])); }, + }); + init?.signal?.addEventListener("abort", () => + bodyController.error(new Error("idle-aborted")), ); + return { ok: true, status: 200, body } as unknown as Response; + }, + ); - const chunks = createMockChunks(1); - - await expect( - processor.processChunks( - chunks, - createMockSchema(), - abortController.signal, - ), - ).rejects.toThrow("Statement was canceled"); - }); - - test("should abort in-flight requests when signal fires", async () => { - const abortController = new AbortController(); - let fetchStarted = false; - - globalThis.fetch = vi - .fn() - .mockImplementation( - async (_url: string, options?: { signal?: AbortSignal }) => { - fetchStarted = true; - - // Simulate slow request that checks signal - return new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - resolve({ - ok: true, - arrayBuffer: async () => new Uint8Array([1]).buffer, - }); - }, 1000); + const p = new ArrowStreamProcessor({ timeout: 50, retries: 1 }); + const gen = p.streamChunks(mockChunks(1)); - options?.signal?.addEventListener("abort", () => { - clearTimeout(timeout); - reject(new DOMException("Aborted", "AbortError")); - }); - }); - }, - ); + const first = await gen.next(); + expect([...(first.value as Uint8Array)]).toEqual([1]); - const chunks = createMockChunks(1); - const promise = processor.processChunks( - chunks, - createMockSchema(), - abortController.signal, - ); + // Consumer stalls well past the 50ms idle timeout while suspended at yield. + await new Promise((r) => setTimeout(r, 120)); - // Wait for fetch to start, then abort - await vi.waitFor(() => expect(fetchStarted).toBe(true)); - abortController.abort(); + // The download is still healthy — the rest streams through. + bodyController.enqueue(new Uint8Array([2])); + bodyController.close(); - await expect(promise).rejects.toThrow("Statement was canceled"); - }); + const second = await gen.next(); + expect(second.done).toBe(false); + expect([...(second.value as Uint8Array)]).toEqual([2]); + expect((await gen.next()).done).toBe(true); }); - describe("buffer concatenation", () => { - test("should concatenate buffers from multiple chunks", async () => { - // Custom fetch that returns distinct byte patterns - globalThis.fetch = vi.fn().mockImplementation(async (url: string) => { - const match = url.match(/chunk-(\d+)/); - const index = match ? parseInt(match[1], 10) : 0; - // Each chunk returns a unique byte pattern - return { - ok: true, - arrayBuffer: async () => new Uint8Array([100 + index]).buffer, - }; - }); - - const chunks = createMockChunks(3); - const result = await processor.processChunks(chunks, createMockSchema()); + test("downloads lazily — one fetch per pulled chunk, not all upfront", async () => { + const fetchMock = globalThis.fetch as ReturnType; + const iterator = processor.streamChunks(mockChunks(3)); - // Should be [100, 101, 102] - expect(result.data).toEqual(new Uint8Array([100, 101, 102])); - }); + // Draining the first chunk's body fetches exactly one chunk. + await iterator.next(); // first body piece of chunk 0 + await iterator.next(); // chunk 0 done → advances to chunk 1's fetch + expect(fetchMock).toHaveBeenCalledTimes(2); + }); - test("should return single buffer without unnecessary allocation", async () => { - globalThis.fetch = vi.fn().mockResolvedValue({ + test("retries establishing the response, then streams the body", async () => { + const fetchMock = vi + .fn() + .mockRejectedValueOnce(new Error("connection reset")) + .mockResolvedValueOnce({ ok: true, - arrayBuffer: async () => new Uint8Array([42, 43, 44]).buffer, - }); - - const chunks = createMockChunks(1); - const result = await processor.processChunks(chunks, createMockSchema()); - - // Single chunk should return as-is - expect(result.data).toEqual(new Uint8Array([42, 43, 44])); - }); + status: 200, + body: streamOf(new Uint8Array([7, 8, 9])), + } as unknown as Response); + globalThis.fetch = fetchMock; + + const bytes = await drain(processor.streamChunks(mockChunks(1))); + + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(bytes).toEqual([7, 8, 9]); + }, 10000); + + test("throws after exhausting retries on a non-2xx response", async () => { + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 1 }); + globalThis.fetch = vi.fn( + async () => + ({ ok: false, status: 500, statusText: "Server Error" }) as Response, + ); + await expect(drain(p.streamChunks(mockChunks(1)))).rejects.toThrow( + /chunk 0/, + ); }); - describe("missing external_link handling", () => { - test("should log error and continue retrying on missing external_link", async () => { - // Create chunk without external_link - const chunks = [ - { chunk_index: 0, external_link: undefined }, - { chunk_index: 1, external_link: "https://example.com/chunk-1" }, - ] as any; - - const consoleSpy = vi - .spyOn(console, "error") - .mockImplementation(() => {}); - - // This will fail because chunk 0 has no external_link and retries will be exhausted - await expect( - processor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(); - - // Logger uses util.format, so the message is pre-formatted - expect(consoleSpy).toHaveBeenCalledWith( - "[appkit:stream:arrow]", - expect.stringContaining("External link is required for chunk:"), - ); - - consoleSpy.mockRestore(); - }); + test("throws immediately when a chunk has no external_link", async () => { + const chunks = [{ chunk_index: 0 }] as any; + await expect(drain(processor.streamChunks(chunks))).rejects.toThrow( + /External link missing/, + ); }); - describe("error messages", () => { - test("should include chunk index in error message", async () => { - globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network error")); - - const chunks = createMockChunks(1); - - await expect( - processor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/chunk 0/); - }); - - test("should include HTTP status in error message for failed responses", async () => { - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: false, - status: 403, - statusText: "Forbidden", - }); + test("aborts during retry backoff instead of waiting it out", async () => { + const controller = new AbortController(); + // Every attempt fails to connect, forcing the retry backoff between them. + const fetchMock = vi.fn().mockRejectedValue(new Error("connection reset")); + globalThis.fetch = fetchMock; - const singleRetryProcessor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 1, - timeout: 5000, - retries: 1, - }); + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const drained = drain(p.streamChunks(mockChunks(1), controller.signal)); - const chunks = createMockChunks(1); - - await expect( - singleRetryProcessor.processChunks(chunks, createMockSchema()), - ).rejects.toThrow(/403 Forbidden/); - }); - }); -}); - -describe("Semaphore (via ArrowStreamProcessor)", () => { - let originalFetch: typeof globalThis.fetch; - - beforeEach(() => { - originalFetch = globalThis.fetch; - vi.clearAllMocks(); - }); + // Let the first attempt fail and enter the (~1s) backoff, then abort. + await new Promise((r) => setTimeout(r, 50)); + controller.abort(); - afterEach(() => { - globalThis.fetch = originalFetch; + await expect(drained).rejects.toThrow(); + // Aborting mid-backoff cancels immediately — no second fetch attempt. + expect(fetchMock).toHaveBeenCalledTimes(1); }); - test("should properly queue and release permits", async () => { - const processor = new ArrowStreamProcessor({ - maxConcurrentDownloads: 1, - timeout: 5000, - retries: 1, - }); - - const order: number[] = []; - - globalThis.fetch = vi.fn().mockImplementation(async (url: string) => { - const match = url.match(/chunk-(\d+)/); - const index = match ? parseInt(match[1], 10) : 0; - order.push(index); - - // Simulate varying response times - await new Promise((resolve) => setTimeout(resolve, 5)); - - return { - ok: true, - arrayBuffer: async () => new Uint8Array([index + 1]).buffer, - }; - }); - - const chunks = [ - { chunk_index: 0, external_link: "https://example.com/chunk-0" }, - { chunk_index: 1, external_link: "https://example.com/chunk-1" }, - { chunk_index: 2, external_link: "https://example.com/chunk-2" }, - ]; - - await processor.processChunks(chunks as any, { columns: [] }); - - // With concurrency of 1, they should complete in order - expect(order).toHaveLength(3); - expect(globalThis.fetch).toHaveBeenCalledTimes(3); + test("an already-aborted signal cancels before fetching", async () => { + const controller = new AbortController(); + controller.abort(); + const fetchMock = globalThis.fetch as ReturnType; + await expect( + drain(processor.streamChunks(mockChunks(1), controller.signal)), + ).rejects.toThrow(); + expect(fetchMock).not.toHaveBeenCalled(); }); }); diff --git a/packages/appkit/src/stream/types.ts b/packages/appkit/src/stream/types.ts index 154d388d..78b14e35 100644 --- a/packages/appkit/src/stream/types.ts +++ b/packages/appkit/src/stream/types.ts @@ -26,7 +26,7 @@ export interface SSEError { error: string; code: SSEErrorCode; /** - * Upstream-domain structured code (e.g. `INLINE_ARROW_STASH_EXHAUSTED`, + * Upstream-domain structured code (e.g. `RESULT_TOO_LARGE_FOR_JSON_FALLBACK`, * `NOT_IMPLEMENTED`). UI code should branch on this instead of parsing * the human-readable `error` string. */ diff --git a/packages/shared/src/sse/analytics.test.ts b/packages/shared/src/sse/analytics.test.ts index f66437c3..cdc8a70e 100644 --- a/packages/shared/src/sse/analytics.test.ts +++ b/packages/shared/src/sse/analytics.test.ts @@ -1,9 +1,5 @@ import { describe, expect, test } from "vitest"; -import { - AnalyticsSseMessage, - makeArrowMessage, - makeResultMessage, -} from "./analytics"; +import { AnalyticsSseMessage, makeResultMessage } from "./analytics"; describe("AnalyticsSseMessage schema", () => { test("accepts a result message with rows", () => { @@ -18,42 +14,11 @@ describe("AnalyticsSseMessage schema", () => { expect(() => AnalyticsSseMessage.parse({ type: "result" })).not.toThrow(); }); - test("accepts an arrow message with warehouse statement_id", () => { - const parsed = AnalyticsSseMessage.parse({ - type: "arrow", - statement_id: "stmt-1", - }); - expect(parsed.type).toBe("arrow"); - }); - - test("accepts an arrow message with synthetic inline- id", () => { - // Inline Arrow payloads are stashed server-side and surfaced through the - // same `arrow` message variant — the `inline-` prefix tells the - // /arrow-result handler to drain the stash instead of hitting the - // warehouse. The schema must accept both id shapes transparently. - const parsed = AnalyticsSseMessage.parse({ - type: "arrow", - statement_id: "inline-abc-123", - }); - expect(parsed.statement_id).toBe("inline-abc-123"); - }); - - test("rejects an arrow message with empty statement_id", () => { - expect(() => - AnalyticsSseMessage.parse({ type: "arrow", statement_id: "" }), - ).toThrow(); - }); - - test("rejects an arrow message with no statement_id", () => { - expect(() => AnalyticsSseMessage.parse({ type: "arrow" })).toThrow(); - }); - - test("rejects the retired arrow_inline message type", () => { - // arrow_inline was the prior wire shape (base64 payload on the SSE - // channel). The current protocol routes all Arrow payloads through - // /arrow-result; the type must no longer parse. + test("rejects a retired arrow message — ARROW_STREAM no longer uses SSE", () => { + // Arrow bytes now stream on the query response body, not the SSE channel, + // so there is no `arrow` message type to parse. expect(() => - AnalyticsSseMessage.parse({ type: "arrow_inline", attachment: "AQID" }), + AnalyticsSseMessage.parse({ type: "arrow", statement_id: "stmt-1" }), ).toThrow(); }); @@ -69,19 +34,9 @@ describe("AnalyticsSseMessage schema", () => { }); }); -describe("typed builders", () => { +describe("typed builder", () => { test("makeResultMessage roundtrips through the schema", () => { const msg = makeResultMessage([{ id: 1 }], { statement_id: "s-1" }); expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); }); - - test("makeArrowMessage roundtrips through the schema", () => { - const msg = makeArrowMessage("stmt-2"); - expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); - }); - - test("makeArrowMessage accepts synthetic inline- ids", () => { - const msg = makeArrowMessage("inline-some-uuid"); - expect(() => AnalyticsSseMessage.parse(msg)).not.toThrow(); - }); }); diff --git a/packages/shared/src/sse/analytics.ts b/packages/shared/src/sse/analytics.ts index eba7597e..41022672 100644 --- a/packages/shared/src/sse/analytics.ts +++ b/packages/shared/src/sse/analytics.ts @@ -3,25 +3,21 @@ import { z } from "zod"; /** * Wire protocol for analytics SSE messages emitted by `/api/analytics/query`. * - * These schemas are the single source of truth for the contract between the - * server (`AnalyticsPlugin._handleQueryRoute`) and the client - * (`useAnalyticsQuery`). Both sides validate with the same schema: - * - * - Server uses the typed builders (`makeResultMessage`, `makeArrowMessage`) - * to construct messages with compile-time guarantees that all required - * fields are present. - * - Client calls `AnalyticsSseMessage.parse(JSON.parse(event.data))` to fail - * loudly on a malformed payload instead of silently treating an undefined - * field as data. + * The SSE channel carries only the JSON_ARRAY path (warehouse-readiness + * events + a `result` message of rows). ARROW_STREAM does NOT use SSE — + * the server streams the raw Arrow IPC bytes back on the query response body + * (`_handleArrowStreamQuery`) and the client reads them directly + * (`fetchArrowDirect`), so there is no `arrow` message type here. * - * Arrow payloads — inline or external-links — never traverse the SSE control - * channel; both flow through `/api/analytics/arrow-result/:jobId` and are - * differentiated by an `inline-` prefix on the job id (see - * `InlineArrowStash`). The wire shape from the client's perspective is - * therefore uniform: an `arrow` message carries an id, the client fetches. + * These schemas are the single source of truth for the JSON contract between + * the server (`AnalyticsPlugin._handleQueryRoute`) and the client + * (`useAnalyticsQuery`). Both sides validate with the same schema: * - * Adding a new message variant requires a schema update here, which keeps - * server and client in lockstep. + * - Server uses the typed builder (`makeResultMessage`) to construct messages + * with compile-time guarantees that all required fields are present. + * - Client calls `AnalyticsSseMessage.safeParse(JSON.parse(event.data))` to + * fail loudly on a malformed payload instead of silently treating an + * undefined field as data. */ /** Successful row-shaped result (JSON_ARRAY format, or empty results). */ @@ -60,30 +56,16 @@ export interface AnalyticsResultMessage { } /** - * ARROW_STREAM result delivered via /arrow-result/:jobId. The id is either: - * - the warehouse-issued `statement_id` for EXTERNAL_LINKS responses, or - * - a synthetic `inline-` id pointing at the server-side - * `InlineArrowStash` for INLINE responses. - * - * Both shapes are fetched the same way; the prefix tells the route handler - * which path to take. + * Every message the analytics SSE stream may emit. Currently only the + * row-shaped `result` message (JSON_ARRAY path); `warehouse_status` and + * `error` events are handled off-schema by the client. ARROW_STREAM never + * uses SSE. */ -export const AnalyticsArrowMessage = z.object({ - type: z.literal("arrow"), - statement_id: z.string().min(1), - status: z.unknown().optional(), -}); -export type AnalyticsArrowMessage = z.infer; - -/** Discriminated union of every message the analytics SSE stream may emit. */ -export const AnalyticsSseMessage = z.discriminatedUnion("type", [ - AnalyticsResultMessage, - AnalyticsArrowMessage, -]); +export const AnalyticsSseMessage = AnalyticsResultMessage; export type AnalyticsSseMessage = z.infer; // ──────────────────────────────────────────────────────────────────────────── -// Typed builders — call from the server route handler. The compiler enforces +// Typed builder — call from the server route handler. The compiler enforces // that every required field is supplied, and the return type narrows so // downstream code (executeStream / SSE writer) keeps full type information. // ──────────────────────────────────────────────────────────────────────────── @@ -94,10 +76,3 @@ export function makeResultMessage( ): AnalyticsResultMessage { return { type: "result", data, ...extras }; } - -export function makeArrowMessage( - statement_id: string, - extras: { status?: unknown } = {}, -): AnalyticsArrowMessage { - return { type: "arrow", statement_id, ...extras }; -} diff --git a/tools/test-helpers.ts b/tools/test-helpers.ts index f15630d7..9161a0f6 100644 --- a/tools/test-helpers.ts +++ b/tools/test-helpers.ts @@ -140,9 +140,16 @@ export function createMockResponse() { const eventListeners: Record void>> = {}; const res = { + // Flips to true once headers/body have gone out — mirrors Express so + // streaming handlers can branch between a JSON error (pre-headers) and + // aborting the socket (mid-stream). + headersSent: false, status: vi.fn().mockReturnThis(), json: vi.fn().mockReturnThis(), - send: vi.fn().mockReturnThis(), + send: vi.fn(function (this: any) { + this.headersSent = true; + return this; + }), sendStatus: vi.fn().mockReturnThis(), end: vi.fn(function (this: any) { this.writableEnded = true; @@ -154,9 +161,16 @@ export function createMockResponse() { } return this; }), - write: vi.fn().mockReturnThis(), - setHeader: vi.fn().mockReturnThis(), + write: vi.fn(function (this: any) { + this.headersSent = true; + return this; + }), + setHeader: vi.fn(function (this: any) { + this.headersSent = true; + return this; + }), flushHeaders: vi.fn().mockReturnThis(), + destroy: vi.fn().mockReturnThis(), on: vi.fn(function ( this: any, event: string, @@ -168,6 +182,18 @@ export function createMockResponse() { eventListeners[event].push(handler); return this; }), + off: vi.fn(function ( + this: any, + event: string, + handler: (...args: any[]) => void, + ) { + if (eventListeners[event]) { + eventListeners[event] = eventListeners[event].filter( + (h) => h !== handler, + ); + } + return this; + }), writableEnded: false, }; return res;