Merge pinned main primary-key repair into invitation candidate

Preserve invitation migration 0069 and admission semantics while integrating main events_pkey reconciliation and schema-parity coverage.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: seiler <seiler@squareup.com>

.github/workflows/ci.yml

@@ -53,6 +53,8 @@ jobs:
               - '.config/nextest.toml'
               - 'scripts/postgres-test-*.sh'
               - 'scripts/reconcile-schema-after-pgschema.sql'
+              - 'scripts/repair-events-pkey.sql'
+              - 'scripts/start-relay-for-tests.sh'
               - 'bin/pgschema'
               - 'bin/.pgschema-*.pkg'
               - 'scripts/check-postgres-test-discovery.py'

crates/buzz-db/src/runtime/migration.rs

@@ -1601,6 +1601,32 @@ mod postgres_tests {
         );
     }
 
+    /// The reconcile script repairs `events_pkey` after every apply, and
+    /// `start-relay-for-tests.sh` runs the same repair before apply so a reused
+    /// database reconciled before the repair existed can still be applied.
+    #[test]
+    fn events_pkey_repair_runs_before_reused_apply_and_matches_reconcile() {
+        let repair = include_str!("../../../../scripts/repair-events-pkey.sql");
+        let block = &repair[repair.find("DO $$").expect("repair script has a DO block")..];
+        assert!(
+            include_str!("../../../../scripts/reconcile-schema-after-pgschema.sql").contains(block),
+            "scripts/reconcile-schema-after-pgschema.sql must carry scripts/repair-events-pkey.sql's \
+             DO block verbatim"
+        );
+
+        let launcher = include_str!("../../../../scripts/start-relay-for-tests.sh");
+        let repair_at = launcher
+            .find("scripts/repair-events-pkey.sql")
+            .expect("start-relay-for-tests.sh runs scripts/repair-events-pkey.sql");
+        let apply_at = launcher
+            .find("./bin/pgschema apply")
+            .expect("start-relay-for-tests.sh runs pgschema apply");
+        assert!(
+            repair_at < apply_at,
+            "start-relay-for-tests.sh must repair events_pkey before pgschema apply"
+        );
+    }
+
     /// `.pgschemaignore` hides dated monthlies from pgschema so a re-apply
     /// cannot drop the ones the relay created at runtime. Each monthly
     /// `schema/partitions/` declares must be exempt, or pgschema would stop
@@ -2735,6 +2761,10 @@ mod postgres_tests {
     /// by name, not ordinal, because migrations append via `ALTER TABLE` while
     /// `schema.sql` declares them inline — positions legitimately differ, shapes
     /// must not. Entity-intake constraints are also compared in the live catalog.
+    /// `events` and its partitions must carry the same primary key on both
+    /// probes, so the reconcile script's `events_pkey` rebuild cannot drift from
+    /// the migrations. Partitions are compared as a set of key definitions
+    /// because the two paths create different months.
     ///
     /// Driving the real binary is load-bearing: `pgschema` 1.7.4 discards
     /// per-key `NULLS FIRST`/`NULLS LAST` when it re-emits an index, so a naive
@@ -2807,6 +2837,23 @@ mod postgres_tests {
             .expect("read constraint definitions")
         }
 
+        // Distinct (is the parent, key definition) pairs across events and its
+        // partitions.
+        async fn events_key_defs(pool: &PgPool) -> Vec<(bool, String)> {
+            sqlx::query_as(
+                "SELECT DISTINCT conrelid = 'events'::regclass, pg_get_constraintdef(oid) \
+                 FROM pg_constraint \
+                 WHERE contype = 'p' AND conrelid IN ( \
+                     SELECT 'events'::regclass \
+                     UNION ALL \
+                     SELECT inhrelid FROM pg_inherits WHERE inhparent = 'events'::regclass) \
+                 ORDER BY 1, 2",
+            )
+            .fetch_all(pool)
+            .await
+            .expect("read events key definitions")
+        }
+
         let base_url = std::env::var("BUZZ_TEST_DATABASE_URL")
             .or_else(|_| std::env::var("DATABASE_URL"))
             .unwrap_or_else(|_| TEST_DB_URL.to_owned());
@@ -2827,8 +2874,8 @@ mod postgres_tests {
             .await
             .expect("create migrated probe database");
 
-        // Bootstrap the desired-state probe through the real pgschema binary, the
-        // same invocation the test-relay launchers use.
+        // Bootstrap the desired-state probe the way the test-relay launchers do:
+        // the real pgschema binary, then the post-apply reconcile script.
         let apply = pgschema(&conn, &desired_db, &["apply", "--auto-approve"]);
         assert!(
             apply.status.success(),
@@ -2875,6 +2922,13 @@ mod postgres_tests {
                  migration and schema.sql must both use a representable shape."
             );
         }
+        assert_eq!(
+            events_key_defs(&desired).await,
+            events_key_defs(&migrated).await,
+            "events primary-key parity mismatch: the pgschema-bootstrapped, reconciled \
+             desired state has drifted from the migrations; update the events_pkey block in \
+             scripts/reconcile-schema-after-pgschema.sql and scripts/repair-events-pkey.sql"
+        );
 
         assert_eq!(
             constraints(&desired, "entity_deletion_requests").await,
@@ -2894,12 +2948,19 @@ mod postgres_tests {
         }
     }
 
-    /// Re-applying `schema/schema.sql` to a database a relay has used plans
-    /// `DROP TABLE IF EXISTS ... CASCADE` for every monthly the partition
-    /// manager created at runtime, and their rows go with them.
-    /// `.pgschemaignore` keeps those monthlies out of pgschema's view, and
-    /// pgschema does not compare partition bounds, so the advanced catch-all
-    /// is left alone too.
+    /// Re-applying `schema/schema.sql` to a database bootstrapped the way every
+    /// caller does it (pgschema apply, then the reconcile script) and then used
+    /// by a relay plans nothing. pgschema does not compare partition bounds, so
+    /// the advanced catch-alls do not show up in the plan. Two things keep the
+    /// rest of it empty:
+    ///
+    /// - `.pgschemaignore` hides the monthlies the partition manager created
+    ///   at runtime, which pgschema would otherwise drop, rows included, with
+    ///   `DROP TABLE IF EXISTS ... CASCADE`.
+    /// - The reconcile script rebuilds `events_pkey` in its declared order.
+    ///   pgschema creates it with the partition key first, then plans to drop
+    ///   and re-add it and every declared child's key, which fails once the
+    ///   children are attached.
     #[tokio::test]
     #[ignore = "requires Postgres"]
     async fn pgschema_reapply_keeps_runtime_partitions() {
@@ -2920,24 +2981,27 @@ mod postgres_tests {
             .await
             .expect("create probe database");
 
-        let apply = pgschema(&conn, &database, &["apply", "--auto-approve"]);
-        assert!(
-            apply.status.success(),
-            "pgschema apply failed: {}\n{}",
-            String::from_utf8_lossy(&apply.stdout),
-            String::from_utf8_lossy(&apply.stderr),
-        );
-
         let pool = PgPool::connect(&format!("{base_prefix}/{database}"))
             .await
             .expect("connect probe database");
-        // Every pgschema caller attaches the declared children this way.
-        sqlx::raw_sql(include_str!(
-            "../../../../scripts/reconcile-schema-after-pgschema.sql"
-        ))
-        .execute(&pool)
-        .await
-        .expect("reconcile schema after pgschema");
+        // The bootstrap every pgschema caller runs; the reconcile script
+        // attaches the declared children.
+        let bootstrap = |round: &str| {
+            let apply = pgschema(&conn, &database, &["apply", "--auto-approve"]);
+            assert!(
+                apply.status.success(),
+                "{round} pgschema apply failed: {}\n{}",
+                String::from_utf8_lossy(&apply.stdout),
+                String::from_utf8_lossy(&apply.stderr),
+            );
+            sqlx::raw_sql(include_str!(
+                "../../../../scripts/reconcile-schema-after-pgschema.sql"
+            ))
+            .execute(&pool)
+        };
+        bootstrap("first")
+            .await
+            .expect("first post-apply reconciliation");
         async fn catch_all_bounds(pool: &PgPool) -> Vec<(String, Option<String>)> {
             sqlx::query_as(
                 "SELECT relname::text, pg_get_expr(relpartbound, oid) FROM pg_class \
@@ -2967,6 +3031,10 @@ mod postgres_tests {
         .await
         .expect("maintain partitions");
         let advanced_bounds = catch_all_bounds(&pool).await;
+        assert!(
+            advanced_bounds.len() == 2 && advanced_bounds.iter().all(|(_, bound)| bound.is_some()),
+            "maintenance dropped or detached a catch-all: {advanced_bounds:?}"
+        );
         for (declared, advanced) in declared_bounds.iter().zip(&advanced_bounds) {
             assert_ne!(declared, advanced, "catch-all did not advance");
         }
@@ -2979,7 +3047,6 @@ mod postgres_tests {
         .fetch_all(&pool)
         .await
         .expect("list partitions");
-        pool.close().await;
         let declared = declared_monthlies();
         let runtime: Vec<_> = children
             .iter()
@@ -2987,9 +3054,18 @@ mod postgres_tests {
                 child.starts_with(&format!("{parent}_p2")) && !declared.contains(child)
             })
             .collect();
-        for parent in ["events", "delivery_log"] {
-            // The first month the declared catch-all served.
-            let first = format!("{parent}_p2026_07");
+        for (catch_all, bound) in &declared_bounds {
+            // The first month the declared catch-all served, from its
+            // `FOR VALUES FROM ('YYYY-MM-01 ...')` lower bound.
+            let parent = catch_all
+                .strip_suffix("_p_future")
+                .expect("catch-all is named <parent>_p_future");
+            let month = bound
+                .as_deref()
+                .and_then(|bound| bound.split_once("FROM ('"))
+                .and_then(|(_, rest)| rest.get(..7))
+                .expect("declared catch-all has a dated lower bound");
+            let first = format!("{parent}_p{}", month.replace('-', "_"));
             assert!(
                 runtime.iter().any(|(_, child)| *child == first),
                 "maintenance did not advance into {first}: {children:?}"
@@ -3004,15 +3080,14 @@ mod postgres_tests {
             String::from_utf8_lossy(&plan.stderr),
         );
         assert!(
-            !plan_sql.contains("DROP TABLE"),
-            "re-applying schema.sql would drop tables:\n{plan_sql}"
+            plan_sql.trim().is_empty(),
+            "re-applying schema.sql to a reconciled database must plan nothing; \
+             runtime monthlies {runtime:?}:\n{plan_sql}"
         );
-        for (_, child) in &runtime {
-            assert!(
-                !plan_sql.contains(child.as_str()),
-                "re-applying schema.sql would change runtime monthly {child}:\n{plan_sql}"
-            );
-        }
+        bootstrap("second")
+            .await
+            .expect("second post-apply reconciliation");
+        pool.close().await;
 
         sqlx::query(AssertSqlSafe(format!(
             "DROP DATABASE {database} WITH (FORCE)"

scripts/reconcile-schema-after-pgschema.sql

@@ -181,6 +181,55 @@ BEGIN
     END IF;
 END $$;
 
+-- pgschema moves the partition key to the front of a partitioned table's
+-- primary key when it creates the table, so events gets
+-- PRIMARY KEY (created_at, community_id, id) instead of the declared
+-- (community_id, created_at, id). Every later plan then drops and re-adds
+-- events_pkey and each declared child's key, which fails once the children are
+-- attached because their keys are inherited. Rebuild the key on the parent in
+-- the declared order; PostgreSQL recreates each partition's key with it.
+-- delivery_log already declares its partition key first. This block is a
+-- verbatim copy of scripts/repair-events-pkey.sql; keep the two identical.
+DO $$
+DECLARE
+    declared CONSTANT text := 'PRIMARY KEY (community_id, created_at, id)';
+BEGIN
+    -- Also runs before pgschema apply, when events may not exist yet.
+    IF to_regclass('events') IS NULL THEN
+        RETURN;
+    END IF;
+
+    IF (
+        SELECT pg_get_constraintdef(oid) FROM pg_constraint
+        WHERE conrelid = 'events'::regclass AND contype = 'p'
+    ) IS DISTINCT FROM declared THEN
+        -- The rebuild locks events and every partition ACCESS EXCLUSIVE. Fail
+        -- fast instead of queueing behind a relay that still uses them.
+        BEGIN
+            PERFORM set_config('lock_timeout', '10s', true);
+            ALTER TABLE events DROP CONSTRAINT IF EXISTS events_pkey;
+            ALTER TABLE events
+                ADD CONSTRAINT events_pkey PRIMARY KEY (community_id, created_at, id);
+        EXCEPTION WHEN lock_not_available THEN
+            RAISE EXCEPTION 'timed out locking events to rebuild events_pkey'
+                USING HINT = 'Stop any relay using this database, then rerun.';
+        END;
+    END IF;
+
+    IF EXISTS (
+        SELECT 1 FROM pg_constraint
+        WHERE contype = 'p'
+          AND conrelid IN (
+              SELECT 'events'::regclass
+              UNION ALL
+              SELECT inhrelid FROM pg_inherits WHERE inhparent = 'events'::regclass
+          )
+          AND pg_get_constraintdef(oid) <> declared
+    ) THEN
+        RAISE EXCEPTION 'events and its partitions must have %', declared;
+    END IF;
+END $$;
+
 -- pgschema drops multi-column CHECK constraints. Restore the direct-action
 -- shape CHECK from schema/schema.sql (and migration 0055) verbatim.
 DO $$

scripts/repair-events-pkey.sql

@@ -0,0 +1,56 @@
+-- Rebuild events_pkey in its declared column order.
+--
+-- pgschema creates events with the partition key moved to the front of its
+-- primary key, and every later plan then tries to drop and re-add the inherited
+-- key on each partition, which fails (SQLSTATE 42P16).
+-- scripts/reconcile-schema-after-pgschema.sql repairs that after every apply
+-- and carries a verbatim copy of the block below. A reused database reconciled
+-- before that repair existed still fails its next apply, before the reconcile
+-- script can run, so scripts/start-relay-for-tests.sh, the one launcher that
+-- reuses its database, runs this file before pgschema apply. It does nothing on
+-- an empty database or one whose key is already in order.
+--
+-- Both copies are temporary. The reconcile block, and the test that keeps the
+-- copies identical, can go once pgschema creates the key in declared order
+-- (1.13.3 does). This file and the launcher step can go once no local database
+-- reconciled before the reconcile block existed remains.
+
+DO $$
+DECLARE
+    declared CONSTANT text := 'PRIMARY KEY (community_id, created_at, id)';
+BEGIN
+    -- Also runs before pgschema apply, when events may not exist yet.
+    IF to_regclass('events') IS NULL THEN
+        RETURN;
+    END IF;
+
+    IF (
+        SELECT pg_get_constraintdef(oid) FROM pg_constraint
+        WHERE conrelid = 'events'::regclass AND contype = 'p'
+    ) IS DISTINCT FROM declared THEN
+        -- The rebuild locks events and every partition ACCESS EXCLUSIVE. Fail
+        -- fast instead of queueing behind a relay that still uses them.
+        BEGIN
+            PERFORM set_config('lock_timeout', '10s', true);
+            ALTER TABLE events DROP CONSTRAINT IF EXISTS events_pkey;
+            ALTER TABLE events
+                ADD CONSTRAINT events_pkey PRIMARY KEY (community_id, created_at, id);
+        EXCEPTION WHEN lock_not_available THEN
+            RAISE EXCEPTION 'timed out locking events to rebuild events_pkey'
+                USING HINT = 'Stop any relay using this database, then rerun.';
+        END;
+    END IF;
+
+    IF EXISTS (
+        SELECT 1 FROM pg_constraint
+        WHERE contype = 'p'
+          AND conrelid IN (
+              SELECT 'events'::regclass
+              UNION ALL
+              SELECT inhrelid FROM pg_inherits WHERE inhparent = 'events'::regclass
+          )
+          AND pg_get_constraintdef(oid) <> declared
+    ) THEN
+        RAISE EXCEPTION 'events and its partitions must have %', declared;
+    END IF;
+END $$;

scripts/start-relay-for-tests.sh

@@ -103,6 +103,10 @@ export PGSCHEMA_PLAN_DB=buzz
 export PGSCHEMA_PLAN_USER=buzz
 export PGSCHEMA_PLAN_PASSWORD=buzz_dev
 
+# This database persists across runs. One reconciled before the events_pkey
+# repair existed would fail apply, so repair it first (see the script).
+docker exec -i -e PGPASSWORD="${PGPASSWORD}" buzz-postgres \
+  psql -U "${PGUSER}" -d "${PGDATABASE}" -v ON_ERROR_STOP=1 < scripts/repair-events-pkey.sql
 ./bin/pgschema apply --file schema/schema.sql --auto-approve
 docker exec -i -e PGPASSWORD="${PGPASSWORD}" buzz-postgres \
   psql -U "${PGUSER}" -d "${PGDATABASE}" -v ON_ERROR_STOP=1 < scripts/reconcile-schema-after-pgschema.sql
@@ -168,7 +172,7 @@ if [[ "${BUZZ_REQUIRE_RELAY_MEMBERSHIP:-}" == "true" ]]; then
 fi
 
 nohup env \
-  DATABASE_URL=postgres://buzz:buzz_dev@localhost:5432/buzz \
+  DATABASE_URL="postgres://${PGUSER}:${PGPASSWORD}@${PGHOST}:${PGPORT}/${PGDATABASE}" \
   REDIS_URL=redis://localhost:6379 \
   RELAY_URL=ws://localhost:3000 \
   BUZZ_BIND_ADDR=0.0.0.0:3000 \