Fix Apex run provenance and migration lock findings

Preserve existing PR history and verified downstream privacy rules.

Co-Authored-By: Paperclip <noreply@paperclip.ing>

doc/DATABASE.md

@@ -226,6 +226,14 @@ When authoring migrations or one-time backfills:
 - Split schema changes, index creation, and data backfill into separate phases so each step has clear locking and rollback behavior.
 - Treat the `check:migrations` CI gate as the enforcement backstop for these rules. If it flags a migration, rewrite the migration or add a suppression comment with the indexed predicate, batch bound, and reason the remaining scan is safe.
 
+Private-task migrations `0311` and `0312` are explicitly allowlisted in the
+Paperclip executor to run outside a file-wide transaction. Their idempotent
+keyset batches commit every 1,000 rows, and migration history is recorded only
+when all batches finish. The executor repairs invalid concurrent indexes on
+retry. Bootstrap uses the same executor; other migrations remain transactional
+per file. Apply these migrations through `pnpm db:migrate` using a direct
+connection, before enabling the new server and UI.
+
 ## Migration snapshots
 
 `drizzle-kit generate` diffs `packages/db/src/schema/` against the newest snapshot in `packages/db/src/migrations/meta/`. That snapshot must describe the schema that every migration produces when they run in order. A snapshot that drifts from the schema makes the *next* migration wrong, because `generate` folds the drift into it. The drift can add a column that an earlier migration already created, which makes that migration fail on a fresh database. It can also drop a column that the schema still uses.

packages/db/src/client.ts

@@ -1,6 +1,5 @@
 import { createHash } from "node:crypto";
 import { drizzle as drizzlePg } from "drizzle-orm/postgres-js";
-import { migrate as migratePg } from "drizzle-orm/postgres-js/migrator";
 import { readFile, readdir } from "node:fs/promises";
 import { fileURLToPath } from "node:url";
 import postgres from "postgres";
@@ -463,6 +462,14 @@ async function recordMigrationHistoryEntry(
   );
 }
 
+// These idempotent privacy migrations commit keyset batches inside their DO
+// blocks. Execute each statement at top level so DDL locks and completed batches
+// are released before the next batch. History is recorded only after completion.
+const MIGRATIONS_WITH_BATCH_COMMITS = new Set([
+  "0311_private_task_access.sql",
+  "0312_private_task_draft_assets.sql",
+]);
+
 async function applyPendingMigrationsManually(
   url: string,
   pendingMigrations: string[],
@@ -492,8 +499,24 @@ async function applyPendingMigrationsManually(
       );
       if (existingEntry) continue;
 
-      await runInTransaction(sql, async () => {
+      const applyMigration = async () => {
         for (const statement of splitMigrationStatements(migrationContent)) {
+          if (MIGRATIONS_WITH_BATCH_COMMITS.has(migrationFile)) {
+            // A cancelled concurrent build leaves an invalid index. IF NOT
+            // EXISTS alone would skip it on retry and journal an incomplete index.
+            const index = statement.replace(/^\s*--.*$/gm, "").trim()
+              .match(/^CREATE (?:UNIQUE )?INDEX CONCURRENTLY IF NOT EXISTS "([A-Za-z_][A-Za-z0-9_]*)"/i);
+            if (index) {
+              const invalid = await sql<{ invalid: boolean }[]>`
+                SELECT NOT i.indisvalid AS invalid FROM pg_index i
+                JOIN pg_class c ON c.oid = i.indexrelid
+                JOIN pg_namespace n ON n.oid = c.relnamespace
+                WHERE n.nspname = 'public' AND c.relname = ${index[1]}`;
+              if (invalid[0]?.invalid) {
+                await sql.unsafe(`DROP INDEX CONCURRENTLY ${quoteIdentifier("public")}.${quoteIdentifier(index[1])}`);
+              }
+            }
+          }
           await sql.unsafe(statement);
         }
 
@@ -505,7 +528,12 @@ async function applyPendingMigrationsManually(
           hash,
           folderMillisByFileName.get(migrationFile) ?? Date.now(),
         );
-      });
+      };
+      if (MIGRATIONS_WITH_BATCH_COMMITS.has(migrationFile)) {
+        await applyMigration();
+      } else {
+        await runInTransaction(sql, applyMigration);
+      }
     }
   } finally {
     await sql.end();
@@ -1033,26 +1061,10 @@ export async function applyPendingMigrations(url: string): Promise<void> {
   if (initialState.status === "upToDate") return;
 
   if (initialState.reason === "no-migration-journal-empty-db") {
-    const sql = createUtilitySql(url);
-    try {
-      const db = drizzlePg(sql);
-      await migratePg(db, { migrationsFolder: MIGRATIONS_FOLDER });
-    } finally {
-      await sql.end();
-    }
-
-    let bootstrappedState = await inspectMigrations(url);
-    if (bootstrappedState.status === "upToDate") return;
-    if (bootstrappedState.reason === "pending-migrations") {
-      const repair = await reconcilePendingMigrationHistory(url);
-      if (repair.repairedMigrations.length > 0) {
-        bootstrappedState = await inspectMigrations(url);
-      }
-      if (bootstrappedState.status === "needsMigrations" && bootstrappedState.reason === "pending-migrations") {
-        await applyPendingMigrationsManually(url, bootstrappedState.pendingMigrations);
-        bootstrappedState = await inspectMigrations(url);
-      }
-    }
+    // Use the same per-file executor for bootstrap and upgrades: Drizzle wraps
+    // the whole journal in one transaction and cannot run committing backfills.
+    await applyPendingMigrationsManually(url, initialState.pendingMigrations);
+    const bootstrappedState = await inspectMigrations(url);
     if (bootstrappedState.status === "upToDate") return;
     throw new Error(
       `Failed to bootstrap migrations: ${bootstrappedState.pendingMigrations.join(", ")}`,
@@ -1116,8 +1128,7 @@ export async function migratePostgresIfEmpty(url: string): Promise<MigrationBoot
       return { migrated: false, reason: "not-empty-no-migration-journal", tableCount };
     }
 
-    const db = drizzlePg(sql);
-    await migratePg(db, { migrationsFolder: MIGRATIONS_FOLDER });
+    await applyPendingMigrationsManually(url, await listMigrationFiles());
 
     return { migrated: true, reason: "migrated-empty-db", tableCount: 0 };
   } finally {

packages/db/src/heartbeat-run-issue-privacy-migration.test.ts

@@ -60,7 +60,7 @@ describeEmbeddedPostgres("heartbeat run issue privacy migration", () => {
       )
       VALUES
         (${firstRunId}, ${companyId}, ${agentId}, 'succeeded', NULL, ${sql.json({ issueId: firstIssueId })}),
-        (${secondRunId}, ${companyId}, ${agentId}, 'failed', NULL, ${sql.json({ issueId: secondIssueId })}),
+        (${secondRunId}, ${companyId}, ${agentId}, 'failed', NULL, ${sql.json({ taskId: secondIssueId })}),
         (${maintenanceRunId}, ${companyId}, ${agentId}, 'succeeded', NULL, ${sql.json({ wakeReason: "heartbeat_timer" })})
     `;
 
@@ -94,13 +94,22 @@ describeEmbeddedPostgres("heartbeat run issue privacy migration", () => {
 
     const [parity] = await sql<{ expected_count: number; populated_count: number }[]>`
       SELECT
-        count(*) FILTER (WHERE "context_snapshot" ? 'issueId')::int AS "expected_count",
-        count(*) FILTER (WHERE "context_snapshot" ? 'issueId' AND "issue_id" IS NOT NULL)::int AS "populated_count"
+        count(*) FILTER (WHERE "context_snapshot" ? 'issueId' OR "context_snapshot" ? 'taskId')::int AS "expected_count",
+        count(*) FILTER (WHERE ("context_snapshot" ? 'issueId' OR "context_snapshot" ? 'taskId') AND "issue_id" IS NOT NULL)::int AS "populated_count"
       FROM "heartbeat_runs"
       WHERE "id" IN (${firstRunId}, ${secondRunId}, ${maintenanceRunId})
     `;
     expect(parity).toEqual({ expected_count: 2, populated_count: 2 });
 
+    const [newRun] = await sql<{ id: string; issue_id: string | null; scope_kind: string }[]>`
+      INSERT INTO heartbeat_runs (company_id, agent_id, status, context_snapshot)
+      VALUES (${companyId}, ${agentId}, 'succeeded', ${sql.json({ taskId: secondIssueId })})
+      RETURNING id, issue_id, scope_kind`;
+    expect(newRun).toMatchObject({ issue_id: secondIssueId, scope_kind: "issue" });
+    await sql`UPDATE heartbeat_runs SET context_snapshot = '{}'::jsonb, scope_kind = 'company' WHERE id = ${newRun.id}`;
+    expect(await sql`SELECT issue_id, scope_kind FROM heartbeat_runs WHERE id = ${newRun.id}`)
+      .toEqual([{ issue_id: secondIssueId, scope_kind: "issue" }]);
+
     const indexes = await sql<{ indexname: string }[]>`
       SELECT "indexname"
       FROM "pg_indexes"
@@ -142,16 +151,83 @@ describeEmbeddedPostgres("heartbeat run issue privacy migration", () => {
       VALUES
         (${companyId}, ${agentId}, 'succeeded', ${sql.json({ issueId: "not-a-uuid" })}),
         (${companyId}, ${agentId}, 'succeeded', ${sql.json({ issueId: randomUUID() })}),
-        (${companyId}, ${agentId}, 'succeeded', ${sql.json({ issueId: foreignIssueId })})
+        (${companyId}, ${agentId}, 'succeeded', ${sql.json({ issueId: foreignIssueId })}),
+        (${companyId}, ${agentId}, 'succeeded', ${sql.json({ taskId: "not-a-uuid" })}),
+        (${companyId}, ${agentId}, 'succeeded', ${sql.json({ taskId: randomUUID() })}),
+        (${companyId}, ${agentId}, 'succeeded', ${sql.json({ taskId: foreignIssueId })}),
+        (${companyId}, ${agentId}, 'succeeded', ${sql.json({ taskId: null })})
     `;
 
     await applyPendingMigrations(database.connectionString);
     const rows = await sql<{ issue_id: string | null; scope_kind: string }[]>`
       SELECT issue_id, scope_kind FROM heartbeat_runs WHERE company_id = ${companyId}`;
-    expect(rows).toHaveLength(3);
+    expect(rows).toHaveLength(7);
     expect(rows.every(row => row.scope_kind === "issue" && row.issue_id === null)).toBe(true);
     // Re-running the migration preserves fail-closed tombstones and constraints.
     await sql`DELETE FROM "drizzle"."__drizzle_migrations" WHERE "hash" = ${await migrationHash()}`;
     await applyPendingMigrations(database.connectionString);
   }, 30_000);
+
+  it("releases schema and completed-batch row locks and resumes after interruption", async () => {
+    const database = await startEmbeddedPostgresTestDatabase("paperclip-privacy-batch-commits-");
+    cleanups.push(database.cleanup);
+    const sql = postgres(database.connectionString, { max: 1, onnotice: () => {} });
+    cleanups.push(async () => sql.end());
+    const companyId = randomUUID(), agentId = randomUUID(), issueId = randomUUID();
+    const runId = (n: number) => `00000000-0000-0000-0000-${String(n).padStart(12, "0")}`;
+    await sql`INSERT INTO companies (id, name, issue_prefix) VALUES (${companyId}, 'Batch commits', 'BCH')`;
+    await sql`INSERT INTO agents (id, company_id, name, adapter_type) VALUES (${agentId}, ${companyId}, 'Runner', 'process')`;
+    await sql`INSERT INTO issues (id, company_id, title, visibility) VALUES (${issueId}, ${companyId}, 'Private task', 'private')`;
+    await sql`DROP TRIGGER heartbeat_runs_set_scope_kind ON heartbeat_runs`;
+    await sql`INSERT INTO heartbeat_runs (id, company_id, agent_id, status, context_snapshot)
+      SELECT ('00000000-0000-0000-0000-' || lpad(n::text, 12, '0'))::uuid,
+        ${companyId}::uuid, ${agentId}::uuid, 'succeeded', ${sql.json({ taskId: issueId })}::jsonb
+      FROM generate_series(1, 1001) n`;
+    // A failed concurrent build must be repaired rather than skipped by
+    // CREATE INDEX IF NOT EXISTS when the migration is retried.
+    await sql`DROP INDEX heartbeat_runs_company_issue_created_idx`;
+    await expect(sql.unsafe("CREATE UNIQUE INDEX CONCURRENTLY heartbeat_runs_company_issue_created_idx ON heartbeat_runs (company_id)"))
+      .rejects.toThrow();
+    await sql.unsafe(`CREATE FUNCTION interrupt_privacy_batch_fixture() RETURNS trigger LANGUAGE plpgsql AS $$
+      BEGIN
+        PERFORM pg_advisory_xact_lock(736319937);
+        RAISE EXCEPTION 'fixture interrupted second batch';
+      END $$`);
+    await sql.unsafe(`CREATE TRIGGER zzz_interrupt_privacy_batch_fixture BEFORE UPDATE ON heartbeat_runs
+      FOR EACH ROW WHEN (OLD.scope_kind = 'company' AND NEW.id = '${runId(1001)}'::uuid)
+      EXECUTE FUNCTION interrupt_privacy_batch_fixture()`);
+    const hash = await migrationHash();
+    await sql`DELETE FROM drizzle.__drizzle_migrations WHERE hash = ${hash}`;
+    await sql`SELECT pg_advisory_lock(736319937)`;
+    const migration = applyPendingMigrations(database.connectionString).then(() => null, error => error);
+    try {
+      // The second batch is blocked while another session can see the first
+      // committed batch and write a row whose lock has already been released.
+      await expect.poll(async () => {
+        const [row] = await sql`SELECT count(*)::int AS n FROM heartbeat_runs
+          WHERE company_id = ${companyId} AND scope_kind = 'issue'`;
+        return row.n;
+      }, { timeout: 10_000 }).toBe(1000);
+      await sql`SET lock_timeout = '500ms'`;
+      await sql`UPDATE heartbeat_runs SET context_snapshot = context_snapshot || '{"probe":"visible"}'::jsonb
+        WHERE id = ${runId(1)}`;
+      expect(await sql`SELECT context_snapshot->>'probe' AS probe FROM heartbeat_runs WHERE id = ${runId(1)}`)
+        .toEqual([{ probe: "visible" }]);
+    } finally {
+      await sql`SELECT pg_advisory_unlock(736319937)`;
+      await sql`RESET lock_timeout`;
+    }
+    expect(await migration).toMatchObject({ message: "fixture interrupted second batch" });
+    expect(await sql`SELECT hash FROM drizzle.__drizzle_migrations WHERE hash = ${hash}`).toHaveLength(0);
+    await sql`DROP TRIGGER zzz_interrupt_privacy_batch_fixture ON heartbeat_runs`;
+    await sql`DROP FUNCTION interrupt_privacy_batch_fixture()`;
+    await applyPendingMigrations(database.connectionString);
+    const [finished] = await sql`SELECT count(*)::int AS n FROM heartbeat_runs
+      WHERE company_id = ${companyId} AND scope_kind = 'issue' AND issue_id = ${issueId}`;
+    expect(finished.n).toBe(1001);
+    expect(await sql`SELECT indisvalid FROM pg_index WHERE indexrelid = 'heartbeat_runs_company_issue_created_idx'::regclass`)
+      .toEqual([{ indisvalid: true }]);
+    expect(await sql`SELECT hash FROM drizzle.__drizzle_migrations WHERE hash = ${hash}`).toHaveLength(1);
+  }, 30_000);
+
 });

packages/db/src/migrations/0311_private_task_access.sql

@@ -1,3 +1,5 @@
+-- Runs through the Paperclip migration executor outside a file-wide transaction.
+-- Each idempotent keyset batch commits before advancing; history is recorded last.
 CREATE TABLE IF NOT EXISTS "issue_access_grants" (
 	"id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL,
 	"issue_id" uuid NOT NULL,
@@ -82,10 +84,10 @@ DO $$ BEGIN
     ALTER TABLE "workspace_operations" ADD CONSTRAINT "workspace_operations_issue_id_issues_id_fk" FOREIGN KEY ("issue_id") REFERENCES "public"."issues"("id") ON DELETE set null ON UPDATE no action;
   END IF;
 END $$;--> statement-breakpoint
-CREATE INDEX IF NOT EXISTS "heartbeat_runs_company_issue_created_idx" ON "heartbeat_runs" USING btree ("company_id","issue_id","created_at");--> statement-breakpoint
-CREATE INDEX IF NOT EXISTS "issues_company_privacy_parent_idx" ON "issues" USING btree ("company_id","privacy_parent_issue_id");--> statement-breakpoint
-CREATE INDEX IF NOT EXISTS "issues_company_privacy_root_idx" ON "issues" USING btree ("company_id","privacy_root_issue_id");--> statement-breakpoint
-CREATE UNIQUE INDEX IF NOT EXISTS "projects_company_personal_owner_uq" ON "projects" USING btree ("company_id","personal_owner_user_id") WHERE "projects"."personal_owner_user_id" is not null;--> statement-breakpoint
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "heartbeat_runs_company_issue_created_idx" ON "heartbeat_runs" USING btree ("company_id","issue_id","created_at");--> statement-breakpoint
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "issues_company_privacy_parent_idx" ON "issues" USING btree ("company_id","privacy_parent_issue_id");--> statement-breakpoint
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "issues_company_privacy_root_idx" ON "issues" USING btree ("company_id","privacy_root_issue_id");--> statement-breakpoint
+CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS "projects_company_personal_owner_uq" ON "projects" USING btree ("company_id","personal_owner_user_id") WHERE "projects"."personal_owner_user_id" is not null;--> statement-breakpoint
 DO $$ BEGIN
   IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'heartbeat_runs_scope_binding_check' AND conrelid = 'public.heartbeat_runs'::regclass) THEN
     ALTER TABLE "heartbeat_runs" ADD CONSTRAINT "heartbeat_runs_scope_binding_check" CHECK (("heartbeat_runs"."scope_kind" = 'company' AND "heartbeat_runs"."issue_id" IS NULL)
@@ -118,10 +120,10 @@ BEGIN
     RETURN NEW;
   END IF;
   candidate := coalesce(NEW.native_issue_id, NEW.issue_id);
-  IF candidate IS NULL AND NEW.context_snapshot->>'issueId' ~* '^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$' THEN
-    candidate := (NEW.context_snapshot->>'issueId')::uuid;
+  IF candidate IS NULL AND coalesce(NEW.context_snapshot->>'issueId', NEW.context_snapshot->>'taskId') ~* '^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$' THEN
+    candidate := coalesce(NEW.context_snapshot->>'issueId', NEW.context_snapshot->>'taskId')::uuid;
   END IF;
-  IF candidate IS NOT NULL OR NEW.context_snapshot ? 'issueId' OR NEW.scope_kind = 'issue' THEN
+  IF candidate IS NOT NULL OR NEW.context_snapshot ? 'issueId' OR NEW.context_snapshot ? 'taskId' OR NEW.scope_kind = 'issue' THEN
     NEW.scope_kind := 'issue';
     SELECT id INTO NEW.issue_id FROM issues WHERE id = candidate AND company_id = NEW.company_id;
   ELSE
@@ -144,6 +146,7 @@ BEGIN LOOP
   SELECT id INTO next_id FROM changed ORDER BY id DESC LIMIT 1;
   EXIT WHEN next_id IS NULL;
   cursor_id := next_id;
+  COMMIT;
 END LOOP; END $$;
 --> statement-breakpoint
 -- Existing structural children acquire downward inheritance without widening
@@ -156,6 +159,7 @@ BEGIN LOOP
   SELECT id INTO next_id FROM changed ORDER BY id DESC LIMIT 1;
   EXIT WHEN next_id IS NULL;
   cursor_id := next_id;
+  COMMIT;
 END LOOP; END $$;
 --> statement-breakpoint
 -- Keep workspace provenance after reassignment/deletion. A previous task may
@@ -202,6 +206,7 @@ BEGIN LOOP
   SELECT id INTO next_id FROM changed ORDER BY id DESC LIMIT 1;
   EXIT WHEN next_id IS NULL;
   cursor_id := next_id;
+  COMMIT;
 END LOOP; END $$;
 
 --> statement-breakpoint
@@ -237,6 +242,7 @@ BEGIN
     EXIT WHEN batch_ids IS NULL;
     UPDATE workspace_operations SET metadata = metadata WHERE id = ANY(batch_ids);
     cursor_id := batch_ids[array_length(batch_ids, 1)];
+    COMMIT;
   END LOOP;
 END $$;
 
@@ -263,5 +269,6 @@ BEGIN
       WHERE id = ANY(batch_ids) AND visibility = 'private' AND privacy_owner_user_id IS NOT NULL
       ON CONFLICT (project_id, subject_type, subject_id) DO NOTHING;
     cursor_id := batch_ids[array_length(batch_ids, 1)];
+    COMMIT;
   END LOOP;
 END $$;