From e36969d910ee812ccbcf57beb193ab9ed6b8b8e4 Mon Sep 17 00:00:00 2001
From: Shantanu Singh <shsingh@broadinstitute.org>
Date: Sat, 26 Sep 2026 17:35:16 -0400
Subject: [PATCH] Scope person fact migration checks to affected tables

---
 internal/store/migrate_person_brief_ledger.go | 24 +++----
 .../person_brief_ledger_migration_test.go     | 66 +++++++++++++++++++
 2 files changed, 76 insertions(+), 14 deletions(-)

diff --git a/internal/store/migrate_person_brief_ledger.go b/internal/store/migrate_person_brief_ledger.go
index d1cf072aa..c2a4717f2 100644
--- a/internal/store/migrate_person_brief_ledger.go
+++ b/internal/store/migrate_person_brief_ledger.go
@@ -101,7 +101,7 @@ func (s *Store) migratePersonFactClaimOriginBriefSQLite(ctx context.Context) (er
 			return fmt.Errorf("rebuild person fact claim origins: %w", err)
 		}
 	}
-	violations, err := countSQLiteForeignKeyViolations(ctx, tx)
+	violations, err := countPersonFactClaimForeignKeyViolations(ctx, tx)
 	if err != nil {
 		return err
 	}
@@ -254,21 +254,17 @@ func sqliteColumnPresent(ctx context.Context, tx *loggedTx, table, column string
 	return false, nil
 }
 
-// countSQLiteForeignKeyViolations reports how many references the rebuilt
-// schema leaves dangling, so a rebuild that lost a row fails the upgrade rather
-// than committing a corrupt ledger.
-func countSQLiteForeignKeyViolations(ctx context.Context, tx *sql.Tx) (int, error) {
-	rows, err := tx.QueryContext(ctx, `PRAGMA foreign_key_check`)
+// countPersonFactClaimForeignKeyViolations checks the rebuilt table and its
+// referencing tables. Unrelated legacy violations must not block this migration.
+func countPersonFactClaimForeignKeyViolations(ctx context.Context, tx *sql.Tx) (int, error) {
+	var violations int
+	err := tx.QueryRowContext(ctx, `SELECT
+		(SELECT COUNT(*) FROM pragma_foreign_key_check('person_fact_claims')) +
+		(SELECT COUNT(*) FROM pragma_foreign_key_check('person_fact_claim_evidence')) +
+		(SELECT COUNT(*) FROM pragma_foreign_key_check('person_fact_decisions'))
+	`).Scan(&violations)
 	if err != nil {
 		return 0, fmt.Errorf("check person fact claim references: %w", err)
 	}
-	defer func() { _ = rows.Close() }()
-	violations := 0
-	for rows.Next() {
-		violations++
-	}
-	if err := rows.Err(); err != nil {
-		return 0, fmt.Errorf("check person fact claim references: %w", err)
-	}
 	return violations, nil
 }
diff --git a/internal/store/person_brief_ledger_migration_test.go b/internal/store/person_brief_ledger_migration_test.go
index dd278f5f3..fa78fe4bd 100644
--- a/internal/store/person_brief_ledger_migration_test.go
+++ b/internal/store/person_brief_ledger_migration_test.go
@@ -234,3 +234,69 @@ func TestPersonSweepAttemptBriefFailureMigrationAddsColumnToExistingArchives(t *
 		`SELECT COUNT(*) FROM person_sweep_attempts WHERE id = ?`), f.attemptID).Scan(&attempts))
 	checks.Equal(1, attempts)
 }
+
+func TestPersonFactClaimOriginMigrationScopesForeignKeyChecks(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		corruptSQL string
+		wantError  bool
+	}{
+		{
+			name: "unrelated legacy references",
+			corruptSQL: `CREATE TABLE legacy_parent (id INTEGER PRIMARY KEY);
+				CREATE TABLE legacy_child (parent_id INTEGER REFERENCES legacy_parent(id));
+				INSERT INTO legacy_child VALUES (1)`,
+		},
+		{
+			name:       "claim parent reference",
+			corruptSQL: `UPDATE person_fact_claims SET generation_id = 1000000`,
+			wantError:  true,
+		},
+		{
+			name:       "evidence claim reference",
+			corruptSQL: `UPDATE person_fact_claim_evidence SET claim_id = 1000000`,
+			wantError:  true,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			require := require.New(t)
+			assert := assert.New(t)
+			f := newPersonSweepBudgetFixture(t, "scoped-claim-migration")
+			if f.store.IsPostgreSQL() {
+				t.Skip("SQLite table rebuild")
+			}
+			seedPersonFactClaimLedger(t, f.store, f.personID, "extraction-claim")
+			installLegacyPersonFactClaimOrigin(t, f.store)
+			conn, err := f.store.DB().Conn(t.Context())
+			require.NoError(err)
+			_, err = conn.ExecContext(t.Context(), `PRAGMA foreign_keys = OFF`)
+			require.NoError(err)
+			_, err = conn.ExecContext(t.Context(), tc.corruptSQL)
+			require.NoError(err)
+			_, err = conn.ExecContext(t.Context(), `PRAGMA foreign_keys = ON`)
+			require.NoError(err)
+			require.NoError(conn.Close())
+			_, err = f.store.DB().ExecContext(t.Context(),
+				`DELETE FROM applied_migrations WHERE name = 'person_fact_claim_origin_brief_v1'`)
+			require.NoError(err)
+
+			err = f.store.InitSchema()
+			if tc.wantError {
+				require.ErrorContains(err, "dangling references")
+				var definition string
+				require.NoError(f.store.DB().QueryRowContext(t.Context(),
+					`SELECT sql FROM sqlite_master WHERE name = 'person_fact_claims'`).Scan(&definition))
+				assert.NotContains(definition, "'brief'", "failed rebuild must roll back")
+			} else {
+				require.NoError(err)
+				require.NoError(insertBriefOriginClaim(t.Context(), f.store, f.personID, "brief-after"))
+				var violations int
+				require.NoError(f.store.DB().QueryRowContext(t.Context(),
+					`SELECT COUNT(*) FROM pragma_foreign_key_check('legacy_child')`).Scan(&violations))
+				assert.Equal(1, violations, "unrelated data must remain unchanged")
+			}
+			_, err = f.store.DB().ExecContext(t.Context(), st1000ClaimEvidenceInsert(f.store))
+			require.Error(err, "foreign key enforcement must be restored")
+		})
+	}
+}