diff --git a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts index af5bfbb18..82b88c96e 100644 --- a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts +++ b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts @@ -1,7 +1,7 @@ import { Binary, Bool, - type DataType, + DataType, DateDay, Decimal, DurationMicrosecond, @@ -16,14 +16,22 @@ import { List, Map_, Null, + RecordBatch, Schema, Struct, Table, TimestampMicrosecond, tableToIPC, + Type, Utf8, + makeData, + makeVector, + vectorFromArray, } from "apache-arrow"; +import type { ColumnInfo } from "../../workspace-client"; +import { hasJsonArrayCodec, parseJsonArrayCell } from "./json-array-cells"; + /** * Parse a Databricks SQL type text (the value returned by the Statement * Execution API in `ColumnInfo.type_text`) into an Apache Arrow DataType. @@ -51,30 +59,136 @@ export function parseDatabricksType(typeText: string): DataType { * deliver a real Arrow Table to the client, preserving the hook's typed * contract. */ -export function buildEmptyArrowIPCBase64( - columns: Array<{ - name?: string; - typeText?: string; - typeName?: string; - }>, -): string { - const fields = columns.map((col, index) => { - const typeText = col.typeText ?? col.typeName ?? "STRING"; - let dataType: DataType; - try { - dataType = parseDatabricksType(typeText); - } catch { - dataType = new Utf8(); - } - const name = col.name && col.name.length > 0 ? col.name : `column_${index}`; - return new Field(name, dataType, true); - }); +export function buildEmptyArrowIPCBase64(columns: ColumnInfo[]): string { + const fields = columns.map( + (col, index) => + new Field(manifestColumnName(col, index), columnArrowType(col), true), + ); const schema = new Schema(fields); const table = new Table(schema); const ipc = tableToIPC(table, "stream"); return Buffer.from(ipc).toString("base64"); } +/** A manifest column's name, or `column_N` when it has none. */ +export function manifestColumnName( + col: ColumnInfo | undefined, + index: number, +): string { + return col?.name || `column_${index}`; +} + +function columnTypeText(col: ColumnInfo): string { + return col.typeText ?? col.typeName ?? "STRING"; +} + +/** The Arrow type for a manifest column; unparseable type text becomes Utf8. */ +function columnArrowType(col: ColumnInfo): DataType { + try { + return parseDatabricksType(columnTypeText(col)); + } catch { + return new Utf8(); + } +} + +/** + * Rebuild the `ARROW_STREAM` IPC bytes from the positional rows of an `INLINE` + * + `JSON_ARRAY` result, keeping the manifest column names (duplicates + * included). Columns without a {@link parseJsonArrayCell} codec stay Utf8, and + * a cell that does not parse as its column type throws. Zoned TIMESTAMP columns + * are labelled UTC, since cells carry an offset rather than the session zone. + */ +export function buildArrowIPCFromRows( + columns: ColumnInfo[], + rows: Array>, +): Uint8Array { + const vectors = columns.map((col, index) => + vectorFromCells( + col, + rows.map((row) => row[index] ?? null), + ), + ); + const fields = vectors.map( + (vector, index) => + new Field(manifestColumnName(columns[index], index), vector.type, true), + ); + const batch = new RecordBatch( + new Schema(fields), + makeData({ + type: new Struct(fields), + length: rows.length, + nullCount: 0, + children: vectors.map((vector) => vector.data[0]), + }), + ); + return tableToIPC(new Table(batch), "stream"); +} + +function vectorFromCells(col: ColumnInfo, cells: Array) { + const type = rebuildType(col); + const values = cells.map((cell) => + cell === null ? null : parseJsonArrayCell(cell, type), + ); + if (DataType.isDecimal(type)) { + return decimalVector(values as Array, type); + } + return vectorFromArray(values, type); +} + +const STRING_ONLY_TYPES = new Set([ + Type.List, + Type.Struct, + Type.Map, + Type.Interval, + Type.Duration, +]); + +function rebuildType(col: ColumnInfo): DataType { + const type = columnArrowType(col); + 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 && + columnTypeText(col).toUpperCase() !== "BINARY" + ) { + return new Utf8(); + } + if (type.typeId === Type.Decimal && col.typePrecision != null) { + return new Decimal(col.typeScale ?? 0, col.typePrecision, 128); + } + if (type.typeId !== Type.Null && !hasJsonArrayCodec(type)) return new Utf8(); + return type; +} + +/** `vectorFromArray` cannot build Decimal128, so write the 128-bit words directly. */ +function decimalVector(unscaled: Array, type: Decimal) { + const words = new Uint32Array(unscaled.length * 4); + const validity = new Uint8Array(Math.ceil(unscaled.length / 8)); + let nullCount = 0; + unscaled.forEach((value, i) => { + if (value === null) { + nullCount++; + return; + } + validity[i >> 3] |= 1 << (i & 7); + const twosComplement = value < 0n ? (1n << 128n) + value : value; + for (let w = 0; w < 4; w++) { + words[i * 4 + w] = Number( + (twosComplement >> BigInt(32 * w)) & 0xffffffffn, + ); + } + }); + return makeVector( + makeData({ + type: new Decimal(type.scale, type.precision, 128), + length: unscaled.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..1c60a210b 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -31,7 +31,7 @@ import type { StatementStatus, WorkspaceClient, } from "../../workspace-client"; -import { buildEmptyArrowIPCBase64 } from "./arrow-schema"; +import { buildEmptyArrowIPCBase64, manifestColumnName } from "./arrow-schema"; import { executeStatementDefaults } from "./defaults"; import { WarehousePollBackoff } from "./warehouse-poll-backoff"; import { WarehouseStatusEmitter } from "./warehouse-status-emitter"; @@ -48,15 +48,12 @@ const logger = createLogger("connectors:sql-warehouse"); function arrowColumnNames(response: StatementResponse): string[] | undefined { const cols = response.manifest?.schema?.columns; if (!cols || cols.length === 0) return undefined; - return cols.map((c, i) => - c.name && c.name.length > 0 ? c.name : `column_${i}`, - ); + return cols.map(manifestColumnName); } /** * Coerce the modular SDK's `bigint` row/byte counts back to `number` (the type - * the legacy SDK used). AppKit never does arithmetic on these — they are purely - * informational — but a stray `bigint` makes `JSON.stringify` throw ("Do not + * the legacy SDK used). A stray `bigint` makes `JSON.stringify` throw ("Do not * know how to serialize a BigInt") the instant the result is cached or written * to an SSE frame. Reyden's INLINE + ARROW_STREAM result — which the analytics * arrow path caches — carries them on `result`/`manifest`, so normalize every @@ -102,7 +99,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 +243,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 +356,7 @@ export class SQLWarehouseConnector { statementId, this.config.timeout, signal, + input.rawRows, ); break; case "SUCCEEDED": @@ -360,6 +364,7 @@ export class SQLWarehouseConnector { response, workspaceClient, signal, + input.rawRows, ); break; case "FAILED": @@ -830,6 +835,7 @@ export class SQLWarehouseConnector { statementId: string, timeout = executeStatementDefaults.timeout, signal?: AbortSignal, + rawRows?: boolean, ) { return this.telemetry.startActiveSpan( "sql.poll", @@ -908,6 +914,7 @@ export class SQLWarehouseConnector { response, workspaceClient, signal, + rawRows, ); case "FAILED": throw ExecutionError.statementFailed( @@ -949,6 +956,7 @@ export class SQLWarehouseConnector { response: StatementResponse, workspaceClient: WorkspaceClient, signal?: AbortSignal, + rawRows?: boolean, ) { if (response.manifest?.format === "ARROW_STREAM") { const result = response.result as @@ -998,17 +1006,28 @@ 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 = {}; row.forEach((value, index) => { const column = columns[index]; - const columnName = column?.name || `column_${index}`; + const columnName = manifestColumnName(column, index); // attempt to parse JSON strings for string columns if ( @@ -1108,6 +1127,7 @@ export class SQLWarehouseConnector { statement_id: string; status: StatementStatus; columnNames?: string[]; + totalByteCount?: number; external_links?: ExternalLink[]; refreshChunkLink?: RefreshChunkLink; }; @@ -1121,6 +1141,10 @@ export class SQLWarehouseConnector { error: response.status?.error, } as StatementStatus, columnNames: arrowColumnNames(response), + totalByteCount: + response.manifest?.totalByteCount === undefined + ? undefined + : Number(response.manifest.totalByteCount), // 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,14 @@ export class SQLWarehouseConnector { chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, + fullHeaderWait?: boolean, ): AsyncGenerator { - return this.arrowProcessor.streamChunks(chunks, signal, refresh); + return this.arrowProcessor.streamChunks( + chunks, + signal, + refresh, + fullHeaderWait, + ); } /** diff --git a/packages/appkit/src/connectors/sql-warehouse/json-array-cells.ts b/packages/appkit/src/connectors/sql-warehouse/json-array-cells.ts new file mode 100644 index 000000000..dc6728b76 --- /dev/null +++ b/packages/appkit/src/connectors/sql-warehouse/json-array-cells.ts @@ -0,0 +1,169 @@ +import { + type DataType, + type Decimal, + type Int, + type Timestamp, + Type, +} from "apache-arrow"; + +/** + * How one scalar type is written in an `INLINE` + `JSON_ARRAY` cell. `parse` + * returns what `vectorFromArray` accepts: an unscaled `bigint` for DECIMAL, + * epoch milliseconds for DATE and TIMESTAMP. + */ +interface CellCodec { + /** Defaults to `String(value)`, or base64 for bytes. */ + format?(value: unknown, type: DataType): string; + parse(cell: string, type: DataType): unknown; +} + +const DECIMAL_TEXT = /^(-?)(\d+)(?:\.(\d+))?$/; +const FLOAT_TEXT = + /^-?(?:\d+\.?\d*|\.\d+)(?:[eE][+-]?\d+)?$|^(?:NaN|-?Infinity)$/; +const BASE64_TEXT = /^[A-Za-z0-9+/]*={0,2}$/; +const DATE_TEXT = /^\d{4}-\d{2}-\d{2}$/; +/** `TIMESTAMP` cells carry the session offset (or `Z`); `TIMESTAMP_NTZ` cells carry none. */ +const ZONED_TIMESTAMP_TEXT = + /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:\d{2})$/; +const LOCAL_TIMESTAMP_TEXT = /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?$/; + +function invalid(kind: string, cell: string): Error { + return new Error(`Invalid ${kind} "${cell}"`); +} + +/** `Date.parse` rolls 2026-02-31 over to March, so the day must survive a round trip. */ +function epochMillis(text: string, kind: string, cell: string): number { + const day = cell.slice(0, 10); + const ms = Date.parse(text); + const dayMs = Date.parse(`${day}T00:00:00Z`); + if ( + Number.isNaN(ms) || + Number.isNaN(dayMs) || + !new Date(dayMs).toISOString().startsWith(day) + ) { + throw invalid(kind, cell); + } + return ms; +} + +const hasTimezone = (type: Timestamp) => type.timezone != null; + +const CODECS: Partial> = { + [Type.Int]: { + parse(cell, type: Int) { + const { bitWidth } = type; + if (bitWidth === 64) { + let n: bigint; + try { + n = BigInt(cell); + } catch { + throw invalid("integer", cell); + } + if (String(n) !== cell) throw invalid("integer", cell); + return n; + } + const n = Number(cell); + const limit = 2 ** (bitWidth - 1); + if ( + String(n) !== cell || + !Number.isInteger(n) || + n < -limit || + n >= limit + ) { + throw invalid("integer", cell); + } + return n; + }, + }, + [Type.Float]: { + parse(cell) { + if (!FLOAT_TEXT.test(cell)) throw invalid("number", cell); + return Number(cell); + }, + }, + [Type.Decimal]: { + format(value, type: Decimal) { + const { scale } = type; + const unscaled = String(value); + if (scale <= 0) return unscaled; + const negative = unscaled.startsWith("-"); + let digits = negative ? unscaled.slice(1) : unscaled; + if (digits.length <= scale) digits = digits.padStart(scale + 1, "0"); + const point = digits.length - scale; + const out = `${digits.slice(0, point)}.${digits.slice(point)}`; + return negative ? `-${out}` : out; + }, + parse(cell, type: Decimal) { + const { scale, precision } = type; + const match = DECIMAL_TEXT.exec(cell); + if (!match) throw invalid("decimal", cell); + const [, sign, whole, fraction = ""] = match; + const digits = `${whole}${fraction.padEnd(scale, "0")}`; + if ( + fraction.length > scale || + digits.replace(/^0+/, "").length > precision + ) { + throw invalid("decimal", cell); + } + const unscaled = BigInt(digits); + return sign === "-" ? -unscaled : unscaled; + }, + }, + [Type.Date]: { + format: (value) => new Date(Number(value)).toISOString().slice(0, 10), + parse(cell) { + if (!DATE_TEXT.test(cell)) throw invalid("date", cell); + return epochMillis(`${cell}T00:00:00Z`, "date", cell); + }, + }, + [Type.Timestamp]: { + format(value, type: Timestamp) { + const iso = new Date(Number(value)).toISOString(); + return hasTimezone(type) ? iso : iso.slice(0, -1); + }, + parse(cell, type: Timestamp) { + if (hasTimezone(type)) { + if (!ZONED_TIMESTAMP_TEXT.test(cell)) throw invalid("timestamp", cell); + return epochMillis(cell, "timestamp", cell); + } + if (!LOCAL_TIMESTAMP_TEXT.test(cell)) throw invalid("timestamp", cell); + return epochMillis(`${cell}Z`, "timestamp", cell); + }, + }, + [Type.Bool]: { + parse(cell) { + if (cell !== "true" && cell !== "false") throw invalid("boolean", cell); + return cell === "true"; + }, + }, + [Type.Binary]: { + parse(cell) { + if (!BASE64_TEXT.test(cell)) throw invalid("binary", cell); + return new Uint8Array(Buffer.from(cell, "base64")); + }, + }, + [Type.Utf8]: { + parse: (cell) => cell, + }, +}; + +/** Whether {@link parseJsonArrayCell} can read cells of this Arrow type. */ +export function hasJsonArrayCodec(type: DataType): boolean { + return CODECS[type.typeId] !== undefined; +} + +/** Render a scalar Arrow value the way the warehouse writes it under `JSON_ARRAY`. */ +export function formatJsonArrayCell(value: unknown, type: DataType): string { + const format = CODECS[type.typeId]?.format; + if (format) return format(value, type); + if (value instanceof Uint8Array) return Buffer.from(value).toString("base64"); + if (value instanceof Date) return value.toISOString(); + return String(value); +} + +/** Read a `JSON_ARRAY` cell as a value of `type`. Throws when the text does not fit the type. */ +export function parseJsonArrayCell(cell: string, type: DataType): unknown { + const codec = CODECS[type.typeId]; + if (!codec) throw new Error(`No JSON_ARRAY codec for ${String(type)}`); + return codec.parse(cell, type); +} 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..4eeb52517 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 @@ -1,7 +1,9 @@ +import { readFileSync } from "node:fs"; + import { Binary, Bool, - type DataType, + DataType, DateDay, Decimal, DurationMicrosecond, @@ -23,7 +25,12 @@ import { } from "apache-arrow"; import { describe, expect, test } from "vitest"; -import { buildEmptyArrowIPCBase64, parseDatabricksType } from "../arrow-schema"; +import type { ColumnInfo } from "../../../workspace-client"; +import { + buildArrowIPCFromRows, + buildEmptyArrowIPCBase64, + parseDatabricksType, +} from "../arrow-schema"; // ============================================================================ // Helpers @@ -513,3 +520,297 @@ 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)); + 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 instanceof Uint8Array ? [...v] : v), + }, + 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(column.name); + expect(field.typeId).toBe(native.typeId); + if (DataType.isDecimal(field.type)) { + expect([field.type.precision, field.type.scale]).toEqual([ + column.typePrecision, + column.typeScale, + ]); + } else { + expect(String(field.type)).toBe(String(native)); + } + const values = [...(table.getChildAt(0) ?? [])]; + expect(values.map((v) => (read ? read(v) : v))).toEqual(expected); + }, + ); + + test("duplicate column names and JSON-looking strings come through 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([ + "x", + "x", + "payload", + ]); + expect(table.getChildAt(0)?.get(0)).toBe(1); + expect(table.getChildAt(1)?.get(0)).toBe(2); + expect(table.getChildAt(2)?.get(0)).toBe(payload); + }); + + test("an unnamed column gets the column_N name the header carries", () => { + const ipc = buildArrowIPCFromRows([{ name: "", typeText: "INT" }], [["1"]]); + expect(tableFromIPC(ipc).schema.fields[0].name).toBe("column_0"); + }); + + 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({ + a: row[0], + m: row[1], + s: row[2], + i: row[3], + g: 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/); + }); +}); + +describe("buildArrowIPCFromRows against a captured warehouse response", () => { + const fixture = JSON.parse( + readFileSync( + new URL("./fixtures/typed-values.json", import.meta.url), + "utf8", + ), + ) as { + columns: ColumnInfo[]; + dataArray: Array>; + arrowIpcBase64: string; + }; + const native = tableFromIPC(Buffer.from(fixture.arrowIpcBase64, "base64")); + const rebuilt = tableFromIPC( + buildArrowIPCFromRows(fixture.columns, fixture.dataArray), + ); + const comparable = (value: unknown) => + value !== null && typeof value === "object" ? String(value) : value; + const values = (table: typeof native, index: number) => + [...(table.getChildAt(index) ?? [])].map(comparable); + const SESSION_ZONE = "America/Los_Angeles"; + + test("has the same rows and column names, duplicates included", () => { + expect(rebuilt.numRows).toBe(native.numRows); + expect(rebuilt.schema.fields.map((f) => f.name)).toEqual( + native.schema.fields.map((f) => f.name), + ); + }); + + test.each( + fixture.columns + .map((column, index) => [column.name ?? "", index] as const) + .filter(([name]) => !["ts", "ts_micro", "arr"].includes(name)), + )("column %s (%i) has the native Arrow type and values", (_name, index) => { + expect(String(rebuilt.schema.fields[index].type)).toBe( + String(native.schema.fields[index].type), + ); + expect(values(rebuilt, index)).toEqual(values(native, index)); + }); + + test("TIMESTAMP keeps every instant but is labeled UTC instead of the session zone", () => { + const index = fixture.columns.findIndex((c) => c.name === "ts"); + expect(native.schema.fields[index].type.timezone).toBe(SESSION_ZONE); + expect(rebuilt.schema.fields[index].type.timezone).toBe("UTC"); + expect(values(rebuilt, index)).toEqual(values(native, index)); + }); + + test("TIMESTAMP keeps milliseconds, because the JSON text has no microseconds", () => { + const index = fixture.columns.findIndex((c) => c.name === "ts_micro"); + expect(fixture.dataArray[0][index]).toMatch(/\.\d{3}-\d{2}:\d{2}$/); + expect(values(rebuilt, index)).toEqual( + values(native, index).map((ms) => + ms === null ? null : Math.trunc(Number(ms)), + ), + ); + }); + + test("ARRAY arrives as JSON text and stays a string column", () => { + const index = fixture.columns.findIndex((c) => c.name === "arr"); + expect(native.schema.fields[index].type.typeId).toBe(Type.List); + expect(rebuilt.schema.fields[index].type.typeId).toBe(Type.Utf8); + expect(values(rebuilt, index)).toEqual( + fixture.dataArray.map((row) => row[index]), + ); + }); +}); diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/fixtures/typed-values.json b/packages/appkit/src/connectors/sql-warehouse/tests/fixtures/typed-values.json new file mode 100644 index 000000000..e3872788a --- /dev/null +++ b/packages/appkit/src/connectors/sql-warehouse/tests/fixtures/typed-values.json @@ -0,0 +1,683 @@ +{ + "source": "Statement Execution API on a serverless PRO SQL warehouse (session time zone America/Los_Angeles), captured 2026-11-10", + "statement": "SELECT\n CAST(id AS INT) AS id,\n IF(id % 5 = 4, NULL, CAST(id AS BIGINT) * 1000000000007) AS big,\n IF(id % 5 = 4, NULL, CAST(id % 100 AS TINYINT)) AS tiny,\n IF(id % 5 = 4, NULL, CAST(id AS SMALLINT)) AS small,\n IF(id % 5 = 4, NULL, CAST(id * 1.5 AS DOUBLE)) AS amount,\n IF(id % 5 = 4, NULL, CAST(id * 0.25 AS FLOAT)) AS ratio,\n IF(id % 5 = 4, NULL, CAST((id * 101 - 50) / 100.0 AS DECIMAL(10, 2))) AS price,\n IF(id % 5 = 4, NULL, DATE_ADD(DATE'2026-01-01', CAST(id AS INT))) AS day,\n IF(id % 5 = 4, NULL, TIMESTAMPADD(HOUR, id, TIMESTAMP'2026-01-01 10:00:00')) AS ts,\n IF(id % 5 = 4, NULL, TIMESTAMPADD(HOUR, id, TIMESTAMP_NTZ'2026-01-01 10:00:00')) AS ts_ntz,\n IF(id % 5 = 4, NULL, TIMESTAMPADD(MICROSECOND, id, TIMESTAMP'2026-01-01 10:00:00.123456')) AS ts_micro,\n IF(id % 5 = 4, NULL, CONCAT('{\"order_id\": 1234567890123456789', CAST(id % 10 AS STRING), ', \"tags\": [1, 2]}')) AS payload,\n IF(id % 5 = 4, NULL, id % 2 = 0) AS flag,\n CASE WHEN id % 7 = 3 THEN NULL ELSE CONCAT('row-', CAST(id AS STRING)) END AS label,\n IF(id % 5 = 4, NULL, ARRAY(id, id + 1)) AS arr,\n IF(id % 5 = 4, NULL, id) AS dup,\n IF(id % 5 = 4, NULL, id + 1) AS dup\nFROM range(0, 30) AS t(id)\nORDER BY id", + "columns": [ + { + "name": "id", + "typeName": "INT", + "typeText": "INT", + "position": 0 + }, + { + "name": "big", + "typeName": "LONG", + "typeText": "BIGINT", + "position": 1 + }, + { + "name": "tiny", + "typeName": "BYTE", + "typeText": "TINYINT", + "position": 2 + }, + { + "name": "small", + "typeName": "SHORT", + "typeText": "SMALLINT", + "position": 3 + }, + { + "name": "amount", + "typeName": "DOUBLE", + "typeText": "DOUBLE", + "position": 4 + }, + { + "name": "ratio", + "typeName": "FLOAT", + "typeText": "FLOAT", + "position": 5 + }, + { + "name": "price", + "typeName": "DECIMAL", + "typeText": "DECIMAL(10,2)", + "position": 6, + "typePrecision": 10, + "typeScale": 2 + }, + { + "name": "day", + "typeName": "DATE", + "typeText": "DATE", + "position": 7 + }, + { + "name": "ts", + "typeName": "TIMESTAMP", + "typeText": "TIMESTAMP", + "position": 8 + }, + { + "name": "ts_ntz", + "typeName": "TIMESTAMP_NTZ", + "typeText": "TIMESTAMP_NTZ", + "position": 9 + }, + { + "name": "ts_micro", + "typeName": "TIMESTAMP", + "typeText": "TIMESTAMP", + "position": 10 + }, + { + "name": "payload", + "typeName": "STRING", + "typeText": "STRING", + "position": 11 + }, + { + "name": "flag", + "typeName": "BOOLEAN", + "typeText": "BOOLEAN", + "position": 12 + }, + { + "name": "label", + "typeName": "STRING", + "typeText": "STRING", + "position": 13 + }, + { + "name": "arr", + "typeName": "ARRAY", + "typeText": "ARRAY", + "position": 14 + }, + { + "name": "dup", + "typeName": "LONG", + "typeText": "BIGINT", + "position": 15 + }, + { + "name": "dup", + "typeName": "LONG", + "typeText": "BIGINT", + "position": 16 + } + ], + "dataArray": [ + [ + "0", + "0", + "0", + "0", + "0.0", + "0.0", + "-0.50", + "2026-01-01", + "2026-01-01T10:00:00.000-08:00", + "2026-01-01T10:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567890, \"tags\": [1, 2]}", + "true", + "row-0", + "[\"0\",\"1\"]", + "0", + "1" + ], + [ + "1", + "1000000000007", + "1", + "1", + "1.5", + "0.25", + "0.51", + "2026-01-02", + "2026-01-01T11:00:00.000-08:00", + "2026-01-01T11:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567891, \"tags\": [1, 2]}", + "false", + "row-1", + "[\"1\",\"2\"]", + "1", + "2" + ], + [ + "2", + "2000000000014", + "2", + "2", + "3.0", + "0.5", + "1.52", + "2026-01-03", + "2026-01-01T12:00:00.000-08:00", + "2026-01-01T12:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567892, \"tags\": [1, 2]}", + "true", + "row-2", + "[\"2\",\"3\"]", + "2", + "3" + ], + [ + "3", + "3000000000021", + "3", + "3", + "4.5", + "0.75", + "2.53", + "2026-01-04", + "2026-01-01T13:00:00.000-08:00", + "2026-01-01T13:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567893, \"tags\": [1, 2]}", + "false", + null, + "[\"3\",\"4\"]", + "3", + "4" + ], + [ + "4", + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + "row-4", + null, + null, + null + ], + [ + "5", + "5000000000035", + "5", + "5", + "7.5", + "1.25", + "4.55", + "2026-01-06", + "2026-01-01T15:00:00.000-08:00", + "2026-01-01T15:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567895, \"tags\": [1, 2]}", + "false", + "row-5", + "[\"5\",\"6\"]", + "5", + "6" + ], + [ + "6", + "6000000000042", + "6", + "6", + "9.0", + "1.5", + "5.56", + "2026-01-07", + "2026-01-01T16:00:00.000-08:00", + "2026-01-01T16:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567896, \"tags\": [1, 2]}", + "true", + "row-6", + "[\"6\",\"7\"]", + "6", + "7" + ], + [ + "7", + "7000000000049", + "7", + "7", + "10.5", + "1.75", + "6.57", + "2026-01-08", + "2026-01-01T17:00:00.000-08:00", + "2026-01-01T17:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567897, \"tags\": [1, 2]}", + "false", + "row-7", + "[\"7\",\"8\"]", + "7", + "8" + ], + [ + "8", + "8000000000056", + "8", + "8", + "12.0", + "2.0", + "7.58", + "2026-01-09", + "2026-01-01T18:00:00.000-08:00", + "2026-01-01T18:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567898, \"tags\": [1, 2]}", + "true", + "row-8", + "[\"8\",\"9\"]", + "8", + "9" + ], + [ + "9", + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + "row-9", + null, + null, + null + ], + [ + "10", + "10000000000070", + "10", + "10", + "15.0", + "2.5", + "9.60", + "2026-01-11", + "2026-01-01T20:00:00.000-08:00", + "2026-01-01T20:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567890, \"tags\": [1, 2]}", + "true", + null, + "[\"10\",\"11\"]", + "10", + "11" + ], + [ + "11", + "11000000000077", + "11", + "11", + "16.5", + "2.75", + "10.61", + "2026-01-12", + "2026-01-01T21:00:00.000-08:00", + "2026-01-01T21:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567891, \"tags\": [1, 2]}", + "false", + "row-11", + "[\"11\",\"12\"]", + "11", + "12" + ], + [ + "12", + "12000000000084", + "12", + "12", + "18.0", + "3.0", + "11.62", + "2026-01-13", + "2026-01-01T22:00:00.000-08:00", + "2026-01-01T22:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567892, \"tags\": [1, 2]}", + "true", + "row-12", + "[\"12\",\"13\"]", + "12", + "13" + ], + [ + "13", + "13000000000091", + "13", + "13", + "19.5", + "3.25", + "12.63", + "2026-01-14", + "2026-01-01T23:00:00.000-08:00", + "2026-01-01T23:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567893, \"tags\": [1, 2]}", + "false", + "row-13", + "[\"13\",\"14\"]", + "13", + "14" + ], + [ + "14", + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + "row-14", + null, + null, + null + ], + [ + "15", + "15000000000105", + "15", + "15", + "22.5", + "3.75", + "14.65", + "2026-01-16", + "2026-01-02T01:00:00.000-08:00", + "2026-01-02T01:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567895, \"tags\": [1, 2]}", + "false", + "row-15", + "[\"15\",\"16\"]", + "15", + "16" + ], + [ + "16", + "16000000000112", + "16", + "16", + "24.0", + "4.0", + "15.66", + "2026-01-17", + "2026-01-02T02:00:00.000-08:00", + "2026-01-02T02:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567896, \"tags\": [1, 2]}", + "true", + "row-16", + "[\"16\",\"17\"]", + "16", + "17" + ], + [ + "17", + "17000000000119", + "17", + "17", + "25.5", + "4.25", + "16.67", + "2026-01-18", + "2026-01-02T03:00:00.000-08:00", + "2026-01-02T03:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567897, \"tags\": [1, 2]}", + "false", + null, + "[\"17\",\"18\"]", + "17", + "18" + ], + [ + "18", + "18000000000126", + "18", + "18", + "27.0", + "4.5", + "17.68", + "2026-01-19", + "2026-01-02T04:00:00.000-08:00", + "2026-01-02T04:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567898, \"tags\": [1, 2]}", + "true", + "row-18", + "[\"18\",\"19\"]", + "18", + "19" + ], + [ + "19", + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + "row-19", + null, + null, + null + ], + [ + "20", + "20000000000140", + "20", + "20", + "30.0", + "5.0", + "19.70", + "2026-01-21", + "2026-01-02T06:00:00.000-08:00", + "2026-01-02T06:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567890, \"tags\": [1, 2]}", + "true", + "row-20", + "[\"20\",\"21\"]", + "20", + "21" + ], + [ + "21", + "21000000000147", + "21", + "21", + "31.5", + "5.25", + "20.71", + "2026-01-22", + "2026-01-02T07:00:00.000-08:00", + "2026-01-02T07:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567891, \"tags\": [1, 2]}", + "false", + "row-21", + "[\"21\",\"22\"]", + "21", + "22" + ], + [ + "22", + "22000000000154", + "22", + "22", + "33.0", + "5.5", + "21.72", + "2026-01-23", + "2026-01-02T08:00:00.000-08:00", + "2026-01-02T08:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567892, \"tags\": [1, 2]}", + "true", + "row-22", + "[\"22\",\"23\"]", + "22", + "23" + ], + [ + "23", + "23000000000161", + "23", + "23", + "34.5", + "5.75", + "22.73", + "2026-01-24", + "2026-01-02T09:00:00.000-08:00", + "2026-01-02T09:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567893, \"tags\": [1, 2]}", + "false", + "row-23", + "[\"23\",\"24\"]", + "23", + "24" + ], + [ + "24", + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null + ], + [ + "25", + "25000000000175", + "25", + "25", + "37.5", + "6.25", + "24.75", + "2026-01-26", + "2026-01-02T11:00:00.000-08:00", + "2026-01-02T11:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567895, \"tags\": [1, 2]}", + "false", + "row-25", + "[\"25\",\"26\"]", + "25", + "26" + ], + [ + "26", + "26000000000182", + "26", + "26", + "39.0", + "6.5", + "25.76", + "2026-01-27", + "2026-01-02T12:00:00.000-08:00", + "2026-01-02T12:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567896, \"tags\": [1, 2]}", + "true", + "row-26", + "[\"26\",\"27\"]", + "26", + "27" + ], + [ + "27", + "27000000000189", + "27", + "27", + "40.5", + "6.75", + "26.77", + "2026-01-28", + "2026-01-02T13:00:00.000-08:00", + "2026-01-02T13:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567897, \"tags\": [1, 2]}", + "false", + "row-27", + "[\"27\",\"28\"]", + "27", + "28" + ], + [ + "28", + "28000000000196", + "28", + "28", + "42.0", + "7.0", + "27.78", + "2026-01-29", + "2026-01-02T14:00:00.000-08:00", + "2026-01-02T14:00:00.000", + "2026-01-01T10:00:00.123-08:00", + "{\"order_id\": 12345678901234567898, \"tags\": [1, 2]}", + "true", + "row-28", + "[\"28\",\"29\"]", + "28", + "29" + ], + [ + "29", + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + null, + "row-29", + null, + null, + null + ] + ], + "arrowIpcBase64": "/////6gEAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAABEAAAAsBAAA4AMAAJQDAABYAwAAGAMAAOACAACcAgAAaAIAABACAADYAQAAgAEAAEgBAAAUAQAA4AAAAHQAAAA8AAAABAAAAHL8//8UAAAAFAAAABQAAAAAAAIBGAAAAAAAAAAAAAAAGPz//wAAAAFAAAAAAwAAAGR1cACm/P//FAAAABQAAAAUAAAAAAACARgAAAAAAAAAAAAAAEz8//8AAAABQAAAAAMAAABkdXAA2vz//xQAAAAUAAAAUAAAAAAADAFMAAAAAAAAAAEAAAAEAAAAtvz//xQAAAAUAAAAFAAAAAAAAAIYAAAAAAAAAAAAAACk/P//AAAAAUAAAAAHAAAAZWxlbWVudAB0////AwAAAGFycgBC/f//FAAAABQAAAAUAAAAAAAFARAAAAAAAAAAAAAAAKD///8FAAAAbGFiZWwAAABy/f//FAAAABQAAAAUAAAAAAAGARAAAAAAAAAAAAAAAND///8EAAAAZmxhZwAAAACi/f//FAAAABQAAAAYAAAAAAAFARQAAAAAAAAAAAAAAAQABAAEAAAABwAAAHBheWxvYWQA1v3//xQAAAAUAAAAFAAAAAAACgEwAAAAAAAAAAAAAAB4////CAAAAAAAAgATAAAAQW1lcmljYS9Mb3NfQW5nZWxlcwAIAAAAdHNfbWljcm8AAAAAKv7//xQAAAAUAAAAFAAAAAAACgEUAAAAAAAAAAAAAADS/v//AAACAAYAAAB0c19udHoAAF7+//8UAAAAFAAAABwAAAAAAAoBOAAAAAAAAAAAAAAACAAMAAoABAAIAAAACAAAAAAAAgATAAAAQW1lcmljYS9Mb3NfQW5nZWxlcwACAAAAdHMAALL+//8UAAAAFAAAABQAAAAAAAgBFAAAAAAAAAAAAAAAWv///wAAAAADAAAAZGF5AOL+//8UAAAAFAAAABwAAAAAAAcBIAAAAAAAAAAAAAAACAAMAAgABAAIAAAAAgAAAAoAAAAFAAAAcHJpY2UAAAAi////FAAAABQAAAAUAAAAAAADARQAAAAAAAAAAAAAAMr///8AAAEABQAAAHJhdGlvAAAAVv///xQAAAAUAAAAHAAAAAAAAwEcAAAAAAAAAAAAAAAAAAYACAAGAAYAAAAAAAIABgAAAGFtb3VudAAAkv///xQAAAAUAAAAFAAAAAAAAgEYAAAAAAAAAAAAAAA4////AAAAARAAAAAFAAAAc21hbGwAAADK////FAAAABQAAAAUAAAAAAACARgAAAAAAAAAAAAAAHD///8AAAABCAAAAAQAAAB0aW55AAASABgAFAATABIADAAAAAgABAASAAAAFAAAABQAAAAUAAAAAAACARgAAAAAAAAAAAAAALj///8AAAABQAAAAAMAAABiaWcAAAASABgAFAAAABMADAAAAAgABAASAAAAFAAAABQAAAAcAAAAAAAAAiAAAAAAAAAAAAAAAAgADAAIAAcACAAAAAAAAAEgAAAAAgAAAGlkAAAAAAAA/////9gDAAAUAAAAAAAAAAwAFgAGAAUACAAMAAwAAAAAAwQAGAAAABgTAAAAAAAAAAAKABgADAAEAAgACgAAAHwCAAAQAAAAHgAAAAAAAAAAAAAAJgAAAAAAAAAAAAAABAAAAAAAAAAIAAAAAAAAAHgAAAAAAAAAgAAAAAAAAAAEAAAAAAAAAIgAAAAAAAAA8AAAAAAAAAB4AQAAAAAAAAQAAAAAAAAAgAEAAAAAAAAeAAAAAAAAAKABAAAAAAAABAAAAAAAAACoAQAAAAAAADwAAAAAAAAA6AEAAAAAAAAEAAAAAAAAAPABAAAAAAAA8AAAAAAAAADgAgAAAAAAAAQAAAAAAAAA6AIAAAAAAAB4AAAAAAAAAGADAAAAAAAABAAAAAAAAABoAwAAAAAAAOABAAAAAAAASAUAAAAAAAAEAAAAAAAAAFAFAAAAAAAAeAAAAAAAAADIBQAAAAAAAAQAAAAAAAAA0AUAAAAAAADwAAAAAAAAAMAGAAAAAAAABAAAAAAAAADIBgAAAAAAAPAAAAAAAAAAuAcAAAAAAAAEAAAAAAAAAMAHAAAAAAAA8AAAAAAAAACwCAAAAAAAAAQAAAAAAAAAuAgAAAAAAAB8AAAAAAAAADgJAAAAAAAAsAQAAAAAAADoDQAAAAAAAAQAAAAAAAAA8A0AAAAAAAAEAAAAAAAAAPgNAAAAAAAABAAAAAAAAAAADgAAAAAAAHwAAAAAAAAAgA4AAAAAAACTAAAAAAAAABgPAAAAAAAABAAAAAAAAAAgDwAAAAAAAHwAAAAAAAAAoA8AAAAAAAAGAAAAAAAAAKgPAAAAAAAAgAEAAAAAAAAoEQAAAAAAAAQAAAAAAAAAMBEAAAAAAADwAAAAAAAAACASAAAAAAAABAAAAAAAAAAoEgAAAAAAAPAAAAAAAAAAAAAAABIAAAAeAAAAAAAAAAAAAAAAAAAAHgAAAAAAAAAGAAAAAAAAAB4AAAAAAAAABgAAAAAAAAAeAAAAAAAAAAYAAAAAAAAAHgAAAAAAAAAGAAAAAAAAAB4AAAAAAAAABgAAAAAAAAAeAAAAAAAAAAYAAAAAAAAAHgAAAAAAAAAGAAAAAAAAAB4AAAAAAAAABgAAAAAAAAAeAAAAAAAAAAYAAAAAAAAAHgAAAAAAAAAGAAAAAAAAAB4AAAAAAAAABgAAAAAAAAAeAAAAAAAAAAYAAAAAAAAAHgAAAAAAAAAEAAAAAAAAAB4AAAAAAAAABgAAAAAAAAAwAAAAAAAAAAAAAAAAAAAAHgAAAAAAAAAGAAAAAAAAAB4AAAAAAAAABgAAAAAAAAD///8/AAAAAAAAAAABAAAAAgAAAAMAAAAEAAAABQAAAAYAAAAHAAAACAAAAAkAAAAKAAAACwAAAAwAAAANAAAADgAAAA8AAAAQAAAAEQAAABIAAAATAAAAFAAAABUAAAAWAAAAFwAAABgAAAAZAAAAGgAAABsAAAAcAAAAHQAAAO+99x4AAAAAAAAAAAAAAAAHEKXU6AAAAA4gSqnRAQAAFTDvfboCAAAAAAAAAAAAACNQOSeMBAAAKmDe+3QFAAAxcIPQXQYAADiAKKVGBwAAAAAAAAAAAABGoHJOGAkAAE2wFyMBCgAAVMC89+kKAABb0GHM0gsAAAAAAAAAAAAAafCrdaQNAABwAFFKjQ4AAHcQ9h52DwAAfiCb814QAAAAAAAAAAAAAIxA5ZwwEgAAk1CKcRkTAACaYC9GAhQAAKFw1BrrFAAAAAAAAAAAAACvkB7EvBYAALagw5ilFwAAvbBobY4YAADEwA1CdxkAAAAAAAAAAAAA7733HgAAAAAAAQIDAAUGBwgACgsMDQAPEBESABQVFhcAGRobHAAAAO+99x4AAAAAAAABAAIAAwAAAAUABgAHAAgAAAAKAAsADAANAAAADwAQABEAEgAAABQAFQAWABcAAAAZABoAGwAcAAAAAAAAAO+99x4AAAAAAAAAAAAAAAAAAAAAAAD4PwAAAAAAAAhAAAAAAAAAEkAAAAAAAAAAAAAAAAAAAB5AAAAAAAAAIkAAAAAAAAAlQAAAAAAAAChAAAAAAAAAAAAAAAAAAAAuQAAAAAAAgDBAAAAAAAAAMkAAAAAAAIAzQAAAAAAAAAAAAAAAAACANkAAAAAAAAA4QAAAAAAAgDlAAAAAAAAAO0AAAAAAAAAAAAAAAAAAAD5AAAAAAACAP0AAAAAAAIBAQAAAAAAAQEFAAAAAAAAAAAAAAAAAAMBCQAAAAAAAgENAAAAAAABAREAAAAAAAABFQAAAAAAAAAAA7733HgAAAAAAAAAAAACAPgAAAD8AAEA/AAAAAAAAoD8AAMA/AADgPwAAAEAAAAAAAAAgQAAAMEAAAEBAAABQQAAAAAAAAHBAAACAQAAAiEAAAJBAAAAAAAAAoEAAAKhAAACwQAAAuEAAAAAAAADIQAAA0EAAANhAAADgQAAAAADvvfceAAAAAM7///////////////////8zAAAAAAAAAAAAAAAAAAAAmAAAAAAAAAAAAAAAAAAAAP0AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAxwEAAAAAAAAAAAAAAAAAACwCAAAAAAAAAAAAAAAAAACRAgAAAAAAAAAAAAAAAAAA9gIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAADAAwAAAAAAAAAAAAAAAAAAJQQAAAAAAAAAAAAAAAAAAIoEAAAAAAAAAAAAAAAAAADvBAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAALkFAAAAAAAAAAAAAAAAAAAeBgAAAAAAAAAAAAAAAAAAgwYAAAAAAAAAAAAAAAAAAOgGAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAsgcAAAAAAAAAAAAAAAAAABcIAAAAAAAAAAAAAAAAAAB8CAAAAAAAAAAAAAAAAAAA4QgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACrCQAAAAAAAAAAAAAAAAAAEAoAAAAAAAAAAAAAAAAAAHUKAAAAAAAAAAAAAAAAAADaCgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAO+99x4AAAAA5k8AAOdPAADoTwAA6U8AAAAAAADrTwAA7E8AAO1PAADuTwAAAAAAAPBPAADxTwAA8k8AAPNPAAAAAAAA9U8AAPZPAAD3TwAA+E8AAAAAAAD6TwAA+08AAPxPAAD9TwAAAAAAAP9PAAAAUAAAAVAAAAJQAAAAAAAA7733HgAAAAAAyIFcV0cGAABsFTNYRwYAABCpCVlHBgAAtDzgWUcGAAAAAAAAAAAAAPxjjVtHBgAAoPdjXEcGAABEizpdRwYAAOgeEV5HBgAAAAAAAAAAAAAwRr5fRwYAANTZlGBHBgAAeG1rYUcGAAAcAUJiRwYAAAAAAAAAAAAAZCjvY0cGAAAIvMVkRwYAAKxPnGVHBgAAUONyZkcGAAAAAAAAAAAAAJgKIGhHBgAAPJ72aEcGAADgMc1pRwYAAITFo2pHBgAAAAAAAAAAAADM7FBsRwYAAHCAJ21HBgAAFBT+bUcGAAC4p9RuRwYAAAAAAAAAAADvvfceAAAAAACo5KdQRwYAAEx4flFHBgAA8AtVUkcGAACUnytTRwYAAAAAAAAAAAAA3MbYVEcGAACAWq9VRwYAACTuhVZHBgAAyIFcV0cGAAAAAAAAAAAAABCpCVlHBgAAtDzgWUcGAABY0LZaRwYAAPxjjVtHBgAAAAAAAAAAAABEizpdRwYAAOgeEV5HBgAAjLLnXkcGAAAwRr5fRwYAAAAAAAAAAAAAeG1rYUcGAAAcAUJiRwYAAMCUGGNHBgAAZCjvY0cGAAAAAAAAAAAAAKxPnGVHBgAAUONyZkcGAAD0dklnRwYAAJgKIGhHBgAAAAAAAAAAAO+99x4AAAAAQKqDXFdHBgBBqoNcV0cGAEKqg1xXRwYAQ6qDXFdHBgAAAAAAAAAAAEWqg1xXRwYARqqDXFdHBgBHqoNcV0cGAEiqg1xXRwYAAAAAAAAAAABKqoNcV0cGAEuqg1xXRwYATKqDXFdHBgBNqoNcV0cGAAAAAAAAAAAAT6qDXFdHBgBQqoNcV0cGAFGqg1xXRwYAUqqDXFdHBgAAAAAAAAAAAFSqg1xXRwYAVaqDXFdHBgBWqoNcV0cGAFeqg1xXRwYAAAAAAAAAAABZqoNcV0cGAFqqg1xXRwYAW6qDXFdHBgBcqoNcV0cGAAAAAAAAAAAA7733HgAAAAAAAAAAMgAAAGQAAACWAAAAyAAAAMgAAAD6AAAALAEAAF4BAACQAQAAkAEAAMIBAAD0AQAAJgIAAFgCAABYAgAAigIAALwCAADuAgAAIAMAACADAABSAwAAhAMAALYDAADoAwAA6AMAABoEAABMBAAAfgQAALAEAACwBAAAAAAAAHsib3JkZXJfaWQiOiAxMjM0NTY3ODkwMTIzNDU2Nzg5MCwgInRhZ3MiOiBbMSwgMl19eyJvcmRlcl9pZCI6IDEyMzQ1Njc4OTAxMjM0NTY3ODkxLCAidGFncyI6IFsxLCAyXX17Im9yZGVyX2lkIjogMTIzNDU2Nzg5MDEyMzQ1Njc4OTIsICJ0YWdzIjogWzEsIDJdfXsib3JkZXJfaWQiOiAxMjM0NTY3ODkwMTIzNDU2Nzg5MywgInRhZ3MiOiBbMSwgMl19eyJvcmRlcl9pZCI6IDEyMzQ1Njc4OTAxMjM0NTY3ODk1LCAidGFncyI6IFsxLCAyXX17Im9yZGVyX2lkIjogMTIzNDU2Nzg5MDEyMzQ1Njc4OTYsICJ0YWdzIjogWzEsIDJdfXsib3JkZXJfaWQiOiAxMjM0NTY3ODkwMTIzNDU2Nzg5NywgInRhZ3MiOiBbMSwgMl19eyJvcmRlcl9pZCI6IDEyMzQ1Njc4OTAxMjM0NTY3ODk4LCAidGFncyI6IFsxLCAyXX17Im9yZGVyX2lkIjogMTIzNDU2Nzg5MDEyMzQ1Njc4OTAsICJ0YWdzIjogWzEsIDJdfXsib3JkZXJfaWQiOiAxMjM0NTY3ODkwMTIzNDU2Nzg5MSwgInRhZ3MiOiBbMSwgMl19eyJvcmRlcl9pZCI6IDEyMzQ1Njc4OTAxMjM0NTY3ODkyLCAidGFncyI6IFsxLCAyXX17Im9yZGVyX2lkIjogMTIzNDU2Nzg5MDEyMzQ1Njc4OTMsICJ0YWdzIjogWzEsIDJdfXsib3JkZXJfaWQiOiAxMjM0NTY3ODkwMTIzNDU2Nzg5NSwgInRhZ3MiOiBbMSwgMl19eyJvcmRlcl9pZCI6IDEyMzQ1Njc4OTAxMjM0NTY3ODk2LCAidGFncyI6IFsxLCAyXX17Im9yZGVyX2lkIjogMTIzNDU2Nzg5MDEyMzQ1Njc4OTcsICJ0YWdzIjogWzEsIDJdfXsib3JkZXJfaWQiOiAxMjM0NTY3ODkwMTIzNDU2Nzg5OCwgInRhZ3MiOiBbMSwgMl19eyJvcmRlcl9pZCI6IDEyMzQ1Njc4OTAxMjM0NTY3ODkwLCAidGFncyI6IFsxLCAyXX17Im9yZGVyX2lkIjogMTIzNDU2Nzg5MDEyMzQ1Njc4OTEsICJ0YWdzIjogWzEsIDJdfXsib3JkZXJfaWQiOiAxMjM0NTY3ODkwMTIzNDU2Nzg5MiwgInRhZ3MiOiBbMSwgMl19eyJvcmRlcl9pZCI6IDEyMzQ1Njc4OTAxMjM0NTY3ODkzLCAidGFncyI6IFsxLCAyXX17Im9yZGVyX2lkIjogMTIzNDU2Nzg5MDEyMzQ1Njc4OTUsICJ0YWdzIjogWzEsIDJdfXsib3JkZXJfaWQiOiAxMjM0NTY3ODkwMTIzNDU2Nzg5NiwgInRhZ3MiOiBbMSwgMl19eyJvcmRlcl9pZCI6IDEyMzQ1Njc4OTAxMjM0NTY3ODk3LCAidGFncyI6IFsxLCAyXX17Im9yZGVyX2lkIjogMTIzNDU2Nzg5MDEyMzQ1Njc4OTgsICJ0YWdzIjogWzEsIDJdfe+99x4AAAAARRVVFAAAAAD3+/0+AAAAAAAAAAAFAAAACgAAAA8AAAAPAAAAFAAAABkAAAAeAAAAIwAAACgAAAAtAAAALQAAADMAAAA5AAAAPwAAAEUAAABLAAAAUQAAAFEAAABXAAAAXQAAAGMAAABpAAAAbwAAAHUAAAB1AAAAewAAAIEAAACHAAAAjQAAAJMAAAAAAAAAcm93LTByb3ctMXJvdy0ycm93LTRyb3ctNXJvdy02cm93LTdyb3ctOHJvdy05cm93LTExcm93LTEycm93LTEzcm93LTE0cm93LTE1cm93LTE2cm93LTE4cm93LTE5cm93LTIwcm93LTIxcm93LTIycm93LTIzcm93LTI1cm93LTI2cm93LTI3cm93LTI4cm93LTI5AAAAAADvvfceAAAAAAAAAAACAAAABAAAAAYAAAAIAAAACAAAAAoAAAAMAAAADgAAABAAAAAQAAAAEgAAABQAAAAWAAAAGAAAABgAAAAaAAAAHAAAAB4AAAAgAAAAIAAAACIAAAAkAAAAJgAAACgAAAAoAAAAKgAAACwAAAAuAAAAMAAAADAAAAAAAAAA////////AAAAAAAAAAAAAAEAAAAAAAAAAQAAAAAAAAACAAAAAAAAAAIAAAAAAAAAAwAAAAAAAAADAAAAAAAAAAQAAAAAAAAABQAAAAAAAAAGAAAAAAAAAAYAAAAAAAAABwAAAAAAAAAHAAAAAAAAAAgAAAAAAAAACAAAAAAAAAAJAAAAAAAAAAoAAAAAAAAACwAAAAAAAAALAAAAAAAAAAwAAAAAAAAADAAAAAAAAAANAAAAAAAAAA0AAAAAAAAADgAAAAAAAAAPAAAAAAAAABAAAAAAAAAAEAAAAAAAAAARAAAAAAAAABEAAAAAAAAAEgAAAAAAAAASAAAAAAAAABMAAAAAAAAAFAAAAAAAAAAVAAAAAAAAABUAAAAAAAAAFgAAAAAAAAAWAAAAAAAAABcAAAAAAAAAFwAAAAAAAAAYAAAAAAAAABkAAAAAAAAAGgAAAAAAAAAaAAAAAAAAABsAAAAAAAAAGwAAAAAAAAAcAAAAAAAAABwAAAAAAAAAHQAAAAAAAADvvfceAAAAAAAAAAAAAAAAAQAAAAAAAAACAAAAAAAAAAMAAAAAAAAAAAAAAAAAAAAFAAAAAAAAAAYAAAAAAAAABwAAAAAAAAAIAAAAAAAAAAAAAAAAAAAACgAAAAAAAAALAAAAAAAAAAwAAAAAAAAADQAAAAAAAAAAAAAAAAAAAA8AAAAAAAAAEAAAAAAAAAARAAAAAAAAABIAAAAAAAAAAAAAAAAAAAAUAAAAAAAAABUAAAAAAAAAFgAAAAAAAAAXAAAAAAAAAAAAAAAAAAAAGQAAAAAAAAAaAAAAAAAAABsAAAAAAAAAHAAAAAAAAAAAAAAAAAAAAO+99x4AAAAAAQAAAAAAAAACAAAAAAAAAAMAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAYAAAAAAAAABwAAAAAAAAAIAAAAAAAAAAkAAAAAAAAAAAAAAAAAAAALAAAAAAAAAAwAAAAAAAAADQAAAAAAAAAOAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAARAAAAAAAAABIAAAAAAAAAEwAAAAAAAAAAAAAAAAAAABUAAAAAAAAAFgAAAAAAAAAXAAAAAAAAABgAAAAAAAAAAAAAAAAAAAAaAAAAAAAAABsAAAAAAAAAHAAAAAAAAAAdAAAAAAAAAAAAAAAAAAAA/////wAAAAA=" +} diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/json-array-cells.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/json-array-cells.test.ts new file mode 100644 index 000000000..6e4910395 --- /dev/null +++ b/packages/appkit/src/connectors/sql-warehouse/tests/json-array-cells.test.ts @@ -0,0 +1,233 @@ +import { + Binary, + Bool, + type DataType, + DateDay, + Decimal, + Float32, + Float64, + Int8, + Int16, + Int32, + Int64, + makeData, + makeVector, + Table, + TimestampMicrosecond, + tableFromIPC, + tableToIPC, + Utf8, + type Vector, + vectorFromArray, +} from "apache-arrow"; +import { describe, expect, test } from "vitest"; + +import { formatJsonArrayCell, parseJsonArrayCell } from "../json-array-cells"; + +function decimalVector(unscaled: bigint[], scale: number, precision: number) { + const words = new Uint32Array(unscaled.length * 4); + unscaled.forEach((value, i) => { + const bits = value < 0n ? (1n << 128n) + value : value; + for (let w = 0; w < 4; w++) { + words[i * 4 + w] = Number((bits >> BigInt(32 * w)) & 0xffffffffn); + } + }); + return makeVector( + makeData({ + type: new Decimal(scale, precision, 128), + length: unscaled.length, + data: words, + }), + ); +} + +const DAY = 86_400_000; + +const cases: Array<[string, () => Vector]> = [ + ["TINYINT", () => vectorFromArray([-128, -1, 0, 127], new Int8())], + ["SMALLINT", () => vectorFromArray([-32768, 32767], new Int16())], + ["INT", () => vectorFromArray([-2147483648, 0, 42, 2147483647], new Int32())], + [ + "BIGINT", + () => + vectorFromArray( + [-9223372036854775808n, 0n, 9007199254740993n, 9223372036854775807n], + new Int64(), + ), + ], + ["FLOAT", () => vectorFromArray([0, -0.25, 0.1, 3.4e38], new Float32())], + [ + "DOUBLE", + () => + vectorFromArray( + [0, 1.1, -1e-300, Number.MAX_VALUE, Number.NaN], + new Float64(), + ), + ], + [ + "DECIMAL(10,2)", + () => decimalVector([-50n, 0n, 51n, 5n, 9999999999n, -9999999999n], 2, 10), + ], + ["DECIMAL(38,0)", () => decimalVector([10n ** 37n, -1n], 0, 38)], + ["DECIMAL(5,5)", () => decimalVector([12345n, -1n], 5, 5)], + [ + "DATE", + () => + vectorFromArray( + [-DAY, Date.UTC(2024, 1, 29), Date.UTC(2026, 0, 1)], + new DateDay(), + ), + ], + [ + "TIMESTAMP", + () => + vectorFromArray( + [ + 0, + Date.UTC(2026, 0, 1, 10, 0, 0, 123), + Date.UTC(1969, 11, 31, 23, 59, 59, 999), + ], + new TimestampMicrosecond("UTC"), + ), + ], + [ + "TIMESTAMP_NTZ", + () => + vectorFromArray( + [0, Date.UTC(2026, 0, 1, 10, 0, 0, 123)], + new TimestampMicrosecond(), + ), + ], + ["BOOLEAN", () => vectorFromArray([true, false], new Bool())], + [ + "STRING", + () => vectorFromArray(["", "a", "ünïcödé ✓", '{"x": 1}'], new Utf8()), + ], + [ + "BINARY", + () => + vectorFromArray( + [new Uint8Array([]), new Uint8Array([0, 255, 1])], + new Binary(), + ), + ], +]; + +/** Values as an Arrow consumer reads them, comparable across decimal representations. */ +const comparable = (value: unknown) => + value instanceof Uint8Array + ? [...value] + : typeof value === "object" || typeof value === "bigint" + ? String(value) + : value; + +describe("JSON_ARRAY cell codecs", () => { + test.each(cases)( + "%s: parsing the formatted text gives back the native Arrow value", + (_name, build) => { + const table = tableFromIPC( + tableToIPC(new Table({ c: build() }), "stream"), + ); + const { type } = table.schema.fields[0]; + for (const native of table.getChildAt(0) ?? []) { + const parsed = parseJsonArrayCell( + formatJsonArrayCell(native, type), + type, + ); + expect(comparable(parsed)).toEqual(comparable(native)); + } + }, + ); + + test("formats match what a warehouse writes", () => { + const formatted = (vector: Vector) => { + const table = tableFromIPC( + tableToIPC(new Table({ c: vector }), "stream"), + ); + const { type } = table.schema.fields[0]; + return [...(table.getChildAt(0) ?? [])].map((v) => + formatJsonArrayCell(v, type), + ); + }; + expect(formatted(decimalVector([-50n, 51n], 2, 10))).toEqual([ + "-0.50", + "0.51", + ]); + expect( + formatted( + vectorFromArray([Date.UTC(2026, 0, 1, 10)], new TimestampMicrosecond()), + ), + ).toEqual(["2026-01-01T10:00:00.000"]); + expect( + formatted(vectorFromArray([Date.UTC(2026, 0, 1)], new DateDay())), + ).toEqual(["2026-01-01"]); + }); + + test("a zoned TIMESTAMP keeps its instant when the warehouse writes the session offset", () => { + const type = new TimestampMicrosecond("UTC"); + expect(parseJsonArrayCell("2026-01-01T10:00:00.000-08:00", type)).toBe( + Date.UTC(2026, 0, 1, 18), + ); + }); + + test.each<[string, DataType, string]>([ + ["INT with a fraction", new Int32(), "1.5"], + ["INT in exponent form", new Int32(), "1e3"], + ["INT with padding", new Int32(), " 1"], + ["INT out of range", new Int32(), "2147483648"], + ["TINYINT out of range", new Int8(), "128"], + ["empty BIGINT", new Int64(), ""], + ["BIGINT with a fraction", new Int64(), "1.0"], + ["BIGINT with a plus sign", new Int64(), "+1"], + ["BIGINT with a leading zero", new Int64(), "01"], + ["empty DOUBLE", new Float64(), ""], + ["DOUBLE that is not a number", new Float64(), "abc"], + ["DOUBLE in hex", new Float64(), "0x1A"], + ["DOUBLE with padding", new Float64(), " 1 "], + [ + "DECIMAL with more digits than its scale", + new Decimal(2, 10, 128), + "1.234", + ], + ["DECIMAL with two points", new Decimal(2, 10, 128), "1.2.3"], + ["DECIMAL with a plus sign", new Decimal(2, 10, 128), "+1.00"], + [ + "DECIMAL with more digits than its precision", + new Decimal(2, 5, 128), + "123456789.00", + ], + ["BINARY that is not base64", new Binary(), "not base64!!"], + ["DATE that is not on the calendar", new DateDay(), "2026-02-31"], + ["DATE without zero padding", new DateDay(), "2026-1-1"], + ["DATE with a time", new DateDay(), "2026-01-01T00:00:00Z"], + [ + "zoned TIMESTAMP without an offset", + new TimestampMicrosecond("UTC"), + "2026-01-01T10:00:00.000", + ], + [ + "zoned TIMESTAMP that is not a date", + new TimestampMicrosecond("UTC"), + "2026-13-01T10:00:00Z", + ], + [ + "TIMESTAMP_NTZ that is not on the calendar", + new TimestampMicrosecond(), + "2026-02-30T10:00:00.000", + ], + [ + "TIMESTAMP_NTZ with Z", + new TimestampMicrosecond(), + "2026-01-01T10:00:00.000Z", + ], + [ + "TIMESTAMP_NTZ with an offset", + new TimestampMicrosecond(), + "2026-01-01T10:00:00.000-08:00", + ], + ["BOOLEAN in upper case", new Bool(), "TRUE"], + ["BOOLEAN as a digit", new Bool(), "1"], + ])("rejects %s", (_name, type, cell) => { + expect(() => parseJsonArrayCell(cell, type)).toThrow(/Invalid/); + }); +}); diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 1d18aae6d..c6b1bca07 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"; @@ -129,9 +130,13 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { * on every query, so once learned we skip that doomed probe; Reyden stays * `"inline"`. Capability is a property of the warehouse, not the user, so it * is not user-scoped. Bounded by the number of distinct warehouses a process - * talks to (effectively one). + * talks to (effectively one). `"links-denied"` reads as `"external"` once + * `until` passes. */ - private _arrowCapability = new Map(); + private _arrowCapability = new Map< + string, + { capability: ArrowCapability; until?: number } + >(); constructor(config: IAnalyticsConfig) { super(config); @@ -861,10 +866,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 +890,15 @@ 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), + capabilityHint: this._arrowCapabilityHint(warehouseId), onCapabilityResolved: (capability) => - this._arrowCapability.set(warehouseId, capability), + this._arrowCapability.set(warehouseId, { + capability, + until: + capability === "links-denied" + ? Date.now() + LINKS_DENIED_TTL_MS + : undefined, + }), }, ); const first = await bytes.next(); @@ -964,6 +976,16 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } } + private _arrowCapabilityHint( + warehouseId: string, + ): ArrowCapability | undefined { + const entry = this._arrowCapability.get(warehouseId); + if (entry?.until !== undefined && Date.now() >= entry.until) { + return "external"; + } + return entry?.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 +1016,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 +1040,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 +1059,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 +1238,9 @@ export function writeChunk( */ const DEFAULT_ARROW_FIRST_BYTE_TIMEOUT_MS = 120_000; +/** How long `"links-denied"` is remembered before Cloud Fetch 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..c08bf492a 100644 --- a/packages/appkit/src/plugins/analytics/result-delivery.ts +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -1,10 +1,17 @@ import { type DataType, type Field, Type, tableFromIPC } from "apache-arrow"; import type { SQLTypeMarker } from "shared"; +import { buildArrowIPCFromRows } from "../../connectors/sql-warehouse/arrow-schema"; +import { MAX_INLINE_ATTACHMENT_BYTES } from "../../connectors/sql-warehouse/client"; +import { formatJsonArrayCell } from "../../connectors/sql-warehouse/json-array-cells"; import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; -import type { RefreshChunkLink } from "../../stream/arrow-stream-processor"; -import type { ExternalLink } from "../../workspace-client"; +import { + type DownloadFailureKind, + downloadFailureKind, + type RefreshChunkLink, +} from "../../stream/arrow-stream-processor"; +import type { ColumnInfo, ExternalLink } from "../../workspace-client"; /** * Centralized disposition/format fallback for analytics result delivery. @@ -33,14 +40,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?: ColumnInfo[]; external_links?: ExternalLink[]; columnNames?: string[]; + totalByteCount?: number; + nextChunkIndex?: number; + nextChunkInternalLink?: string; statement_id?: string; status?: unknown; refreshChunkLink?: RefreshChunkLink; @@ -55,11 +72,26 @@ export interface ArrowChunkStreamer { chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, + fullHeaderWait?: 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"`: + * storage refused or never answered Cloud Fetch, so inline `JSON_ARRAY` goes first. + */ +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(fullHeaderWait: boolean): AsyncGenerator; + }; +} /** Optional per-warehouse capability memo, avoiding a doomed probe per query. */ interface ArrowDeliveryOptions { @@ -190,7 +222,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 +303,190 @@ 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 = (fullHeaderWait: boolean) => + streamer.streamExternalLinks( + links, + signal, + ext.refreshChunkLink, + fullHeaderWait, + ); + 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 bytes Cloud Fetch would have served. */ +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 }; + } +} + +/** Download refused or unanswered, and the result cannot be sent inline. */ +function linksDenied(cause: unknown): ExecutionError { + return new ExecutionError( + `Cloud Fetch download failed and the result is too large to send inline: ${errMessage(cause)}`, + { + errorCode: "EXTERNAL_LINKS_DENIED", + cause: cause instanceof Error ? cause : undefined, + clientMessage: + "The query result could not be downloaded from workspace storage and is too large to send inline. Check the workspace storage firewall and the app's network policy.", + }, + ); +} + +type Download = + | { stream: AsyncGenerator } + | { error: unknown; kind: DownloadFailureKind }; + +/** Pull the first chunk, so a failed download surfaces before the route sends headers. */ +async function startDownload( + bytes: AsyncGenerator, + signal?: AbortSignal, +): Promise { + let first: IteratorResult; + try { + first = await bytes.next(); + } catch (error: unknown) { + const kind = signal?.aborted ? undefined : downloadFailureKind(error); + if (!kind) throw error; + return { error, kind }; + } + return { + stream: (async function* () { + if (!first.done) yield first.value; + yield* bytes; + })(), + }; +} + +/** + * {@link deliverArrowBytes}, falling back to an inline `JSON_ARRAY` run when the + * Cloud Fetch download fails before its first byte. + * + * A refusal (400/403 after the re-mint) reports `"links-denied"` at once; a + * download that never answers reports it after the fallback, since blocked + * hosts hang instead of refusing. Other download errors fall back without being + * remembered. A result too large for inline gets one Cloud Fetch retry with the + * full header wait, and `EXTERNAL_LINKS_DENIED` only when storage refused or + * never answered. With the `"links-denied"` hint, inline runs first and only a + * result too large for it goes back to Cloud Fetch. + */ +export async function* deliverArrowBytesWithFallback( + executor: QueryExecutor, + streamer: ArrowChunkStreamer, + query: string, + processedParams: Record | undefined, + out: ArrowDeliveryOut, + signal?: AbortSignal, + opts?: ArrowDeliveryOptions, +): AsyncGenerator { + const inline = () => + inlineJsonAsArrow(executor, query, processedParams, out, signal); + const remember = (capability: ArrowCapability) => + opts?.onCapabilityResolved?.(capability); + const hinted = opts?.capabilityHint === "links-denied"; + let tooLarge = false; + if (hinted) { + const attempt = await inline(); + if ("bytes" in attempt) { + yield attempt.bytes; + return; + } + if (!(attempt.error instanceof InlineDoesNotFit)) throw attempt.error; + tooLarge = true; + } + + const download = await startDownload( + deliverArrowBytes( + executor, + streamer, + query, + processedParams, + out, + signal, + hinted ? { capabilityHint: "external" } : opts, + ), signal, - ext.refreshChunkLink, ); + if ("stream" in download) { + if (hinted) remember("external"); + yield* download.stream; + return; + } + + const { kind } = download; + if (kind === "refused") remember("links-denied"); + const size = out.cloudFetch?.totalByteCount; + tooLarge ||= size !== undefined && size > MAX_INLINE_ATTACHMENT_BYTES; + if (!tooLarge) { + logger.warn( + "Cloud Fetch download failed (%s), serving the result inline: %s", + kind, + errMessage(download.error), + ); + const attempt = await inline(); + if ("bytes" in attempt) { + if (kind === "unanswered") remember("links-denied"); + yield attempt.bytes; + return; + } + if (!(attempt.error instanceof InlineDoesNotFit)) throw attempt.error; + } + if (kind === "refused" || !out.cloudFetch) throw linksDenied(download.error); + + logger.warn( + "Cloud Fetch download failed (%s) on a result too large for inline, retrying with the full header wait", + kind, + ); + const retry = await startDownload(out.cloudFetch.stream(true), signal); + if ("stream" in retry) { + if (hinted) remember("external"); + yield* retry.stream; + return; + } + const denied = retry.kind !== "other"; + if (denied || kind === "unanswered") remember("links-denied"); + throw denied ? linksDenied(retry.error) : retry.error; } /** JSON row result plus the metadata the SSE `result` message forwards. */ @@ -355,70 +566,11 @@ export async function deliverJsonResult( const JSON_ARRAY_FALLBACK_MAX_ROWS = 100_000; const JSON_ARRAY_FALLBACK_MAX_BYTES = 64 * 1024 * 1024; -/** Render an apache-arrow Decimal cell to a fixed-point string. */ -function formatDecimalCell( - value: { toString(): string }, - scale: number, -): string { - const unscaled = value.toString(); - if (scale <= 0) return unscaled; - const negative = unscaled.startsWith("-"); - let digits = negative ? unscaled.slice(1) : unscaled; - if (digits.length <= scale) digits = digits.padStart(scale + 1, "0"); - const point = digits.length - scale; - const out = `${digits.slice(0, point)}.${digits.slice(point)}`; - return negative ? `-${out}` : out; -} - -/** - * Render an apache-arrow Timestamp cell to ISO-8601 ms precision — `Z` for - * zoned columns, no `Z` for TIMESTAMP_NTZ — matching native JSON_ARRAY. - */ -function formatTimestampCell(epochMs: number, hasTimezone: boolean): string { - const iso = new Date(epochMs).toISOString(); - return hasTimezone ? iso : iso.slice(0, -1); -} - -/** Render an apache-arrow Date cell to `yyyy-MM-dd` (UTC). */ -function formatDateCell(epochMs: number): string { - return new Date(epochMs).toISOString().slice(0, 10); -} - -/** - * Render a single scalar Arrow value to the string form the warehouse emits - * under native JSON_ARRAY. Used for both top-level cells and (recursively, via - * {@link renderNestedValue}) scalar leaves inside List/Struct/Map columns, - * keyed on the Arrow field type so a nested `DECIMAL` keeps its scale and a - * nested `TIMESTAMP` becomes ISO-8601 instead of raw epoch-ms. - */ -function formatScalarCell(value: unknown, type: DataType): string { - switch (type.typeId) { - case Type.Decimal: - return formatDecimalCell( - value as { toString(): string }, - (type as DataType & { scale: number }).scale, - ); - case Type.Timestamp: - return formatTimestampCell( - Number(value), - (type as DataType & { timezone?: string | null }).timezone != null, - ); - case Type.Date: - return formatDateCell(Number(value)); - default: - if (value instanceof Uint8Array) { - return Buffer.from(value).toString("base64"); - } - if (value instanceof Date) return value.toISOString(); - return String(value); - } -} - /** * Recursively render a nested Arrow value (List / Struct / Map, or a scalar * leaf) to the plain JS shape the warehouse nests inside a JSON_ARRAY cell: * structs and maps become objects, lists become arrays, and every scalar leaf - * is stringified per {@link formatScalarCell}. The caller `JSON.stringify`s the + * is stringified per {@link formatJsonArrayCell}. The caller `JSON.stringify`s the * result so nested columns match the warehouse's native `data_array` exactly. */ function renderNestedValue(value: unknown, type: DataType): unknown { @@ -458,10 +610,17 @@ function renderNestedValue(value: unknown, type: DataType): unknown { return obj; } default: - return formatScalarCell(value, type); + return formatJsonArrayCell(value, type); } } +const NESTED_TYPES = new Set([ + Type.List, + Type.FixedSizeList, + Type.Struct, + Type.Map, +]); + /** Parse a STRING cell that looks like JSON into an object/array. */ function maybeParseJsonString(value: string): unknown { if (value && (value[0] === "{" || value[0] === "[")) { @@ -519,37 +678,14 @@ export function decodeArrowAttachmentToRows( const v = col?.get(i); if (v == null) { row[name] = null; - continue; - } - switch (type.typeId) { - case Type.Decimal: - case Type.Timestamp: - case Type.Date: - row[name] = formatScalarCell(v, type); - continue; - // Nested columns (STRUCT / ARRAY / MAP) arrive on the native - // JSON_ARRAY wire as a JSON string with every scalar leaf stringified - // by type; render the same shape and stringify so callers cannot tell - // the fallback path from the native one. - case Type.List: - case Type.FixedSizeList: - case Type.Struct: - case Type.Map: - row[name] = JSON.stringify(renderNestedValue(v, type)); - continue; - } - if ( - typeof v === "number" || - typeof v === "bigint" || - typeof v === "boolean" - ) { - row[name] = String(v); + } else if (NESTED_TYPES.has(type.typeId)) { + // The native JSON_ARRAY wire sends nested columns as JSON text with + // every scalar leaf stringified by type. + row[name] = JSON.stringify(renderNestedValue(v, type)); } else if (typeof v === "string") { row[name] = maybeParseJsonString(v); - } else if (v instanceof Uint8Array) { - row[name] = Buffer.from(v).toString("base64"); } else { - row[name] = formatScalarCell(v, type); + row[name] = formatJsonArrayCell(v, type); } } rows.push(row); 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..7e32bc0eb --- /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({ + id: vectorFromArray(id, new Int32()), + big: vectorFromArray(big, new Int64()), + amount: vectorFromArray(amount, new Float64()), + price: decimalVector(price, 2), + day: vectorFromArray(day, new DateDay()), + ts: vectorFromArray(ts, new TimestampMicrosecond("UTC")), + label: vectorFromArray(label, new Utf8()), + flag: vectorFromArray(flag, new Bool()), + }); + + const text = (v: T | null, render: (x: T) => string) => + v === null ? null : render(v); + 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, index) => + [...(table.getChildAt(index) ?? [])].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..4360be4ec 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,339 @@ 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?: () => Awaited>; + 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(); + 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, fullHeaderWait) => + processor.streamChunks(links, signal, refresh, fullHeaderWait), + }; + 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) { + if (!(e instanceof ExecutionError)) throw e; + error = e; + } + 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|storage account/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 retry with the full header wait 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 full-header-wait retry", 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.each([ + ["never answered", { kind: "hang" }, true], + ["reset after headers", { kind: "reset-after-headers" }, false], + ] as const)( + "a large result whose download %s and whose retry then fails otherwise keeps the download error", + async (_name, storage, remembered) => { + const ctx = await setup({ storage, totalByteCount: LARGE }); + const streamExternalLinks = ctx.streamer.streamExternalLinks; + ctx.streamer.streamExternalLinks = (links, signal, refresh, fullWait) => { + if (fullWait) ctx.storage.setMode({ kind: "status", status: 404 }); + return streamExternalLinks(links, signal, refresh, fullWait); + }; + const { error, reported } = await run(ctx); + + expect(error?.errorCode).toBe("EXTERNAL_LINK_DOWNLOAD_FAILED"); + expect(ctx.storage.gets()).toBe(2); + expect(reported.includes("links-denied")).toBe(remembered); + }, + ); + + 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 surfaces without running the query again", 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"); + expect(ctx.executions).toEqual(["INLINE+JSON_ARRAY"]); + expect(ctx.storage.gets()).toBe(0); + }); + + 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..d91c01020 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -18,14 +18,72 @@ 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 header wait with `fullHeaderWait`. */ 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. */ +const EXTERNAL_LINK_DOWNLOAD_FAILED = "EXTERNAL_LINK_DOWNLOAD_FAILED"; + const BACKOFF_MULTIPLIER = 1000; +type UsableLink = ExternalLink & { externalLink: string }; + +const hasUrl = (link: ExternalLink | undefined): link is UsableLink => + Boolean(link?.externalLink); + +/** A 4xx won't change on retry, except a request timeout (408) or a rate limit (429). */ +const isFinalStatus = (status: number) => + status >= 400 && status < 500 && status !== 408 && status !== 429; + +interface DownloadFailure { + status?: number; + error: unknown; + timedOut: boolean; +} + +type Connection = + | { + ok: true; + response: Response; + controller: AbortController; + detach: () => void; + } + | ({ ok: false } & DownloadFailure); + +function downloadFailed( + message: string, + { status, error, timedOut }: DownloadFailure, +): ExecutionError { + return new ExecutionError(message, { + errorCode: EXTERNAL_LINK_DOWNLOAD_FAILED, + context: { httpStatus: status, timedOut }, + cause: error instanceof Error ? error : undefined, + }); +} + +/** Storage refused the download (400/403), never answered it, or it failed some other way. */ +export type DownloadFailureKind = "refused" | "unanswered" | "other"; + +/** The {@link DownloadFailureKind} of `err`, or `undefined` when it is not a chunk download failure. */ +export function downloadFailureKind( + err: unknown, +): DownloadFailureKind | undefined { + if ( + !(err instanceof ExecutionError) || + err.errorCode !== EXTERNAL_LINK_DOWNLOAD_FAILED + ) { + return undefined; + } + const status = err.context?.httpStatus; + if (status === 400 || status === 403) return "refused"; + return err.context?.timedOut === true ? "unanswered" : "other"; +} + /** * 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 +94,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 +102,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 +115,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. + * + * `fullHeaderWait` 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, + fullHeaderWait = 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, fullHeaderWait); } } @@ -75,89 +142,85 @@ 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, since a pre-signed URL may have expired. Other + * 4xx (except 408/429) and header timeouts stop: a blocked host hangs instead + * of refusing, so a retry only repeats the wait. */ private async *streamChunkBody( chunk: ExternalLink, - signal?: AbortSignal, - refresh?: RefreshChunkLink, + signal: AbortSignal | undefined, + refresh: RefreshChunkLink | undefined, + fullHeaderWait: boolean, ): AsyncGenerator { - let externalLink = chunk.externalLink; - if (!externalLink) { + if (!hasUrl(chunk)) { // A missing link cannot be fixed by retrying — fail loudly. throw ExecutionError.statementFailed( `External link missing for chunk ${chunk.chunkIndex}`, ); } + let link: UsableLink = chunk; - let response: Response | undefined; - let controller: AbortController | undefined; - let onOuterAbort: (() => void) | undefined; - let lastError: unknown; + const { timeout, retries, responseTimeoutMs } = this.options; + const headerWaitMs = fullHeaderWait + ? timeout + : Math.min(responseTimeoutMs, timeout); + let connection: Extract | undefined; + let failure: DownloadFailure = { + error: new Error("no download attempts"), + timedOut: false, + }; + 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 (hasUrl(fresh)) 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; - 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, - ); - } - } - } + continue; + } + if (status !== undefined && isFinalStatus(status)) break; + 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, ); } + const { response, controller, detach } = connection; const body = response.body; let idleTimer: ReturnType | undefined; try { @@ -170,10 +233,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 +257,52 @@ export class ArrowStreamProcessor { ); throw error instanceof ExecutionError ? error - : ExecutionError.statementFailed( + : downloadFailed( `Failed streaming chunk ${chunk.chunkIndex}: ${ error instanceof Error ? error.message : String(error) }`, + { error, timedOut: false }, ); } finally { clearTimeout(idleTimer); - if (onOuterAbort) signal?.removeEventListener("abort", onOuterAbort); + detach(); + } + } + + /** One GET; only the wait for response headers is timed. */ + private async connect( + link: UsableLink, + 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, { + 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..822b3946e 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,268 @@ describe("ArrowStreamProcessor.streamChunks", () => { expect(bytes).toEqual([7]); }, 10000); }); + +const response = (status: number, statusText = "") => + new Response("denied", { status, statusText }); +const okResponse = (...pieces: number[]) => + new Response(streamOf(new Uint8Array(pieces))); + +async function failureOf(gen: AsyncGenerator) { + try { + await drain(gen); + } catch (error) { + if (error instanceof ExecutionError) return error; + throw error; + } + 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("fullHeaderWait 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/); + }); +});