@query-farm/vgi 0.6.0

package/src/functions/scalar.ts

@@ -0,0 +1,356 @@
+// Copyright 2025, 2026 Query Farm LLC - https://query.farm
+// Scalar function implementation.
+// Scalar functions receive columnar input and produce a single output column
+// with 1:1 row mapping.
+
+import {
+  type VgiSchema,
+  type VgiField,
+  type VgiBatch,
+  type VgiDataType,
+  type ValueFor,
+  type Repr,
+  schema as makeSchema,
+  field,
+  nullType,
+  isNull,
+  isBatch,
+} from "../arrow/index.js";
+import type { OutputCollector, AuthContext } from "@query-farm/vgi-rpc";
+import {
+  FunctionType,
+  FunctionStability,
+  NullHandling,
+  DEFAULT_MAX_WORKERS,
+} from "../types.js";
+import { RowCountMismatchError } from "../errors.js";
+import type {
+  BindRequest,
+  BindResponse,
+  InitRequest,
+  GlobalInitResponse,
+} from "../protocol/types.js";
+import type {
+  VgiFunction,
+  FunctionMeta,
+  StreamHandlers,
+  FunctionExample,
+} from "./types.js";
+import type { ArgumentSpec } from "../arguments/argument-spec.js";
+import { argumentSpecsToSchema } from "../arguments/argument-spec.js";
+import { batchToScalarDict, batchToSecretDict, batchFromColumns, projectBatch, safeNumber } from "../util/arrow/index.js";
+
+// ============================================================================
+// Scalar Bind Parameters
+// ============================================================================
+
+export interface ScalarBindParameters {
+  /** Constant argument values (resolved at bind time) */
+  constArgs: Record<string, any>;
+  /** Schema of all arguments including columnar params */
+  argumentsSchema: VgiSchema;
+  /** Settings from DuckDB */
+  settings: Record<string, any>;
+  /** Secrets from DuckDB */
+  secrets: Record<string, Record<string, any>>;
+  /** Original bind request — exposes attach_opaque_data, transaction_opaque_data, function_type, etc. */
+  bindCall: BindRequest;
+}
+
+// ============================================================================
+// Functional API
+// ============================================================================
+
+/** Ordered parameter definition for scalar functions. */
+export interface ScalarParameterDef {
+  name: string;
+  type: VgiDataType;
+  /** If true, this is a constant parameter (scalar value resolved at bind time). */
+  const?: boolean;
+  /** If true, this parameter accepts variable number of arguments. */
+  varargs?: boolean;
+  /** Human-readable per-argument description (surfaced as `vgi_doc`). */
+  doc?: string;
+}
+
+/**
+ * The typed compute row for a scalar function: one property per declared
+ * columnar param, valued in the representation (`rich` / `raw`) selected by the
+ * function's `repr`. (Provided as a convenience type for authors who build rows
+ * out of the input batch.)
+ */
+export type ScalarComputeRow<
+  P extends Record<string, VgiDataType>,
+  M extends Repr,
+> = { [K in keyof P]: ValueFor<P[K], M> | null };
+
+/** A single output value for a scalar function under representation `M`. */
+export type ScalarOutputValue<R extends VgiDataType, M extends Repr> =
+  ValueFor<R, M> | null;
+
+/**
+ * Accepted compute return shapes: an array (or iterable) of output values in
+ * the declared representation, or a pre-built VgiBatch. The representation is
+ * enforced statically — returning a `Date` under `repr: 'raw'` (which expects a
+ * branded `Date32`/`Date64Ms`), or a branded value under `repr: 'rich'` (which
+ * expects a `Date`), is a COMPILE error.
+ */
+export type ScalarComputeResult<R extends VgiDataType, M extends Repr> =
+  | Array<ScalarOutputValue<R, M>>
+  | Iterable<ScalarOutputValue<R, M>>
+  | VgiBatch;
+
+export interface ScalarFunctionConfig<
+  P extends Record<string, VgiDataType> = Record<string, VgiDataType>,
+  R extends VgiDataType = VgiDataType,
+  M extends Repr = "rich",
+> {
+  name: string;
+  description?: string;
+  /** Columnar params (receive Arrow arrays at process time) */
+  params?: P;
+  /** Constant params (receive scalar values resolved at bind time) */
+  constParams?: Record<string, VgiDataType>;
+  /**
+   * Per-argument descriptions keyed by param name, surfaced as `vgi_doc` field
+   * metadata (and via the extension's `vgi_function_arguments()`). Applies to
+   * `params` and `constParams`; for the ordered `parameters` list use each
+   * entry's `doc` field instead.
+   */
+  argDocs?: Record<string, string>;
+  /**
+   * Ordered parameter list. When provided, overrides `params` and `constParams`.
+   * Use this when parameters need non-default ordering (e.g. const, param, const).
+   */
+  parameters?: ScalarParameterDef[];
+  /** Output type (static) */
+  returns?: R;
+  /** Dynamic output type at bind time */
+  outputType?: (params: ScalarBindParameters) => VgiDataType | Promise<VgiDataType>;
+  /**
+   * Value representation for compute I/O. `'rich'` (default) uses JS `Date` for
+   * date32/date64 and plain number/bigint elsewhere. `'raw'` uses the branded
+   * unit-carrying aliases (Date32, TimestampMicros, UnscaledDecimal, …). The
+   * choice flows into compute()'s statically-checked return type and selects
+   * the runtime converter used to build the output column.
+   */
+  repr?: M;
+  /**
+   * Process: receives the columnar input batch + const values, returns the
+   * output column as an array/iterable of values (statically typed from
+   * `returns` and `repr`) or a pre-built VgiBatch.
+   */
+  compute: (
+    batch: VgiBatch,
+    consts: Record<string, any>,
+    info: {
+      settings: Record<string, any>;
+      secrets: Record<string, Record<string, any>>;
+      auth: AuthContext;
+    }
+    // NoInfer pins R/M from `returns`/`repr` so the compute return type is
+    // CHECKED against them rather than widening them to fit a wrong value.
+  ) => ScalarComputeResult<NoInfer<R>, NoInfer<M>>;
+  // Metadata
+  stability?: FunctionStability;
+  nullHandling?: NullHandling;
+  examples?: FunctionExample[];
+  categories?: string[];
+  tags?: Record<string, string>;
+  maxWorkers?: number;
+  requiredSettings?: string[];
+  requiredSecrets?: string[];
+}
+
+export function defineScalarFunction<
+  P extends Record<string, VgiDataType> = Record<string, VgiDataType>,
+  R extends VgiDataType = VgiDataType,
+  M extends Repr = "rich",
+>(config: ScalarFunctionConfig<P, R, M>): VgiFunction {
+  const repr: Repr = config.repr ?? "rich";
+  // Build argument specs
+  const specs: ArgumentSpec[] = [];
+
+  if (config.parameters) {
+    // Ordered parameter list — supports arbitrary interleaving of const/non-const
+    for (const [idx, p] of config.parameters.entries()) {
+      const isAny = isNull(p.type);
+      specs.push({
+        name: p.name,
+        position: idx,
+        arrowType: isAny ? nullType() : p.type,
+        isAnyType: isAny,
+        isConst: p.const,
+        isVarargs: p.varargs,
+        doc: p.doc,
+      });
+    }
+  } else {
+    // Legacy: params first, then constParams
+    let posIdx = 0;
+
+    if (config.params) {
+      for (const [name, type] of Object.entries(config.params)) {
+        const isAny = isNull(type);
+        specs.push({
+          name,
+          position: posIdx++,
+          arrowType: isAny ? nullType() : type,
+          isAnyType: isAny,
+          doc: config.argDocs?.[name],
+        });
+      }
+    }
+
+    if (config.constParams) {
+      for (const [name, type] of Object.entries(config.constParams)) {
+        specs.push({
+          name,
+          position: posIdx++,
+          arrowType: type,
+          isConst: true,
+          doc: config.argDocs?.[name],
+        });
+      }
+    }
+  }
+
+  const meta: FunctionMeta = {
+    name: config.name,
+    description: config.description,
+    stability: config.stability,
+    nullHandling: config.nullHandling,
+    examples: config.examples,
+    categories: config.categories,
+    tags: config.tags,
+    maxWorkers: config.maxWorkers,
+    requiredSettings: config.requiredSettings,
+    requiredSecrets: config.requiredSecrets,
+  };
+
+  // Default output schema for catalog registration.
+  // Static returns: use the declared type.
+  // Dynamic outputType: use null with vgi:any metadata (matches Python convention).
+  const defaultOutputSchema = config.returns
+    ? makeSchema([field("result", config.returns, true)])
+    : makeSchema([field("result", nullType(), true, new Map([["vgi:any", "true"]]))]);
+
+  return {
+    kind: "scalar",
+    meta,
+    argumentSpecs: specs,
+    defaultOutputSchema,
+
+    async bind(request: BindRequest): Promise<BindResponse> {
+      const settings = batchToScalarDict(request.settings);
+      const secrets = batchToSecretDict(request.secrets);
+
+      // Extract constant args from the Arguments.
+      // DuckDB only sends const arg values in Arguments (column args are in inputSchema).
+      // Arguments are indexed sequentially (0, 1, ...) for const args only.
+      const constArgs: Record<string, any> = {};
+      {
+        let constIdx = 0;
+        for (const spec of specs) {
+          if (spec.isConst) {
+            constArgs[spec.name] = request.arguments.get(constIdx, undefined);
+            constIdx++;
+          }
+        }
+      }
+
+      // Determine output type
+      let outputType: VgiDataType;
+      if (config.outputType) {
+        outputType = await config.outputType({
+          constArgs,
+          argumentsSchema: request.input_schema ?? makeSchema([]),
+          settings,
+          secrets,
+          bindCall: request,
+        });
+      } else if (config.returns) {
+        outputType = config.returns;
+      } else {
+        throw new Error(
+          `Scalar function '${config.name}' must specify either 'returns' or 'outputType'`
+        );
+      }
+
+      const outputSchema = makeSchema([
+        field("result", outputType, true),
+      ]);
+
+      return { output_schema: outputSchema, opaque_data: null };
+    },
+
+    globalInit(request: InitRequest): GlobalInitResponse {
+      const executionId = new Uint8Array(16);
+      crypto.getRandomValues(executionId);
+      return {
+        max_workers: config.maxWorkers ?? DEFAULT_MAX_WORKERS,
+        execution_id: executionId,
+        opaque_data: null,
+      };
+    },
+
+    createStreamHandlers(
+      request: InitRequest,
+      response: GlobalInitResponse
+    ): StreamHandlers {
+      const outputSchema = request.output_schema;
+      const settings = batchToScalarDict(request.bind_call.settings);
+      const secrets = batchToSecretDict(request.bind_call.secrets);
+
+      // Extract const args (DuckDB only sends const values, indexed sequentially)
+      const constArgs: Record<string, any> = {};
+      let constIdx = 0;
+      for (const spec of specs) {
+        if (spec.isConst) {
+          constArgs[spec.name] = request.bind_call.arguments.get(constIdx, undefined);
+          constIdx++;
+        }
+      }
+
+      const inputSchema = request.bind_call.input_schema;
+
+      return {
+        outputSchema,
+        inputSchema: inputSchema ?? undefined,
+        exchangeInit: () => ({}),
+        exchangeFn: (state: any, input: VgiBatch, out: OutputCollector) => {
+          const result = config.compute(input as any, constArgs, { settings, secrets, auth: out.auth });
+
+          // Build output batch
+          let outputBatch: VgiBatch;
+          if (isBatch(result)) {
+            outputBatch = result as VgiBatch;
+          } else if (Array.isArray(result)) {
+            // Array of values -> single column. Route through batchFromColumns
+            // so complex types (Decimal, BigInt-backed Timestamp/Duration,
+            // List, Map, Struct) are built via our buildColumnData. The `repr`
+            // selects raw<->canonical (branded) vs rich<->canonical conversion.
+            const fieldName = outputSchema.fields[0].name;
+            outputBatch = batchFromColumns({ [fieldName]: result }, outputSchema, repr);
+          } else {
+            // User returned an Arrow Vector / typed-array-shaped column. Iterate
+            // it into a JS array and re-route through batchFromColumns. Adds an
+            // O(n) copy but matches both backends' contract.
+            const fieldName = outputSchema.fields[0].name;
+            outputBatch = batchFromColumns({ [fieldName]: [...(result as Iterable<any>)] }, outputSchema, repr);
+          }
+
+          // Validate row count
+          if (outputBatch.numRows !== input.numRows) {
+            throw new RowCountMismatchError(
+              input.numRows,
+              outputBatch.numRows
+            );
+          }
+
+          out.emit(outputBatch);
+        },
+      };
+    },
+  };
+}

package/src/functions/storage-cf-do.ts

@@ -0,0 +1,303 @@
+// Copyright 2025, 2026 Query Farm LLC - https://query.farm
+// Cloudflare Durable Object storage for VGI function state.
+//
+// Port of vgi-python's vgi/function_storage_cf_do.py. Talks the Cloudflare
+// Worker + Durable Object's unified state_* / queue_* protocol
+// (vgi-cloudflare-durable-object-storage/src/index.ts). Every request carries
+// the per-attach `shard_key` (set via forShard); destructive ops carry a fresh
+// 32-hex `attempt_id`. The DO is single-threaded SQLite, so ops are atomic.
+//
+// Usage:
+//   Set `VGI_WORKER_SHARED_STORAGE=cloudflare-do` plus `VGI_CF_DO_URL`.
+//   Optionally set `VGI_CF_DO_TOKEN` for bearer auth.
+
+import type { FunctionStorage } from "./storage.js";
+import { UnknownInvocationError } from "./storage.js";
+
+/** Minimal fetch signature satisfied by both the global `fetch` and a
+ *  Cloudflare service-binding Fetcher (`env.<BINDING>`). */
+export type FetchLike = (
+  input: string,
+  init?: { method?: string; headers?: Record<string, string>; body?: string },
+) => Promise<Response>;
+
+// Namespaces under the unified state_* table for the legacy worker/queue
+// families the DO no longer has dedicated endpoints for.
+const NS_WORKER = new TextEncoder().encode("worker");
+
+export class FunctionStorageCfDo implements FunctionStorage {
+  private readonly _baseUrl: string;
+  private readonly _token: string | null;
+  /** Per-attach routing key (att-<hex uuid>); "" until pinned via forShard. */
+  private readonly _shardKey: string;
+  /** Fetch implementation. Defaults to the global `fetch`; pass a Cloudflare
+   *  service-binding Fetcher (env.<BINDING>) when the storage worker lives on
+   *  the same account/zone — a same-zone public-URL fetch is rejected (CF
+   *  error 1042), but a service binding routes worker-to-worker directly. */
+  private readonly _fetch: FetchLike;
+
+  constructor(opts: {
+    url: string;
+    token?: string | null;
+    shardKey?: string;
+    fetch?: FetchLike;
+  }) {
+    // Strip trailing slash so endpoint paths can be appended unconditionally.
+    this._baseUrl = opts.url.replace(/\/+$/, "");
+    this._token = opts.token ?? null;
+    this._shardKey = opts.shardKey ?? "";
+    this._fetch = opts.fetch ?? ((input, init) => fetch(input as any, init));
+  }
+
+  /** Build an instance from environment variables. */
+  static fromEnv(): FunctionStorageCfDo {
+    const url = process.env.VGI_CF_DO_URL;
+    if (!url) {
+      throw new Error(
+        "VGI_CF_DO_URL environment variable is required when " +
+          "VGI_WORKER_SHARED_STORAGE=cloudflare-do",
+      );
+    }
+    return new FunctionStorageCfDo({
+      url,
+      token: process.env.VGI_CF_DO_TOKEN || null,
+    });
+  }
+
+  /**
+   * Return a view of this backend pinned to one shard key, so callers can route
+   * per logical ATTACH. Shares the URL/token; only the shard key differs.
+   */
+  forShard(shardKey: string): FunctionStorageCfDo {
+    return new FunctionStorageCfDo({
+      url: this._baseUrl,
+      token: this._token,
+      shardKey,
+      fetch: this._fetch,
+    });
+  }
+
+  // --- Worker state → ns=worker, key = int64(worker_id) ---
+
+  async workerPut(executionId: Uint8Array, workerId: number, state: Uint8Array): Promise<void> {
+    await this._statePutMany(executionId, NS_WORKER, [[int64Key(workerId), state]]);
+  }
+
+  async workerCollect(executionId: Uint8Array): Promise<Uint8Array[]> {
+    const rows = await this._statePaged("state_drain", executionId, NS_WORKER, newAttemptId());
+    return rows.map(([, v]) => v);
+  }
+
+  async workerScan(executionId: Uint8Array): Promise<Array<[number, Uint8Array]>> {
+    const rows = await this._statePaged("state_scan", executionId, NS_WORKER, null);
+    return rows.map(([k, v]) => [int64FromKey(k), v]);
+  }
+
+  // --- Work queue ---
+
+  async queuePush(executionId: Uint8Array, items: Uint8Array[]): Promise<number> {
+    const data = await this._post<{ count: number }>("queue_push", {
+      execution_id: bytesToB64(executionId),
+      items: items.map(bytesToB64),
+      attempt_id: newAttemptId(),
+    });
+    return Number(data.count);
+  }
+
+  async queuePop(executionId: Uint8Array): Promise<Uint8Array | null> {
+    const data = await this._post<{ item: string | null }>("queue_pop", {
+      execution_id: bytesToB64(executionId),
+      attempt_id: newAttemptId(),
+    });
+    return data.item ? b64ToBytes(data.item) : null;
+  }
+
+  async queueClear(executionId: Uint8Array): Promise<number> {
+    const data = await this._post<{ cleared: number }>("queue_clear", {
+      execution_id: bytesToB64(executionId),
+      attempt_id: newAttemptId(),
+    });
+    return Number(data.cleared);
+  }
+
+  // --- Namespaced key/value state ---
+
+  async stateGet(scopeId: Uint8Array, ns: Uint8Array, key: Uint8Array): Promise<Uint8Array | null> {
+    const data = await this._post<{ rows: Array<{ value: string } | null> }>("state_get_many", {
+      scope_id: bytesToB64(scopeId),
+      ns: bytesToB64(ns),
+      keys: [bytesToB64(key)],
+    });
+    const row = (data.rows ?? [])[0];
+    return row ? b64ToBytes(row.value) : null;
+  }
+
+  async statePut(scopeId: Uint8Array, ns: Uint8Array, key: Uint8Array, value: Uint8Array): Promise<void> {
+    await this._statePutMany(scopeId, ns, [[key, value]]);
+  }
+
+  // --- Append-only log ---
+
+  async stateAppend(scopeId: Uint8Array, ns: Uint8Array, key: Uint8Array, item: Uint8Array): Promise<number> {
+    const data = await this._post<{ ordinal: number }>("state_append", {
+      scope_id: bytesToB64(scopeId),
+      ns: bytesToB64(ns),
+      key: bytesToB64(key),
+      item: bytesToB64(item),
+      attempt_id: newAttemptId(),
+    });
+    return Number(data.ordinal);
+  }
+
+  async stateLogScan(
+    scopeId: Uint8Array,
+    ns: Uint8Array,
+    key: Uint8Array,
+    afterId: number = -1,
+    limit: number | null = null,
+  ): Promise<Array<[number, Uint8Array]>> {
+    const body: Record<string, unknown> = {
+      scope_id: bytesToB64(scopeId),
+      ns: bytesToB64(ns),
+      key: bytesToB64(key),
+      after_id: afterId,
+    };
+    if (limit != null && limit > 0) body.limit = limit;
+    const data = await this._post<{ rows: Array<{ id: number; value: string }> }>("state_log_scan", body);
+    return (data.rows ?? []).map((r) => [Number(r.id), b64ToBytes(r.value)]);
+  }
+
+  async executionClear(scopeId: Uint8Array): Promise<number> {
+    const data = await this._post<{ deleted: number }>("execution_clear", {
+      scope_id: bytesToB64(scopeId),
+      attempt_id: newAttemptId(),
+    });
+    return Number(data.deleted);
+  }
+
+  // --- unified state_* helpers ---
+
+  private async _statePutMany(
+    scopeId: Uint8Array,
+    ns: Uint8Array,
+    items: Array<[Uint8Array, Uint8Array]>,
+  ): Promise<void> {
+    await this._post("state_put_many", {
+      scope_id: bytesToB64(scopeId),
+      ns: bytesToB64(ns),
+      items: items.map(([k, v]) => ({ key: bytesToB64(k), value: bytesToB64(v) })),
+      attempt_id: newAttemptId(),
+    });
+  }
+
+  /** Drive state_scan / state_drain across pages. attemptId is null for scan. */
+  private async _statePaged(
+    endpoint: string,
+    scopeId: Uint8Array,
+    ns: Uint8Array,
+    attemptId: string | null,
+  ): Promise<Array<[Uint8Array, Uint8Array]>> {
+    const out: Array<[Uint8Array, Uint8Array]> = [];
+    let afterKey: string | undefined;
+    for (;;) {
+      const body: Record<string, unknown> = { scope_id: bytesToB64(scopeId), ns: bytesToB64(ns) };
+      if (afterKey != null) body.after_key = afterKey;
+      if (attemptId != null) body.attempt_id = attemptId;
+      const data = await this._post<{ rows: Array<{ key: string; value: string }>; next_after?: string }>(
+        endpoint,
+        body,
+      );
+      for (const r of data.rows ?? []) out.push([b64ToBytes(r.key), b64ToBytes(r.value)]);
+      if (!data.next_after) break;
+      afterKey = data.next_after;
+    }
+    return out;
+  }
+
+  // --- HTTP plumbing ---
+
+  private async _post<T = Record<string, unknown>>(
+    endpoint: string,
+    body: Record<string, unknown>,
+  ): Promise<T> {
+    const url = `${this._baseUrl}/${endpoint}`;
+    const headers: Record<string, string> = { "Content-Type": "application/json" };
+    if (this._token) headers["Authorization"] = `Bearer ${this._token}`;
+    // The Worker routes on shard_key (idFromName) and rejects requests without
+    // one — always splice it in.
+    const payload = JSON.stringify({ ...body, shard_key: this._shardKey });
+
+    let lastErr: unknown;
+    // One retry on transport-level failures, matching the python client.
+    for (let attempt = 0; attempt < 2; attempt++) {
+      let resp: Response;
+      try {
+        resp = await this._fetch(url, { method: "POST", headers, body: payload });
+      } catch (err) {
+        lastErr = err;
+        continue;
+      }
+
+      let data: any;
+      try {
+        data = await resp.json();
+      } catch {
+        if (!resp.ok) {
+          throw new Error(`CF DO storage error ${resp.status} on ${endpoint}: <non-json body>`);
+        }
+        return {} as T;
+      }
+
+      if (resp.status === 404 && data?.error === "unknown_invocation") {
+        throw new UnknownInvocationError(
+          data.message ??
+            "Invocation is not registered. Call queuePush first to register the invocation.",
+        );
+      }
+      if (resp.status === 401) {
+        throw new Error(`Authentication failed: ${data?.error ?? "unauthorized"}`);
+      }
+      if (!resp.ok) {
+        throw new Error(`CF DO storage error ${resp.status} on ${endpoint}: ${JSON.stringify(data)}`);
+      }
+      return data as T;
+    }
+    throw lastErr instanceof Error
+      ? lastErr
+      : new Error(`CF DO storage transport error on ${endpoint}: ${String(lastErr)}`);
+  }
+}
+
+/** Fresh 32-char lowercase-hex idempotency token (the DO's attempt_id shape). */
+function newAttemptId(): string {
+  const b = new Uint8Array(16);
+  crypto.getRandomValues(b);
+  let s = "";
+  for (let i = 0; i < 16; i++) s += b[i].toString(16).padStart(2, "0");
+  return s;
+}
+
+/** Encode an int64 worker/group id as an 8-byte big-endian state key. */
+function int64Key(v: number): Uint8Array {
+  const b = new Uint8Array(8);
+  new DataView(b.buffer).setBigInt64(0, BigInt(v), false);
+  return b;
+}
+
+function int64FromKey(b: Uint8Array): number {
+  if (b.length !== 8) return 0;
+  return Number(new DataView(b.buffer, b.byteOffset, 8).getBigInt64(0, false));
+}
+
+function bytesToB64(buf: Uint8Array): string {
+  let binary = "";
+  for (let i = 0; i < buf.length; i++) binary += String.fromCharCode(buf[i]);
+  return btoa(binary);
+}
+
+function b64ToBytes(b64: string): Uint8Array {
+  const binary = atob(b64);
+  const out = new Uint8Array(binary.length);
+  for (let i = 0; i < binary.length; i++) out[i] = binary.charCodeAt(i);
+  return out;
+}