fix(crypto): walk every encrypted column in bounded pages

Key rotation and the rotation status read every column that was not
flagged as a blob with one `findMany` over the whole table.
`Measurement.notesEncrypted` and `MoodEntry.noteEncrypted` are such
columns, so rotation loaded every measurement of the instance: 1.25
million rows on one large account (#1031), and the rotation status page
did the same on every view.

Every column is now walked with the id-cursor pages the blob columns
already used: 25 rows a page for blobs, 5 000 for everything else. The
maintenance script's own per-column walks page the same way. On the seeded
account with 1 460 notes to rotate, peak heap fell from 272 MB to 153 MB,
and it no longer grows with the table.

`scanned` now counts the rows that hold ciphertext for every column, as
it already did for the blob columns, rather than every row of the table
for some columns and only the ciphertext rows for others. The runbook
says so.

docs/ops/encryption-key-rotation.md

@@ -60,8 +60,10 @@ under `v1` (the synthetic id assigned to the existing `ENCRYPTION_KEY`).
 
    Read three things off the output before going further:
 
-   - the per-column line, `scanned` / `rotated` / `errors` / `dropped`. Treat
-     `errors > 0` as a hard failure and re-run after fixing the cause. A
+   - the per-column line, `scanned` / `rotated` / `errors` / `dropped`.
+     `scanned` counts the rows that hold ciphertext in that column, not every
+     row of the table. Treat `errors > 0` as a hard failure and re-run after
+     fixing the cause. A
      `dropped` count is only ever a cache row the run could not read and
      deleted rather than leave unreadable.
    - `Columns walked: N/M registered`. The two numbers must match.

scripts/rotate-encryption-key.ts

@@ -108,6 +108,49 @@ function shouldRotate(value: string | null): boolean {
   return id !== getActiveKeyId();
 }
 
+/** The part of a Prisma delegate the per-column walks below need. */
+interface PagedDelegate<V> {
+  findMany: (args: {
+    select: Record<string, true>;
+    orderBy: { id: "asc" };
+    take: number;
+    cursor?: { id: string };
+    skip?: number;
+  }) => Promise<Array<Record<string, unknown>>>;
+  update: (args: {
+    where: { id: string };
+    data: Record<string, V>;
+  }) => Promise<unknown>;
+}
+
+/** Rows per page when walking a column. */
+const PAGE_SIZE = 5_000;
+
+/**
+ * Every row of one column, `id` and value only, a page at a time in id order.
+ * The walks used to read a column in one `findMany`, which for
+ * `Measurement.notesEncrypted` is every measurement of the instance: 1.25
+ * million rows on one large account (#1031).
+ */
+async function* pagedRows(
+  delegate: PagedDelegate<unknown>,
+  field: string,
+): AsyncGenerator<Record<string, unknown>> {
+  let cursor: string | null = null;
+  for (;;) {
+    const rows = await delegate.findMany({
+      select: { id: true, [field]: true },
+      orderBy: { id: "asc" },
+      take: PAGE_SIZE,
+      ...(cursor ? { cursor: { id: cursor }, skip: 1 } : {}),
+    });
+    if (rows.length === 0) return;
+    yield* rows;
+    if (rows.length < PAGE_SIZE) return;
+    cursor = rows[rows.length - 1]!.id as string;
+  }
+}
+
 /**
  * Rotate one `String` ciphertext column on a Prisma model. `delegate` is the
  * `prisma.<model>` accessor; `field` is the column. Reads only `id` + the
@@ -116,31 +159,28 @@ function shouldRotate(value: string | null): boolean {
 async function rotateStringColumn(
   table: string,
   field: string,
-  delegate: {
-    findMany: (args: {
-      select: Record<string, true>;
-    }) => Promise<Array<Record<string, unknown>>>;
-    update: (args: {
-      where: { id: string };
-      data: Record<string, string>;
-    }) => Promise<unknown>;
-  },
+  delegate: PagedDelegate<string>,
 ): Promise<RotationResult> {
-  const rows = await delegate.findMany({ select: { id: true, [field]: true } });
   const result: RotationResult = {
     table,
     field,
-    scanned: rows.length,
+    scanned: 0,
     rotated: 0,
     errors: 0,
     dropped: 0,
   };
-  for (const row of rows) {
+  for await (const row of pagedRows(
+    delegate as PagedDelegate<unknown>,
+    field,
+  )) {
     const v = row[field] as string | null;
+    // `scanned` counts rows that hold ciphertext, as the in-app rotation does.
+    if (!v) continue;
+    result.scanned++;
     if (!shouldRotate(v)) continue;
     const id = row.id as string;
     try {
-      const re = encrypt(decrypt(v as string));
+      const re = encrypt(decrypt(v));
       await delegate.update({ where: { id }, data: { [field]: re } });
       result.rotated++;
     } catch (err) {
@@ -159,28 +199,23 @@ async function rotateStringColumn(
 async function rotateBytesColumn(
   table: string,
   field: string,
-  delegate: {
-    findMany: (args: {
-      select: Record<string, true>;
-    }) => Promise<Array<Record<string, unknown>>>;
-    update: (args: {
-      where: { id: string };
-      data: Record<string, Uint8Array>;
-    }) => Promise<unknown>;
-  },
+  delegate: PagedDelegate<Uint8Array>,
 ): Promise<RotationResult> {
-  const rows = await delegate.findMany({ select: { id: true, [field]: true } });
   const result: RotationResult = {
     table,
     field,
-    scanned: rows.length,
+    scanned: 0,
     rotated: 0,
     errors: 0,
     dropped: 0,
   };
-  for (const row of rows) {
+  for await (const row of pagedRows(
+    delegate as PagedDelegate<unknown>,
+    field,
+  )) {
     const buf = row[field] as Uint8Array | null;
     if (!buf || buf.byteLength === 0) continue;
+    result.scanned++;
     const asString = Buffer.from(buf).toString("utf8");
     if (!shouldRotate(asString)) continue;
     const id = row.id as string;

src/lib/crypto/__tests__/encryption-corpus.test.ts

@@ -214,4 +214,45 @@ describe("encryption-corpus scan + rotate", () => {
     expect(col.legacy).toBe(1);
     expect(col.byKeyId[LEGACY_BUCKET]).toBe(1);
   });
+
+  it("reads every column in bounded pages and counts only rows that hold ciphertext (#1031)", async () => {
+    // A notes column is read for every row of its table, 1.25 million on a
+    // large account, and was read in one findMany. 12 000 measurements, one
+    // in ten with a note written under the old key.
+    const measurements: Row[] = Array.from({ length: 12_000 }, (_, i) => ({
+      id: `m-${String(i).padStart(5, "0")}`,
+      notesEncrypted:
+        i % 10 === 0
+          ? new Uint8Array(Buffer.from(activeUnder("v1", `note ${i}`), "utf8"))
+          : null,
+    }));
+    const { client } = makeClient({ Measurement: measurements });
+    const takes: Array<number | undefined> = [];
+    const measurement = client.measurement!;
+    const findMany = measurement.findMany;
+    measurement.findMany = (args) => {
+      takes.push(args.take);
+      return findMany(args);
+    };
+
+    const scan = await scanCorpus(client);
+    const col = scan.columns.find(
+      (c) => c.model === "Measurement" && c.field === "notesEncrypted",
+    )!;
+    expect(col.total).toBe(1_200);
+    expect(col.byKeyId.v1).toBe(1_200);
+
+    const rotation = await rotateCorpus(client);
+    const result = rotation.results.find(
+      (r) => r.model === "Measurement" && r.field === "notesEncrypted",
+    )!;
+    expect(result.scanned).toBe(1_200);
+    expect(result.rotated).toBe(1_200);
+
+    expect(takes.length).toBeGreaterThan(2);
+    for (const take of takes) {
+      expect(take).toBeDefined();
+      expect(take!).toBeLessThanOrEqual(5_000);
+    }
+  });
 });

src/lib/crypto/encryption-corpus.ts

@@ -22,7 +22,7 @@
  *    `ENCRYPTION_KEYS` / `ENCRYPTION_ACTIVE_KEY_ID`. The operator's env key map
  *    is the only place keys live; a key drop stays a deliberate env + redeploy
  *    act, never a button.
- *  - IDEMPOTENT. `shouldRotate()` skips rows already on the active key, so a
+ *  - IDEMPOTENT. The walk skips rows already on the active key, so a
  *    second pass (or two racing workers) re-encrypts zero rows.
  *  - FAIL-CLOSED. A row written under a key id that is no longer configured
  *    throws on decrypt (counted as an error, the row is left untouched) rather
@@ -54,6 +54,15 @@ export const LEGACY_BUCKET = "legacy";
  */
 export const BLOB_ROTATION_BATCH_SIZE = 25;
 
+/**
+ * Batch size for every other column. These rows are short strings, but a
+ * column can have one per measurement: `Measurement.notesEncrypted` is read
+ * for every row of the table, 1.25 million on a large account (#1031). The
+ * walk used to read such a column in one `findMany`; it is id-cursor paged
+ * like the blob columns, only with larger pages.
+ */
+export const ROW_ROTATION_BATCH_SIZE = 5_000;
+
 /** Minimal Prisma delegate shape this module needs. */
 interface ColumnDelegate {
   findMany: (args: {
@@ -114,10 +123,6 @@ function fromCiphertext(
   return next;
 }
 
-function shouldRotate(ciphertext: string): boolean {
-  return extractKeyId(ciphertext) !== getActiveKeyId();
-}
-
 // ─── Codec-dispatched blob columns (document vault) ─────────────────────────
 
 /**
@@ -152,15 +157,16 @@ function reencryptBlob(value: Uint8Array, codec: string): Uint8Array {
 }
 
 /**
- * Walk a blob column in bounded id-cursor batches, invoking `onRow` per
- * non-empty row. At most `BLOB_ROTATION_BATCH_SIZE` rows are in memory per
- * step; an interrupted run resumes safely on re-invocation because processing
- * is idempotent (already-active rows are skipped by the callers).
+ * Walk a column in bounded id-cursor batches, invoking `onRow` per non-empty
+ * row. At most one batch is in memory per step (`BLOB_ROTATION_BATCH_SIZE`
+ * rows of a blob column, `ROW_ROTATION_BATCH_SIZE` of any other); an
+ * interrupted run resumes safely on re-invocation because processing is
+ * idempotent (already-active rows are skipped by the callers).
  *
  * `codec` is the row's own codec label for a codec-dispatched column and null
  * for a plain `batched` one — a backup blob has one layout, just a large one.
  */
-async function walkBlobColumn(
+async function walkColumn(
   delegate: ColumnDelegate,
   col: EncryptedColumn,
   onRow: (row: {
@@ -171,6 +177,9 @@ async function walkBlobColumn(
 ): Promise<void> {
   const codecField = col.codecField;
   const pk = pkField(col);
+  const batchSize = isBlobColumn(col)
+    ? BLOB_ROTATION_BATCH_SIZE
+    : ROW_ROTATION_BATCH_SIZE;
   let cursor: string | null = null;
   for (;;) {
     const rows = await delegate.findMany({
@@ -180,7 +189,7 @@ async function walkBlobColumn(
         ...(codecField ? { [codecField]: true } : {}),
       },
       orderBy: { [pk]: "asc" },
-      take: BLOB_ROTATION_BATCH_SIZE,
+      take: batchSize,
       ...(cursor ? { cursor: { [pk]: cursor }, skip: 1 } : {}),
     });
     if (rows.length === 0) break;
@@ -196,7 +205,7 @@ async function walkBlobColumn(
       });
     }
     cursor = rows[rows.length - 1]![pk] as string;
-    if (rows.length < BLOB_ROTATION_BATCH_SIZE) break;
+    if (rows.length < batchSize) break;
   }
 }
 
@@ -205,8 +214,8 @@ function pkField(col: EncryptedColumn): string {
   return col.pkField ?? "id";
 }
 
-/** True when the column's rows must be walked in bounded batches. */
-function isBatched(col: EncryptedColumn): boolean {
+/** True when the column's rows are blobs, walked in the small batches. */
+function isBlobColumn(col: EncryptedColumn): boolean {
   return Boolean(col.codecField ?? col.batched);
 }
 
@@ -242,25 +251,12 @@ export async function scanColumn(
   const byKeyId: Record<string, number> = {};
   let total = 0;
 
-  if (isBatched(col)) {
-    // Blob column: bounded batches, per-row codec where the column has one.
-    await walkBlobColumn(delegate, col, ({ value, codec }) => {
-      total += 1;
-      const id = walkedKeyId(value, codec, col.kind) ?? LEGACY_BUCKET;
-      byKeyId[id] = (byKeyId[id] ?? 0) + 1;
-    });
-  } else {
-    const rows = await delegate.findMany({
-      select: { [pkField(col)]: true, [col.field]: true },
-    });
-    for (const row of rows) {
-      const ciphertext = toCiphertext(row[col.field], col.kind);
-      if (ciphertext == null) continue;
-      total += 1;
-      const id = extractKeyId(ciphertext) ?? LEGACY_BUCKET;
-      byKeyId[id] = (byKeyId[id] ?? 0) + 1;
-    }
-  }
+  // Bounded batches for every column, per-row codec where the column has one.
+  await walkColumn(delegate, col, ({ value, codec }) => {
+    total += 1;
+    const id = walkedKeyId(value, codec, col.kind) ?? LEGACY_BUCKET;
+    byKeyId[id] = (byKeyId[id] ?? 0) + 1;
+  });
   return {
     model: col.model,
     field: col.field,
@@ -356,57 +352,33 @@ export async function rotateColumn(
     }
   };
 
-  if (isBatched(col)) {
-    // Blob column: bounded id-cursor batches (never an unbounded findMany),
-    // re-encrypted under each row's OWN codec where it has one. Idempotent —
-    // rows already on the active key are skipped, so an interrupted run
-    // resumes cleanly on the next invocation.
-    await walkBlobColumn(delegate, col, async ({ id, value, codec }) => {
-      result.scanned += 1;
-      if (walkedKeyId(value, codec, col.kind) === getActiveKeyId()) return;
-      try {
-        // The plaintext is never inspected, only re-sealed, so every envelope
-        // a stored value can legitimately carry survives rotation untouched.
-        const next =
-          codec !== null
-            ? reencryptBlob(value as Uint8Array, codec)
-            : fromCiphertext(
-                encrypt(decrypt(toCiphertext(value, col.kind)!)),
-                col.kind,
-              );
-        await delegate.update({
-          where: { [pkField(col)]: id },
-          data: { [col.field]: next },
-        });
-        result.rotated += 1;
-      } catch {
-        await onUnreadable(id);
-      }
-    });
-    return result;
-  }
-
-  const pk = pkField(col);
-  const rows = await delegate.findMany({
-    select: { [pk]: true, [col.field]: true },
-  });
-  result.scanned = rows.length;
-  for (const row of rows) {
-    const ciphertext = toCiphertext(row[col.field], col.kind);
-    if (ciphertext == null || !shouldRotate(ciphertext)) continue;
-    const id = row[pk] as string;
+  // Bounded id-cursor batches for every column (never an unbounded
+  // findMany), re-encrypted under each row's OWN codec where it has one.
+  // Idempotent — rows already on the active key are skipped, so an
+  // interrupted run resumes cleanly on the next invocation. `scanned` counts
+  // the rows that hold ciphertext, whatever the column.
+  await walkColumn(delegate, col, async ({ id, value, codec }) => {
+    result.scanned += 1;
+    if (walkedKeyId(value, codec, col.kind) === getActiveKeyId()) return;
     try {
-      // ACTIVE-KEY-ONLY: encrypt() always writes the active key id.
-      const reencrypted = encrypt(decrypt(ciphertext));
+      // The plaintext is never inspected, only re-sealed, so every envelope
+      // a stored value can legitimately carry survives rotation untouched.
+      const next =
+        codec !== null
+          ? reencryptBlob(value as Uint8Array, codec)
+          : fromCiphertext(
+              encrypt(decrypt(toCiphertext(value, col.kind)!)),
+              col.kind,
+            );
       await delegate.update({
-        where: { [pk]: id },
-        data: { [col.field]: fromCiphertext(reencrypted, col.kind) },
+        where: { [pkField(col)]: id },
+        data: { [col.field]: next },
       });
       result.rotated += 1;
     } catch {
       await onUnreadable(id);
     }
-  }
+  });
   return result;
 }