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

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<Array<string | null | undefined>>,
+): 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<string | null>) {
+  const type = rebuildType(col);
+  const values = cells.map((cell) =>
+    cell === null ? null : parseJsonArrayCell(cell, type),
+  );
+  if (DataType.isDecimal(type)) {
+    return decimalVector(values as Array<bigint | null>, type);
+  }
+  return vectorFromArray(values, type);
+}
+
+const STRING_ONLY_TYPES = new Set<Type>([
+  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<bigint | null>, 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
 // ============================================================================

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<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 +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,13 +356,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 +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<string, unknown> = {};
       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<Uint8Array, void, unknown> {
-    return this.arrowProcessor.streamChunks(chunks, signal, refresh);
+    return this.arrowProcessor.streamChunks(
+      chunks,
+      signal,
+      refresh,
+      fullHeaderWait,
+    );
   }
 
   /**