@intentius/chant-lexicon-sql 0.0.0-stage → 0.112.1

package/src/postgres/migrate/batches.e2e.test.ts

@@ -0,0 +1,163 @@
+/**
+ * The backfill's batches for a primary key that is not one integer column
+ * (#3322), against the pinned server: a uuid key's boundaries are walked and
+ * recorded once, and a resumed backfill finds the same batches and skips
+ * the ones its receipts vouch for; a composite key and a text key run the
+ * Op through its gates to the end.
+ *
+ * Needs Docker; skips cleanly without it.
+ */
+
+import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterAll, beforeAll, describe, expect, test } from "vitest";
+import { dockerAvailable, startTestPostgres, type TestPostgres } from "../testing/server";
+import { ApprovingLedger, runMigrationOp, runOutcome } from "../testing/migration";
+import type { PostgresClient } from "../live/client";
+import { planPgAgainstServer } from "../plan/commands";
+import { postgresApply } from "../../op/activities/postgres-apply";
+import * as m from "../../op/activities/postgres-migration";
+import type { PostgresMigrationArgs, PostgresMigrationOpConfig } from "./op";
+import { POSTGRES_RECEIPTS_TABLE } from "./receipts";
+
+const enabled = await dockerAvailable();
+let server: TestPostgres | undefined;
+let admin: PostgresClient | undefined;
+const dir = mkdtempSync(join(tmpdir(), "chant-pg-batches-"));
+
+beforeAll(async () => {
+  if (!enabled) return;
+  server = await startTestPostgres();
+  admin = await server.connect();
+}, 600_000);
+
+afterAll(async () => {
+  await admin?.end();
+  await server?.stop();
+  rmSync(dir, { recursive: true, force: true });
+});
+
+interface Obj {
+  export: string;
+  type: string;
+  ddl: string;
+  dependsOn?: string[];
+}
+
+function writeBuild(file: string, objects: Obj[]): string {
+  writeFileSync(join(dir, file), JSON.stringify({ dialect: "postgres", postgresMajor: 18, applyOrder: objects.map((o) => o.export), objects: objects.map((o) => ({ dependsOn: [], ...o })) }));
+  return file;
+}
+
+const MARKER = { stack: "e2e", env: "test" };
+const profile = () => ({
+  config: { ownership: MARKER, sql: { profiles: { e2e: { url: server!.endpoint().url, password: { env: "PG_E2E_PASSWORD" } } } } },
+  env: { PG_E2E_PASSWORD: server!.endpoint().password },
+});
+const deps = (extra: Partial<m.PostgresMigrationDeps> = {}): m.PostgresMigrationDeps => ({ ...profile(), log: () => undefined, ...extra });
+const apply = (file: string) => postgresApply({ buildPath: join(dir, file), environment: "e2e" }, undefined, { ...profile(), log: () => undefined });
+const plan = (file: string) => planPgAgainstServer("e2e", join(dir, file), profile());
+const one = async <T>(sql: string): Promise<T> => (await admin!.query<T>(sql))[0]!;
+const schemaObj = (name: string): Obj => ({ export: name, type: "Postgres::Schema", ddl: `CREATE SCHEMA ${name}` });
+const config = (over: Partial<PostgresMigrationOpConfig> & Pick<PostgresMigrationOpConfig, "name" | "table" | "column" | "output">): PostgresMigrationOpConfig => ({
+  env: "e2e",
+  build: false,
+  path: dir,
+  retain: "0s",
+  stack: MARKER.stack,
+  ownershipEnv: MARKER.env,
+  ...over,
+});
+
+/** Run the Op until it is done, approving each gate it stops at. */
+async function runToEnd(c: PostgresMigrationOpConfig) {
+  const ledger = new ApprovingLedger();
+  const runs = [];
+  for (let i = 0; i < 4; i++) {
+    const r = await runMigrationOp(c, ledger.port, () => deps());
+    runs.push(r);
+    if (r.status !== "gated") break;
+    ledger.approveLast();
+  }
+  return runs;
+}
+
+describe.skipIf(!enabled)("a uuid primary key", () => {
+  const UU = schemaObj("uu");
+  const tokens = (score: string): Obj => ({ export: "tokens", type: "Postgres::Table", dependsOn: ["uu"], ddl: `CREATE TABLE uu.tokens (id uuid PRIMARY KEY, score ${score})` });
+  const args = (): PostgresMigrationArgs => ({ table: "uu.tokens", column: "score", buildPath: "tokens-v2.json", environment: "e2e", batchSize: 1000, stack: MARKER.stack, ownershipEnv: MARKER.env, cwd: dir });
+
+  test("boundaries walked and recorded once; a resumed backfill finds the same batches and updates no row twice", async () => {
+    await apply(writeBuild("tokens-v1.json", [UU, tokens("text")]));
+    await admin!.query("INSERT INTO uu.tokens SELECT gen_random_uuid(), (g % 50)::text FROM generate_series(1, 2500) g");
+    await admin!.query("CREATE SCHEMA audit");
+    await admin!.query("CREATE TABLE audit.updates (id uuid)");
+    await admin!.query("CREATE FUNCTION audit.count_update() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN INSERT INTO audit.updates VALUES (NEW.id); RETURN NEW; END $$");
+    await admin!.query("CREATE TRIGGER audit_update AFTER UPDATE ON uu.tokens FOR EACH ROW EXECUTE FUNCTION audit.count_update()");
+    writeBuild("tokens-v2.json", [UU, tokens("integer")]);
+    expect((await plan("tokens-v2.json")).changes.map((c) => c.rule)).toEqual(["SQLPG208"]);
+
+    expect(await m.postgresMigrationPlan(args(), undefined, deps())).toMatchObject({ state: "migrate", summary: expect.stringContaining("batched by id") });
+    await m.postgresMigrationExpand(args(), undefined, deps());
+    await m.postgresMigrationDualWrite(args(), undefined, deps());
+
+    // Interrupted after the first batch.
+    const stop = new AbortController();
+    const interrupted = await m.postgresMigrationBackfill(args(), stop.signal, deps({ backfill: { afterBatch: () => stop.abort() } })).then(
+      () => undefined,
+      (e: unknown) => e,
+    );
+    expect(interrupted).toBeDefined();
+    const record = await one<{ e: string }>(`SELECT expectation AS e FROM uu.${POSTGRES_RECEIPTS_TABLE} WHERE address LIKE '%/migrate/uu.tokens.score/batches'`);
+    const recorded = JSON.parse(record.e) as { keys: string[]; size: number; bounds: string[][] };
+    expect(recorded).toMatchObject({ keys: ["id"], size: 1000 });
+    expect(recorded.bounds).toHaveLength(3);
+    expect(await one("SELECT count(*)::int AS n FROM uu.tokens WHERE score__chant_new IS NOT NULL")).toEqual({ n: 1000 });
+
+    // Rows written meanwhile land anywhere in the key's order; the dual write fills them.
+    await admin!.query("INSERT INTO uu.tokens SELECT gen_random_uuid(), '7' FROM generate_series(1, 500)");
+
+    const resumed = await m.postgresMigrationBackfill(args(), undefined, deps());
+    expect(resumed).toMatchObject({ batches: 3, skipped: 1, filled: 2, rows: 1500 });
+    expect(await one("SELECT count(*)::int AS n FROM uu.tokens WHERE score__chant_new IS NULL")).toEqual({ n: 0 });
+    // The recorded boundaries were read back, not walked again over 3000 rows.
+    expect(await one<{ e: string }>(`SELECT expectation AS e FROM uu.${POSTGRES_RECEIPTS_TABLE} WHERE address LIKE '%/batches'`)).toEqual(record);
+    expect(await one("SELECT count(*)::int AS rows, max(n)::int AS most FROM (SELECT id, count(*) AS n FROM audit.updates GROUP BY id) s")).toEqual({ rows: 2500, most: 1 });
+    expect(await m.postgresMigrationVerify(args(), undefined, deps())).toMatchObject({ rows: 3000, mismatched: 0 });
+
+    const again = await m.postgresMigrationBackfill(args(), undefined, deps());
+    expect(again).toMatchObject({ batches: 3, skipped: 3, filled: 0 });
+    const dropped = await m.postgresMigrationCompensate(args(), undefined, deps());
+    expect(dropped.dropped.at(-1)).toBe("4 receipt(s)");
+  }, 300_000);
+});
+
+describe.skipIf(!enabled)("a composite and a text primary key, through the gates", () => {
+  test("(tenant_id, id): a varchar narrowed in batches of the key's order", async () => {
+    const MT = schemaObj("mt");
+    const events = (kind: string): Obj => ({ export: "events", type: "Postgres::Table", dependsOn: ["mt"], ddl: `CREATE TABLE mt.events (tenant_id integer, id bigint, kind ${kind} NOT NULL, PRIMARY KEY (tenant_id, id))` });
+    await apply(writeBuild("events-v1.json", [MT, events("varchar(20)")]));
+    await admin!.query("INSERT INTO mt.events SELECT g % 4, g, CASE WHEN g % 2 = 0 THEN 'click' ELSE 'view' END FROM generate_series(1, 1000) g");
+    writeBuild("events-v2.json", [MT, events("varchar(10)")]);
+    const runs = await runToEnd(config({ name: "migrate-events-kind", table: "mt.events", column: "kind", output: "events-v2.json", batchSize: 300 }));
+    expect(runs.map((r) => r.status)).toEqual(["gated", "gated", "ok"]);
+    expect(runOutcome(runs[0]!, "Filled")).toBe(4);
+    expect(runOutcome(runs[0]!, "BackfilledRows")).toBe(1000);
+    expect(await one("SELECT pg_catalog.format_type(atttypid, atttypmod) AS t FROM pg_catalog.pg_attribute WHERE attrelid = 'mt.events'::regclass AND attname = 'kind'")).toEqual({ t: "character varying(10)" });
+    expect((await plan("events-v2.json")).changes).toEqual([]);
+  }, 300_000);
+
+  test("a text key: a rename", async () => {
+    const TX = schemaObj("tx");
+    const codes = (col: string): Obj => ({ export: "codes", type: "Postgres::Table", dependsOn: ["tx"], ddl: `CREATE TABLE tx.codes (code text PRIMARY KEY,\n  ${col}\n)` });
+    await apply(writeBuild("codes-v1.json", [TX, codes("label text")]));
+    await admin!.query("INSERT INTO tx.codes SELECT 'c' || g, 'L' || g FROM generate_series(1, 700) g");
+    writeBuild("codes-v2.json", [TX, codes("title text -- previously: label")]);
+    const runs = await runToEnd(config({ name: "migrate-codes-title", table: "tx.codes", column: "title", output: "codes-v2.json", batchSize: 250 }));
+    expect(runs.map((r) => r.status)).toEqual(["gated", "gated", "ok"]);
+    expect(runOutcome(runs[0]!, "Filled")).toBe(3);
+    expect(await one("SELECT count(*)::int AS n FROM tx.codes WHERE title = 'L' || substr(code, 2)")).toEqual({ n: 700 });
+    expect((await plan("codes-v2.json")).changes).toEqual([]);
+  }, 300_000);
+});

package/src/postgres/migrate/batches.test.ts

@@ -0,0 +1,36 @@
+/**
+ * The backfill's batches (#3322), the parts that need no server: which keys
+ * batch by arithmetic, and the conditions between recorded boundaries.
+ * `batches.e2e.test.ts` runs them against the pinned server.
+ */
+
+import { describe, expect, test } from "vitest";
+import { boundaryBatches, isIntegerKey, keyText } from "./batches";
+
+describe("batches", () => {
+  test("one integer column batches by arithmetic; anything else by recorded boundaries", () => {
+    expect(isIntegerKey([{ name: "id", type: "bigint" }])).toBe(true);
+    expect(isIntegerKey([{ name: "id", type: "uuid" }])).toBe(false);
+    expect(isIntegerKey([{ name: "tenant_id", type: "integer" }, { name: "id", type: "integer" }])).toBe(false);
+    expect(keyText([{ name: "tenant_id", type: "integer" }, { name: "id", type: "bigint" }])).toBe("(tenant_id, id)");
+  });
+
+  test("between boundaries: the first batch has no lower bound and the last no upper one, in the key's own order", () => {
+    const keys = [
+      { name: "tenant_id", type: "integer" },
+      { name: "id", type: "uuid" },
+    ];
+    const batches = boundaryBatches(keys, [
+      ["1", "00000000-0000-0000-0000-000000000001"],
+      ["2", "00000000-0000-0000-0000-000000000005"],
+      ["3", "00000000-0000-0000-0000-000000000002"],
+    ]);
+    expect(batches.map((b) => [b.id, b.where, b.params])).toEqual([
+      ["k0", "(tenant_id, id) < ($1::integer, $2::uuid)", ["2", "00000000-0000-0000-0000-000000000005"]],
+      ["k1", "(tenant_id, id) >= ($1::integer, $2::uuid) AND (tenant_id, id) < ($3::integer, $4::uuid)", ["2", "00000000-0000-0000-0000-000000000005", "3", "00000000-0000-0000-0000-000000000002"]],
+      ["k2", "(tenant_id, id) >= ($1::integer, $2::uuid)", ["3", "00000000-0000-0000-0000-000000000002"]],
+    ]);
+    expect(boundaryBatches([{ name: "code", type: "text" }], [["a"]]).map((b) => [b.where, b.params])).toEqual([["true", []]]);
+    expect(boundaryBatches([{ name: "code", type: "text" }], [])).toEqual([]);
+  });
+});

package/src/postgres/migrate/batches.ts

@@ -0,0 +1,159 @@
+/**
+ * How the backfill cuts a table into batches (#3281, #3322). A batch must be
+ * the same rows on every run, or a receipt written for it by one run would
+ * vouch for different rows in the next.
+ *
+ * - A primary key of one smallint, integer or bigint column: ranges of its
+ *   value, `[b * size, (b + 1) * size)`. Arithmetic, so nothing needs to be
+ *   remembered, and a batch is the same range however many rows were
+ *   written since.
+ * - Any other primary key (uuid, text, a timestamp, several columns):
+ *   ranges between recorded boundaries in the key's order. The first run
+ *   walks the key's index once, taking every `size`-th key as a boundary
+ *   (`WHERE key >= boundary ORDER BY key OFFSET size LIMIT 1`, so the walk
+ *   reads each key once and sorts nothing), and records the boundaries in
+ *   the receipt store, beside the batches' receipts, before any batch runs.
+ *   Every later run reads them back instead of walking again, so batch `i`
+ *   is `boundary[i] <= key < boundary[i + 1]` in every run. The first batch
+ *   has no lower bound and the last no upper one, so a row written since,
+ *   anywhere in the key's order, falls in exactly one batch (and the dual
+ *   write has already filled it). Each batch is a range scan of the primary
+ *   key's index, with a row comparison (`(a, b) >= ($1, $2)`) in the index's
+ *   own column order.
+ *
+ * Hash buckets (`hash(key) % n`) were the other option: they need nothing
+ * recorded but the bucket count, but no index can serve them, so every batch
+ * would read the whole table. Ranges recorded once read each row once.
+ *
+ * The record is bound to the table's oid and the new column's attribute
+ * number, like the batches' receipts: a migration started again after
+ * onFailure (which drops the new column and the receipts) walks the key
+ * again. A record left from another batch size is kept, since its batches'
+ * receipts are what the run resumes from.
+ */
+
+import { col } from "./names";
+import type { PostgresClient } from "../live/client";
+import type { PostgresReceiptStore } from "./receipts";
+
+/** One primary-key column, in the key's order. */
+export interface KeyColumn {
+  name: string;
+  /** As `format_type` prints it: what a boundary's text is cast back to. */
+  type: string;
+}
+
+/** One batch: its id (the receipt's address suffix), its rows as a condition with parameters from `$1`, and what its receipt binds. */
+export interface Batch {
+  id: string;
+  where: string;
+  params: string[];
+  inputs: Record<string, unknown>;
+}
+
+export const INTEGER_KEY_TYPES: ReadonlySet<string> = new Set(["smallint", "integer", "bigint"]);
+
+/** Whether the key is batched by arithmetic ranges of its one integer column. */
+export const isIntegerKey = (keys: readonly KeyColumn[]): boolean => keys.length === 1 && INTEGER_KEY_TYPES.has(keys[0]!.type);
+
+/** `id`, or `(tenant_id, id)`. */
+export const keyText = (keys: readonly KeyColumn[]): string => (keys.length === 1 ? keys[0]!.name : `(${keys.map((k) => k.name).join(", ")})`);
+
+const tuple = (keys: readonly KeyColumn[]): string => (keys.length === 1 ? col(keys[0]!.name) : `(${keys.map((k) => col(k.name)).join(", ")})`);
+const placeholders = (keys: readonly KeyColumn[], from: number): string => {
+  const p = keys.map((k, i) => `$${from + i}::${k.type}`);
+  return keys.length === 1 ? p[0]! : `(${p.join(", ")})`;
+};
+
+/** The batches of an integer key: one per range that holds rows. */
+export async function integerBatches(client: PostgresClient, table: string, key: KeyColumn, size: number): Promise<Batch[]> {
+  const k = col(key.name);
+  const rows = await client.query<{ b: string }>(`SELECT pg_catalog.floor(${k}::numeric / ${size})::bigint::text AS b FROM ${table} GROUP BY 1 ORDER BY 1`);
+  return rows.map(({ b }) => {
+    const lo = (BigInt(b) * BigInt(size)).toString();
+    const hi = ((BigInt(b) + 1n) * BigInt(size)).toString();
+    return { id: b, where: `${k} >= $1 AND ${k} < $2`, params: [lo, hi], inputs: { batch: b, size, key: key.name } };
+  });
+}
+
+/** The batches between recorded boundaries: `[undefined, b1)`, `[b1, b2)`, ..., `[bn, undefined)`, where `b0` is only where the walk started. */
+export function boundaryBatches(keys: readonly KeyColumn[], bounds: readonly string[][]): Batch[] {
+  const k = tuple(keys);
+  return bounds.map((lo, i) => {
+    const hi = bounds[i + 1];
+    const where: string[] = [];
+    const params: string[] = [];
+    if (i > 0) {
+      where.push(`${k} >= ${placeholders(keys, 1)}`);
+      params.push(...lo);
+    }
+    if (hi) {
+      where.push(`${k} < ${placeholders(keys, params.length + 1)}`);
+      params.push(...hi);
+    }
+    return {
+      id: `k${i}`,
+      where: where.join(" AND ") || "true",
+      params,
+      inputs: { batch: `k${i}`, key: keys.map((c) => c.name), lo: i > 0 ? lo : null, hi: hi ?? null },
+    };
+  });
+}
+
+/** Walk the key's order and take every `size`-th key: the boundaries the batches lie between. */
+export async function walkBoundaries(client: PostgresClient, table: string, keys: readonly KeyColumn[], size: number, signal?: AbortSignal): Promise<string[][]> {
+  const select = keys.map((k, i) => `${col(k.name)}::text AS k${i}`).join(", ");
+  const order = keys.map((k) => col(k.name)).join(", ");
+  const values = (row: Record<string, unknown> | undefined): string[] | undefined => (row ? keys.map((_, i) => String(row[`k${i}`])) : undefined);
+  const bounds: string[][] = [];
+  let at = values((await client.query(`SELECT ${select} FROM ${table} ORDER BY ${order} LIMIT 1`))[0]);
+  while (at) {
+    signal?.throwIfAborted();
+    bounds.push(at);
+    at = values((await client.query(`SELECT ${select} FROM ${table} WHERE ${tuple(keys)} >= ${placeholders(keys, 1)} ORDER BY ${order} OFFSET ${size} LIMIT 1`, at))[0]);
+  }
+  return bounds;
+}
+
+/** The batching a run recorded, as kept in the receipt store. */
+export interface BatchRecord {
+  v: 1;
+  /** The table's oid and the new column's attribute number, which the record is for. */
+  table: string;
+  column: number;
+  keys: string[];
+  size: number;
+  bounds: string[][];
+}
+
+/**
+ * The boundary batches of a key that is not one integer column: read from
+ * the record in the receipt store, or walked and recorded when there is
+ * none for this table and new column.
+ */
+export async function recordedBatches(
+  client: PostgresClient,
+  input: { table: string; oid: string; attnum: number; keys: readonly KeyColumn[]; size: number; effect: string; receipts: PostgresReceiptStore; log: (line: string) => void; signal?: AbortSignal },
+): Promise<{ batches: Batch[]; record: BatchRecord; walked: boolean }> {
+  const ref = { name: input.effect, effect: input.effect, flavor: "hash" as const, inputs: {} };
+  const keys = input.keys.map((k) => k.name);
+  const raw = await input.receipts.read(ref);
+  let record: BatchRecord | undefined;
+  try {
+    const parsed = raw === undefined ? undefined : (JSON.parse(raw) as BatchRecord);
+    if (parsed?.v === 1 && parsed.table === input.oid && parsed.column === input.attnum && JSON.stringify(parsed.keys) === JSON.stringify(keys)) record = parsed;
+  } catch {
+    record = undefined;
+  }
+  let walked = false;
+  if (!record) {
+    const bounds = await walkBoundaries(client, input.table, input.keys, input.size, input.signal);
+    record = { v: 1, table: input.oid, column: input.attnum, keys, size: input.size, bounds };
+    await input.receipts.write(ref, JSON.stringify(record));
+    walked = true;
+    input.log(`-- batches of ${keyText(input.keys)}: ${bounds.length} boundary(ies) every ${input.size} key(s), recorded`);
+  } else if (record.size !== input.size) {
+    input.log(`-- batches of ${keyText(input.keys)}: kept as recorded every ${record.size} key(s); batchSize ${input.size} applies to a migration started again`);
+  }
+  return { batches: boundaryBatches(input.keys, record.bounds), record, walked };
+}