feat: Allow adding unique constraints to existing tables

crates/core/src/db/relational_db.rs

@@ -48,7 +48,7 @@ use spacetimedb_sats::{AlgebraicType, AlgebraicValue, ProductType, ProductValue}
 use spacetimedb_schema::def::{ModuleDef, TableDef, ViewDef};
 use spacetimedb_schema::reducer_name::ReducerName;
 use spacetimedb_schema::schema::{
-    ColumnSchema, IndexSchema, RowLevelSecuritySchema, Schema, SequenceSchema, TableSchema,
+    ColumnSchema, ConstraintSchema, IndexSchema, RowLevelSecuritySchema, Schema, SequenceSchema, TableSchema,
 };
 use spacetimedb_schema::table_name::TableName;
 use spacetimedb_snapshot::{ReconstructedSnapshot, SnapshotError, SnapshotRepository};
@@ -1483,6 +1483,15 @@ impl RelationalDB {
         Ok(self.inner.drop_sequence_mut_tx(tx, seq_id)?)
     }
 
+    /// Creates a new constraint in the database instance.
+    pub fn create_constraint(
+        &self,
+        tx: &mut MutTx,
+        constraint: ConstraintSchema,
+    ) -> Result<ConstraintId, DBError> {
+        Ok(self.inner.create_constraint_mut_tx(tx, constraint)?)
+    }
+
     ///Removes the [Constraints] from database instance
     pub fn drop_constraint(&self, tx: &mut MutTx, constraint_id: ConstraintId) -> Result<(), DBError> {
         Ok(self.inner.drop_constraint_mut_tx(tx, constraint_id)?)

crates/core/src/db/update.rs

@@ -6,6 +6,7 @@ use spacetimedb_lib::db::auth::StTableType;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::AlgebraicValue;
 use spacetimedb_primitives::{ColSet, TableId};
+use spacetimedb_schema::schema::ConstraintSchema;
 use spacetimedb_schema::auto_migrate::{AutoMigratePlan, ManualMigratePlan, MigratePlan};
 use spacetimedb_schema::def::{TableDef, ViewDef};
 use spacetimedb_schema::schema::{column_schemas_from_defs, IndexSchema, Schema, SequenceSchema, TableSchema};
@@ -233,6 +234,29 @@ fn auto_migrate_database(
                 );
                 stdb.drop_constraint(tx, constraint_schema.constraint_id)?;
             }
+            spacetimedb_schema::auto_migrate::AutoMigrateStep::AddConstraint(constraint_name) => {
+                let table_def = plan
+                    .new
+                    .stored_in_table_def(constraint_name)
+                    .expect("AddConstraint references a table that should exist in the new module def");
+                let constraint_def = &table_def.constraints[constraint_name];
+                let table_id = stdb
+                    .table_id_from_name_mut(tx, &table_def.name)?
+                    .expect("table should exist in the database for AddConstraint");
+                let constraint_schema = ConstraintSchema::from_module_def(
+                    plan.new,
+                    constraint_def,
+                    table_id,
+                    spacetimedb_primitives::ConstraintId::SENTINEL,
+                );
+                log!(
+                    logger,
+                    "Adding constraint `{}` on table `{}`",
+                    constraint_name,
+                    table_def.name
+                );
+                stdb.create_constraint(tx, constraint_schema)?;
+            }
             spacetimedb_schema::auto_migrate::AutoMigrateStep::AddSequence(sequence_name) => {
                 let table_def = plan.new.stored_in_table_def(sequence_name).unwrap();
                 let sequence_def = table_def.sequences.get(sequence_name).unwrap();

crates/datastore/src/locking_tx_datastore/committed_state.rs

@@ -758,14 +758,38 @@ impl CommittedState {
                     .unwrap_or_else(|e| match e {});
             }
             // A constraint was removed. Add it back.
-            ConstraintRemoved(table_id, constraint_schema) => {
+            ConstraintRemoved(table_id, constraint_schema, index_ids) => {
                 let table = self.tables.get_mut(&table_id)?;
                 table.with_mut_schema(|s| s.update_constraint(constraint_schema));
+                // If the constraint had unique indices, make them unique again.
+                for index_id in index_ids {
+                    if let Some(idx) = table.indexes.get_mut(&index_id) {
+                        idx.make_unique().expect("rollback: index should have no duplicates");
+                    }
+                }
+                // Forward `drop_constraint` calls `Table::make_index_non_unique`, which
+                // rebuilds the pointer map when no unique index remained. Whatever the
+                // forward path did, the table invariant "pointer map is present iff no
+                // unique index exists" (see `table.rs`) is about to be re-established
+                // by the `make_unique` calls above — drop any rebuilt map now.
+                // `take_pointer_map` is idempotent: it returns `None` when the map is
+                // already absent, so it is safe to call unconditionally.
+                table.take_pointer_map();
             }
             // A constraint was added. Remove it.
-            ConstraintAdded(table_id, constraint_id) => {
+            ConstraintAdded(table_id, constraint_id, index_ids, pointer_map) => {
                 let table = self.tables.get_mut(&table_id)?;
                 table.with_mut_schema(|s| s.remove_constraint(constraint_id));
+                // If the constraint made indices unique, revert them to non-unique.
+                for index_id in index_ids {
+                    if let Some(idx) = table.indexes.get_mut(&index_id) {
+                        idx.make_non_unique();
+                    }
+                }
+                // Restore the pointer map if it was taken.
+                if let Some(pm) = pointer_map {
+                    table.restore_pointer_map(pm);
+                }
             }
             // A sequence was removed. Add it back.
             SequenceRemoved(table_id, seq, schema) => {