src/node.ts

Functions

Imports

It imports src/drill.ts, src/types.ts. Imported by src/cli.ts.

Source

1import { existsSync, mkdtempSync, rmSync } from "node:fs";
2import { readFile, readdir } from "node:fs/promises";
3import { tmpdir } from "node:os";
4import { join, relative, sep } from "node:path";
5import { DatabaseSync, type SQLInputValue } from "node:sqlite";
6import { SCRATCH_PREFIX, type Scratch } from "./drill.js";
7import type { D1Like, D1Statement, R2Like } from "./types.js";
8
9// A downloaded dump as an R2Like, read only, for a drill run from Node. `root` holds the
10// keys as paths: <root>/backups/json/<run id>/<table>.json. Not exported from the package
11// entry, which has to stay free of node: imports so a Worker can bundle it.
12export function directoryBucket(root: string): R2Like {
13  const readOnly = async (): Promise<never> => {
14    throw new Error("d1-dump: a directory bucket is read only");
15  };
16  const keysUnder = async (dir: string): Promise<string[]> => {
17    const out: string[] = [];
18    for (const entry of await readdir(dir, { withFileTypes: true }).catch(() => [])) {
19      const path = join(dir, entry.name);
20      if (entry.isDirectory()) out.push(...(await keysUnder(path)));
21      else out.push(relative(root, path).split(sep).join("/"));
22    }
23    return out;
24  };
25  return {
26    get: async (key) => {
27      const path = join(root, ...key.split("/"));
28      if (relative(root, path).startsWith("..")) throw new Error(`d1-dump: ${key} is outside ${root}`);
29      const text = await readFile(path, "utf8").catch((err: NodeJS.ErrnoException) => {
30        if (err.code === "ENOENT") return null;
31        throw err;
32      });
33      return text === null ? null : { text: async () => text };
34    },
35    list: async ({ prefix = "" }) => ({ objects: (await keysUnder(root)).filter((k) => k.startsWith(prefix)).sort().map((key) => ({ key })), truncated: false }),
36    put: readOnly,
37    delete: readOnly,
38    createMultipartUpload: readOnly,
39  };
40}
41
42type Backend = "sqlite" | "d1";
43
44// sqlite unless Cloudflare credentials are set; either can be asked for. Says why, so the
45// drill's output shows which one ran. asking for d1 without credentials is an error.
46export function pickBackend(asked: string | undefined, credentials: boolean): { backend: Backend; why: string } | { error: string } {
47  if (asked !== undefined && asked !== "sqlite" && asked !== "d1") return { error: `--backend is sqlite or d1, not ${JSON.stringify(asked)}` };
48  const backend: Backend = asked ?? (credentials ? "d1" : "sqlite");
49  if (backend === "d1" && !credentials) return { error: "--backend d1 needs CLOUDFLARE_ACCOUNT_ID and CLOUDFLARE_API_TOKEN in the environment" };
50  return { backend, why: asked ? "asked for" : credentials ? "Cloudflare credentials are set" : "no Cloudflare credentials are set" };
51}
52
53// A scratch database in a local SQLite file, for a drill that must hold no Cloudflare
54// credential. node:sqlite is built into Node (unflagged since 22.13; this package
55// already needs 24.14.1) and carries FTS5. The file lives in a directory of its own,
56// named restore-drill-<date>-<pid>, and delete() closes it and removes the directory,
57// then checks it is gone.
58export function sqliteScratch(): (name: string) => Promise<Scratch> {
59  return async (name) => {
60    if (!name.startsWith(SCRATCH_PREFIX)) throw new Error(`d1-dump: refusing to create ${JSON.stringify(name)}: a scratch database starts with ${SCRATCH_PREFIX}`);
61    const dir = mkdtempSync(join(tmpdir(), `${name}-`));
62    const file = join(dir, `${name}.sqlite`);
63    const sqlite = new DatabaseSync(file);
64
65    // D1 binds a BLOB as an array of byte values and refuses undefined; SQLite takes bytes.
66    const sqlValues = (values: unknown[]): SQLInputValue[] =>
67      values.map((v) => (Array.isArray(v) ? Uint8Array.from(v as number[]) : (v as SQLInputValue)));
68    const rowsOf = (sql: string, values: unknown[]) =>
69      sqlite
70        .prepare(sql)
71        .all(...sqlValues(values))
72        .map((r) => ({ ...r }));
73
74    const statement = (sql: string, values: unknown[] = []): D1Statement & { exec(): { results: unknown[] } } => ({
75      bind: (...next) => statement(sql, next),
76      all: (async <T>() => ({ results: rowsOf(sql, values) as T[] })) as D1Statement["all"],
77      run: async () => sqlite.prepare(sql).run(...sqlValues(values)),
78      exec: () => ({ results: rowsOf(sql, values) }),
79    });
80    const db: D1Like = {
81      prepare: (sql) => statement(sql),
82      // One batch is one transaction, as D1's is: all of it lands or none of it does.
83      batch: async <T>(statements: D1Statement[]) => {
84        sqlite.exec("BEGIN");
85        try {
86          const out = statements.map((s) => (s as ReturnType<typeof statement>).exec() as { results: T[] });
87          sqlite.exec("COMMIT");
88          return out;
89        } catch (err) {
90          sqlite.exec("ROLLBACK");
91          throw err;
92        }
93      },
94    };
95    return {
96      db,
97      delete: async () => {
98        sqlite.close();
99        rmSync(dir, { recursive: true, force: true });
100        if (existsSync(dir)) throw new Error(`${dir} is still there after the delete`);
101      },
102    };
103  };
104}