@specific.dev/spectest 0.4.0

package/src/components/postgres.ts

@@ -0,0 +1,281 @@
+import type { ServiceDefinition } from "../index.js";
+import { recordDb, reserveEvent, safeSerialize } from "../recorder.js";
+import { wrap } from "../inspect.js";
+import type { Wrapped } from "../inspect.js";
+
+export interface PostgresOptions {
+  /** Image tag for the official `postgres` image. Default `"18-alpine"`. */
+  version?: string;
+  /** Database created on first boot. */
+  database: string;
+  /** Superuser name created on first boot. */
+  user: string;
+  /** Superuser password. */
+  password: string;
+  /**
+   * Whether postgres data persists across container restarts via a bind
+   * mount under `.spectest/volumes/<name>/`. Default `true`.
+   * Snapshots/forks always preserve state regardless of this flag — the
+   * volume just means a plain `docker rm`/`docker run` cycle keeps data.
+   */
+  persistent?: boolean;
+  /** TCP port the container listens on. Default `5432`. */
+  port?: number;
+  /** Extra environment variables forwarded to the container. */
+  env?: Record<string, string>;
+}
+
+/**
+ * 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`.
+ *
+ * `SqlClient` is callable as a tagged template (`sql\`...\``) and has an
+ * `unsafe(text, params?)` escape hatch. Both return a thenable that
+ * executes lazily; we wrap them so each settle is recorded as a DbEvent.
+ *
+ * Both forms are generic in the row type and resolve to an array of rows,
+ * so callers name the shape once at the call site instead of casting the
+ * result: `` await ctx.svc.db.client<TodoRow>`SELECT * FROM todos` `` is
+ * `Promise<TodoRow[]>`. `T` defaults to `unknown` (the row shape is the
+ * caller's to assert) — pass it to drop the cast.
+ */
+interface SqlClientBase {
+  close?(opts?: { timeout?: number }): Promise<void>;
+  /** Close the connection / drain the pool (Bun's `SQL.end`). */
+  end(opts?: { timeout?: number }): Promise<void>;
+}
+
+export interface SqlClient extends SqlClientBase {
+  // Callable as a tagged template; resolves to the result rows.
+  <T = unknown>(
+    strings: TemplateStringsArray,
+    ...values: unknown[]
+  ): Promise<T[]>;
+  unsafe<T = unknown>(text: string, params?: unknown[]): Promise<T[]>;
+}
+
+/**
+ * The recorder-instrumented client wired onto `ctx.svc.<name>.client`. Same
+ * surface as {@link SqlClient}, but each query resolves to a **wrapped** row
+ * array ({@link Wrapped}) rather than the raw rows: every settle is recorded
+ * as a DbEvent, 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 `===`. The raw escape-hatch {@link sqlClient}
+ * stays un-instrumented (plain {@link SqlClient}), so its results are genuinely
+ * raw and belong with `expectRaw`.
+ */
+export interface RecordingSqlClient extends SqlClientBase {
+  <T = unknown>(
+    strings: TemplateStringsArray,
+    ...values: unknown[]
+  ): Promise<Wrapped<T[]>>;
+  unsafe<T = unknown>(text: string, params?: unknown[]): Promise<Wrapped<T[]>>;
+}
+
+interface BunGlobal {
+  SQL: new (url: string) => SqlClient;
+}
+
+/**
+ * Open a raw Bun SQL client against `url`, typed as {@link SqlClient}.
+ *
+ * Unlike the pool `postgres(...)` wires onto `ctx.svc.<name>.client`, this
+ * is **not** recorder-instrumented — reach for it when you only learn a
+ * connection string at runtime (e.g. one a fake handed back after
+ * provisioning a database via `ctx.startService(...)`), where the boot-time
+ * component isn't an option. Queries are still generic in the row type:
+ * `` await sqlClient(url)<{ name: string }>`SELECT name FROM t` ``.
+ *
+ * Requires the Bun runtime (it runs inside the spectest daemon), so it
+ * replaces the `(globalThis as any).Bun` reach-around with a typed entry
+ * point.
+ */
+export function sqlClient(url: string): SqlClient {
+  const bun = (globalThis as { Bun?: BunGlobal }).Bun;
+  if (!bun?.SQL) {
+    throw new Error(
+      "sqlClient(url) requires Bun.SQL (Bun >= 1.2) — it runs inside the spectest daemon.",
+    );
+  }
+  return new bun.SQL(url);
+}
+
+/** Helpers a `postgres(...)` service exposes on `ctx.svc.<name>`. */
+export interface PostgresHelpers {
+  /** A Bun SQL pool, wrapped so each query lands on the test event
+   * log alongside http/exec/assertion events — its rows come back
+   * inspect-wrapped, so `expect(...)` on them links to the query. */
+  client: RecordingSqlClient;
+}
+
+/**
+ * A ready-to-use Postgres service. Drop into `environment.services`:
+ *
+ * ```ts
+ * services: {
+ *   db: postgres({ database: "todos", user: "todos", password: "todos" }),
+ *   ...
+ * }
+ * ```
+ *
+ * Peer services reach it at `<key>:<port>` (e.g. `db:5432` for the
+ * example above). Tests get a wired-up Bun SQL client at
+ * `ctx.svc.<key>.client` — `await ctx.svc.db.client\`SELECT 1\`` — with
+ * every query recorded on the test event log.
+ */
+export function postgres(opts: PostgresOptions) {
+  const version = opts.version ?? "18-alpine";
+  const port = opts.port ?? 5432;
+  const persistent = opts.persistent ?? true;
+  // `satisfies` (instead of a return-type annotation) preserves the
+  // literal type — in particular, the `helpers` factory's return shape.
+  // The mapped type in `ServiceHandlesFor<S>` reads that to type
+  // `ctx.svc.<name>` as `{ client: SqlClient }`.
+  return {
+    image: { type: "registry", reference: `postgres:${version}` },
+    env: {
+      POSTGRES_DB: opts.database,
+      POSTGRES_USER: opts.user,
+      POSTGRES_PASSWORD: opts.password,
+      // PGDATA lives in a subdir so postgres can initialise inside a bind
+      // mount that may not be empty on first boot.
+      PGDATA: "/var/lib/postgresql/data/pgdata",
+      ...(opts.env ?? {}),
+    },
+    volumes: persistent ? [{ target: "/var/lib/postgresql/data" }] : [],
+    ports: [port],
+    readyCheck: { type: "tcp" as const, port, timeoutSecs: 60 },
+    helpers: ({ name }: { name: string }): PostgresHelpers => {
+      const bun = (globalThis as { Bun?: BunGlobal }).Bun;
+      if (!bun?.SQL) {
+        throw new Error(
+          "postgres component requires Bun.SQL (Bun >= 1.2). " +
+            "This helpers factory is meant to run inside the spectest daemon.",
+        );
+      }
+      const url =
+        `postgres://${encodeURIComponent(opts.user)}` +
+        `:${encodeURIComponent(opts.password)}` +
+        `@${name}:${port}/${encodeURIComponent(opts.database)}`;
+      return { client: wrapSqlForRecording(new bun.SQL(url), name) };
+    },
+  } satisfies ServiceDefinition<PostgresHelpers>;
+}
+
+/**
+ * Proxy a Bun.SQL instance so each tagged-template call and each
+ * `unsafe(...)` call emits a DbEvent into the active recorder when its
+ * promise settles. Other Bun.SQL methods (`.transaction`, `.array`,
+ * `.file`, …) pass through unwrapped.
+ */
+function wrapSqlForRecording(raw: SqlClient, service: string): RecordingSqlClient {
+  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,
+          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,
+          query,
+          params: params.length > 0 ? params.map(safeSerialize) : undefined,
+          durationMs: Date.now() - started,
+          error: e?.message ?? String(err),
+        }, resv);
+        throw err;
+      },
+    );
+  };
+
+  const handler: ProxyHandler<SqlClient> = {
+    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 RecordingSqlClient;
+}
+
+/** 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,
+  };
+}