diff --git a/packages/db/src/backup-lib.test.ts b/packages/db/src/backup-lib.test.ts index 775226957b9..1e527182e9a 100644 --- a/packages/db/src/backup-lib.test.ts +++ b/packages/db/src/backup-lib.test.ts @@ -140,6 +140,21 @@ describeEmbeddedPostgres("runDatabaseBackup", () => { "created_at" timestamptz NOT NULL DEFAULT now() ); `); + await sourceSql.unsafe(` + CREATE FUNCTION "public"."backup_test_mark_done"() + RETURNS trigger + LANGUAGE plpgsql + AS $$ + BEGIN + NEW."state" := 'done'; + RETURN NEW; + END; + $$; + CREATE TRIGGER "backup_test_mark_done_trigger" + BEFORE UPDATE OF "title" ON "public"."backup_test_records" + FOR EACH ROW + EXECUTE FUNCTION "public"."backup_test_mark_done"(); + `); const payload = "x".repeat(8192); for (let index = 0; index < 160; index += 1) { @@ -213,6 +228,18 @@ describeEmbeddedPostgres("runDatabaseBackup", () => { metadata: { index: 159, even: false }, }, ]); + + await restoreSql.unsafe(` + UPDATE "public"."backup_test_records" + SET "title" = 'triggered' + WHERE "title" = 'row-0' + `); + const triggeredRows = await restoreSql.unsafe<{ state: string }[]>(` + SELECT "state"::text AS "state" + FROM "public"."backup_test_records" + WHERE "title" = 'triggered' + `); + expect(triggeredRows).toEqual([{ state: "done" }]); } finally { await sourceSql.end(); await restoreSql.end(); diff --git a/packages/db/src/backup-lib.ts b/packages/db/src/backup-lib.ts index 9e3760c7588..3b83ef40d96 100644 --- a/packages/db/src/backup-lib.ts +++ b/packages/db/src/backup-lib.ts @@ -883,6 +883,33 @@ export async function runDatabaseBackup(opts: RunDatabaseBackupOptions): Promise emit(""); } + // JavaScript backups are used when a worktree seed filters or transforms + // table data. Preserve user-defined routines so a restored schema retains + // behavior implemented by triggers as well as its tables and indexes. + const routines = await sql<{ definition: string }[]>` + SELECT pg_get_functiondef(p.oid) AS definition + FROM pg_proc p + JOIN pg_namespace n ON n.oid = p.pronamespace + WHERE ${sql.unsafe(nonSystemSchemaPredicate("n.nspname"))} + AND p.prokind IN ('f', 'p') + AND NOT EXISTS ( + SELECT 1 + FROM pg_depend d + WHERE d.classid = 'pg_proc'::regclass + AND d.objid = p.oid + AND d.deptype = 'e' + ) + ORDER BY n.nspname, p.proname, pg_get_function_identity_arguments(p.oid) + `; + if (routines.length > 0) { + emit("-- Functions and procedures"); + for (const routine of routines) { + const definition = routine.definition.trimEnd(); + emitStatement(definition.endsWith(";") ? definition : `${definition};`); + } + emit(""); + } + // Dump data for each table for (const { schema_name, tablename } of tables) { const currentTableKey = tableKey(schema_name, tablename); @@ -938,6 +965,33 @@ export async function runDatabaseBackup(opts: RunDatabaseBackupOptions): Promise emit(""); } + const allTriggers = await sql<{ + schema_name: string; + tablename: string; + definition: string; + }[]>` + SELECT + n.nspname AS schema_name, + c.relname AS tablename, + pg_get_triggerdef(t.oid, true) AS definition + FROM pg_trigger t + JOIN pg_class c ON c.oid = t.tgrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE NOT t.tgisinternal + AND ${sql.unsafe(nonSystemSchemaPredicate("n.nspname"))} + ORDER BY n.nspname, c.relname, t.tgname + `; + const triggers = allTriggers.filter((entry) => ( + includedTableNames.has(tableKey(entry.schema_name, entry.tablename)) + )); + if (triggers.length > 0) { + emit("-- Triggers"); + for (const trigger of triggers) { + emitStatement(`${trigger.definition};`); + } + emit(""); + } + // Sequence values if (sequences.length > 0) { emit("-- Sequence values"); diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index c0337620491..45e50f962a8 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -411,6 +411,25 @@ async function columnExists( return rows[0]?.exists ?? false; } +async function columnHasDataType( + sql: ReturnType, + tableName: string, + columnName: string, + dataType: string, +): Promise { + const rows = await sql<{ dataType: string; udtName: string }[]>` + SELECT data_type AS "dataType", udt_name AS "udtName" + FROM information_schema.columns + WHERE table_schema = 'public' + AND table_name = ${tableName} + AND column_name = ${columnName} + `; + const expected = dataType.toLowerCase(); + return rows.some((row) => ( + row.dataType.toLowerCase() === expected || row.udtName.toLowerCase() === expected + )); +} + async function indexExists( sql: ReturnType, indexName: string, @@ -444,11 +463,79 @@ async function constraintExists( return rows[0]?.exists ?? false; } +async function functionExists( + sql: ReturnType, + functionName: string, +): Promise { + const rows = await sql<{ exists: boolean }[]>` + SELECT EXISTS ( + SELECT 1 + FROM pg_proc p + JOIN pg_namespace n ON n.oid = p.pronamespace + WHERE n.nspname = 'public' + AND p.proname = ${functionName} + ) AS exists + `; + return rows[0]?.exists ?? false; +} + +async function triggerExists( + sql: ReturnType, + triggerName: string, +): Promise { + const rows = await sql<{ exists: boolean }[]>` + SELECT EXISTS ( + SELECT 1 + FROM pg_trigger t + JOIN pg_class c ON c.oid = t.tgrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = 'public' + AND t.tgname = ${triggerName} + AND NOT t.tgisinternal + ) AS exists + `; + return rows[0]?.exists ?? false; +} + +async function heartbeatEventSequencesAreUnique( + sql: ReturnType, +): Promise { + const rows = await sql<{ unique: boolean }[]>` + SELECT NOT EXISTS ( + SELECT 1 + FROM heartbeat_run_events + GROUP BY run_id, seq + HAVING count(*) > 1 + ) AS unique + `; + return rows[0]?.unique ?? false; +} + +async function heartbeatNextEventSequencesAreCurrent( + sql: ReturnType, +): Promise { + const rows = await sql<{ current: boolean }[]>` + SELECT NOT EXISTS ( + SELECT 1 + FROM heartbeat_runs run + WHERE run.next_event_seq IS DISTINCT FROM COALESCE(( + SELECT max(event.seq) + 1 + FROM heartbeat_run_events event + WHERE event.run_id = run.id + ), 1) + ) AS current + `; + return rows[0]?.current ?? false; +} + async function migrationStatementAlreadyApplied( sql: ReturnType, statement: string, ): Promise { - const normalized = statement.replace(/\s+/g, " ").trim(); + const normalized = statement + .replace(/^\s*--.*$/gm, "") + .replace(/\s+/g, " ") + .trim(); const createTableMatch = normalized.match(/^CREATE TABLE(?: IF NOT EXISTS)? "([^"]+)"/i); if (createTableMatch) { @@ -462,6 +549,18 @@ async function migrationStatementAlreadyApplied( return columnExists(sql, addColumnMatch[1], addColumnMatch[2]); } + const alterColumnTypeMatch = normalized.match( + /^ALTER TABLE "([^"]+)" ALTER COLUMN "([^"]+)" SET DATA TYPE ([A-Za-z0-9_]+)/i, + ); + if (alterColumnTypeMatch) { + return columnHasDataType( + sql, + alterColumnTypeMatch[1], + alterColumnTypeMatch[2], + alterColumnTypeMatch[3], + ); + } + const createIndexMatch = normalized.match(/^CREATE (?:UNIQUE )?INDEX(?: IF NOT EXISTS)? "([^"]+)"/i); if (createIndexMatch) { return indexExists(sql, createIndexMatch[1]); @@ -472,6 +571,36 @@ async function migrationStatementAlreadyApplied( return constraintExists(sql, addConstraintMatch[2]); } + const createFunctionMatch = normalized.match( + /^CREATE OR REPLACE FUNCTION "?([A-Za-z_][A-Za-z0-9_]*)"?\s*\(/i, + ); + if (createFunctionMatch) { + return functionExists(sql, createFunctionMatch[1]); + } + + const createTriggerMatch = normalized.match( + /^CREATE TRIGGER "?([A-Za-z_][A-Za-z0-9_]*)"?/i, + ); + if (createTriggerMatch) { + return triggerExists(sql, createTriggerMatch[1]); + } + + // These audited Runner D data repairs are followed by persistent database + // invariants. Verify their postconditions instead of replaying them when a + // restored database is missing only the migration-history row. + if ( + normalized.startsWith("WITH ranked AS (") + && normalized.includes('UPDATE "heartbeat_run_events" AS event') + ) { + return heartbeatEventSequencesAreUnique(sql); + } + if ( + normalized.startsWith('UPDATE "heartbeat_runs" AS run') + && normalized.includes('SET "next_event_seq" = COALESCE') + ) { + return heartbeatNextEventSequencesAreCurrent(sql); + } + // If we cannot reason about a statement safely, require manual migration. return false; }