import type { Config } from "@libsql/client"; import { parseArgs } from "@std/cli/parse-args"; import { load } from "@std/dotenv"; import { basename, dirname, join, resolve, toFileUrl } from "@std/path"; import * as v from "valibot"; import { withCleanup } from "#scripts/cleanup.ts"; import { lstatOrNull, statNumberOrNull, statOrNull, } from "#scripts/not-found.ts"; import { secureUrlCheck } from "#scripts/secure-url.ts"; import { getEnv } from "#shared/env.ts"; import { parseOrThrow } from "#shared/validation/parse.ts"; export const SNAPSHOT_USAGE = "Usage: deno task snapshot --out "; export interface SnapshotOptions { outputPath: string; } export interface SnapshotRequest extends SnapshotOptions { dbToken: string; dbUrl: string; } export interface SnapshotQueryResult { rows: readonly Readonly>[]; } export interface SnapshotQueryCheck { sql: string; verify: (result: SnapshotQueryResult) => void; } export interface SnapshotClient { close(): void; execute(sql: string): Promise; sync(): Promise; } export type SnapshotClientFactory = (config: Config) => SnapshotClient; export type SnapshotEnvReader = (key: string) => string | undefined; export type SnapshotEnvFileLoader = ( path: string, ) => Promise>; export const SNAPSHOT_PROGRESS = { checking: "[1/4] Checking destination", publishing: "[4/4] Publishing standalone SQLite file", syncing: "[2/4] Syncing remote database", verifying: "[3/4] Checkpointing and checking integrity", } as const; export type SnapshotProgress = (typeof SNAPSHOT_PROGRESS)[keyof typeof SNAPSHOT_PROGRESS]; export type SnapshotProgressWriter = (message: SnapshotProgress) => void; const ignoreSnapshotProgress: SnapshotProgressWriter = () => {}; const OutputValuesSchema = v.strictTuple([ v.pipe( v.string(), v.check((path) => path.trim().length > 0), ), ]); const EmptyStringsSchema = v.strictTuple([]); const snapshotArgsSchema = < Help extends boolean, OutSchema extends v.GenericSchema, >( help: Help, out: OutSchema, ) => v.strictObject({ _: EmptyStringsSchema, h: v.literal(help), help: v.literal(help), out, }); const SnapshotArgsSchema = snapshotArgsSchema(false, OutputValuesSchema); const SnapshotHelpArgsSchema = snapshotArgsSchema(true, EmptyStringsSchema); const CheckpointRowsSchema = v.pipe( v.array( v.strictObject({ busy: v.literal(0), checkpointed: v.literal(0), log: v.literal(0), }), ), v.length(1), ); const IntegrityRowsSchema = v.pipe( v.array(v.strictObject({ integrity_check: v.literal("ok") })), v.length(1), ); const SECURE_DATABASE_PROTOCOLS = new Set(["https:", "libsql:"]); const invalidUsage = (): never => { throw new Error(SNAPSHOT_USAGE); }; export const parseSnapshotArgs = (args: string[]): SnapshotOptions | null => { const parsed = parseArgs(args, { alias: { h: "help" }, boolean: ["help"], collect: ["out"], string: ["out"], }); if (parsed.help) { if (!v.safeParse(SnapshotHelpArgsSchema, parsed).success) invalidUsage(); return null; } const result = v.safeParse(SnapshotArgsSchema, parsed); if (!result.success) return invalidUsage(); return { outputPath: result.output.out[0] }; }; const requiredEnv = (key: string, readEnv: SnapshotEnvReader): string => { const value = readEnv(key); if (!value || value.trim().length === 0) { throw new Error(`${key} environment variable is required`); } return value; }; const requireRemoteDatabaseUrl = (value: string): string => secureUrlCheck(SECURE_DATABASE_PROTOCOLS)(value, "DB_URL"); export const readSnapshotRequest = ( options: SnapshotOptions, readEnv: SnapshotEnvReader = getEnv, ): SnapshotRequest => ({ dbToken: requiredEnv("DB_TOKEN", readEnv), dbUrl: requireRemoteDatabaseUrl(requiredEnv("DB_URL", readEnv)), outputPath: options.outputPath, }); const loadSnapshotEnvFile: SnapshotEnvFileLoader = (path) => load({ envPath: path }); export const readSnapshotRequestFromEnvFile = async ( options: SnapshotOptions, path = ".env", readEnv: SnapshotEnvReader = getEnv, loadFile: SnapshotEnvFileLoader = loadSnapshotEnvFile, ): Promise => { const fileEnv = await loadFile(path); return readSnapshotRequest(options, (key) => fileEnv[key] ?? readEnv(key)); }; const outputAlreadyExists = (path: string): Error => new Error(`Output already exists: ${path}`); const requireOutputDirectory = async (path: string): Promise => { const info = await statOrNull(path); if (info === null) throw new Error(`Output directory does not exist: ${path}`); if (!info.isDirectory) throw new Error(`Output directory is not a directory: ${path}`); }; const outputFiles = (path: string): string[] => [ path, `${path}-wal`, `${path}-shm`, ]; const requireMissingOutputFiles = async (path: string): Promise => { const existingPath = ( await Promise.all( outputFiles(path).map(async (filePath) => (await lstatOrNull(filePath)) === null ? null : filePath, ), ) ).find((filePath) => filePath !== null); if (existingPath !== undefined) throw outputAlreadyExists(existingPath); }; const withSnapshotClient = ( factory: SnapshotClientFactory, config: Config, run: (client: SnapshotClient) => Promise, ): Promise => { const client = factory(config); return withCleanup(() => run(client), [() => client.close()]); }; const waitForOrAbort = async ( operation: () => Promise, signal?: AbortSignal, ): Promise => { if (signal === undefined) return await operation(); signal.throwIfAborted(); const interrupted = Promise.withResolvers(); const stop = (): void => interrupted.reject(signal.reason); signal.addEventListener("abort", stop, { once: true }); try { return await Promise.race([operation(), interrupted.promise]); } finally { signal.removeEventListener("abort", stop); } }; export const checkLocalSnapshot = ( path: string, factory: SnapshotClientFactory, checks: SnapshotQueryCheck[], signal?: AbortSignal, ): Promise => withSnapshotClient(factory, { url: toFileUrl(path).href }, async (client) => { for (const check of checks) { signal?.throwIfAborted(); check.verify( await waitForOrAbort(() => client.execute(check.sql), signal), ); } }); const syncReplica = ( path: string, request: SnapshotRequest, factory: SnapshotClientFactory, signal?: AbortSignal, ): Promise => withSnapshotClient( factory, { authToken: request.dbToken, syncUrl: request.dbUrl, url: toFileUrl(path).href, }, (client) => waitForOrAbort(() => client.sync(), signal), ); const verifyRows = ( schema: v.GenericSchema, result: SnapshotQueryResult, message: string, ): void => { parseOrThrow(schema, result.rows, () => new Error(message)); }; const snapshotQueryCheck = ( schema: v.GenericSchema, sql: string, message: string, ): SnapshotQueryCheck => ({ sql, verify: (result) => verifyRows(schema, result, message), }); const SNAPSHOT_QUERY_CHECKS = [ snapshotQueryCheck( CheckpointRowsSchema, "PRAGMA wal_checkpoint(TRUNCATE)", "Database snapshot checkpoint did not empty the WAL", ), snapshotQueryCheck( IntegrityRowsSchema, "PRAGMA integrity_check", "Database snapshot integrity check failed", ), ]; // SQLite can either remove an empty WAL or leave a zero-byte file. const fileSizeOrNull = statNumberOrNull((info) => info.size); const requireEmptyWal = async (path: string): Promise => { const size = await fileSizeOrNull(`${path}-wal`); if (size !== null && size !== 0) { throw new Error("Database snapshot WAL is not empty"); } }; const publishSnapshot = async ( temporaryPath: string, outputPath: string, ): Promise => { await Deno.chmod(temporaryPath, 0o600); await requireMissingOutputFiles(outputPath); try { await Deno.link(temporaryPath, outputPath); } catch (error) { if (error instanceof Deno.errors.AlreadyExists) { throw outputAlreadyExists(outputPath); } throw error; } }; export const createDatabaseSnapshot = async ( request: SnapshotRequest, factory: SnapshotClientFactory, writeProgress: SnapshotProgressWriter = ignoreSnapshotProgress, signal?: AbortSignal, ): Promise => { const outputPath = resolve(request.outputPath); const outputDirectory = dirname(outputPath); writeProgress(SNAPSHOT_PROGRESS.checking); await requireOutputDirectory(outputDirectory); await requireMissingOutputFiles(outputPath); const temporaryDirectory = await Deno.makeTempDir({ dir: outputDirectory, prefix: `.${basename(outputPath)}-snapshot-`, }); return await withCleanup(async () => { const temporaryPath = join(temporaryDirectory, "snapshot.sqlite"); writeProgress(SNAPSHOT_PROGRESS.syncing); await syncReplica(temporaryPath, request, factory, signal); writeProgress(SNAPSHOT_PROGRESS.verifying); await checkLocalSnapshot( temporaryPath, factory, SNAPSHOT_QUERY_CHECKS, signal, ); await requireEmptyWal(temporaryPath); writeProgress(SNAPSHOT_PROGRESS.publishing); await publishSnapshot(temporaryPath, outputPath); return outputPath; }, [() => Deno.remove(temporaryDirectory, { recursive: true })]); };