Fix primary key migration causing stale schema (#3934) (#4666)

## Summary

Fixes #3934. Removing or changing a `#[primary_key]` annotation succeeds
on the first re-publish, but the stored schema's `table_primary_key`
field in `st_table` is never updated. On the **next** publish,
`check_compatible` fails with:

```
Primary key mismatch: self.primary_key: Some(ColId(0)), def.primary_key: None
```

## Root Cause

`auto_migrate_table` handles removing the PK's index and unique
constraint, but there was no `AutoMigrateStep` to update the
`table_primary_key` field in the system table.

## Fix

- Add `AutoMigrateStep::ChangePrimaryKey` variant to the auto-migration
planner
- Detect `old.primary_key != new.primary_key` in `auto_migrate_table`
and emit the step
- Add `alter_table_primary_key` to the datastore layer (`mut_tx`,
`datastore`, `relational_db`) with proper rollback support via
`PendingSchemaChange::TableAlterPrimaryKey`
- Handle the step in `auto_migrate_database`
- Add migration plan formatter support (termcolor output)

## Repro Script

```bash
# 1. Publish with primary key
spacetime publish repro --yes  # table has #[primary_key] on name

# 2. Remove primary key — succeeds
spacetime publish repro --yes  # removed #[primary_key]

# 3. Any change — CRASHES
spacetime publish repro --yes  # Primary key mismatch error
```

## Tests

- **Unit test** in `crates/core/src/db/update.rs`: Reproduces the exact
three-publish sequence from the issue (create table with PK → remove PK
→ trivial change)
- **Smoketest** in
`crates/smoketests/tests/smoketests/auto_migration.rs`: Full end-to-end
publish flow exercising the same scenario

## Files Changed

| File | Change |
|------|--------|
| `crates/schema/src/auto_migrate.rs` | Add `ChangePrimaryKey` variant +
detection |
| `crates/schema/src/auto_migrate/formatter.rs` | Format the new step |
| `crates/schema/src/auto_migrate/termcolor_formatter.rs` | Colored
output |
| `crates/datastore/src/locking_tx_datastore/tx_state.rs` |
`TableAlterPrimaryKey` pending change |
| `crates/datastore/src/locking_tx_datastore/mut_tx.rs` |
`alter_table_primary_key` |
| `crates/datastore/src/locking_tx_datastore/datastore.rs` | Expose
through datastore |
| `crates/datastore/src/locking_tx_datastore/committed_state.rs` |
Rollback support |
| `crates/core/src/db/relational_db.rs` | Expose through RelationalDB |
| `crates/core/src/db/update.rs` | Handle step + unit test |
| `crates/smoketests/.../auto_migration.rs` | Smoketest |

---------

Co-authored-by: clockwork-labs-bot <clockwork-labs-bot@users.noreply.github.com>

Author: clockwork-labs-bot

crates/core/src/db/relational_db.rs

@@ -989,6 +989,15 @@ impl RelationalDB {
         Ok(self.inner.alter_table_access_mut_tx(tx, name, access)?)
     }
 
+    pub(crate) fn alter_table_primary_key(
+        &self,
+        tx: &mut MutTx,
+        name: &str,
+        primary_key: Option<ColId>,
+    ) -> Result<(), DBError> {
+        Ok(self.inner.alter_table_primary_key_mut_tx(tx, name, primary_key)?)
+    }
+
     pub(crate) fn alter_table_row_type(
         &self,
         tx: &mut MutTx,

crates/core/src/db/update.rs

@@ -282,6 +282,11 @@ fn auto_migrate_database(
                 let table_def = plan.new.stored_in_table_def(&table_name.clone().into()).unwrap();
                 stdb.alter_table_access(tx, table_name, table_def.table_access.into())?;
             }
+            spacetimedb_schema::auto_migrate::AutoMigrateStep::ChangePrimaryKey(table_name) => {
+                let table_def = plan.new.stored_in_table_def(&table_name.clone().into()).unwrap();
+                log!(logger, "Changing primary key for table `{table_name}`");
+                stdb.alter_table_primary_key(tx, table_name, table_def.primary_key)?;
+            }
             spacetimedb_schema::auto_migrate::AutoMigrateStep::AddSchedule(_) => {
                 anyhow::bail!("Adding schedules is not yet implemented");
             }
@@ -336,7 +341,7 @@ mod test {
     };
     use spacetimedb_datastore::locking_tx_datastore::PendingSchemaChange;
     use spacetimedb_lib::db::raw_def::v9::{btree, RawModuleDefV9Builder, TableAccess};
-    use spacetimedb_sats::{product, AlgebraicType::U64};
+    use spacetimedb_sats::{product, AlgebraicType, AlgebraicType::U64};
     use spacetimedb_schema::{auto_migrate::ponder_migrate, def::ModuleDef};
 
     struct TestLogger;
@@ -419,6 +424,77 @@ mod test {
         Ok(())
     }
 
+    /// Regression test for #3934: removing a primary key annotation and then
+    /// re-publishing causes "Primary key mismatch" on the NEXT publish.
+    #[test]
+    fn update_db_remove_primary_key_issue_3934() -> anyhow::Result<()> {
+        let auth_ctx = AuthCtx::for_testing();
+        let stdb = TestDB::durable()?;
+
+        // Step 1: Table with a primary key (requires unique constraint + index).
+        let module_v1 = {
+            let mut builder = RawModuleDefV9Builder::new();
+            builder
+                .build_table_with_new_type("person", [("name", AlgebraicType::String)], true)
+                .with_primary_key(0)
+                .with_unique_constraint(0)
+                .with_index(btree(0), "person_name_idx")
+                .with_access(TableAccess::Public)
+                .finish();
+            let raw: ModuleDef = builder.finish().try_into().expect("valid module def");
+            raw
+        };
+
+        // Step 2: Same table, but primary key removed.
+        let module_v2 = {
+            let mut builder = RawModuleDefV9Builder::new();
+            builder
+                .build_table_with_new_type("person", [("name", AlgebraicType::String)], true)
+                .with_access(TableAccess::Public)
+                .finish();
+            let raw: ModuleDef = builder.finish().try_into().expect("valid module def");
+            raw
+        };
+
+        // Step 3: Trivially different module (same as v2, simulates "change anything").
+        let module_v3 = {
+            let mut builder = RawModuleDefV9Builder::new();
+            builder
+                .build_table_with_new_type("person", [("name", AlgebraicType::String)], true)
+                .with_access(TableAccess::Public)
+                .finish();
+            builder.add_reducer("noop", spacetimedb_sats::ProductType::unit(), None);
+            let raw: ModuleDef = builder.finish().try_into().expect("valid module def");
+            raw
+        };
+
+        // Publish v1.
+        let mut tx = begin_mut_tx(&stdb);
+        for def in module_v1.tables() {
+            create_table_from_def(&stdb, &mut tx, &module_v1, def)?;
+        }
+        stdb.commit_tx(tx)?;
+
+        // Migrate v1 → v2 (remove primary key). Should succeed.
+        let mut tx = begin_mut_tx(&stdb);
+        let plan = ponder_migrate(&module_v1, &module_v2)?;
+        let res = update_database(&stdb, &mut tx, auth_ctx.clone(), plan, &TestLogger)?;
+        assert!(matches!(res, UpdateResult::Success), "v1 → v2 migration failed");
+        stdb.commit_tx(tx)?;
+
+        // Migrate v2 → v3 (trivial change). This is where #3934 crashes.
+        let mut tx = begin_mut_tx(&stdb);
+        let plan = ponder_migrate(&module_v2, &module_v3)?;
+        let res = update_database(&stdb, &mut tx, auth_ctx, plan, &TestLogger)?;
+        assert!(
+            matches!(res, UpdateResult::Success),
+            "v2 → v3 migration failed (issue #3934)"
+        );
+        stdb.commit_tx(tx)?;
+
+        Ok(())
+    }
+
     fn empty_module() -> ModuleDef {
         RawModuleDefV9Builder::new()
             .finish()
@@ -461,7 +537,6 @@ mod test {
             "removing a table should disconnect clients"
         );
         assert!(stdb.table_id_from_name_mut(&tx, "droppable")?.is_none());
-        // `drop_table` records a `TableRemoved` schema change for the subscription system.
         assert!(
             tx.pending_schema_changes()
                 .iter()
@@ -498,7 +573,6 @@ mod test {
             err.to_string().contains("table contains data"),
             "error should mention that the table contains data, got: {err}"
         );
-        // The migration failed, so no schema changes should be pending.
         assert!(
             tx.pending_schema_changes().is_empty(),
             "failed migration should leave no pending schema changes: {:?}",

crates/datastore/src/locking_tx_datastore/committed_state.rs

@@ -1295,6 +1295,11 @@ impl CommittedState {
                 let table = self.tables.get_mut(&table_id)?;
                 table.with_mut_schema(|s| s.table_access = access);
             }
+            // A table's primary key was changed. Change back to the old one.
+            TableAlterPrimaryKey(table_id, old_pk) => {
+                let table = self.tables.get_mut(&table_id)?;
+                table.with_mut_schema(|s| s.primary_key = old_pk.and_then(|cl| cl.as_singleton()));
+            }
             // A table's row type was changed. Change back to the old one.
             // The row representation of old rows hasn't changed,
             // so it's safe to not rewrite the rows and merely change the type back.

crates/datastore/src/locking_tx_datastore/datastore.rs

@@ -34,7 +34,7 @@ use spacetimedb_durability::TxOffset;
 use spacetimedb_lib::{db::auth::StAccess, metrics::ExecutionMetrics};
 use spacetimedb_lib::{ConnectionId, Identity};
 use spacetimedb_paths::server::SnapshotDirPath;
-use spacetimedb_primitives::{ColList, ConstraintId, IndexId, SequenceId, TableId, ViewId};
+use spacetimedb_primitives::{ColId, ColList, ConstraintId, IndexId, SequenceId, TableId, ViewId};
 use spacetimedb_sats::{
     algebraic_value::de::ValueDeserializer, bsatn, buffer::BufReader, AlgebraicValue, ProductValue,
 };
@@ -330,6 +330,19 @@ impl Locking {
         tx.alter_table_access(table_id, access)
     }
 
+    pub fn alter_table_primary_key_mut_tx(
+        &self,
+        tx: &mut MutTxId,
+        name: &str,
+        primary_key: Option<ColId>,
+    ) -> Result<()> {
+        let table_id = self
+            .table_id_from_name_mut_tx(tx, name)?
+            .ok_or_else(|| TableError::NotFound(name.into()))?;
+
+        tx.alter_table_primary_key(table_id, primary_key)
+    }
+
     pub fn alter_table_row_type_mut_tx(
         &self,
         tx: &mut MutTxId,

crates/datastore/src/locking_tx_datastore/mut_tx.rs

@@ -1082,6 +1082,26 @@ impl MutTxId {
         Ok(())
     }
 
+    /// Change the primary key of the table identified by `table_id`.
+    ///
+    /// Updates both the in-memory schema and the `st_table` system table.
+    /// See: <https://github.com/clockworklabs/SpacetimeDB/issues/3934>
+    pub(crate) fn alter_table_primary_key(&mut self, table_id: TableId, new_primary_key: Option<ColId>) -> Result<()> {
+        // Write to the table in the tx state.
+        let ((tx_table, ..), (commit_table, ..)) = self.get_or_create_insert_table_mut(table_id)?;
+        tx_table.with_mut_schema_and_clone(commit_table, |s| s.primary_key = new_primary_key);
+
+        // Update system tables.
+        let new_pk_col_list = new_primary_key.map(|col| col.into());
+        let old_pk =
+            self.update_st_table_row(table_id, |st| mem::replace(&mut st.table_primary_key, new_pk_col_list))?;
+
+        // Remember the pending change so we can undo if necessary.
+        self.push_schema_change(PendingSchemaChange::TableAlterPrimaryKey(table_id, old_pk));
+
+        Ok(())
+    }
+
     /// Change the row type of the table identified by `table_id`.
     ///
     /// In practice, this should not error,

crates/datastore/src/locking_tx_datastore/tx_state.rs

@@ -123,6 +123,9 @@ pub enum PendingSchemaChange {
     /// Only non-representational row-type changes are allowed here,
     /// so existing rows in the table will be compatible with the new row type.
     TableAlterRowType(TableId, Vec<ColumnSchema>),
+    /// The primary key of the table with [`TableId`] was changed.
+    /// The old primary key was stored.
+    TableAlterPrimaryKey(TableId, Option<ColList>),
     /// The constraint with [`ConstraintSchema`] was added to the table with [`TableId`].
     ConstraintRemoved(TableId, ConstraintSchema),
     /// The constraint with [`ConstraintId`] was added to the table with [`TableId`].
@@ -146,6 +149,7 @@ impl MemoryUsage for PendingSchemaChange {
             Self::TableAdded(table_id) => table_id.heap_usage(),
             Self::TableAlterAccess(table_id, st_access) => table_id.heap_usage() + st_access.heap_usage(),
             Self::TableAlterRowType(table_id, column_schemas) => table_id.heap_usage() + column_schemas.heap_usage(),
+            Self::TableAlterPrimaryKey(table_id, pk) => table_id.heap_usage() + pk.heap_usage(),
             Self::ConstraintRemoved(table_id, constraint_schema) => {
                 table_id.heap_usage() + constraint_schema.heap_usage()
             }

crates/schema/src/auto_migrate.rs

@@ -301,6 +301,15 @@ pub enum AutoMigrateStep<'def> {
     /// Change the access of a table.
     ChangeAccess(<TableDef as ModuleDefLookup>::Key<'def>),
 
+    /// Change the primary key of a table.
+    ///
+    /// This updates the `table_primary_key` field in `st_table`
+    /// to match the new module definition.
+    /// Without this step, a stale primary key in the stored schema
+    /// causes `check_compatible` to fail on the next publish.
+    /// See: <https://github.com/clockworklabs/SpacetimeDB/issues/3934>
+    ChangePrimaryKey(<TableDef as ModuleDefLookup>::Key<'def>),
+
     /// Recompute a view, update its backing table, and push updates to clients
     UpdateView(<ViewDef as ModuleDefLookup>::Key<'def>),
 
@@ -676,6 +685,9 @@ fn auto_migrate_table<'def>(plan: &mut AutoMigratePlan<'def>, old: &'def TableDe
     if old.table_access != new.table_access {
         plan.steps.push(AutoMigrateStep::ChangeAccess(key));
     }
+    if old.primary_key != new.primary_key {
+        plan.steps.push(AutoMigrateStep::ChangePrimaryKey(key));
+    }
     if old.schedule != new.schedule {
         // Note: this handles the case where there's an altered ScheduleDef for some reason.
         if let Some(old_schedule) = old.schedule.as_ref() {

crates/schema/src/auto_migrate/formatter.rs

@@ -10,6 +10,7 @@ use crate::{
 };
 use itertools::Itertools;
 use spacetimedb_lib::db::raw_def::v9::{RawRowLevelSecurityDefV9, TableAccess, TableType};
+use spacetimedb_primitives::ColId;
 use spacetimedb_sats::raw_identifier::RawIdentifier;
 use spacetimedb_sats::{AlgebraicType, AlgebraicValue, WithTypespace};
 use thiserror::Error;
@@ -74,6 +75,11 @@ fn format_step<F: MigrationFormatter>(
             let access_info = extract_access_change_info(*table, plan)?;
             f.format_change_access(&access_info)
         }
+        AutoMigrateStep::ChangePrimaryKey(table) => {
+            let old_table = plan.old.lookup_expect::<TableDef>(*table);
+            let new_table = plan.new.lookup_expect::<TableDef>(*table);
+            f.format_change_primary_key(table, old_table.primary_key, new_table.primary_key)
+        }
         AutoMigrateStep::AddSchedule(schedule) => {
             let schedule_info = extract_schedule_info(*schedule, plan.new)?;
             f.format_schedule(&schedule_info, Action::Created)
@@ -150,6 +156,12 @@ pub trait MigrationFormatter {
     fn format_constraint(&mut self, constraint_info: &ConstraintInfo, action: Action) -> io::Result<()>;
     fn format_sequence(&mut self, sequence_info: &SequenceInfo, action: Action) -> io::Result<()>;
     fn format_change_access(&mut self, access_info: &AccessChangeInfo) -> io::Result<()>;
+    fn format_change_primary_key(
+        &mut self,
+        table_name: &str,
+        old_pk: Option<ColId>,
+        new_pk: Option<ColId>,
+    ) -> io::Result<()>;
     fn format_schedule(&mut self, schedule_info: &ScheduleInfo, action: Action) -> io::Result<()>;
     fn format_rls(&mut self, rls_info: &RlsInfo, action: Action) -> io::Result<()>;
     fn format_change_columns(&mut self, column_changes: &ColumnChanges) -> io::Result<()>;

crates/schema/src/auto_migrate/termcolor_formatter.rs

@@ -2,6 +2,7 @@ use std::io::Write;
 use std::{fmt, io};
 
 use spacetimedb_lib::{db::raw_def::v9::TableAccess, AlgebraicType};
+use spacetimedb_primitives::ColId;
 use spacetimedb_sats::algebraic_type::fmt::fmt_algebraic_type;
 use termcolor::{Buffer, Color, ColorChoice, ColorSpec, WriteColor};
 
@@ -324,6 +325,24 @@ impl MigrationFormatter for TermColorFormatter {
         self.buffer.write_all(b")\n")
     }
 
+    fn format_change_primary_key(
+        &mut self,
+        table_name: &str,
+        old_pk: Option<ColId>,
+        new_pk: Option<ColId>,
+    ) -> io::Result<()> {
+        let description = match (old_pk, new_pk) {
+            (Some(_), None) => "removed".to_string(),
+            (None, Some(col)) => format!("added on column {col}"),
+            (Some(_), Some(col)) => format!("changed to column {col}"),
+            (None, None) => return Ok(()),
+        };
+        self.write_action_prefix(&Action::Changed)?;
+        self.buffer.write_all(b" primary key on table ")?;
+        self.write_colored(table_name, Some(self.colors.table_name), true)?;
+        self.buffer.write_all(format!(" ({description})\n").as_bytes())
+    }
+
     fn format_schedule(&mut self, s: &ScheduleInfo, action: Action) -> io::Result<()> {
         self.write_action_prefix(&action)?;
         self.buffer.write_all(b" schedule for table ")?;

crates/smoketests/tests/smoketests/auto_migration.rs

@@ -409,3 +409,82 @@ fn test_add_table_columns() {
         logs2
     );
 }
+
+// --- Issue #3934: Removing a primary key breaks subsequent publishes ---
+
+const MODULE_CODE_WITH_PK: &str = r#"
+use spacetimedb::{ReducerContext, Table};
+
+#[spacetimedb::table(accessor = person, public)]
+pub struct Person {
+    #[primary_key]
+    name: String,
+}
+
+#[spacetimedb::reducer]
+pub fn add(ctx: &ReducerContext, name: String) {
+    ctx.db.person().insert(Person { name });
+}
+"#;
+
+const MODULE_CODE_WITHOUT_PK: &str = r#"
+use spacetimedb::{ReducerContext, Table};
+
+#[spacetimedb::table(accessor = person, public)]
+pub struct Person {
+    name: String,
+}
+
+#[spacetimedb::reducer]
+pub fn add(ctx: &ReducerContext, name: String) {
+    ctx.db.person().insert(Person { name });
+}
+"#;
+
+const MODULE_CODE_WITHOUT_PK_V2: &str = r#"
+use spacetimedb::{ReducerContext, Table};
+
+#[spacetimedb::table(accessor = person, public)]
+pub struct Person {
+    name: String,
+}
+
+#[spacetimedb::reducer]
+pub fn add(ctx: &ReducerContext, name: String) {
+    ctx.db.person().insert(Person { name });
+}
+
+#[spacetimedb::reducer]
+pub fn noop(_ctx: &ReducerContext) {}
+"#;
+
+/// Regression test for <https://github.com/clockworklabs/SpacetimeDB/issues/3934>.
+///
+/// Removing a `#[primary_key]` annotation and re-publishing succeeds,
+/// but the stored schema retains the stale primary key. On the *next*
+/// publish, `check_compatible` sees the mismatch and fails with:
+///
+///   "Primary key mismatch: self.primary_key: Some(ColId(0)), def.primary_key: None"
+///
+/// The fix adds a `ChangePrimaryKey` auto-migration step that updates
+/// `table_primary_key` in `st_table`.
+#[test]
+fn test_remove_primary_key_issue_3934() {
+    let mut test = Smoketest::builder().module_code(MODULE_CODE_WITH_PK).build();
+
+    // Step 1: Publish with primary key.
+    let identity = test
+        .database_identity
+        .clone()
+        .expect("database should be published after build");
+
+    // Step 2: Remove primary key. Should succeed.
+    test.write_module_code(MODULE_CODE_WITHOUT_PK).unwrap();
+    test.publish_module_with_options(&identity, false, true)
+        .expect("Removing primary key should succeed");
+
+    // Step 3: Trivial change (add a reducer). This is where #3934 crashes.
+    test.write_module_code(MODULE_CODE_WITHOUT_PK_V2).unwrap();
+    test.publish_module_with_options(&identity, false, true)
+        .expect("Publish after PK removal should succeed (issue #3934)");
+}