@specific.dev/spectest 0.9.0 → 0.11.0

package/src/sql.ts

@@ -0,0 +1,223 @@
+// `SQL` — a drop-in for `Bun.SQL` that records every query on the test event
+// log and returns its rows inspect-wrapped, so `expect(rows[0]!.id)` links the
+// assertion back to the query in the timeline (same provenance mechanism as
+// `fetch`). Swap `new Bun.SQL(url)` → `new SQL(url)` and nothing else changes
+// at the call site; query *results* come back `Wrapped<T[]>` instead of raw
+// `T[]` (that wrapper is what carries provenance), so `.unwrap()` before feeding
+// a row value to another real client or a `===` — exactly the rule that already
+// holds for `ctx.svc.<postgres>.client`.
+//
+// This is the core primitive; the `postgres(...)` component is a thin user of
+// it (its `helpers.client` is `new SQL(url, { label })`). Open one yourself
+// whenever you stand up a database without the component — a Bun-native image
+// like timescaledb/postgis, or a connection string a fake handed back after
+// provisioning a DB at runtime.
+
+import { recordDb, reserveEvent, safeSerialize } from "./recorder.js";
+import { wrap } from "./inspect.js";
+import type { Wrapped } from "./inspect.js";
+
+/**
+ * Minimal structural type for the parts of `Bun.SQL` we use. Declared locally
+ * so user projects don't need `@types/bun` for type-checking to follow the SDK
+ * through to the daemon. The real type comes from Bun at runtime via
+ * `globalThis.Bun.SQL`.
+ */
+interface RawSqlClient {
+  <T = unknown>(strings: TemplateStringsArray, ...values: unknown[]): Promise<T[]>;
+  unsafe<T = unknown>(text: string, params?: unknown[]): Promise<T[]>;
+  close?(opts?: { timeout?: number }): Promise<void>;
+  end(opts?: { timeout?: number }): Promise<void>;
+}
+
+interface SqlClientBase {
+  close?(opts?: { timeout?: number }): Promise<void>;
+  /** Close the connection / drain the pool (Bun's `SQL.end`). */
+  end(opts?: { timeout?: number }): Promise<void>;
+}
+
+/**
+ * The instrumented SQL surface. Same shape as `Bun.SQL` — callable as a tagged
+ * template (`` sql`SELECT 1` ``) with an `unsafe(text, params?)` escape hatch —
+ * but each query resolves to a **wrapped** row array ({@link Wrapped}): every
+ * settle is recorded as a `db` event and the wrapper carries provenance back to
+ * it. That's why `expect(rows)` / `expect(rows[0]!.id)` link under the query in
+ * the timeline with no cast — and why you `.unwrap()` before feeding a row value
+ * to a real client or a `===`.
+ *
+ * Generic in the row type, so callers name the shape once at the call site
+ * instead of casting: `` await client<TodoRow>`SELECT * FROM todos` `` is
+ * `Promise<Wrapped<TodoRow[]>>`.
+ */
+export interface SqlClient extends SqlClientBase {
+  <T = unknown>(strings: TemplateStringsArray, ...values: unknown[]): Promise<Wrapped<T[]>>;
+  unsafe<T = unknown>(text: string, params?: unknown[]): Promise<Wrapped<T[]>>;
+}
+
+/** Options accepted by the {@link SQL} constructor beyond Bun's own. */
+export interface SqlOptions {
+  /**
+   * Label shown on each recorded `db` event (the timeline step's service tag).
+   * Defaults to the connection URL's host — e.g. `db` for
+   * `postgres://…@db:5432/app` — which is usually the service name.
+   */
+  label?: string;
+}
+
+interface BunGlobal {
+  SQL: new (url: string) => RawSqlClient;
+}
+
+interface SqlConstructor {
+  new (url: string, opts?: SqlOptions): SqlClient;
+  (url: string, opts?: SqlOptions): SqlClient;
+}
+
+/**
+ * Open an instrumented SQL client against `url`. Usable with or without `new`
+ * (`new SQL(url)` mirrors `new Bun.SQL(url)`; a constructor that returns an
+ * object yields that object). Requires the Bun runtime — it runs inside the
+ * spectest daemon.
+ */
+export const SQL = function SQL(url: string, opts?: SqlOptions): SqlClient {
+  const bun = (globalThis as { Bun?: BunGlobal }).Bun;
+  if (!bun?.SQL) {
+    throw new Error(
+      "SQL(url) requires the Bun runtime (Bun >= 1.2) — it runs inside the spectest daemon.",
+    );
+  }
+  return instrumentSql(new bun.SQL(url), opts?.label ?? hostLabel(url));
+} as unknown as SqlConstructor;
+
+/** Derive a default event label from a connection URL's host. */
+function hostLabel(url: string): string {
+  try {
+    return new URL(url).hostname || "db";
+  } catch {
+    return "db";
+  }
+}
+
+/**
+ * Proxy a `Bun.SQL` instance so each tagged-template call and each
+ * `unsafe(...)` call emits a `db` event into the active recorder when its
+ * promise settles, and resolves to a {@link Wrapped} result. Other `Bun.SQL`
+ * methods (`.transaction`, `.array`, `.file`, …) pass through unwrapped.
+ *
+ * Exported so a client built some other way (e.g. a pool you already hold) can
+ * opt into the same instrumentation without going through {@link SQL}.
+ */
+export function instrumentSql(raw: RawSqlClient, label: string): SqlClient {
+  const wrapResult = <T>(
+    result: PromiseLike<T>,
+    query: string,
+    params: unknown[],
+  ): Promise<T> => {
+    const started = Date.now();
+    const resv = reserveEvent();
+    return Promise.resolve(result).then(
+      (value) => {
+        const { rows, rowsTruncated, columns } = captureRows(value);
+        const seq = recordDb(
+          {
+            service: label,
+            query,
+            params: params.length > 0 ? params.map(safeSerialize) : undefined,
+            rowCount: rowCountOf(value),
+            rows,
+            rowsTruncated,
+            columns,
+            durationMs: Date.now() - started,
+          },
+          resv,
+        );
+        return wrap(value, seq) as T;
+      },
+      (err) => {
+        const e = err as Error;
+        recordDb(
+          {
+            service: label,
+            query,
+            params: params.length > 0 ? params.map(safeSerialize) : undefined,
+            durationMs: Date.now() - started,
+            error: e?.message ?? String(err),
+          },
+          resv,
+        );
+        throw err;
+      },
+    );
+  };
+
+  const handler: ProxyHandler<RawSqlClient> = {
+    apply(target, thisArg, args) {
+      const [strings, ...values] = args as [TemplateStringsArray, ...unknown[]];
+      const query = reconstructSqlTemplate(strings, values.length);
+      const result = Reflect.apply(
+        target as unknown as (...a: unknown[]) => PromiseLike<unknown>,
+        thisArg,
+        args,
+      );
+      return wrapResult(result, query, values);
+    },
+    get(target, prop, receiver) {
+      if (prop === "unsafe") {
+        return (text: string, params?: unknown[]) => {
+          const result = (
+            target.unsafe as (t: string, p?: unknown[]) => PromiseLike<unknown>
+          ).call(target, text, params);
+          return wrapResult(result, text, params ?? []);
+        };
+      }
+      const v = Reflect.get(target, prop, receiver);
+      return typeof v === "function" ? v.bind(target) : v;
+    },
+  };
+  return new Proxy(raw, handler) as unknown as SqlClient;
+}
+
+/** Rebuild a parameterized SQL string from a tagged-template's pieces.
+ * Bun.SQL substitutes `${value}` for `$1`, `$2`, …; we do the same so
+ * the recorded query matches what the server sees. */
+function reconstructSqlTemplate(
+  strings: TemplateStringsArray,
+  valueCount: number,
+): string {
+  let out = strings[0] ?? "";
+  for (let i = 0; i < valueCount; i++) {
+    out += `$${i + 1}` + (strings[i + 1] ?? "");
+  }
+  return out.trim();
+}
+
+function rowCountOf(value: unknown): number | undefined {
+  if (Array.isArray(value)) return value.length;
+  if (value && typeof value === "object") {
+    const obj = value as { count?: unknown; rowCount?: unknown };
+    if (typeof obj.count === "number") return obj.count;
+    if (typeof obj.rowCount === "number") return obj.rowCount;
+  }
+  return undefined;
+}
+
+const MAX_DB_ROWS = 50;
+
+function captureRows(value: unknown): {
+  rows?: unknown[];
+  rowsTruncated?: boolean;
+  columns?: string[];
+} {
+  if (!Array.isArray(value) || value.length === 0) return {};
+  const slice = value.slice(0, MAX_DB_ROWS).map(safeSerialize);
+  const first = slice[0];
+  const columns =
+    first && typeof first === "object" && !Array.isArray(first)
+      ? Object.keys(first as Record<string, unknown>)
+      : undefined;
+  return {
+    rows: slice,
+    rowsTruncated: value.length > MAX_DB_ROWS ? true : undefined,
+    columns,
+  };
+}