src/drill.ts

Functions

Imports

It imports src/dump.ts, src/restore.ts, src/runs.ts, src/schema.ts, src/types.ts. Imported by src/cli.ts, src/cloudflare.ts, src/index.ts, src/node.ts.

Source

1import { COMPLETE_MARKER, SCHEMA_SIDECAR, type CompleteMarker, type SchemaSidecar } from "./dump.js";
2import { restoreDump, type RestoreOptions, type RestoreResult } from "./restore.js";
3import { DEFAULT_STALE_HOURS, latestDump } from "./runs.js";
4import { quoteIdent } from "./schema.js";
5import type { D1Like, D1Statement, R2Like } from "./types.js";
6
7// D1 refuses one SQL statement over 100,000 bytes (developers.cloudflare.com/d1/platform/limits).
8export const D1_MAX_STATEMENT_BYTES = 100_000;
9export const SCRATCH_PREFIX = "restore-drill-";
10
11export interface StatementStats {
12  statements: number;
13  batches: number;
14  // The longest SQL text prepared, in UTF-8 bytes. Bound values are not part of it.
15  maxStatementBytes: number;
16  // The largest single bound text or BLOB, in characters or bytes.
17  maxBoundValue: number;
18  // The most text and BLOB the statements of one batch bound together.
19  maxBatchBound: number;
20}
21
22function boundOf(values: unknown[]): number {
23  return values.reduce<number>((n, v) => n + (typeof v === "string" || Array.isArray(v) ? v.length : 0), 0);
24}
25
26// Wraps a database so every statement is measured before it is sent. A statement over
27// D1's cap throws here and is never sent, so a drill that finishes is a drill in which
28// none was.
29export function guardStatements(db: D1Like): { db: D1Like; stats: StatementStats } {
30  const stats: StatementStats = { statements: 0, batches: 0, maxStatementBytes: 0, maxBoundValue: 0, maxBatchBound: 0 };
31  // Each wrapper maps back to the statement it stands for, and to the size it bound.
32  const real = new WeakMap<D1Statement, { inner: D1Statement; bound: number }>();
33  const wrap = (inner: D1Statement, bound: number): D1Statement => {
34    const stmt: D1Statement = {
35      bind: (...values) => {
36        for (const v of values) if (typeof v === "string" || Array.isArray(v)) stats.maxBoundValue = Math.max(stats.maxBoundValue, v.length);
37        return wrap(inner.bind(...values), boundOf(values));
38      },
39      all: ((): Promise<{ results: never[] }> => {
40        stats.statements++;
41        return inner.all() as Promise<{ results: never[] }>;
42      }) as D1Statement["all"],
43      run: () => {
44        stats.statements++;
45        return inner.run();
46      },
47    };
48    real.set(stmt, { inner, bound });
49    return stmt;
50  };
51  const guarded: D1Like = {
52    prepare: (sql) => {
53      const bytes = new TextEncoder().encode(sql).length;
54      if (bytes > D1_MAX_STATEMENT_BYTES) throw new Error(`d1-dump: a statement of ${bytes} bytes is over D1's ${D1_MAX_STATEMENT_BYTES} byte cap and was not sent`);
55      stats.maxStatementBytes = Math.max(stats.maxStatementBytes, bytes);
56      return wrap(db.prepare(sql), 0);
57    },
58    batch: (statements) => {
59      const parts = statements.map((s) => real.get(s) ?? { inner: s, bound: 0 });
60      stats.batches++;
61      stats.statements += statements.length;
62      stats.maxBatchBound = Math.max(stats.maxBatchBound, parts.reduce((n, p) => n + p.bound, 0));
63      return db.batch(parts.map((p) => p.inner));
64    },
65  };
66  return { db: guarded, stats };
67}
68
69export interface Scratch {
70  db: D1Like;
71  delete(): Promise<void>;
72}
73
74export interface DrillOptions {
75  // Creates an empty scratch database with this name. The name always starts with
76  // restore-drill-, so a leaked one is recognisable and nothing else is ever written.
77  createScratch: (name: string) => Promise<Scratch>;
78  // A pinned run prefix, e.g. backups/json/<run id>/. Defaults to the newest complete run.
79  prefix?: string;
80  now?: Date;
81  pid?: number;
82  // A dump older than this many hours fails the drill. null turns the check off.
83  maxAgeHours?: number | null;
84  restore?: RestoreOptions;
85}
86
87export interface DrillCheck {
88  name: string;
89  ok: boolean;
90  detail: string;
91}
92
93export interface DrillResult {
94  ok: boolean;
95  scratch: string;
96  run: string | null;
97  checks: DrillCheck[];
98  restored: RestoreResult | null;
99  stats: StatementStats | null;
100}
101
102export function scratchName(now: Date, pid: number): string {
103  return `${SCRATCH_PREFIX}${now.toISOString().slice(0, 10)}-${pid}`;
104}
105
106async function readJson<T>(bucket: R2Like, key: string): Promise<T> {
107  const obj = await bucket.get(key);
108  if (!obj) throw new Error(`${key} is missing`);
109  return JSON.parse(await obj.text()) as T;
110}
111
112async function count(db: D1Like, table: string): Promise<number> {
113  const { results } = await db.prepare(`SELECT COUNT(*) AS n FROM ${quoteIdent(table)}`).all<{ n: number }>();
114  return results[0]?.n ?? Number.NaN;
115}
116
117// A word of four or more letters from the first text a table holds, to search for.
118function wordIn(rows: Array<Record<string, unknown>>, column: string): string | null {
119  for (const row of rows) {
120    const v = row[column];
121    const m = typeof v === "string" ? /[A-Za-z]{4,}/.exec(v) : null;
122    if (m) return m[0];
123  }
124  return null;
125}
126
127// Restores the newest complete dump (or a pinned run) into a scratch database it
128// creates, checks the copy, and deletes the scratch database whatever happened. It
129// reports failed checks rather than throwing, so a caller sees every one.
130//
131// Checks, each carrying what it read so "ok" cannot mean "read nothing":
132//   dump      the run is complete and no older than maxAgeHours
133//   restore   restoreDump finished, which itself checks each table against the marker
134//   counts    every table's rows, counted again in the copy, equal the marker's
135//   fts       each external FTS5 index holds as many documents as its content table
136//             (counted through _docsize: COUNT(*) on the FTS table reads through to the
137//             content table and cannot disagree), and each internal one its dumped rows
138//   search    a word taken from a restored row is found through the rebuilt index
139//   statements  no statement sent was over D1's cap
140//   scratch   the scratch database is deleted
141export async function runRestoreDrill(bucket: R2Like, options: DrillOptions): Promise<DrillResult> {
142  const now = options.now ?? new Date();
143  const name = scratchName(now, options.pid ?? (globalThis as { process?: { pid?: number } }).process?.pid ?? 0);
144  const checks: DrillCheck[] = [];
145  const check = (n: string, ok: boolean, detail: string) => checks.push({ name: n, ok, detail });
146  let run: string | null = null;
147  let restored: RestoreResult | null = null;
148  let stats: StatementStats | null = null;
149
150  const latest = options.prefix ? null : await latestDump(bucket);
151  const prefix = options.prefix ?? latest?.prefix ?? null;
152  let marker: CompleteMarker | null = null;
153  if (!prefix) {
154    check("dump", false, "no complete dump found");
155  } else {
156    run = prefix;
157    try {
158      marker = await readJson<CompleteMarker>(bucket, `${prefix}${COMPLETE_MARKER}`);
159      await readJson<SchemaSidecar>(bucket, `${prefix}${SCHEMA_SIDECAR}`);
160      const maxAge = options.maxAgeHours === undefined ? DEFAULT_STALE_HOURS : options.maxAgeHours;
161      const ageHours = (now.getTime() - Date.parse(marker.exported_at)) / 3_600_000;
162      if (maxAge !== null && !(ageHours <= maxAge)) {
163        check("dump", false, `exported ${marker.exported_at}, ${ageHours.toFixed(1)}h before the drill, over ${maxAge}h`);
164      } else {
165        check("dump", true, `${prefix} exported ${marker.exported_at}, ${marker.tables.length} tables`);
166      }
167    } catch (err) {
168      marker = null;
169      check("dump", false, (err as Error).message);
170    }
171  }
172
173  if (prefix && marker) {
174    const scratch = await options.createScratch(name);
175    try {
176      const guarded = guardStatements(scratch.db);
177      stats = guarded.stats;
178      try {
179        restored = await restoreDump(guarded.db, bucket, prefix, options.restore);
180        check("restore", true, `${restored.tables.length} tables, ${restored.rebuilt.length} FTS5 indexes rebuilt, ${restored.empty.length} contentless left empty`);
181      } catch (err) {
182        check("restore", false, (err as Error).message);
183      }
184      if (restored) await verify(guarded.db, bucket, marker, check);
185      check(
186        "statements",
187        stats.statements > 0 && stats.maxStatementBytes <= D1_MAX_STATEMENT_BYTES,
188        `${stats.statements} statements in ${stats.batches} batches, longest ${stats.maxStatementBytes} bytes of ${D1_MAX_STATEMENT_BYTES}; largest bound value ${stats.maxBoundValue}, most bound by one batch ${stats.maxBatchBound}`
189      );
190    } finally {
191      try {
192        await scratch.delete();
193        check("scratch", true, `${name} deleted`);
194      } catch (err) {
195        check("scratch", false, `${name} was NOT deleted, delete it by hand: ${(err as Error).message}`);
196      }
197    }
198  }
199  return { ok: checks.every((c) => c.ok), scratch: name, run, checks, restored, stats };
200}
201
202async function verify(db: D1Like, bucket: R2Like, marker: CompleteMarker, check: (n: string, ok: boolean, detail: string) => void): Promise<void> {
203  let total = 0;
204  const wrong: string[] = [];
205  for (const t of marker.tables) {
206    const n = await count(db, t.name);
207    total += n;
208    if (n !== t.rows) wrong.push(`${t.name} has ${n}, the marker says ${t.rows}`);
209  }
210  check("counts", wrong.length === 0 && total > 0, wrong.length > 0 ? wrong.join("; ") : total === 0 ? "the copy holds no rows at all" : `${marker.tables.length} tables, ${total} rows, all equal to the marker`);
211
212  const behind: string[] = [];
213  let indexed = 0;
214  for (const f of marker.fts) {
215    if (f.mode === "external" && f.content) {
216      const docs = await count(db, `${f.name}_docsize`);
217      const rows = await count(db, f.content);
218      indexed += docs;
219      if (docs !== rows) behind.push(`${f.name} indexes ${docs} documents, ${f.content} holds ${rows}`);
220    } else if (f.mode === "internal") {
221      const dumped = marker.tables.find((t) => t.name === f.name)?.rows;
222      const docs = await count(db, `${f.name}_docsize`);
223      indexed += docs;
224      if (docs !== dumped) behind.push(`${f.name} indexes ${docs} documents, the dump held ${dumped}`);
225    }
226  }
227  check("fts", behind.length === 0, behind.length > 0 ? behind.join("; ") : `${marker.fts.length} FTS5 tables, ${indexed} indexed documents`);
228
229  // One search: the first external or internal index that holds a document.
230  for (const f of marker.fts) {
231    if (f.mode === "contentless") continue;
232    const source = f.mode === "external" ? f.content! : f.name;
233    const file = marker.tables.find((t) => t.name === source);
234    if (!file || file.rows === 0) continue;
235    const { results: cols } = await db.prepare(`PRAGMA table_info(${quoteIdent(f.name)})`).all<{ name: string }>();
236    const { rows } = await readJson<{ rows: Array<Record<string, unknown>> }>(bucket, file.key);
237    for (const col of cols) {
238      const word = wordIn(rows, col.name);
239      if (!word) continue;
240      const { results } = await db.prepare(`SELECT rowid FROM ${quoteIdent(f.name)} WHERE ${quoteIdent(f.name)} MATCH ?1`).bind(`"${word}"`).all();
241      check("search", results.length > 0, results.length > 0 ? `"${word}" found ${results.length} document(s) in ${f.name}.${col.name}` : `"${word}" is in a restored ${source} row but ${f.name} does not find it`);
242      return;
243    }
244  }
245  check("search", false, "no FTS5 table held a word to search for, so no search ran");
246}