From ebed179748caeb9e92fadbde065b6337da288932 Mon Sep 17 00:00:00 2001 From: Pedro Figueiredo Date: Mon, 9 Nov 2026 16:25:46 +0000 Subject: [PATCH 1/3] test(appkit): reproduce refused and stalled Cloud Fetch downloads Charts on the Arrow path fail with HTTP 500 when workspace storage refuses the app server's Cloud Fetch download, and hang until the route's 120 s first-byte timer fires when storage never answers (LA-2169). Add tests for the behavior the fix needs, all failing against v0.91.0: - Processor: link headers are sent, a 400 is not retried, a 403 re-mints the link once, a wait for headers is bounded and counts as a denial, every download failure carries EXTERNAL_LINK_DOWNLOAD_FAILED. - Delivery: a failed first chunk is replaced by an inline JSON_ARRAY run rebuilt as Arrow, and a result too large for inline is never stranded. - Route: over real HTTP against a loopback storage host, charts render, a definite refusal is remembered for ten minutes, and errors say why. - Rebuild: Arrow rebuilt from positional rows matches the native schema and values for every scalar type, including nulls. Signed-off-by: Pedro Figueiredo --- .../sql-warehouse/tests/arrow-schema.test.ts | 231 ++++++++++++- .../cloud-fetch-fallback.integration.test.ts | 317 +++++++++++++++++ .../analytics/tests/cloud-fetch-fixtures.ts | 219 ++++++++++++ .../analytics/tests/result-delivery.test.ts | 326 +++++++++++++++++- .../tests/arrow-stream-processor.test.ts | 274 +++++++++++++++ 5 files changed, 1364 insertions(+), 3 deletions(-) create mode 100644 packages/appkit/src/plugins/analytics/tests/cloud-fetch-fallback.integration.test.ts create mode 100644 packages/appkit/src/plugins/analytics/tests/cloud-fetch-fixtures.ts diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts index b7826e87e..65b45a234 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts @@ -23,7 +23,11 @@ import { } from "apache-arrow"; import { describe, expect, test } from "vitest"; -import { buildEmptyArrowIPCBase64, parseDatabricksType } from "../arrow-schema"; +import { + buildArrowIPCFromRows, + buildEmptyArrowIPCBase64, + parseDatabricksType, +} from "../arrow-schema"; // ============================================================================ // Helpers @@ -513,3 +517,228 @@ describe("buildEmptyArrowIPCBase64", () => { expect(table.numCols).toBe(0); }); }); + +// ============================================================================ +// buildArrowIPCFromRows (rebuilds ARROW_STREAM from INLINE JSON_ARRAY rows) +// ============================================================================ + +describe("buildArrowIPCFromRows", () => { + type Case = { + column: Parameters[0][number]; + cells: Array; + expected: unknown[]; + }; + + const decimalUnscaled = (v: unknown) => + v === null ? null : String(v as { toString(): string }); + const cases: Record unknown }> = { + INT: { + column: { name: "c", typeName: "INT", typeText: "INT" }, + cells: ["1", null, "-3"], + expected: [1, null, -3], + }, + TINYINT: { + column: { name: "c", typeName: "BYTE", typeText: "TINYINT" }, + cells: ["7", null, "-128"], + expected: [7, null, -128], + }, + SMALLINT: { + column: { name: "c", typeName: "SHORT", typeText: "SMALLINT" }, + cells: ["32767", null], + expected: [32767, null], + }, + BIGINT: { + column: { name: "c", typeName: "LONG", typeText: "BIGINT" }, + cells: ["9007199254740993", null, "-9223372036854775808"], + expected: [9007199254740993n, null, -9223372036854775808n], + }, + FLOAT: { + column: { name: "c", typeName: "FLOAT", typeText: "FLOAT" }, + cells: ["1.5", null, "-0.25"], + expected: [1.5, null, -0.25], + }, + DOUBLE: { + column: { name: "c", typeName: "DOUBLE", typeText: "DOUBLE" }, + cells: ["1.1", "NaN", null], + expected: [1.1, Number.NaN, null], + }, + DECIMAL: { + column: { + name: "c", + typeName: "DECIMAL", + typeText: "DECIMAL(10,2)", + typePrecision: 10, + typeScale: 2, + }, + cells: ["1.50", "-0.05", null, "12345678.99", "7"], + expected: ["150", "-5", null, "1234567899", "700"], + read: decimalUnscaled, + }, + "DECIMAL without (p,s) in typeText": { + column: { + name: "c", + typeName: "DECIMAL", + typeText: "DECIMAL", + typePrecision: 12, + typeScale: 3, + }, + cells: ["1.5", null], + expected: ["1500", null], + read: decimalUnscaled, + }, + DATE: { + column: { name: "c", typeName: "DATE", typeText: "DATE" }, + cells: ["2026-01-02", null], + expected: [Date.UTC(2026, 0, 2), null], + }, + TIMESTAMP: { + column: { name: "c", typeName: "TIMESTAMP", typeText: "TIMESTAMP" }, + cells: ["2026-01-02T10:00:00.123Z", null], + expected: [Date.UTC(2026, 0, 2, 10, 0, 0, 123), null], + }, + TIMESTAMP_NTZ: { + column: { + name: "c", + typeName: "TIMESTAMP_NTZ", + typeText: "TIMESTAMP_NTZ", + }, + cells: ["2026-01-02T10:00:00.123", null], + expected: [Date.UTC(2026, 0, 2, 10, 0, 0, 123), null], + }, + STRING: { + column: { name: "c", typeName: "STRING", typeText: "STRING" }, + cells: ["a", "", null], + expected: ["a", "", null], + }, + BOOLEAN: { + column: { name: "c", typeName: "BOOLEAN", typeText: "BOOLEAN" }, + cells: ["true", "false", null], + expected: [true, false, null], + }, + BINARY: { + column: { name: "c", typeName: "BINARY", typeText: "BINARY" }, + cells: ["AQI=", null], + expected: [[1, 2], null], + read: (v) => (v === null ? null : [...(v as Uint8Array)]), + }, + VOID: { + column: { name: "c", typeName: "NULL", typeText: "VOID" }, + cells: [null, null], + expected: [null, null], + }, + }; + + test.each(Object.entries(cases))( + "%s keeps the native Arrow type and values, nulls included", + (_name, { column, cells, expected, read }) => { + const ipc = buildArrowIPCFromRows( + [column], + cells.map((cell) => [cell]), + ); + const table = tableFromIPC(ipc); + + const native = parseDatabricksType(column.typeText ?? "STRING"); + const field = table.schema.fields[0]; + expect(field.name).toBe("col_0"); + expect(field.typeId).toBe(native.typeId); + if (column.typeName === "DECIMAL") { + const type = field.type as Decimal; + expect([type.precision, type.scale]).toEqual([ + column.typePrecision, + column.typeScale, + ]); + } else { + expect(String(field.type)).toBe(String(native)); + } + const values = [...(table.getChild("col_0") as Iterable)]; + expect(values.map((v) => (read ? read(v) : v))).toEqual(expected); + }, + ); + + test("positional rows keep duplicate column names and JSON-looking strings verbatim", () => { + const payload = '{"id":9007199254740993, "label":"kept verbatim"}'; + const ipc = buildArrowIPCFromRows( + [ + { name: "x", typeName: "INT", typeText: "INT" }, + { name: "x", typeName: "INT", typeText: "INT" }, + { name: "payload", typeName: "STRING", typeText: "STRING" }, + ], + [["1", "2", payload]], + ); + const table = tableFromIPC(ipc); + expect(table.schema.fields.map((f) => f.name)).toEqual([ + "col_0", + "col_1", + "col_2", + ]); + expect(table.getChild("col_0")?.get(0)).toBe(1); + expect(table.getChild("col_1")?.get(0)).toBe(2); + expect(table.getChild("col_2")?.get(0)).toBe(payload); + }); + + test("nested, interval and geospatial columns stay strings", () => { + const columns = [ + { name: "a", typeText: "ARRAY" }, + { name: "m", typeText: "MAP" }, + { name: "s", typeText: "STRUCT" }, + { name: "i", typeText: "INTERVAL DAY" }, + { name: "g", typeName: "GEOGRAPHY", typeText: "GEOGRAPHY(4326)" }, + ]; + const row = [ + "[1,2]", + '{"k":1}', + '{"f":1}', + "1 00:00:00.000000000", + "POINT(1 2)", + ]; + const table = tableFromIPC(buildArrowIPCFromRows(columns, [row])); + expect(table.schema.fields.map((f) => f.typeId)).toEqual( + columns.map(() => Type.Utf8), + ); + expect(table.toArray()[0].toJSON()).toEqual({ + col_0: row[0], + col_1: row[1], + col_2: row[2], + col_3: row[3], + col_4: row[4], + }); + }); + + test("zero rows still yield the typed schema", () => { + const table = tableFromIPC( + buildArrowIPCFromRows( + [{ name: "n", typeName: "LONG", typeText: "BIGINT" }], + [], + ), + ); + expect(table.numRows).toBe(0); + expect(table.schema.fields[0].typeId).toBe(Type.Int); + }); + + test("a cell that does not parse as its column type throws", () => { + expect(() => + buildArrowIPCFromRows( + [{ typeText: "INT" }, { typeText: "DATE" }], + [["1", "not-a-date"]], + ), + ).toThrow(/Invalid date/); + expect(() => + buildArrowIPCFromRows( + [{ typeText: "DECIMAL(5,2)", typePrecision: 5, typeScale: 2 }], + [["abc"]], + ), + ).toThrow(/Invalid decimal/); + expect(() => + buildArrowIPCFromRows([{ typeText: "INT" }], [["abc"]]), + ).toThrow(/Invalid integer/); + expect(() => + buildArrowIPCFromRows([{ typeText: "BIGINT" }], [["1.5"]]), + ).toThrow(); + expect(() => + buildArrowIPCFromRows([{ typeText: "DOUBLE" }], [["x"]]), + ).toThrow(/Invalid number/); + expect(() => + buildArrowIPCFromRows([{ typeText: "BOOLEAN" }], [["maybe"]]), + ).toThrow(/Invalid boolean/); + }); +}); diff --git a/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fallback.integration.test.ts b/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fallback.integration.test.ts new file mode 100644 index 000000000..39d1c39dc --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fallback.integration.test.ts @@ -0,0 +1,317 @@ +import { + createTestApp, + getMock, + type TestApp, +} from "@databricks/appkit/testing"; +import { tableFromIPC } from "apache-arrow"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { AppManager } from "../../../app"; +import { analytics } from "../index"; +import { + cellsOf, + type Dataset, + type FakeStorage, + schemaOf, + startStorage, + type StorageMode, + typedDataset, +} from "./cloud-fetch-fixtures"; + +const LARGE = 26 * 1024 * 1024; +const TTL_MS = 10 * 60_000; + +const getAppQuerySpy = vi.spyOn(AppManager.prototype, "getAppQuery"); + +type Warehouse = { + executions: string[]; +}; + +describe("Arrow charts when workspace storage refuses the Cloud Fetch download", () => { + let app: TestApp<[ReturnType]>; + let storage: FakeStorage; + let warehouse: Warehouse; + let data: Dataset; + let executeStatement: ReturnType; + let getResultData: ReturnType; + + /** The warehouse behind `statementExecution`, answering like a standard SQL warehouse. */ + function installWarehouse( + options: { + totalByteCount?: number; + inline?: () => never; + } = {}, + ) { + let seq = 0; + executeStatement.mockImplementation(async (body: any) => { + warehouse.executions.push(`${body.disposition}+${body.format}`); + const statementId = `stmt-${++seq}`; + const manifest = (format: string) => ({ + format, + schema: { columns: data.columns }, + totalChunkCount: 1, + totalByteCount: + options.totalByteCount === undefined + ? undefined + : BigInt(options.totalByteCount), + }); + if (body.disposition === "INLINE" && body.format === "ARROW_STREAM") { + throw Object.assign( + new Error( + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + ), + { errorCode: "INVALID_PARAMETER_VALUE" }, + ); + } + if (body.disposition === "EXTERNAL_LINKS") { + return { + statementId, + status: { state: "SUCCEEDED" }, + manifest: manifest("ARROW_STREAM"), + result: { + externalLinks: [ + { + chunkIndex: 0, + externalLink: storage.url("chunk-0"), + httpHeaders: { "x-ms-encryption-key": "k" }, + }, + ], + }, + }; + } + if (options.inline) return options.inline(); + return { + statementId, + status: { state: "SUCCEEDED" }, + manifest: manifest("JSON_ARRAY"), + result: { dataArray: data.dataArray }, + }; + }); + getResultData.mockImplementation(async () => ({ + externalLinks: [ + { + chunkIndex: 0, + externalLink: storage.url("chunk-0-fresh"), + httpHeaders: { "x-ms-encryption-key": "k" }, + }, + ], + })); + } + + async function boot(config: { timeout?: number } = {}) { + app = await createTestApp({ plugins: [analytics(config)] }); + executeStatement = getMock( + app.client, + "statementExecution.executeStatement", + ); + getResultData = getMock(app.client, "statementExecution.getResultData"); + getAppQuerySpy.mockResolvedValue({ + query: "SELECT * FROM sales", + isAsUser: false, + }); + } + + async function setStorage(mode: StorageMode) { + storage = await startStorage(mode); + storage.setBody(data.ipc); + } + + beforeEach(() => { + warehouse = { executions: [] }; + data = typedDataset(); + }); + + afterEach(async () => { + vi.useRealTimers(); + getAppQuerySpy.mockReset(); + await app?.close(); + await storage?.close(); + }); + + const chart = (key: string) => + app.post(`/api/analytics/query/${key}`, { + body: { parameters: {}, format: "ARROW_STREAM" }, + }); + + const columnsHeader = (res: Response) => + JSON.parse( + decodeURIComponent(res.headers.get("X-Appkit-Arrow-Columns") ?? "null"), + ); + + async function expectNativeArrow(res: Response) { + expect(res.status).toBe(200); + expect(res.headers.get("Content-Type")).toBe( + "application/vnd.apache.arrow.stream", + ); + const native = tableFromIPC(data.ipc); + const served = tableFromIPC(new Uint8Array(await res.arrayBuffer())); + expect(schemaOf(served)).toEqual(schemaOf(native)); + expect(cellsOf(served)).toEqual(cellsOf(native)); + expect(columnsHeader(res)).toEqual(data.columns.map((c) => c.name)); + } + + test("healthy storage is untouched by the fallback: same bytes, no inline run, one GET with the link's headers", async () => { + await setStorage({ kind: "ok" }); + await boot(); + installWarehouse(); + + const res = await chart("healthy"); + + expect(res.status).toBe(200); + expect( + Buffer.from(await res.arrayBuffer()).equals(Buffer.from(data.ipc)), + ).toBe(true); + expect(columnsHeader(res)).toEqual(data.columns.map((c) => c.name)); + expect(warehouse.executions).toEqual([ + "INLINE+ARROW_STREAM", + "EXTERNAL_LINKS+ARROW_STREAM", + ]); + expect(storage.gets()).toBe(1); + expect(storage.requestHeaders[0]["x-ms-encryption-key"]).toBe("k"); + }); + + test("a firewall that answers 403 still renders the chart, and the re-minted link is tried once", async () => { + await setStorage({ kind: "status", status: 403 }); + await boot(); + installWarehouse(); + + const res = await chart("firewalled"); + + await expectNativeArrow(res); + expect(warehouse.executions).toEqual([ + "INLINE+ARROW_STREAM", + "EXTERNAL_LINKS+ARROW_STREAM", + "INLINE+JSON_ARRAY", + ]); + expect(storage.gets()).toBe(2); + expect(getResultData).toHaveBeenCalledTimes(1); + }); + + test("a definite refusal is remembered for ten minutes, then storage is tried again", async () => { + vi.useFakeTimers({ toFake: ["Date"] }); + await setStorage({ kind: "status", status: 403 }); + await boot(); + installWarehouse(); + + await expectNativeArrow(await chart("first")); + const getsAfterFirst = storage.gets(); + + warehouse.executions.length = 0; + await expectNativeArrow(await chart("second")); + expect(warehouse.executions).toEqual(["INLINE+JSON_ARRAY"]); + expect(storage.gets()).toBe(getsAfterFirst); + + warehouse.executions.length = 0; + await expectNativeArrow(await chart("second")); + expect(warehouse.executions).toEqual([]); + + vi.setSystemTime(Date.now() + TTL_MS + 1000); + storage.setMode({ kind: "ok" }); + warehouse.executions.length = 0; + const recovered = await chart("third"); + expect( + Buffer.from(await recovered.arrayBuffer()).equals(Buffer.from(data.ipc)), + ).toBe(true); + expect(warehouse.executions).toEqual(["EXTERNAL_LINKS+ARROW_STREAM"]); + expect(storage.gets()).toBe(getsAfterFirst + 1); + }); + + test("storage that never answers falls back, and the next chart tries storage again", async () => { + await setStorage({ kind: "hang" }); + await boot({ timeout: 300 }); + installWarehouse(); + + await expectNativeArrow(await chart("stalled")); + expect(storage.gets()).toBe(1); + + warehouse.executions.length = 0; + await expectNativeArrow(await chart("stalled-again")); + expect(warehouse.executions).toContain("EXTERNAL_LINKS+ARROW_STREAM"); + expect(storage.gets()).toBe(2); + }); + + test("a body that stalls after the headers falls back instead of hanging", async () => { + await setStorage({ kind: "hang-after-headers" }); + await boot({ timeout: 300 }); + installWarehouse(); + + await expectNativeArrow(await chart("stalled-body")); + expect(warehouse.executions.at(-1)).toBe("INLINE+JSON_ARRAY"); + }); + + test("duplicate column names and JSON-looking strings come through the fallback intact", async () => { + const payload = '{"id":9007199254740993, "label":"kept verbatim"}'; + data = { + columns: [ + { name: "x", typeName: "INT", typeText: "INT", position: 0 }, + { name: "x", typeName: "INT", typeText: "INT", position: 1 }, + { + name: "payload", + typeName: "STRING", + typeText: "STRING", + position: 2, + }, + ], + dataArray: [["1", "2", payload]], + ipc: new Uint8Array(), + }; + await setStorage({ kind: "status", status: 403 }); + await boot(); + installWarehouse(); + + const res = await chart("duplicates"); + + expect(res.status).toBe(200); + const served = tableFromIPC(new Uint8Array(await res.arrayBuffer())); + expect(cellsOf(served)).toEqual([["1", "2", payload]]); + expect(columnsHeader(res)).toEqual(["x", "x", "payload"]); + }); + + test("a refused download of a result too large for inline answers with EXTERNAL_LINKS_DENIED and says why", async () => { + await setStorage({ kind: "status", status: 403 }); + await boot(); + installWarehouse({ totalByteCount: LARGE }); + + const res = await chart("large"); + + expect(res.status).toBe(500); + const body = (await res.json()) as { error: string; errorCode: string }; + expect(body.errorCode).toBe("EXTERNAL_LINKS_DENIED"); + expect(body.error).toMatch(/workspace storage/); + expect(body.error).toMatch(/too large to send inline/); + expect(warehouse.executions).not.toContain("INLINE+JSON_ARRAY"); + }); + + test("a large result behind storage that never answers is tried twice, then says why", async () => { + await setStorage({ kind: "hang" }); + await boot({ timeout: 300 }); + installWarehouse({ totalByteCount: LARGE }); + + const res = await chart("large-stalled"); + + expect(res.status).toBe(500); + const body = (await res.json()) as { errorCode: string }; + expect(body.errorCode).toBe("EXTERNAL_LINKS_DENIED"); + expect(storage.gets()).toBe(2); + expect(warehouse.executions).not.toContain("INLINE+JSON_ARRAY"); + }); + + test("when the inline run fails for another reason the user sees that reason, not a storage message", async () => { + await setStorage({ kind: "status", status: 403 }); + await boot(); + installWarehouse({ + inline: () => { + throw Object.assign(new Error("User does not have SELECT"), { + errorCode: "PERMISSION_DENIED", + }); + }, + }); + + const res = await chart("permission"); + + expect(res.status).toBe(500); + const body = (await res.json()) as { error: string; errorCode: string }; + expect(body.errorCode).toBe("PERMISSION_DENIED"); + expect(body.error).not.toMatch(/storage|firewall/i); + }); +}); diff --git a/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fixtures.ts b/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fixtures.ts new file mode 100644 index 000000000..d6431a1c9 --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fixtures.ts @@ -0,0 +1,219 @@ +import { createServer, type IncomingMessage, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; + +import { + Bool, + DateDay, + Decimal, + Float64, + Int32, + Int64, + makeData, + makeVector, + type Table, + Table as ArrowTable, + TimestampMicrosecond, + tableToIPC, + Utf8, + vectorFromArray, +} from "apache-arrow"; + +/** A standard SQL warehouse's view of one result: manifest columns, JSON rows, native Arrow. */ +export interface Dataset { + columns: Array<{ + name: string; + typeName: string; + typeText: string; + position: number; + typePrecision?: number; + typeScale?: number; + }>; + /** Positional string rows, exactly as INLINE + JSON_ARRAY returns them. */ + dataArray: Array>; + /** The IPC stream Cloud Fetch serves for the same result. */ + ipc: Uint8Array; +} + +const DAY = 86_400_000; +const isNull = (row: number, col: number) => (row + col) % 7 === 3; + +function decimalVector(unscaled: Array, scale: number) { + const words = new Uint32Array(unscaled.length * 4); + const validity = new Uint8Array(Math.ceil(unscaled.length / 8)); + let nullCount = 0; + unscaled.forEach((v, i) => { + if (v === null) { + nullCount++; + return; + } + validity[i >> 3] |= 1 << (i & 7); + const twos = v < 0n ? (1n << 128n) + v : v; + for (let w = 0; w < 4; w++) { + words[i * 4 + w] = Number((twos >> BigInt(32 * w)) & 0xffffffffn); + } + }); + return makeVector( + makeData({ + type: new Decimal(scale, 10, 128), + length: unscaled.length, + nullCount, + nullBitmap: validity, + data: words, + }), + ); +} + +const fixed = (v: bigint) => { + const sign = v < 0n ? "-" : ""; + const digits = (v < 0n ? -v : v).toString().padStart(3, "0"); + return `${sign}${digits.slice(0, -2)}.${digits.slice(-2)}`; +}; + +/** Eight scalar types with nulls sprinkled through every column. */ +export function typedDataset(rows = 20): Dataset { + const col = (c: number, f: (i: number) => T) => + Array.from({ length: rows }, (_, i) => (isNull(i, c) ? null : f(i))); + const id = col(0, (i) => i); + const big = col(1, (i) => 9007199254740993n + BigInt(i)); + const amount = col(2, (i) => i * 1.5); + const price = col(3, (i) => BigInt(i * 101 - 50)); + const day = col(4, (i) => Date.UTC(2026, 0, 1) + i * DAY); + const ts = col(5, (i) => Date.UTC(2026, 0, 1, 10) + i * 3_600_000); + const label = col(6, (i) => `row-${i}`); + const flag = col(7, (i) => i % 2 === 0); + + const table = new ArrowTable({ + col_0: vectorFromArray(id, new Int32()), + col_1: vectorFromArray(big, new Int64()), + col_2: vectorFromArray(amount, new Float64()), + col_3: decimalVector(price, 2), + col_4: vectorFromArray(day, new DateDay()), + col_5: vectorFromArray(ts, new TimestampMicrosecond("UTC")), + col_6: vectorFromArray(label, new Utf8()), + col_7: vectorFromArray(flag, new Bool()), + }); + + const text = (v: unknown, render: (x: never) => string) => + v === null ? null : render(v as never); + const dataArray = Array.from({ length: rows }, (_, i) => [ + text(id[i], String), + text(big[i], String), + text(amount[i], String), + text(price[i], fixed), + text(day[i], (x: number) => new Date(x).toISOString().slice(0, 10)), + text(ts[i], (x: number) => new Date(x).toISOString()), + text(label[i], String), + text(flag[i], String), + ]); + + const column = ( + name: string, + typeName: string, + typeText: string, + position: number, + extra: { typePrecision?: number; typeScale?: number } = {}, + ) => ({ name, typeName, typeText, position, ...extra }); + return { + columns: [ + column("id", "INT", "INT", 0), + column("big", "LONG", "BIGINT", 1), + column("amount", "DOUBLE", "DOUBLE", 2), + column("price", "DECIMAL", "DECIMAL(10,2)", 3, { + typePrecision: 10, + typeScale: 2, + }), + column("day", "DATE", "DATE", 4), + column("ts", "TIMESTAMP", "TIMESTAMP", 5), + column("label", "STRING", "STRING", 6), + column("flag", "BOOLEAN", "BOOLEAN", 7), + ], + dataArray, + ipc: tableToIPC(table, "stream"), + }; +} + +/** Every cell as a string, so schemas and values compare across Arrow builds. */ +export function cellsOf(table: Table): Array> { + const columns = table.schema.fields.map((f) => + [...(table.getChild(f.name) as Iterable)].map((v) => + v === null || v === undefined ? null : String(v), + ), + ); + return Array.from({ length: table.numRows }, (_, i) => + columns.map((c) => c[i]), + ); +} + +export function schemaOf(table: Table): string[] { + return table.schema.fields.map((f) => `${f.name}:${String(f.type)}`); +} + +export type StorageMode = + | { kind: "ok"; delayMs?: number } + | { kind: "status"; status: number } + | { kind: "hang" } + | { kind: "hang-after-headers" } + | { kind: "reset-after-headers" }; + +export interface FakeStorage { + /** GETs received, whatever the answer. */ + readonly gets: () => number; + readonly requestHeaders: IncomingMessage["headers"][]; + url(path?: string): string; + setMode(mode: StorageMode): void; + /** Body served on a healthy GET. */ + setBody(bytes: Uint8Array): void; + close(): Promise; +} + +/** A loopback host standing in for workspace storage behind a presigned URL. */ +export async function startStorage( + initial: StorageMode = { kind: "ok" }, +): Promise { + let mode = initial; + let body: Uint8Array = new Uint8Array(); + let gets = 0; + const requestHeaders: IncomingMessage["headers"][] = []; + const server: Server = createServer((req, res) => { + gets++; + requestHeaders.push(req.headers); + const current = mode; + switch (current.kind) { + case "ok": + setTimeout(() => res.end(Buffer.from(body)), current.delayMs ?? 0); + return; + case "status": + res.writeHead(current.status).end("denied"); + return; + case "hang": + return; + case "hang-after-headers": + res.writeHead(200, { "content-length": 100_000 }); + res.flushHeaders(); + return; + case "reset-after-headers": + res.writeHead(200, { "content-length": 100_000 }); + res.flushHeaders(); + setTimeout(() => req.socket.destroy(), 30); + return; + } + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const base = `http://127.0.0.1:${(server.address() as AddressInfo).port}`; + return { + gets: () => gets, + requestHeaders, + url: (path = "chunk") => `${base}/${path}`, + setMode: (next) => { + mode = next; + }, + setBody: (bytes) => { + body = bytes; + }, + close: () => + new Promise((resolve) => { + server.closeAllConnections(); + server.close(() => resolve()); + }), + }; +} diff --git a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts index 06438b5a5..14025074e 100644 --- a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts @@ -1,16 +1,34 @@ -import { Table, tableToIPC, Utf8, vectorFromArray } from "apache-arrow"; -import { describe, expect, test, vi } from "vitest"; +import { + Table, + tableFromIPC, + tableToIPC, + Utf8, + vectorFromArray, +} from "apache-arrow"; +import { afterEach, describe, expect, test, vi } from "vitest"; import { ExecutionError } from "../../../errors"; +import { ArrowStreamProcessor } from "../../../stream/arrow-stream-processor"; import { type ArrowChunkStreamer, + type ArrowDeliveryOut, + type ArrowCapability, arrowDeliveryUnsupported, classifyDispositionRejection, decodeArrowAttachmentToRows, deliverArrowBytes, + deliverArrowBytesWithFallback, deliverJsonResult, type QueryExecutor, } from "../result-delivery"; +import { + cellsOf, + type Dataset, + type FakeStorage, + schemaOf, + startStorage, + typedDataset, +} from "./cloud-fetch-fixtures"; /** Build a base64 Arrow IPC stream (2 cols, given names + rows). */ function arrowBase64(names: [string, string]): string { @@ -170,6 +188,7 @@ describe("deliverArrowBytes — normal warehouse (fallback to EXTERNAL_LINKS)", links, undefined, undefined, + false, ); expect(out.columnNames).toEqual(["a", "b"]); expect(chunks).toEqual([extBytes]); @@ -424,3 +443,306 @@ describe("decodeArrowAttachmentToRows — captured Reyden fixtures", () => { expect(rows).toEqual([REYDEN_NESTED_FIXTURE.expected]); }); }); + +describe("deliverArrowBytesWithFallback", () => { + const LARGE = 26 * 1024 * 1024; + const storages: FakeStorage[] = []; + + afterEach(async () => { + await Promise.all(storages.splice(0).map((st) => st.close())); + }); + + async function setup( + options: { + storage?: Parameters[0]; + totalByteCount?: number; + inline?: () => unknown; + rows?: number; + } = {}, + ) { + const data: Dataset = typedDataset(options.rows); + const storage = await startStorage(options.storage); + storage.setBody(data.ipc); + storages.push(storage); + const executions: string[] = []; + const executor: QueryExecutor = { + query: async (_q, _p, fp) => { + executions.push(`${fp.disposition}+${fp.format}`); + if (fp.disposition === "INLINE" && fp.format === "ARROW_STREAM") { + throw ExecutionError.statementFailed( + "The format field must be JSON_ARRAY when the disposition field is INLINE.", + "INVALID_PARAMETER_VALUE", + ); + } + if (fp.disposition === "EXTERNAL_LINKS") { + return { + statement_id: "stmt-links", + columnNames: data.columns.map((c) => c.name), + totalByteCount: options.totalByteCount, + external_links: [{ chunkIndex: 0, externalLink: storage.url() }], + }; + } + if (options.inline) return options.inline() as never; + return { + statement_id: "stmt-json", + columns: data.columns, + columnNames: data.columns.map((c) => c.name), + dataArray: data.dataArray, + }; + }, + }; + const processor = new ArrowStreamProcessor({ + timeout: 1000, + retries: 2, + responseTimeoutMs: 100, + }); + const streamer: ArrowChunkStreamer = { + streamExternalLinks: (links, signal, refresh, patient) => + processor.streamChunks(links, signal, refresh, patient), + }; + return { data, storage, executions, executor, streamer }; + } + + async function run( + ctx: Awaited>, + capabilityHint?: ArrowCapability, + ) { + const out: ArrowDeliveryOut = {}; + const reported: ArrowCapability[] = []; + const chunks: Uint8Array[] = []; + let error: ExecutionError | undefined; + try { + for await (const bytes of deliverArrowBytesWithFallback( + ctx.executor, + ctx.streamer, + "SELECT 1", + undefined, + out, + undefined, + { + capabilityHint, + onCapabilityResolved: (c) => reported.push(c), + }, + )) { + chunks.push(bytes); + } + } catch (e) { + error = e as ExecutionError; + } + const bytes = Buffer.concat(chunks.map((c) => Buffer.from(c))); + return { bytes, out, reported, error }; + } + + const expectNativeArrow = (ctx: { data: Dataset }, bytes: Uint8Array) => { + const native = tableFromIPC(ctx.data.ipc); + const served = tableFromIPC(new Uint8Array(bytes)); + expect(schemaOf(served)).toEqual(schemaOf(native)); + expect(cellsOf(served)).toEqual(cellsOf(native)); + }; + + test("a healthy workspace streams the download and never runs the inline fallback", async () => { + const ctx = await setup(); + const { bytes, reported, error } = await run(ctx); + + expect(error).toBeUndefined(); + expect(Buffer.from(ctx.data.ipc).equals(bytes)).toBe(true); + expect(ctx.executions).toEqual([ + "INLINE+ARROW_STREAM", + "EXTERNAL_LINKS+ARROW_STREAM", + ]); + expect(ctx.storage.gets()).toBe(1); + expect(reported).toEqual(["external"]); + }); + + test.each([403, 400])( + "a %i from storage serves rebuilt Arrow equal to the native result and reports links-denied", + async (status) => { + const ctx = await setup({ storage: { kind: "status", status } }); + const { bytes, out, reported, error } = await run(ctx); + + expect(error).toBeUndefined(); + expectNativeArrow(ctx, bytes); + expect(out.columnNames).toEqual(ctx.data.columns.map((c) => c.name)); + expect(out.statementId).toBe("stmt-json"); + expect(ctx.executions).toEqual([ + "INLINE+ARROW_STREAM", + "EXTERNAL_LINKS+ARROW_STREAM", + "INLINE+JSON_ARRAY", + ]); + expect(reported.at(-1)).toBe("links-denied"); + }, + ); + + test.each([ + ["no answer", { kind: "hang" }], + ["headers then no body", { kind: "hang-after-headers" }], + ["headers then a reset", { kind: "reset-after-headers" }], + ] as const)( + "storage with %s falls back once and is not remembered as denied", + async (_name, storage) => { + const ctx = await setup({ storage }); + const { bytes, reported, error } = await run(ctx); + + expect(error).toBeUndefined(); + expectNativeArrow(ctx, bytes); + expect(ctx.storage.gets()).toBe(1); + expect(ctx.executions.at(-1)).toBe("INLINE+JSON_ARRAY"); + expect(reported).not.toContain("links-denied"); + }, + ); + + test("with the links-denied hint the query goes inline and storage is not touched", async () => { + const ctx = await setup({ storage: { kind: "status", status: 403 } }); + const { bytes, error } = await run(ctx, "links-denied"); + + expect(error).toBeUndefined(); + expectNativeArrow(ctx, bytes); + expect(ctx.executions).toEqual(["INLINE+JSON_ARRAY"]); + expect(ctx.storage.gets()).toBe(0); + }); + + test("a refusal that is followed by an unrelated inline failure surfaces that failure unchanged", async () => { + const ctx = await setup({ + storage: { kind: "status", status: 403 }, + inline: () => { + throw ExecutionError.statementFailed( + "User does not have SELECT on catalog.schema.table", + "PERMISSION_DENIED", + "You do not have permission to run this query.", + ); + }, + }); + const { error } = await run(ctx); + + expect(error?.errorCode).toBe("PERMISSION_DENIED"); + expect(error?.clientMessage).toBe( + "You do not have permission to run this query.", + ); + }); + + test("a refused download of a result too large for inline fails saying why, without running inline", async () => { + const ctx = await setup({ + storage: { kind: "status", status: 403 }, + totalByteCount: LARGE, + }); + const { error, reported } = await run(ctx); + + expect(error?.errorCode).toBe("EXTERNAL_LINKS_DENIED"); + expect(error?.clientMessage).toMatch(/workspace storage/); + expect(error?.clientMessage).toMatch(/too large to send inline/); + expect(error?.clientMessage).not.toMatch(/azure/i); + expect(ctx.executions).toEqual([ + "INLINE+ARROW_STREAM", + "EXTERNAL_LINKS+ARROW_STREAM", + ]); + expect(reported.at(-1)).toBe("links-denied"); + }); + + test("a slow first byte on a large result gets one patient pass instead of a denial", async () => { + const ctx = await setup({ + storage: { kind: "ok", delayMs: 300 }, + totalByteCount: LARGE, + }); + const { bytes, reported, error } = await run(ctx); + + expect(error).toBeUndefined(); + expect(Buffer.from(ctx.data.ipc).equals(bytes)).toBe(true); + expect(ctx.storage.gets()).toBe(2); + expect(ctx.executions).not.toContain("INLINE+JSON_ARRAY"); + expect(reported).not.toContain("links-denied"); + }); + + test("a large result whose storage never answers fails saying why after the patient pass", async () => { + const ctx = await setup({ + storage: { kind: "hang" }, + totalByteCount: LARGE, + }); + const { error, reported } = await run(ctx); + + expect(error?.errorCode).toBe("EXTERNAL_LINKS_DENIED"); + expect(ctx.storage.gets()).toBe(2); + expect(ctx.executions).not.toContain("INLINE+JSON_ARRAY"); + expect(reported).not.toContain("links-denied"); + }); + + test("an inline result that spans several chunks is not served truncated", async () => { + const ctx = await setup({ + storage: { kind: "status", status: 403 }, + inline: () => ({ + statement_id: "stmt-json", + columns: typedDataset().columns, + dataArray: typedDataset().dataArray.slice(0, 3), + nextChunkIndex: 1, + }), + }); + const { error } = await run(ctx); + + expect(error?.errorCode).toBe("EXTERNAL_LINKS_DENIED"); + }); + + test("with the links-denied hint, a result too big for inline still downloads once storage recovers", async () => { + const ctx = await setup({ + inline: () => ({ + statement_id: "stmt-json", + columns: typedDataset().columns, + dataArray: [], + nextChunkIndex: 1, + }), + }); + const { bytes, reported, error } = await run(ctx, "links-denied"); + + expect(error).toBeUndefined(); + expect(Buffer.from(ctx.data.ipc).equals(bytes)).toBe(true); + expect(ctx.executions).toEqual([ + "INLINE+JSON_ARRAY", + "EXTERNAL_LINKS+ARROW_STREAM", + ]); + expect(ctx.storage.gets()).toBe(1); + expect(reported.at(-1)).toBe("external"); + }); + + test("with the links-denied hint, a too-big result while storage still refuses fails saying why", async () => { + const ctx = await setup({ + storage: { kind: "status", status: 403 }, + inline: () => ({ + statement_id: "stmt-json", + columns: typedDataset().columns, + dataArray: [], + nextChunkIndex: 1, + }), + }); + const { error } = await run(ctx, "links-denied"); + + expect(error?.errorCode).toBe("EXTERNAL_LINKS_DENIED"); + }); + + test("with the links-denied hint, an unrelated inline failure is what the user sees when storage still refuses", async () => { + const ctx = await setup({ + storage: { kind: "status", status: 403 }, + inline: () => { + throw ExecutionError.statementFailed( + "warehouse is restarting", + "TEMPORARILY_UNAVAILABLE", + ); + }, + }); + const { error } = await run(ctx, "links-denied"); + + expect(error?.errorCode).toBe("TEMPORARILY_UNAVAILABLE"); + }); + + test("a result the rebuild cannot parse fails with a code and a message, not a generic error", async () => { + const ctx = await setup({ + storage: { kind: "status", status: 403 }, + inline: () => ({ + statement_id: "stmt-json", + columns: [{ name: "n", typeName: "INT", typeText: "INT" }], + dataArray: [["not-a-number"]], + }), + }); + const { error } = await run(ctx); + + expect(error?.errorCode).toBe("ARROW_REBUILD_FAILED"); + expect(error?.clientMessage).toMatch(/JSON_ARRAY/); + }); +}); 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 84d1ee031..2d2caf623 100644 --- a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts +++ b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts @@ -1,7 +1,13 @@ +import { createServer, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; + import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { ExecutionError } from "../../errors"; import { ArrowStreamProcessor } from "../arrow-stream-processor"; +const EXTERNAL_LINK_DOWNLOAD_FAILED = "EXTERNAL_LINK_DOWNLOAD_FAILED"; + /** A ReadableStream that emits the given pieces in order, then closes. */ function streamOf(...pieces: Uint8Array[]): ReadableStream { return new ReadableStream({ @@ -264,3 +270,271 @@ describe("ArrowStreamProcessor.streamChunks", () => { expect(bytes).toEqual([7]); }, 10000); }); + +const response = (status: number, statusText = "") => + ({ ok: false, status, statusText }) as Response; +const okResponse = (...pieces: number[]) => + ({ + ok: true, + status: 200, + body: streamOf(new Uint8Array(pieces)), + }) as unknown as Response; + +async function failureOf(gen: AsyncGenerator) { + try { + await drain(gen); + } catch (error) { + return error as ExecutionError; + } + throw new Error("expected the download to fail"); +} + +describe("ArrowStreamProcessor — refused downloads", () => { + let originalFetch: typeof globalThis.fetch; + + beforeEach(() => { + originalFetch = globalThis.fetch; + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + vi.restoreAllMocks(); + }); + + test("sends the link's httpHeaders on every GET, and the new link's headers after a re-mint", async () => { + const fetchMock = vi + .fn() + .mockResolvedValueOnce(response(403, "Forbidden")) + .mockResolvedValueOnce(okResponse(1)); + globalThis.fetch = fetchMock; + const refresh = vi.fn(async (chunkIndex: number) => ({ + chunkIndex, + externalLink: "https://example.com/fresh", + httpHeaders: { "x-key": "fresh" }, + })); + + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + await drain( + p.streamChunks( + [ + { + chunkIndex: 0, + externalLink: "https://example.com/stale", + httpHeaders: { "x-key": "stale" }, + }, + ], + undefined, + refresh, + ), + ); + + expect(fetchMock.mock.calls.map(([, init]) => init?.headers)).toEqual([ + { "x-key": "stale" }, + { "x-key": "fresh" }, + ]); + }); + + test.each([400, 404, 410])( + "a %i is not retried, and the error carries the code and status", + async (status) => { + const fetchMock = vi.fn().mockResolvedValue(response(status, "No")); + globalThis.fetch = fetchMock; + const refresh = vi.fn(); + + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const error = await failureOf( + p.streamChunks(mockChunks(1), undefined, refresh), + ); + + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(refresh).not.toHaveBeenCalled(); + expect(error.errorCode).toBe(EXTERNAL_LINK_DOWNLOAD_FAILED); + expect(error.context?.httpStatus).toBe(status); + }, + ); + + test("a 403 re-mints the link once, retries at once, then stops", async () => { + const fetchMock = vi.fn().mockResolvedValue(response(403, "Forbidden")); + globalThis.fetch = fetchMock; + const refresh = vi.fn(async (chunkIndex: number) => ({ + chunkIndex, + externalLink: "https://example.com/fresh", + })); + + // retries: 1 proves the re-mint retry is not charged to the attempt budget. + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 1 }); + const started = Date.now(); + const error = await failureOf( + p.streamChunks(mockChunks(1), undefined, refresh), + ); + + expect(Date.now() - started).toBeLessThan(500); + expect(refresh).toHaveBeenCalledTimes(1); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(error.errorCode).toBe(EXTERNAL_LINK_DOWNLOAD_FAILED); + expect(error.context?.httpStatus).toBe(403); + }); + + test.each([429, 408])( + "a %i is retried after the backoff and can succeed", + async (status) => { + const fetchMock = vi + .fn() + .mockResolvedValueOnce(response(status)) + .mockResolvedValueOnce(okResponse(5)); + globalThis.fetch = fetchMock; + + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const bytes = await drain(p.streamChunks(mockChunks(1))); + + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(bytes).toEqual([5]); + }, + 10000, + ); + + test("a 5xx keeps retrying and ends without a client-error status", async () => { + const fetchMock = vi.fn().mockResolvedValue(response(503, "Unavailable")); + globalThis.fetch = fetchMock; + + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 2 }); + const error = await failureOf(p.streamChunks(mockChunks(1))); + + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(error.errorCode).toBe(EXTERNAL_LINK_DOWNLOAD_FAILED); + expect(error.context?.httpStatus).toBe(503); + }, 10000); +}); + +describe("ArrowStreamProcessor — storage that stalls", () => { + let server: Server; + let gets = 0; + let handler: Parameters[1]; + + const link = () => + `http://127.0.0.1:${(server.address() as AddressInfo).port}/chunk`; + const chunks = () => [{ chunkIndex: 0, externalLink: link() }]; + + beforeEach(async () => { + gets = 0; + server = createServer((req, res) => { + gets++; + handler?.(req, res); + }); + await new Promise((r) => server.listen(0, "127.0.0.1", r)); + }); + + afterEach(async () => { + server.closeAllConnections(); + await new Promise((r) => server.close(r)); + }); + + test("no response headers within responseTimeoutMs is a denial and is not retried", async () => { + handler = () => {}; + + const p = new ArrowStreamProcessor({ + timeout: 5000, + retries: 3, + responseTimeoutMs: 100, + }); + const started = Date.now(); + const error = await failureOf(p.streamChunks(chunks())); + + expect(gets).toBe(1); + expect(Date.now() - started).toBeLessThan(1000); + expect(error.errorCode).toBe(EXTERNAL_LINK_DOWNLOAD_FAILED); + expect(error.context?.httpStatus).toBeUndefined(); + }); + + test("the default wait for response headers is 10 seconds", () => { + expect(ArrowStreamProcessor.DEFAULT_RESPONSE_TIMEOUT).toBe(10_000); + }); + + test("a healthy chunk that streams for longer than responseTimeoutMs finishes", async () => { + handler = (_req, res) => { + res.writeHead(200); + let sent = 0; + const timer = setInterval(() => { + res.write(Buffer.from([sent++])); + if (sent === 8) { + clearInterval(timer); + res.end(); + } + }, 60); + }; + + const p = new ArrowStreamProcessor({ + timeout: 1000, + retries: 1, + responseTimeoutMs: 100, + }); + const bytes = await drain(p.streamChunks(chunks())); + + expect(bytes).toEqual([0, 1, 2, 3, 4, 5, 6, 7]); + }); + + test("a body that stalls after the headers fails with the download code", async () => { + handler = (_req, res) => { + res.writeHead(200, { "content-length": 1000 }); + res.flushHeaders(); + }; + + const p = new ArrowStreamProcessor({ + timeout: 150, + retries: 3, + responseTimeoutMs: 100, + }); + const error = await failureOf(p.streamChunks(chunks())); + + expect(gets).toBe(1); + expect(error.errorCode).toBe(EXTERNAL_LINK_DOWNLOAD_FAILED); + }); + + test("a connection reset after the headers fails with the download code", async () => { + handler = (req, res) => { + res.writeHead(200, { "content-length": 1000 }); + res.flushHeaders(); + setTimeout(() => req.socket.destroy(), 30); + }; + + const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); + const error = await failureOf(p.streamChunks(chunks())); + + expect(error.errorCode).toBe(EXTERNAL_LINK_DOWNLOAD_FAILED); + }); + + test("patient mode waits for headers as long as the attempt timeout", async () => { + handler = (_req, res) => { + setTimeout(() => res.end(Buffer.from([9])), 300); + }; + const p = new ArrowStreamProcessor({ + timeout: 2000, + retries: 1, + responseTimeoutMs: 100, + }); + + const strict = await failureOf(p.streamChunks(chunks())); + expect(strict.errorCode).toBe(EXTERNAL_LINK_DOWNLOAD_FAILED); + + const bytes = await drain( + p.streamChunks(chunks(), undefined, undefined, true), + ); + expect(bytes).toEqual([9]); + }); + + test("an abort from the caller still reads as a cancellation", async () => { + handler = () => {}; + const controller = new AbortController(); + const p = new ArrowStreamProcessor({ + timeout: 5000, + retries: 3, + responseTimeoutMs: 5000, + }); + const failed = failureOf(p.streamChunks(chunks(), controller.signal)); + setTimeout(() => controller.abort(), 50); + + const error = await failed; + expect(error.errorCode).toBeUndefined(); + expect(error.message).toMatch(/canceled/); + }); +}); From e93840c2e129533e1e09cc96a5c955f2f5cd101f Mon Sep 17 00:00:00 2001 From: Pedro Figueiredo Date: Mon, 9 Nov 2026 16:25:46 +0000 Subject: [PATCH 2/3] fix(appkit): fall back to inline results when a Cloud Fetch download is refused A chart on the Arrow path asks for ARROW_STREAM, which a standard SQL warehouse serves only as EXTERNAL_LINKS. When workspace storage refused the app server's download, the route made three attempts and returned HTTP 500 with "Query execution failed". When storage never answered, the route's 120 s first-byte timer returned a 503 that blamed the warehouse (LA-2169). Processor - Send the link's httpHeaders with every GET, and the new link's headers after a re-mint. - Wait at most responseTimeoutMs (10 s) for response headers only. A timeout is not retried. Body reads keep the per-read idle timeout, so a healthy chunk may stream for longer. - A 403 re-mints the link once and retries at once. Any other 4xx except 408 and 429 stops. Network errors and 5xx keep the 1 s and 2 s backoff. - Every download failure, before or after the headers, carries errorCode EXTERNAL_LINK_DOWNLOAD_FAILED and the last HTTP status. Delivery - deliverArrowBytesWithFallback pulls the first chunk before the route sends headers. On a failure before the first byte it runs the query as INLINE + JSON_ARRAY and rebuilds the Arrow IPC stream from the positional rows and manifest columns, typed by parseDatabricksType. - Only a 400 or 403 after the re-mint is remembered as "links-denied", for ten minutes. A timeout, network error, or body failure falls back once and is forgotten. - A result the inline run cannot hold (manifest byte count over 25 MiB, or a multi-chunk inline response) is never served truncated. After an inconclusive failure Cloud Fetch gets one patient pass. After a refusal the error is EXTERNAL_LINKS_DENIED and says why. - Any other inline failure is rethrown unchanged. A rebuild failure is ARROW_REBUILD_FAILED. Connector - executeStatement accepts rawRows to keep dataArray and carry the manifest columns on the result. EXTERNAL_LINKS results now carry totalByteCount. The inline fallback run is cached like the inline Arrow attempt, under its own key. @databricks/appkit-ui is unchanged. Signed-off-by: Pedro Figueiredo --- .../connectors/sql-warehouse/arrow-schema.ts | 181 ++++++++++++++ .../src/connectors/sql-warehouse/client.ts | 37 ++- .../appkit/src/plugins/analytics/analytics.ts | 69 ++++-- .../src/plugins/analytics/result-delivery.ts | 225 +++++++++++++++++- .../src/stream/arrow-stream-processor.ts | 215 +++++++++++------ 5 files changed, 630 insertions(+), 97 deletions(-) diff --git a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts index af5bfbb18..8193a074e 100644 --- a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts +++ b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts @@ -21,7 +21,12 @@ import { Table, TimestampMicrosecond, tableToIPC, + Type, Utf8, + type Vector, + makeData, + makeVector, + vectorFromArray, } from "apache-arrow"; /** @@ -75,6 +80,182 @@ export function buildEmptyArrowIPCBase64( return Buffer.from(ipc).toString("base64"); } +/** Column metadata from `manifest.schema.columns`. */ +export interface ArrowRowColumn { + name?: string; + typeText?: string; + typeName?: string; + typePrecision?: number; + typeScale?: number; +} + +/** + * Rebuild the Arrow IPC stream a warehouse serves for `ARROW_STREAM` from the + * positional string rows of an `INLINE` + `JSON_ARRAY` result. Fields are named + * `col_N` like the native stream, and each type is what + * {@link parseDatabricksType} gives that column, with DECIMAL precision and + * scale taken from `typePrecision` and `typeScale`. + * + * Nested, interval, and geospatial columns arrive as strings on the JSON wire + * and stay strings. A cell that does not parse as its column type throws. + */ +export function buildArrowIPCFromRows( + columns: ArrowRowColumn[], + rows: Array>, +): Uint8Array { + const vectors: Record = {}; + columns.forEach((col, index) => { + const cells = rows.map((row) => row[index] ?? null); + vectors[`col_${index}`] = vectorFromCells(col, cells); + }); + return tableToIPC(new Table(vectors), "stream"); +} + +function vectorFromCells(col: ArrowRowColumn, cells: Array) { + const type = rebuildType(col); + const map = (parse: (cell: string) => T) => + cells.map((cell) => (cell === null ? null : parse(cell))); + switch (type.typeId) { + case Type.Int: { + const { bitWidth } = type as unknown as { bitWidth: number }; + return vectorFromArray( + bitWidth === 64 ? map(parseBigInt) : map(parseInteger), + type as never, + ); + } + case Type.Float: + return vectorFromArray(map(parseFloatCell), type as never); + case Type.Bool: + return vectorFromArray(map(parseBoolean), type as never); + case Type.Date: + return vectorFromArray(map(parseUtcMillis), type as never); + case Type.Timestamp: + return vectorFromArray(map(parseUtcMillis), type as never); + case Type.Binary: + return vectorFromArray( + map((cell) => new Uint8Array(Buffer.from(cell, "base64"))), + type as never, + ); + case Type.Decimal: { + const { scale, precision } = type as unknown as { + scale: number; + precision: number; + }; + return decimalVector(cells, scale, precision); + } + case Type.Null: + return vectorFromArray( + cells.map(() => null), + type as never, + ); + default: + return vectorFromArray(cells, new Utf8()); + } +} + +const STRING_ONLY_TYPES = new Set([ + Type.List, + Type.Struct, + Type.Map, + Type.Interval, + Type.Duration, +]); + +function rebuildType(col: ArrowRowColumn): DataType { + const typeText = col.typeText ?? col.typeName ?? "STRING"; + let type: DataType; + try { + type = parseDatabricksType(typeText); + } catch { + return new Utf8(); + } + if (STRING_ONLY_TYPES.has(type.typeId)) return new Utf8(); + // GEOGRAPHY and GEOMETRY map to Binary but arrive as text on the JSON wire. + if (type.typeId === Type.Binary && typeText.toUpperCase() !== "BINARY") { + return new Utf8(); + } + if (type.typeId === Type.Decimal && col.typePrecision != null) { + return new Decimal(col.typeScale ?? 0, col.typePrecision, 128); + } + return type; +} + +function parseInteger(cell: string): number { + const n = Number(cell); + if (cell.trim() === "" || !Number.isInteger(n)) { + throw new Error(`Invalid integer "${cell}"`); + } + return n; +} + +function parseBigInt(cell: string): bigint { + if (cell.trim() === "") throw new Error(`Invalid integer "${cell}"`); + return BigInt(cell); +} + +function parseFloatCell(cell: string): number { + const n = Number(cell); + if (cell.trim() === "" || (Number.isNaN(n) && cell !== "NaN")) { + throw new Error(`Invalid number "${cell}"`); + } + return n; +} + +function parseBoolean(cell: string): boolean { + if (cell !== "true" && cell !== "false") { + throw new Error(`Invalid boolean "${cell}"`); + } + return cell === "true"; +} + +/** Epoch ms for an ISO date or timestamp; a value with no zone is UTC. */ +function parseUtcMillis(cell: string): number { + const ms = Date.parse( + /(?:Z|[+-]\d\d(?::?\d\d)?)$/i.test(cell) || !cell.includes("T") + ? cell + : `${cell}Z`, + ); + if (Number.isNaN(ms)) throw new Error(`Invalid date or timestamp "${cell}"`); + return ms; +} + +function decimalVector( + cells: Array, + scale: number, + precision: number, +) { + const words = new Uint32Array(cells.length * 4); + const validity = new Uint8Array(Math.ceil(cells.length / 8)); + let nullCount = 0; + cells.forEach((cell, i) => { + if (cell === null) { + nullCount++; + return; + } + validity[i >> 3] |= 1 << (i & 7); + const negative = cell.startsWith("-"); + const [whole, fraction = ""] = cell.replace(/^[-+]/, "").split("."); + const digits = `${whole}${fraction.padEnd(scale, "0").slice(0, scale)}`; + if (!/^\d+$/.test(digits)) throw new Error(`Invalid decimal "${cell}"`); + const unscaled = BigInt(digits); + const twosComplement = negative ? (1n << 128n) - unscaled : unscaled; + for (let w = 0; w < 4; w++) { + words[i * 4 + w] = Number( + (twosComplement >> BigInt(32 * w)) & 0xffffffffn, + ); + } + }); + return makeVector( + makeData({ + type: new Decimal(scale, precision, 128), + length: cells.length, + nullCount, + nullBitmap: validity, + data: words, + }), + ); +} + // ============================================================================ // Recursive-descent parser // ============================================================================ diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index b980a66d8..73d6128a9 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -102,7 +102,7 @@ function normalizeStatementCounts(response: T): T { * 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; +export const MAX_INLINE_ATTACHMENT_BYTES = 25 * 1024 * 1024; /** * Safety cap on how many additional EXTERNAL_LINKS chunks @@ -246,7 +246,13 @@ export class SQLWarehouseConnector { async executeStatement( workspaceClient: WorkspaceClient, - input: ExecuteStatementRequest, + input: ExecuteStatementRequest & { + /** + * Keep the warehouse's positional `dataArray` and carry the manifest + * `columns` on the result, instead of building keyed row objects. + */ + rawRows?: boolean; + }, signal?: AbortSignal, ) { const startTime = Date.now(); @@ -353,6 +359,7 @@ export class SQLWarehouseConnector { statementId, this.config.timeout, signal, + input.rawRows, ); break; case "SUCCEEDED": @@ -360,6 +367,7 @@ export class SQLWarehouseConnector { response, workspaceClient, signal, + input.rawRows, ); break; case "FAILED": @@ -830,6 +838,7 @@ export class SQLWarehouseConnector { statementId: string, timeout = executeStatementDefaults.timeout, signal?: AbortSignal, + rawRows?: boolean, ) { return this.telemetry.startActiveSpan( "sql.poll", @@ -908,6 +917,7 @@ export class SQLWarehouseConnector { response, workspaceClient, signal, + rawRows, ); case "FAILED": throw ExecutionError.statementFailed( @@ -949,6 +959,7 @@ export class SQLWarehouseConnector { response: StatementResponse, workspaceClient: WorkspaceClient, signal?: AbortSignal, + rawRows?: boolean, ) { if (response.manifest?.format === "ARROW_STREAM") { const result = response.result as @@ -998,11 +1009,22 @@ export class SQLWarehouseConnector { // current Databricks warehouses. } - if (!response.result?.dataArray || !response.manifest?.schema?.columns) { - return response; + const columns = response.manifest?.schema?.columns; + if (rawRows && columns) { + return { + ...response, + result: { + ...response.result, + columns, + columnNames: arrowColumnNames(response), + statement_id: response.statementId, + }, + }; } - const columns = response.manifest.schema.columns; + if (!response.result?.dataArray || !columns) { + return response; + } const transformedData = response.result.dataArray.map((row) => { const obj: Record = {}; @@ -1108,6 +1130,7 @@ export class SQLWarehouseConnector { statement_id: string; status: StatementStatus; columnNames?: string[]; + totalByteCount?: number; external_links?: ExternalLink[]; refreshChunkLink?: RefreshChunkLink; }; @@ -1121,6 +1144,7 @@ export class SQLWarehouseConnector { error: response.status?.error, } as StatementStatus, columnNames: arrowColumnNames(response), + totalByteCount: response.manifest?.totalByteCount as number | undefined, // 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 @@ -1227,8 +1251,9 @@ export class SQLWarehouseConnector { chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, + patient?: boolean, ): AsyncGenerator { - return this.arrowProcessor.streamChunks(chunks, signal, refresh); + return this.arrowProcessor.streamChunks(chunks, signal, refresh, patient); } /** diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 1d18aae6d..7834114f2 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -48,7 +48,8 @@ import { import { QueryProcessor } from "./query"; import { type ArrowCapability, - deliverArrowBytes, + type ArrowDeliveryOut, + deliverArrowBytesWithFallback, deliverJsonResult, type QueryExecutor, } from "./result-delivery"; @@ -133,6 +134,9 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { */ private _arrowCapability = new Map(); + /** When each warehouse's `"links-denied"` entry lapses back to `"external"`. */ + private _linksDeniedUntil = new Map(); + constructor(config: IAnalyticsConfig) { super(config); this.config = config; @@ -861,10 +865,11 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { const warehouseId = await getWarehouseId(); - // Populated by `deliverArrowBytes` from the result manifest before the - // first chunk is yielded, so the header below carries the real names. - const columnsRef: { columnNames?: string[]; statementId?: string } = {}; - const bytes = deliverArrowBytes( + // Populated by `deliverArrowBytesWithFallback` from the result manifest + // before the first chunk is yielded, so the header below carries the + // real names. + const columnsRef: ArrowDeliveryOut = {}; + const bytes = deliverArrowBytesWithFallback( // Wrap the executor so the INLINE attempt runs through the interceptor // chain (cache + retry), matching the JSON path. Only inline attachments // are cached; EXTERNAL_LINKS carry expiring pre-signed URLs and are @@ -884,9 +889,16 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { { // Skip the doomed INLINE probe on a warehouse already known to need // EXTERNAL_LINKS; remember the resolved mode for next time. - capabilityHint: this._arrowCapability.get(warehouseId), - onCapabilityResolved: (capability) => - this._arrowCapability.set(warehouseId, capability), + capabilityHint: this._arrowCapabilityHint(warehouseId), + onCapabilityResolved: (capability) => { + this._arrowCapability.set(warehouseId, capability); + if (capability === "links-denied") { + this._linksDeniedUntil.set( + warehouseId, + Date.now() + LINKS_DENIED_TTL_MS, + ); + } + }, }, ); const first = await bytes.next(); @@ -964,6 +976,24 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } } + /** + * The remembered arrow capability for a warehouse. A lapsed `"links-denied"` + * reverts to `"external"`, so the next query tries Cloud Fetch again. + */ + private _arrowCapabilityHint( + warehouseId: string, + ): ArrowCapability | undefined { + const capability = this._arrowCapability.get(warehouseId); + if ( + capability === "links-denied" && + Date.now() >= (this._linksDeniedUntil.get(warehouseId) ?? 0) + ) { + this._arrowCapability.set(warehouseId, "external"); + return "external"; + } + return capability; + } + /** * Await SQL warehouse readiness for the direct-binary Arrow path. There is no * SSE progress channel here — readiness is simply awaited, bounded by the @@ -994,6 +1024,8 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { * * - INLINE results carry a bounded (<=25 MiB) base64 `attachment` with the * same lifecycle as cached JSON rows — safe to cache. + * - The `INLINE + JSON_ARRAY` run behind the Cloud Fetch fallback (the + * `rawRows` request) is bounded the same way, and cached under its own key. * - EXTERNAL_LINKS results carry short-lived pre-signed URLs that expire in * minutes; caching them would serve dead links, so those pass through * uncached (only the tiny link metadata would be cached anyway). @@ -1016,12 +1048,16 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { const ttl = queryDefaults.cache?.ttl; return { query: (q, params, formatParameters, signal) => { - // Only the inline-arrow attempt is cacheable — EXTERNAL_LINKS carry + // Only inline attempts are cacheable — EXTERNAL_LINKS carry // short-lived pre-signed URLs, so those pass straight through. - if ( - formatParameters.disposition !== "INLINE" || - formatParameters.format !== "ARROW_STREAM" - ) { + const inlineArrow = + formatParameters.disposition === "INLINE" && + formatParameters.format === "ARROW_STREAM"; + const inlineRows = + formatParameters.disposition === "INLINE" && + formatParameters.format === "JSON_ARRAY" && + formatParameters.rawRows; + if (!inlineArrow && !inlineRows) { return executor.query(q, params, formatParameters, signal); } // On a standard warehouse this throws a capability rejection — the @@ -1031,7 +1067,9 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { // dedupes concurrent renders (e.g. React StrictMode double-mount). return cache.getOrExecute( [ - "analytics:query:arrow", + inlineArrow + ? "analytics:query:arrow" + : "analytics:query:arrow-rows", query_key, JSON.stringify(parameters), hashedQuery, @@ -1208,6 +1246,9 @@ export function writeChunk( */ const DEFAULT_ARROW_FIRST_BYTE_TIMEOUT_MS = 120_000; +/** How long a definite Cloud Fetch refusal is remembered before storage is tried again. */ +const LINKS_DENIED_TTL_MS = 10 * 60_000; + /** * Byte ceiling for the `X-Appkit-Arrow-Columns` header value. Beyond this (a * very wide schema) the names are served via the `/columns/:statementId` diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts index 3f40439d2..5803a9c0b 100644 --- a/packages/appkit/src/plugins/analytics/result-delivery.ts +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -1,9 +1,17 @@ import { type DataType, type Field, Type, tableFromIPC } from "apache-arrow"; import type { SQLTypeMarker } from "shared"; +import { + type ArrowRowColumn, + buildArrowIPCFromRows, +} from "../../connectors/sql-warehouse/arrow-schema"; +import { MAX_INLINE_ATTACHMENT_BYTES } from "../../connectors/sql-warehouse/client"; import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; -import type { RefreshChunkLink } from "../../stream/arrow-stream-processor"; +import { + EXTERNAL_LINK_DOWNLOAD_FAILED, + type RefreshChunkLink, +} from "../../stream/arrow-stream-processor"; import type { ExternalLink } from "../../workspace-client"; /** @@ -33,14 +41,24 @@ export interface QueryExecutor { query( query: string, parameters: Record | undefined, - formatParameters: { disposition: string; format: string }, + formatParameters: { + disposition: string; + format: string; + /** Ask for positional `dataArray` rows plus manifest `columns`. */ + rawRows?: boolean; + }, signal?: AbortSignal, ): Promise< | { attachment?: string; data?: Record[]; + dataArray?: Array>; + columns?: ArrowRowColumn[]; external_links?: ExternalLink[]; columnNames?: string[]; + totalByteCount?: number; + nextChunkIndex?: number; + nextChunkInternalLink?: string; statement_id?: string; status?: unknown; refreshChunkLink?: RefreshChunkLink; @@ -55,11 +73,27 @@ export interface ArrowChunkStreamer { chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, + patient?: boolean, ): AsyncGenerator; } -/** The arrow delivery mode a warehouse actually supports. */ -export type ArrowCapability = "inline" | "external"; +/** + * The arrow delivery mode a warehouse actually supports. `"links-denied"` means + * storage definitively refused the Cloud Fetch download, so results are served + * from an inline `JSON_ARRAY` run until the caller lets the hint lapse. + */ +export type ArrowCapability = "inline" | "external" | "links-denied"; + +/** Filled in by delivery for the route: column names and, for the fallback, how to re-stream. */ +export interface ArrowDeliveryOut { + columnNames?: string[]; + statementId?: string; + /** Set once EXTERNAL_LINKS resolves. */ + cloudFetch?: { + totalByteCount?: number; + stream(patient: boolean): AsyncGenerator; + }; +} /** Optional per-warehouse capability memo, avoiding a doomed probe per query. */ interface ArrowDeliveryOptions { @@ -190,7 +224,7 @@ export async function* deliverArrowBytes( streamer: ArrowChunkStreamer, query: string, processedParams: Record | undefined, - out: { columnNames?: string[]; statementId?: string }, + out: ArrowDeliveryOut, signal?: AbortSignal, opts?: ArrowDeliveryOptions, ): AsyncGenerator { @@ -271,11 +305,186 @@ export async function* deliverArrowBytes( // need no auth to download, so there is no second `getStatement` under a // mismatched identity. `refreshChunkLink` (also bound to this context) lets // the streamer re-mint a link that expires mid-download. - yield* streamer.streamExternalLinks( - ext.external_links, + const links = ext.external_links; + const stream = (patient: boolean) => + streamer.streamExternalLinks(links, signal, ext.refreshChunkLink, patient); + out.cloudFetch = { totalByteCount: ext.totalByteCount, stream }; + yield* stream(false); +} + +/** The inline result spans several chunks, and only the first is in hand. */ +class InlineDoesNotFit extends Error {} + +type InlineAttempt = { bytes: Uint8Array } | { error: unknown }; + +/** + * Run the query as `INLINE` + `JSON_ARRAY` and rebuild the Arrow IPC stream a + * Cloud Fetch download would have carried. The route sees the same bytes and + * header names either way. + */ +async function inlineJsonAsArrow( + executor: QueryExecutor, + query: string, + processedParams: Record | undefined, + out: ArrowDeliveryOut, + signal?: AbortSignal, +): Promise { + try { + const result = await executor.query( + query, + processedParams, + { disposition: "INLINE", format: "JSON_ARRAY", rawRows: true }, + signal, + ); + if ( + result?.nextChunkIndex != null || + result?.nextChunkInternalLink != null + ) { + return { error: new InlineDoesNotFit() }; + } + if (!result?.columns) throw ExecutionError.missingData("manifest columns"); + let bytes: Uint8Array; + try { + bytes = buildArrowIPCFromRows(result.columns, result.dataArray ?? []); + } catch (err: unknown) { + throw ExecutionError.statementFailed( + `Failed to rebuild Arrow from inline JSON_ARRAY rows: ${errMessage(err)}`, + "ARROW_REBUILD_FAILED", + 'The query result could not be converted to Arrow. Re-run the query with format="JSON_ARRAY".', + ); + } + out.columnNames = result.columnNames; + out.statementId = result.statement_id; + return { bytes }; + } catch (error: unknown) { + if (signal?.aborted) throw error; + return { error }; + } +} + +function isDownloadFailure(err: unknown): err is ExecutionError { + return ( + err instanceof ExecutionError && + err.errorCode === EXTERNAL_LINK_DOWNLOAD_FAILED + ); +} + +/** Download refused (or unanswered) and the result cannot be sent inline. */ +function linksDenied(cause: ExecutionError): ExecutionError { + return new ExecutionError( + `Cloud Fetch download failed and the result is too large to send inline: ${cause.message}`, + { + errorCode: "EXTERNAL_LINKS_DENIED", + cause, + clientMessage: + "The query result could not be downloaded from workspace storage and is too large to send inline. Check the storage account's firewall and the app's network policy.", + }, + ); +} + +/** + * {@link deliverArrowBytes} with a way out when the Cloud Fetch download fails + * before its first byte (storage refuses, resets, or never answers). + * + * The first chunk is pulled here, before the route sends headers, so a failed + * download can still be replaced by `inlineJsonAsArrow`. Only a 400 or 403 that + * survived the link re-mint counts as a definite refusal and reports + * `"links-denied"`; a timeout, network error, or body failure falls back once + * and is forgotten. A result too large for inline is never stranded: after an + * inconclusive failure Cloud Fetch gets one more, patient pass, and after a + * definite refusal the error says why. Any other inline failure is rethrown + * unchanged. + * + * With the `"links-denied"` hint the inline run goes first, and Cloud Fetch is + * tried only if that cannot deliver. + */ +export async function* deliverArrowBytesWithFallback( + executor: QueryExecutor, + streamer: ArrowChunkStreamer, + query: string, + processedParams: Record | undefined, + out: ArrowDeliveryOut, + signal?: AbortSignal, + opts?: ArrowDeliveryOptions, +): AsyncGenerator { + const hinted = opts?.capabilityHint === "links-denied"; + let tooLarge = false; + let inlineError: unknown; + if (hinted) { + const inline = await inlineJsonAsArrow( + executor, + query, + processedParams, + out, + signal, + ); + if ("bytes" in inline) { + yield inline.bytes; + return; + } + if (inline.error instanceof InlineDoesNotFit) tooLarge = true; + else inlineError = inline.error; + } + + const bytes = deliverArrowBytes( + executor, + streamer, + query, + processedParams, + out, signal, - ext.refreshChunkLink, + hinted ? { capabilityHint: "external" } : opts, ); + let first: IteratorResult; + let downloadError: ExecutionError | undefined; + try { + first = await bytes.next(); + } catch (err: unknown) { + if (signal?.aborted || !isDownloadFailure(err)) throw err; + downloadError = err; + first = { done: true, value: undefined }; + } + if (!downloadError) { + if (hinted) opts?.onCapabilityResolved?.("external"); + if (!first.done) yield first.value; + yield* bytes; + return; + } + + const status = downloadError.context?.httpStatus; + const refused = status === 400 || status === 403; + if (refused) opts?.onCapabilityResolved?.("links-denied"); + const size = out.cloudFetch?.totalByteCount; + tooLarge ||= size !== undefined && size > MAX_INLINE_ATTACHMENT_BYTES; + if (!tooLarge && inlineError === undefined) { + const inline = await inlineJsonAsArrow( + executor, + query, + processedParams, + out, + signal, + ); + if ("bytes" in inline) { + yield inline.bytes; + return; + } + if (!(inline.error instanceof InlineDoesNotFit)) throw inline.error; + tooLarge = true; + } + if (!tooLarge) throw inlineError; + if (refused || !out.cloudFetch) throw linksDenied(downloadError); + + const patient = out.cloudFetch.stream(true); + let retried: IteratorResult; + try { + retried = await patient.next(); + } catch (err: unknown) { + if (signal?.aborted || !isDownloadFailure(err)) throw err; + throw linksDenied(err); + } + if (hinted) opts?.onCapabilityResolved?.("external"); + if (!retried.done) yield retried.value; + yield* patient; } /** JSON row result plus the metadata the SSE `result` message forwards. */ diff --git a/packages/appkit/src/stream/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index e063b6abb..822de4a88 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -18,14 +18,40 @@ export type RefreshChunkLink = ( ) => Promise; interface ArrowStreamOptions { - /** Idle timeout (ms) for a single chunk: no progress → abort the download. */ + /** Idle timeout (ms) for a single read, and the response wait in patient mode. */ timeout: number; /** Attempts to establish a chunk's response before any bytes are yielded. */ retries: number; + /** Longest wait (ms) for response headers before the download counts as denied. */ + responseTimeoutMs: number; } +/** `errorCode` on every failure to download a chunk, whether before or after the response headers. */ +export const EXTERNAL_LINK_DOWNLOAD_FAILED = "EXTERNAL_LINK_DOWNLOAD_FAILED"; + const BACKOFF_MULTIPLIER = 1000; +type Connection = + | { + ok: true; + response: Response; + controller: AbortController; + detach: () => void; + } + | { ok: false; status?: number; error: unknown; timedOut: boolean }; + +function downloadFailed( + message: string, + httpStatus?: number, + cause?: unknown, +): ExecutionError { + return new ExecutionError(message, { + errorCode: EXTERNAL_LINK_DOWNLOAD_FAILED, + context: { httpStatus }, + cause: cause instanceof Error ? cause : undefined, + }); +} + /** * 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 @@ -36,6 +62,7 @@ const BACKOFF_MULTIPLIER = 1000; export class ArrowStreamProcessor { static readonly DEFAULT_TIMEOUT = 30000; static readonly DEFAULT_RETRIES = 3; + static readonly DEFAULT_RESPONSE_TIMEOUT = 10000; private options: ArrowStreamOptions; @@ -43,6 +70,9 @@ export class ArrowStreamProcessor { this.options = { timeout: options?.timeout ?? ArrowStreamProcessor.DEFAULT_TIMEOUT, retries: options?.retries ?? ArrowStreamProcessor.DEFAULT_RETRIES, + responseTimeoutMs: + options?.responseTimeoutMs ?? + ArrowStreamProcessor.DEFAULT_RESPONSE_TIMEOUT, }; } @@ -53,18 +83,23 @@ export class ArrowStreamProcessor { * 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. + * + * `patient` waits for response headers as long as the per-attempt `timeout` + * instead of `responseTimeoutMs`, for a second pass over a download that + * only looked stalled. */ async *streamChunks( chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, + patient = false, ): AsyncGenerator { if (chunks.length === 0) { throw ValidationError.missingField("chunks"); } for (const chunk of chunks) { - yield* this.streamChunkBody(chunk, signal, refresh); + yield* this.streamChunkBody(chunk, signal, refresh, patient); } } @@ -75,89 +110,92 @@ export class ArrowStreamProcessor { * 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. + * + * A 403 re-mints the link once and retries at once, since a pre-signed URL + * may have expired. Any other 4xx except 408 and 429 stops, and so does a + * wait for headers that outlasts the limit: a blocked host can hang for + * minutes instead of refusing, and a retry would only repeat the wait. */ private async *streamChunkBody( chunk: ExternalLink, - signal?: AbortSignal, - refresh?: RefreshChunkLink, + signal: AbortSignal | undefined, + refresh: RefreshChunkLink | undefined, + patient: boolean, ): AsyncGenerator { - let externalLink = chunk.externalLink; - if (!externalLink) { + let link = chunk; + if (!link.externalLink) { // A missing link cannot be fixed by retrying — fail loudly. throw ExecutionError.statementFailed( `External link missing for chunk ${chunk.chunkIndex}`, ); } - let response: Response | undefined; - let controller: AbortController | undefined; - let onOuterAbort: (() => void) | undefined; - let lastError: unknown; + const { timeout, retries, responseTimeoutMs } = this.options; + const headerWaitMs = patient + ? timeout + : Math.min(responseTimeoutMs, timeout); + let connection: Extract | undefined; + let failure: Extract | undefined; + let attempts = 0; + let reminted = false; - for (let attempt = 0; attempt < this.options.retries; attempt++) { + while (attempts < retries) { if (signal?.aborted) throw ExecutionError.canceled(); + attempts++; + const result = await this.connect(link, headerWaitMs, signal); + if (result.ok) { + connection = result; + break; + } + failure = result; + if (signal?.aborted) throw ExecutionError.canceled(); + if (result.timedOut) break; - const attemptController = new AbortController(); - const listener = () => attemptController.abort(); - signal?.addEventListener("abort", listener, { once: true }); - const timer = setTimeout( - () => attemptController.abort(), - this.options.timeout, - ); - - try { - const r = await fetch(externalLink, { - signal: attemptController.signal, - }); - clearTimeout(timer); - if (!r.ok) { - throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunkIndex}: ${r.status} ${r.statusText}`, + const { status } = result; + if (status === 403 && refresh && chunk.chunkIndex != null && !reminted) { + reminted = true; + attempts--; + try { + const fresh = await refresh(chunk.chunkIndex, signal); + if (fresh?.externalLink) link = fresh; + } catch (refreshError) { + // Retry the current URL; surface the original error if it fails. + logger.warn( + "Failed to re-resolve link for chunk %s: %O", + chunk.chunkIndex, + refreshError, ); } - // Keep this attempt's controller alive to drive the body read + idle - // timeout below. - response = r; - controller = attemptController; - onOuterAbort = listener; + continue; + } + if ( + status !== undefined && + status >= 400 && + status < 500 && + status !== 408 && + status !== 429 + ) { break; - } catch (error) { - 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, signal); - // Pre-signed URLs expire (<= 15 min). Before retrying, re-mint this - // chunk's link — a stale URL would just 403 again on the same address. - // Only meaningful before any bytes are yielded (below), which is why - // this lives in the establish-response loop. - if (refresh && chunk.chunkIndex != null) { - try { - const fresh = await refresh(chunk.chunkIndex, signal); - if (fresh?.externalLink) externalLink = fresh.externalLink; - } catch (refreshError) { - // Keep retrying the current URL; surface the original error if - // all attempts fail. - logger.warn( - "Failed to re-resolve link for chunk %s: %O", - chunk.chunkIndex, - refreshError, - ); - } - } - } + } + if (attempts < retries) { + await this.delay(2 ** (attempts - 1) * BACKOFF_MULTIPLIER, signal); } } - if (!response || !controller) { - throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunkIndex} after ${this.options.retries} attempts: ${ - lastError instanceof Error ? lastError.message : String(lastError) - }`, + if (!connection) { + const reason = failure?.timedOut + ? `no response within ${headerWaitMs}ms` + : failure?.error instanceof Error + ? failure.error.message + : String(failure?.error); + throw downloadFailed( + `Failed to download chunk ${chunk.chunkIndex}: ${reason}`, + failure?.status, + failure?.error, ); } + const { response, controller, detach } = connection; const body = response.body; let idleTimer: ReturnType | undefined; try { @@ -170,10 +208,7 @@ export class ArrowStreamProcessor { // 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, - ); + idleTimer = setTimeout(() => controller.abort(), timeout); let done: boolean; let value: Uint8Array | undefined; try { @@ -197,14 +232,56 @@ export class ArrowStreamProcessor { ); throw error instanceof ExecutionError ? error - : ExecutionError.statementFailed( + : downloadFailed( `Failed streaming chunk ${chunk.chunkIndex}: ${ error instanceof Error ? error.message : String(error) }`, + undefined, + error, ); } finally { clearTimeout(idleTimer); - if (onOuterAbort) signal?.removeEventListener("abort", onOuterAbort); + detach(); + } + } + + /** + * One GET. Only the wait for response headers is timed here; the timer is + * cleared the moment they arrive, so a healthy chunk may stream for longer. + */ + private async connect( + link: ExternalLink, + headerWaitMs: number, + signal?: AbortSignal, + ): Promise { + const controller = new AbortController(); + const onAbort = () => controller.abort(); + signal?.addEventListener("abort", onAbort, { once: true }); + const detach = () => signal?.removeEventListener("abort", onAbort); + const timer = setTimeout(() => controller.abort(), headerWaitMs); + try { + const response = await fetch(link.externalLink as string, { + signal: controller.signal, + headers: link.httpHeaders, + }); + clearTimeout(timer); + if (response.ok) return { ok: true, response, controller, detach }; + await response.body?.cancel().catch(() => {}); + detach(); + return { + ok: false, + status: response.status, + error: new Error(`${response.status} ${response.statusText}`), + timedOut: false, + }; + } catch (error) { + clearTimeout(timer); + detach(); + return { + ok: false, + error, + timedOut: !signal?.aborted && controller.signal.aborted, + }; } } From dbcad69482b2e8b70bd4c1932d068e4aa0e2484c Mon Sep 17 00:00:00 2001 From: Pedro Figueiredo Date: Mon, 9 Nov 2026 17:51:51 +0000 Subject: [PATCH 3/3] fix(appkit): remember a Cloud Fetch download that never answers Builder pods run with internet off and no SEG, so a download straight to cloud storage hangs instead of refusing. Remember that as links-denied once the inline fallback has served the result, so later charts skip the wait. Signed-off-by: Pedro Figueiredo --- .../src/plugins/analytics/result-delivery.ts | 31 +++++++++++++------ .../cloud-fetch-fallback.integration.test.ts | 6 ++-- .../analytics/tests/result-delivery.test.ts | 14 +++++++-- .../src/stream/arrow-stream-processor.ts | 4 ++- 4 files changed, 39 insertions(+), 16 deletions(-) diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts index 5803a9c0b..e955b6211 100644 --- a/packages/appkit/src/plugins/analytics/result-delivery.ts +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -79,7 +79,7 @@ export interface ArrowChunkStreamer { /** * The arrow delivery mode a warehouse actually supports. `"links-denied"` means - * storage definitively refused the Cloud Fetch download, so results are served + * storage refused or never answered the Cloud Fetch download, so results are served * from an inline `JSON_ARRAY` run until the caller lets the hint lapse. */ export type ArrowCapability = "inline" | "external" | "links-denied"; @@ -369,6 +369,12 @@ function isDownloadFailure(err: unknown): err is ExecutionError { ); } +/** Storage answered 400 or 403 even after the link re-mint. */ +function isRefusal(err: ExecutionError): boolean { + const status = err.context?.httpStatus; + return status === 400 || status === 403; +} + /** Download refused (or unanswered) and the result cannot be sent inline. */ function linksDenied(cause: ExecutionError): ExecutionError { return new ExecutionError( @@ -387,13 +393,14 @@ function linksDenied(cause: ExecutionError): ExecutionError { * before its first byte (storage refuses, resets, or never answers). * * The first chunk is pulled here, before the route sends headers, so a failed - * download can still be replaced by `inlineJsonAsArrow`. Only a 400 or 403 that - * survived the link re-mint counts as a definite refusal and reports - * `"links-denied"`; a timeout, network error, or body failure falls back once - * and is forgotten. A result too large for inline is never stranded: after an - * inconclusive failure Cloud Fetch gets one more, patient pass, and after a - * definite refusal the error says why. Any other inline failure is rethrown - * unchanged. + * download can still be replaced by `inlineJsonAsArrow`. A 400 or 403 that + * survived the link re-mint reports `"links-denied"` at once. No response at all + * reports it once the inline result has been served, or once the patient pass + * also fails, because blocked hosts hang instead of refusing. A network error + * or body failure falls back once and is forgotten. A result too large for + * inline is never stranded: after an inconclusive failure Cloud Fetch gets one + * more, patient pass, and after a definite refusal the error says why. Any + * other inline failure is rethrown unchanged. * * With the `"links-denied"` hint the inline run goes first, and Cloud Fetch is * tried only if that cannot deliver. @@ -451,8 +458,8 @@ export async function* deliverArrowBytesWithFallback( return; } - const status = downloadError.context?.httpStatus; - const refused = status === 400 || status === 403; + const refused = isRefusal(downloadError); + const unanswered = downloadError.context?.timedOut === true; if (refused) opts?.onCapabilityResolved?.("links-denied"); const size = out.cloudFetch?.totalByteCount; tooLarge ||= size !== undefined && size > MAX_INLINE_ATTACHMENT_BYTES; @@ -465,6 +472,7 @@ export async function* deliverArrowBytesWithFallback( signal, ); if ("bytes" in inline) { + if (unanswered) opts?.onCapabilityResolved?.("links-denied"); yield inline.bytes; return; } @@ -480,6 +488,9 @@ export async function* deliverArrowBytesWithFallback( retried = await patient.next(); } catch (err: unknown) { if (signal?.aborted || !isDownloadFailure(err)) throw err; + if (isRefusal(err) || err.context?.timedOut === true) { + opts?.onCapabilityResolved?.("links-denied"); + } throw linksDenied(err); } if (hinted) opts?.onCapabilityResolved?.("external"); diff --git a/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fallback.integration.test.ts b/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fallback.integration.test.ts index 39d1c39dc..22776ce22 100644 --- a/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fallback.integration.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/cloud-fetch-fallback.integration.test.ts @@ -216,7 +216,7 @@ describe("Arrow charts when workspace storage refuses the Cloud Fetch download", expect(storage.gets()).toBe(getsAfterFirst + 1); }); - test("storage that never answers falls back, and the next chart tries storage again", async () => { + test("storage that never answers falls back, and the next chart skips storage", async () => { await setStorage({ kind: "hang" }); await boot({ timeout: 300 }); installWarehouse(); @@ -226,8 +226,8 @@ describe("Arrow charts when workspace storage refuses the Cloud Fetch download", warehouse.executions.length = 0; await expectNativeArrow(await chart("stalled-again")); - expect(warehouse.executions).toContain("EXTERNAL_LINKS+ARROW_STREAM"); - expect(storage.gets()).toBe(2); + expect(warehouse.executions).toEqual(["INLINE+JSON_ARRAY"]); + expect(storage.gets()).toBe(1); }); test("a body that stalls after the headers falls back instead of hanging", async () => { diff --git a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts index 14025074e..07946951d 100644 --- a/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/result-delivery.test.ts @@ -573,8 +573,18 @@ describe("deliverArrowBytesWithFallback", () => { }, ); + test("storage that never answers falls back and is remembered as denied", async () => { + const ctx = await setup({ storage: { kind: "hang" } }); + const { bytes, reported, error } = await run(ctx); + + expect(error).toBeUndefined(); + expectNativeArrow(ctx, bytes); + expect(ctx.storage.gets()).toBe(1); + expect(ctx.executions.at(-1)).toBe("INLINE+JSON_ARRAY"); + expect(reported.at(-1)).toBe("links-denied"); + }); + test.each([ - ["no answer", { kind: "hang" }], ["headers then no body", { kind: "hang-after-headers" }], ["headers then a reset", { kind: "reset-after-headers" }], ] as const)( @@ -662,7 +672,7 @@ describe("deliverArrowBytesWithFallback", () => { expect(error?.errorCode).toBe("EXTERNAL_LINKS_DENIED"); expect(ctx.storage.gets()).toBe(2); expect(ctx.executions).not.toContain("INLINE+JSON_ARRAY"); - expect(reported).not.toContain("links-denied"); + expect(reported.at(-1)).toBe("links-denied"); }); test("an inline result that spans several chunks is not served truncated", async () => { diff --git a/packages/appkit/src/stream/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index 822de4a88..2021eddf9 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -44,10 +44,11 @@ function downloadFailed( message: string, httpStatus?: number, cause?: unknown, + timedOut = false, ): ExecutionError { return new ExecutionError(message, { errorCode: EXTERNAL_LINK_DOWNLOAD_FAILED, - context: { httpStatus }, + context: { httpStatus, timedOut }, cause: cause instanceof Error ? cause : undefined, }); } @@ -192,6 +193,7 @@ export class ArrowStreamProcessor { `Failed to download chunk ${chunk.chunkIndex}: ${reason}`, failure?.status, failure?.error, + failure?.timedOut, ); }