fix(appkit): fall back to inline results when a Cloud Fetch download is refused

packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts

@@ -21,7 +21,12 @@ import {
   Table,
   TimestampMicrosecond,
   tableToIPC,
+  Type,
   Utf8,
+  type Vector,
+  makeData,
+  makeVector,
+  vectorFromArray,
 } from "apache-arrow";
 
 /**
@@ -75,6 +80,182 @@ export function buildEmptyArrowIPCBase64(
   return Buffer.from(ipc).toString("base64");
 }
 
+/** Column metadata from `manifest.schema.columns`. */
+export interface ArrowRowColumn {
+  name?: string;
+  typeText?: string;
+  typeName?: string;
+  typePrecision?: number;
+  typeScale?: number;
+}
+
+/**
+ * Rebuild the Arrow IPC stream a warehouse serves for `ARROW_STREAM` from the
+ * positional string rows of an `INLINE` + `JSON_ARRAY` result. Fields are named
+ * `col_N` like the native stream, and each type is what
+ * {@link parseDatabricksType} gives that column, with DECIMAL precision and
+ * scale taken from `typePrecision` and `typeScale`.
+ *
+ * Nested, interval, and geospatial columns arrive as strings on the JSON wire
+ * and stay strings. A cell that does not parse as its column type throws.
+ */
+export function buildArrowIPCFromRows(
+  columns: ArrowRowColumn[],
+  rows: Array<Array<string | null | undefined>>,
+): Uint8Array {
+  const vectors: Record<string, Vector> = {};
+  columns.forEach((col, index) => {
+    const cells = rows.map((row) => row[index] ?? null);
+    vectors[`col_${index}`] = vectorFromCells(col, cells);
+  });
+  return tableToIPC(new Table(vectors), "stream");
+}
+
+function vectorFromCells(col: ArrowRowColumn, cells: Array<string | null>) {
+  const type = rebuildType(col);
+  const map = <T>(parse: (cell: string) => T) =>
+    cells.map((cell) => (cell === null ? null : parse(cell)));
+  switch (type.typeId) {
+    case Type.Int: {
+      const { bitWidth } = type as unknown as { bitWidth: number };
+      return vectorFromArray(
+        bitWidth === 64 ? map(parseBigInt) : map(parseInteger),
+        type as never,
+      );
+    }
+    case Type.Float:
+      return vectorFromArray(map(parseFloatCell), type as never);
+    case Type.Bool:
+      return vectorFromArray(map(parseBoolean), type as never);
+    case Type.Date:
+      return vectorFromArray(map(parseUtcMillis), type as never);
+    case Type.Timestamp:
+      return vectorFromArray(map(parseUtcMillis), type as never);
+    case Type.Binary:
+      return vectorFromArray(
+        map((cell) => new Uint8Array(Buffer.from(cell, "base64"))),
+        type as never,
+      );
+    case Type.Decimal: {
+      const { scale, precision } = type as unknown as {
+        scale: number;
+        precision: number;
+      };
+      return decimalVector(cells, scale, precision);
+    }
+    case Type.Null:
+      return vectorFromArray(
+        cells.map(() => null),
+        type as never,
+      );
+    default:
+      return vectorFromArray(cells, new Utf8());
+  }
+}
+
+const STRING_ONLY_TYPES = new Set<Type>([
+  Type.List,
+  Type.Struct,
+  Type.Map,
+  Type.Interval,
+  Type.Duration,
+]);
+
+function rebuildType(col: ArrowRowColumn): DataType {
+  const typeText = col.typeText ?? col.typeName ?? "STRING";
+  let type: DataType;
+  try {
+    type = parseDatabricksType(typeText);
+  } catch {
+    return new Utf8();
+  }
+  if (STRING_ONLY_TYPES.has(type.typeId)) return new Utf8();
+  // GEOGRAPHY and GEOMETRY map to Binary but arrive as text on the JSON wire.
+  if (type.typeId === Type.Binary && typeText.toUpperCase() !== "BINARY") {
+    return new Utf8();
+  }
+  if (type.typeId === Type.Decimal && col.typePrecision != null) {
+    return new Decimal(col.typeScale ?? 0, col.typePrecision, 128);
+  }
+  return type;
+}
+
+function parseInteger(cell: string): number {
+  const n = Number(cell);
+  if (cell.trim() === "" || !Number.isInteger(n)) {
+    throw new Error(`Invalid integer "${cell}"`);
+  }
+  return n;
+}
+
+function parseBigInt(cell: string): bigint {
+  if (cell.trim() === "") throw new Error(`Invalid integer "${cell}"`);
+  return BigInt(cell);
+}
+
+function parseFloatCell(cell: string): number {
+  const n = Number(cell);
+  if (cell.trim() === "" || (Number.isNaN(n) && cell !== "NaN")) {
+    throw new Error(`Invalid number "${cell}"`);
+  }
+  return n;
+}
+
+function parseBoolean(cell: string): boolean {
+  if (cell !== "true" && cell !== "false") {
+    throw new Error(`Invalid boolean "${cell}"`);
+  }
+  return cell === "true";
+}
+
+/** Epoch ms for an ISO date or timestamp; a value with no zone is UTC. */
+function parseUtcMillis(cell: string): number {
+  const ms = Date.parse(
+    /(?:Z|[+-]\d\d(?::?\d\d)?)$/i.test(cell) || !cell.includes("T")
+      ? cell
+      : `${cell}Z`,
+  );
+  if (Number.isNaN(ms)) throw new Error(`Invalid date or timestamp "${cell}"`);
+  return ms;
+}
+
+function decimalVector(
+  cells: Array<string | null>,
+  scale: number,
+  precision: number,
+) {
+  const words = new Uint32Array(cells.length * 4);
+  const validity = new Uint8Array(Math.ceil(cells.length / 8));
+  let nullCount = 0;
+  cells.forEach((cell, i) => {
+    if (cell === null) {
+      nullCount++;
+      return;
+    }
+    validity[i >> 3] |= 1 << (i & 7);
+    const negative = cell.startsWith("-");
+    const [whole, fraction = ""] = cell.replace(/^[-+]/, "").split(".");
+    const digits = `${whole}${fraction.padEnd(scale, "0").slice(0, scale)}`;
+    if (!/^\d+$/.test(digits)) throw new Error(`Invalid decimal "${cell}"`);
+    const unscaled = BigInt(digits);
+    const twosComplement = negative ? (1n << 128n) - unscaled : unscaled;
+    for (let w = 0; w < 4; w++) {
+      words[i * 4 + w] = Number(
+        (twosComplement >> BigInt(32 * w)) & 0xffffffffn,
+      );
+    }
+  });
+  return makeVector(
+    makeData({
+      type: new Decimal(scale, precision, 128),
+      length: cells.length,
+      nullCount,
+      nullBitmap: validity,
+      data: words,
+    }),
+  );
+}
+
 // ============================================================================
 // Recursive-descent parser
 // ============================================================================

packages/appkit/src/connectors/sql-warehouse/client.ts

@@ -102,7 +102,7 @@ function normalizeStatementCounts<T extends StatementResponse>(response: T): T {
  * results fall through to `disposition: "EXTERNAL_LINKS"`, which streams
  * chunk-by-chunk and is not bound by this limit.
  */
-const MAX_INLINE_ATTACHMENT_BYTES = 25 * 1024 * 1024;
+export const MAX_INLINE_ATTACHMENT_BYTES = 25 * 1024 * 1024;
 
 /**
  * Safety cap on how many additional EXTERNAL_LINKS chunks
@@ -246,7 +246,13 @@ export class SQLWarehouseConnector {
 
   async executeStatement(
     workspaceClient: WorkspaceClient,
-    input: ExecuteStatementRequest,
+    input: ExecuteStatementRequest & {
+      /**
+       * Keep the warehouse's positional `dataArray` and carry the manifest
+       * `columns` on the result, instead of building keyed row objects.
+       */
+      rawRows?: boolean;
+    },
     signal?: AbortSignal,
   ) {
     const startTime = Date.now();
@@ -353,13 +359,15 @@ export class SQLWarehouseConnector {
                 statementId,
                 this.config.timeout,
                 signal,
+                input.rawRows,
               );
               break;
             case "SUCCEEDED":
               result = await this._transformDataArray(
                 response,
                 workspaceClient,
                 signal,
+                input.rawRows,
               );
               break;
             case "FAILED":
@@ -830,6 +838,7 @@ export class SQLWarehouseConnector {
     statementId: string,
     timeout = executeStatementDefaults.timeout,
     signal?: AbortSignal,
+    rawRows?: boolean,
   ) {
     return this.telemetry.startActiveSpan(
       "sql.poll",
@@ -908,6 +917,7 @@ export class SQLWarehouseConnector {
                   response,
                   workspaceClient,
                   signal,
+                  rawRows,
                 );
               case "FAILED":
                 throw ExecutionError.statementFailed(
@@ -949,6 +959,7 @@ export class SQLWarehouseConnector {
     response: StatementResponse,
     workspaceClient: WorkspaceClient,
     signal?: AbortSignal,
+    rawRows?: boolean,
   ) {
     if (response.manifest?.format === "ARROW_STREAM") {
       const result = response.result as
@@ -998,11 +1009,22 @@ export class SQLWarehouseConnector {
       // current Databricks warehouses.
     }
 
-    if (!response.result?.dataArray || !response.manifest?.schema?.columns) {
-      return response;
+    const columns = response.manifest?.schema?.columns;
+    if (rawRows && columns) {
+      return {
+        ...response,
+        result: {
+          ...response.result,
+          columns,
+          columnNames: arrowColumnNames(response),
+          statement_id: response.statementId,
+        },
+      };
     }
 
-    const columns = response.manifest.schema.columns;
+    if (!response.result?.dataArray || !columns) {
+      return response;
+    }
 
     const transformedData = response.result.dataArray.map((row) => {
       const obj: Record<string, unknown> = {};
@@ -1108,6 +1130,7 @@ export class SQLWarehouseConnector {
       statement_id: string;
       status: StatementStatus;
       columnNames?: string[];
+      totalByteCount?: number;
       external_links?: ExternalLink[];
       refreshChunkLink?: RefreshChunkLink;
     };
@@ -1121,6 +1144,7 @@ export class SQLWarehouseConnector {
           error: response.status?.error,
         } as StatementStatus,
         columnNames: arrowColumnNames(response),
+        totalByteCount: response.manifest?.totalByteCount as number | undefined,
         // Resolve the pre-signed links for EVERY chunk in the caller's own
         // execution context. Streaming these directly (see
         // {@link streamExternalLinks}) avoids a second `getStatement` under
@@ -1227,8 +1251,9 @@ export class SQLWarehouseConnector {
     chunks: ExternalLink[],
     signal?: AbortSignal,
     refresh?: RefreshChunkLink,
+    patient?: boolean,
   ): AsyncGenerator<Uint8Array, void, unknown> {
-    return this.arrowProcessor.streamChunks(chunks, signal, refresh);
+    return this.arrowProcessor.streamChunks(chunks, signal, refresh, patient);
   }
 
   /**