src/restore.ts

Functions

Imports

It imports src/dump.ts, src/schema.ts, src/types.ts. Imported by src/drill.ts, src/index.ts.

Source

1import { COMPLETE_MARKER, SCHEMA_SIDECAR, type CompleteMarker, type SchemaSidecar } from "./dump.js";
2import { quoteIdent } from "./schema.js";
3import type { D1Like, D1Statement, R2Like } from "./types.js";
4
5// D1 binds at most 100 parameters to one statement.
6const D1_MAX_PARAMS = 100;
7const DEFAULT_BATCH_ROWS = 100;
8// A batch is also cut by the size of the values it binds, so a table of large rows is
9// not sent as one huge request. A single row over this is a batch of its own.
10const DEFAULT_BATCH_BYTES = 1_000_000;
11
12export interface RestoreOptions {
13  // Rows inserted per D1 batch. Each batch is one transaction.
14  batchRows?: number;
15  // Bound text and BLOB size after which a batch is sent, whatever its row count.
16  batchBytes?: number;
17}
18
19export interface RestoreResult {
20  tables: Array<{ name: string; rows: number }>;
21  // External-content FTS5 tables, rebuilt from their content tables.
22  rebuilt: string[];
23  // Contentless FTS5 tables, created empty: the dump held no text for them.
24  empty: string[];
25}
26
27// Characters of a text value, bytes of a BLOB (an array of byte values), 8 for the rest.
28// An estimate for batching, not a measurement.
29function boundSize(v: unknown): number {
30  return typeof v === "string" || Array.isArray(v) ? v.length : 8;
31}
32
33async function readJson<T>(bucket: R2Like, key: string): Promise<T> {
34  const obj = await bucket.get(key);
35  if (!obj) throw new Error(`d1-dump: ${key} is missing from the dump`);
36  return JSON.parse(await obj.text()) as T;
37}
38
39// The columns the restored table has, from the CREATE statement in _schema.json that was
40// just run. An internal FTS5 table is dumped with its rowid too.
41async function schemaColumns(db: D1Like, table: string, internalFts: boolean): Promise<string[]> {
42  const { results } = await db.prepare(`PRAGMA table_info(${quoteIdent(table)})`).all<{ name: string }>();
43  const columns = results.map((c) => c.name);
44  if (columns.length === 0) throw new Error(`d1-dump: ${table} has no columns after its CREATE statement ran`);
45  return internalFts ? ["rowid", ...columns] : columns;
46}
47
48// Every row must carry exactly the schema's columns. Comparing to the first row instead
49// would pass a column missing from all of them, which restores as NULL or as a default.
50function checkColumns(table: string, columns: string[], rows: Array<Record<string, unknown>>): void {
51  const expected = new Set(columns);
52  for (let i = 0; i < rows.length; i++) {
53    const keys = Object.keys(rows[i]);
54    if (keys.length === columns.length && keys.every((k) => expected.has(k))) continue;
55    const lacks = columns.filter((c) => !(c in rows[i]));
56    const extra = keys.filter((k) => !expected.has(k));
57    throw new Error(
58      `d1-dump: ${table} row ${i} does not match the schema's columns (lacks: ${lacks.join(", ") || "none"}; not in the schema: ${extra.join(", ") || "none"})`
59    );
60  }
61}
62
63// Rebuilds one complete dump into `db`, which should be empty. Tables and FTS5 tables
64// first, then the rows (parents before children), then each external FTS5 index is
65// rebuilt, then indexes, triggers and views, so no trigger fires on restored rows.
66// Every table's count is checked against the marker, and a mismatch throws.
67export async function restoreDump(db: D1Like, bucket: R2Like, runPrefix: string, options: RestoreOptions = {}): Promise<RestoreResult> {
68  const prefix = runPrefix.endsWith("/") ? runPrefix : `${runPrefix}/`;
69  const marker = await readJson<CompleteMarker>(bucket, `${prefix}${COMPLETE_MARKER}`).catch((err: Error) => {
70    throw new Error(`d1-dump: ${prefix} has no readable ${COMPLETE_MARKER}, so it is not a complete dump (${err.message})`);
71  });
72  const { schema } = await readJson<SchemaSidecar>(bucket, `${prefix}${SCHEMA_SIDECAR}`);
73  const batchRows = options.batchRows ?? DEFAULT_BATCH_ROWS;
74  const batchBytes = options.batchBytes ?? DEFAULT_BATCH_BYTES;
75
76  for (const e of schema.filter((s) => s.type === "table")) await db.prepare(e.sql).run();
77  const internalFts = new Set(marker.fts.filter((f) => f.mode === "internal").map((f) => f.name));
78
79  const restored: RestoreResult["tables"] = [];
80  for (const table of marker.tables) {
81    const { rows } = await readJson<{ rows: Array<Record<string, unknown>> }>(bucket, table.key);
82    if (rows.length !== table.rows) throw new Error(`d1-dump: ${table.key} holds ${rows.length} rows, the marker says ${table.rows}`);
83    if (rows.length > 0) {
84      const columns = await schemaColumns(db, table.name, internalFts.has(table.name));
85      checkColumns(table.name, columns, rows);
86      if (columns.length > D1_MAX_PARAMS) throw new Error(`d1-dump: ${table.name} has ${columns.length} columns, over D1's ${D1_MAX_PARAMS} bound parameters`);
87      const sql = `INSERT INTO ${quoteIdent(table.name)} (${columns.map(quoteIdent).join(", ")}) VALUES (${columns.map((_, i) => `?${i + 1}`).join(", ")})`;
88      let statements: D1Statement[] = [];
89      let bytes = 0;
90      const flush = async () => {
91        if (statements.length === 0) return;
92        await db.batch([db.prepare("PRAGMA defer_foreign_keys = ON"), ...statements]);
93        statements = [];
94        bytes = 0;
95      };
96      for (const row of rows) {
97        const values = columns.map((c) => row[c]);
98        statements.push(db.prepare(sql).bind(...values));
99        bytes += values.reduce<number>((n, v) => n + boundSize(v), 0);
100        if (statements.length >= batchRows || bytes >= batchBytes) await flush();
101      }
102      await flush();
103    }
104    const { results } = await db.prepare(`SELECT COUNT(*) AS n FROM ${quoteIdent(table.name)}`).all<{ n: number }>();
105    if (results[0]?.n !== table.rows) throw new Error(`d1-dump: ${table.name} restored ${results[0]?.n} rows, the dump holds ${table.rows}`);
106    restored.push({ name: table.name, rows: table.rows });
107  }
108
109  const rebuilt: string[] = [];
110  const empty: string[] = [];
111  for (const f of marker.fts) {
112    if (f.mode === "external") {
113      await db.prepare(`INSERT INTO ${quoteIdent(f.name)} (${quoteIdent(f.name)}) VALUES ('rebuild')`).run();
114      rebuilt.push(f.name);
115    } else if (f.mode === "contentless") {
116      empty.push(f.name);
117    }
118  }
119
120  for (const type of ["index", "trigger", "view"] as const) {
121    for (const e of schema.filter((s) => s.type === type)) await db.prepare(e.sql).run();
122  }
123  return { tables: restored, rebuilt, empty };
124}