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/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/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..e955b6211 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 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"; + +/** 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,197 @@ 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 + ); +} + +/** 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( + `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`. 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. + */ +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 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; + if (!tooLarge && inlineError === undefined) { + const inline = await inlineJsonAsArrow( + executor, + query, + processedParams, + out, + signal, + ); + if ("bytes" in inline) { + if (unanswered) opts?.onCapabilityResolved?.("links-denied"); + 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; + if (isRefusal(err) || err.context?.timedOut === true) { + opts?.onCapabilityResolved?.("links-denied"); + } + 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/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..22776ce22 --- /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 skips storage", 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).toEqual(["INLINE+JSON_ARRAY"]); + expect(storage.gets()).toBe(1); + }); + + 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..07946951d 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,316 @@ 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("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([ + ["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.at(-1)).toBe("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/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index e063b6abb..2021eddf9 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -18,14 +18,41 @@ 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, + timedOut = false, +): ExecutionError { + return new ExecutionError(message, { + errorCode: EXTERNAL_LINK_DOWNLOAD_FAILED, + context: { httpStatus, timedOut }, + 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 +63,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 +71,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 +84,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 +111,93 @@ 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, + failure?.timedOut, ); } + const { response, controller, detach } = connection; const body = response.body; let idleTimer: ReturnType | undefined; try { @@ -170,10 +210,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 +234,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, + }; } } 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/); + }); +});