Fix adding enum variants (#3200)

# Description of Changes

First commit reverts disabling adding enum variants
(#3178).
Second commit makes replay recognize inserts / deletes to `st_column`
and triggers a refresh of the in-memory table that was referenced in the
`st_column` change.

# API and ABI breaking changes

None

# Expected complexity level and risk

3, small code bug "deep".

# Testing

I've confirmed manually that a module that couldn't be restarted before
this PR can now be restarted.
We should probably follow up this PR with a smoketest.

Author: Centril

crates/core/src/db/relational_db.rs

@@ -1043,7 +1043,6 @@ impl RelationalDB {
         Ok(self.inner.alter_table_access_mut_tx(tx, name, access)?)
     }
 
-    #[allow(unused)]
     pub(crate) fn alter_table_row_type(
         &self,
         tx: &mut MutTx,

crates/core/src/db/update.rs

@@ -9,7 +9,7 @@ use spacetimedb_lib::AlgebraicValue;
 use spacetimedb_primitives::{ColSet, TableId};
 use spacetimedb_schema::auto_migrate::{AutoMigratePlan, ManualMigratePlan, MigratePlan};
 use spacetimedb_schema::def::TableDef;
-use spacetimedb_schema::schema::{IndexSchema, Schema, SequenceSchema, TableSchema};
+use spacetimedb_schema::schema::{column_schemas_from_defs, IndexSchema, Schema, SequenceSchema, TableSchema};
 use std::sync::Arc;
 
 /// The logger used for by [`update_database`] and friends.
@@ -222,8 +222,14 @@ fn auto_migrate_database(
                 );
                 stdb.drop_sequence(tx, sequence_schema.sequence_id)?;
             }
-            spacetimedb_schema::auto_migrate::AutoMigrateStep::ChangeColumns(_table_name) => {
-                anyhow::bail!("Unsupported: Changing column types");
+            spacetimedb_schema::auto_migrate::AutoMigrateStep::ChangeColumns(table_name) => {
+                let table_def = plan.new.stored_in_table_def(table_name).unwrap();
+                let table_id = stdb.table_id_from_name_mut(tx, table_name).unwrap().unwrap();
+                let column_schemas = column_schemas_from_defs(plan.new, &table_def.columns, table_id);
+
+                log!(logger, "Changing columns of table `{}`", table_name);
+
+                stdb.alter_table_row_type(tx, table_id, column_schemas)?;
             }
             spacetimedb_schema::auto_migrate::AutoMigrateStep::ChangeAccess(table_name) => {
                 let table_def = plan.new.stored_in_table_def(table_name).unwrap();

crates/datastore/src/locking_tx_datastore/committed_state.rs

@@ -8,8 +8,9 @@ use super::{
 };
 use crate::{
     db_metrics::DB_METRICS,
-    error::{IndexError, TableError},
+    error::{DatastoreError, IndexError, TableError},
     execution_context::ExecutionContext,
+    locking_tx_datastore::state_view::iter_st_column_for_table,
     system_tables::{
         system_tables, StColumnRow, StConstraintData, StConstraintRow, StIndexRow, StSequenceRow, StTableFields,
         StTableRow, SystemTable, ST_CLIENT_ID, ST_CLIENT_IDX, ST_COLUMN_ID, ST_COLUMN_IDX, ST_COLUMN_NAME,
@@ -28,8 +29,8 @@ use spacetimedb_lib::{
     db::auth::{StAccess, StTableType},
     Identity,
 };
-use spacetimedb_primitives::{ColList, ColSet, IndexId, TableId};
-use spacetimedb_sats::memory_usage::MemoryUsage;
+use spacetimedb_primitives::{ColId, ColList, ColSet, IndexId, TableId};
+use spacetimedb_sats::{algebraic_value::de::ValueDeserializer, memory_usage::MemoryUsage, Deserialize};
 use spacetimedb_sats::{AlgebraicValue, ProductValue};
 use spacetimedb_schema::{def::IndexAlgorithm, schema::TableSchema};
 use spacetimedb_table::{
@@ -339,6 +340,7 @@ impl CommittedState {
             .delete_equal_row(&self.page_pool, blob_store, rel)
             .map_err(TableError::Bflatn)?
             .ok_or_else(|| anyhow!("Delete for non-existent row when replaying transaction"))?;
+
         Ok(())
     }
 
@@ -349,11 +351,55 @@ impl CommittedState {
         row: &ProductValue,
     ) -> Result<()> {
         let (table, blob_store, pool) = self.get_table_and_blob_store_or_create(table_id, schema);
-        table.insert(pool, blob_store, row).map(drop).map_err(|e| match e {
-            InsertError::Bflatn(e) => TableError::Bflatn(e).into(),
-            InsertError::Duplicate(e) => TableError::Duplicate(e).into(),
-            InsertError::IndexError(e) => IndexError::UniqueConstraintViolation(e).into(),
-        })
+        let (_, row_ref) = table.insert(pool, blob_store, row).map_err(|e| -> DatastoreError {
+            match e {
+                InsertError::Bflatn(e) => TableError::Bflatn(e).into(),
+                InsertError::Duplicate(e) => TableError::Duplicate(e).into(),
+                InsertError::IndexError(e) => IndexError::UniqueConstraintViolation(e).into(),
+            }
+        })?;
+
+        if table_id == ST_COLUMN_ID {
+            // We've made a modification to `st_column`.
+            // The type of a table has changed, so figure out which.
+            // The first column in `StColumnRow` is `table_id`.
+            let row_ptr = row_ref.pointer();
+            self.st_column_changed(row, row_ptr)?;
+        }
+
+        Ok(())
+    }
+
+    /// Refreshes the columns and layout of a table
+    /// when a `row` has been inserted from `st_column`.
+    ///
+    /// The `row_ptr` is a pointer to `row`.
+    fn st_column_changed(&mut self, row: &ProductValue, row_ptr: RowPointer) -> Result<()> {
+        let target_table_id = TableId::deserialize(ValueDeserializer::from_ref(&row.elements[0]))
+            .expect("first field in `st_column` should decode to a `TableId`");
+        let target_col_id = ColId::deserialize(ValueDeserializer::from_ref(&row.elements[1]))
+            .expect("second field in `st_column` should decode to a `ColId`");
+
+        // We're replaying and we don't have unique constraints yet.
+        // Due to replay handling all inserts first and deletes after,
+        // when processing `st_column` insert/deletes,
+        // we may end up with two definitions for the same `col_pos`.
+        // Of those two, we're interested in the one we just inserted
+        // and not the other one, as it is being replaced.
+        let columns = iter_st_column_for_table(self, &target_table_id.into())?
+            .filter_map(|row_ref| {
+                StColumnRow::try_from(row_ref)
+                    .map(|c| (c.col_pos != target_col_id || row_ref.pointer() == row_ptr).then(|| c.into()))
+                    .transpose()
+            })
+            .collect::<Result<Vec<_>>>()?;
+
+        // Update the columns and layout of the the in-memory table.
+        if let Some(table) = self.tables.get_mut(&target_table_id) {
+            table.change_columns_to(columns).map_err(TableError::from)?;
+        }
+
+        Ok(())
     }
 
     pub(super) fn build_sequence_state(&mut self, sequence_state: &mut SequencesState) -> Result<()> {

crates/datastore/src/locking_tx_datastore/state_view.rs

@@ -79,12 +79,8 @@ pub trait StateView {
         let table_primary_key = row.table_primary_key.as_ref().and_then(ColList::as_singleton);
 
         // Look up the columns for the table in question.
-        let mut columns: Vec<ColumnSchema> = self
-            .iter_by_col_eq(ST_COLUMN_ID, StColumnFields::TableId, value_eq)?
-            .map(|row| {
-                let row = StColumnRow::try_from(row)?;
-                Ok(row.into())
-            })
+        let mut columns: Vec<ColumnSchema> = iter_st_column_for_table(self, &table_id.into())?
+            .map(|row| Ok(StColumnRow::try_from(row)?.into()))
             .collect::<Result<Vec<_>>>()?;
         columns.sort_by_key(|col| col.col_pos);
 
@@ -149,6 +145,14 @@ pub trait StateView {
     }
 }
 
+/// Returns an iterator over all `st_column` rows for `table_id`.
+pub(crate) fn iter_st_column_for_table<'a>(
+    this: &'a (impl StateView + ?Sized),
+    table_id: &'a AlgebraicValue,
+) -> Result<impl 'a + Iterator<Item = RowRef<'a>>> {
+    this.iter_by_col_eq(ST_COLUMN_ID, StColumnFields::TableId, table_id)
+}
+
 pub struct IterMutTx<'a> {
     tx_state_ins: Option<(&'a Table, &'a HashMapBlobStore)>,
     stage: ScanStage<'a>,

smoketests/tests/auto_migration.py

@@ -4,23 +4,27 @@
 
 
 class AddTableAutoMigration(Smoketest):
-    MODULE_CODE = """
+    MODULE_CODE_INIT = """
 use spacetimedb::{log, ReducerContext, Table, SpacetimeType};
+use PersonKind::*;
 
 #[spacetimedb::table(name = person, public)]
 pub struct Person {
     name: String,
+    kind: PersonKind,
 }
 
 #[spacetimedb::reducer]
-pub fn add_person(ctx: &ReducerContext, name: String) {
-    ctx.db.person().insert(Person { name });
+pub fn add_person(ctx: &ReducerContext, name: String, kind: String) {
+    let kind = kind_from_string(kind);
+    ctx.db.person().insert(Person { name, kind });
 }
 
 #[spacetimedb::reducer]
 pub fn print_persons(ctx: &ReducerContext, prefix: String) {
     for person in ctx.db.person().iter() {
-        log::info!("{}: {}", prefix, person.name);
+        let kind = kind_to_string(person.kind);
+        log::info!("{prefix}: {} - {kind}", person.name);
     }
 }
 
@@ -39,11 +43,47 @@ class AddTableAutoMigration(Smoketest):
 
 #[spacetimedb::client_visibility_filter]
 const PERSON_VISIBLE: spacetimedb::Filter = spacetimedb::Filter::Sql("SELECT * FROM person");
+"""
+
+    MODULE_CODE = MODULE_CODE_INIT + """
+#[derive(SpacetimeType, Clone, Copy, PartialEq, Eq)]
+pub enum PersonKind {
+    Student,
+}
+
+fn kind_from_string(_: String) -> PersonKind {
+    Student
+}
+
+fn kind_to_string(Student: PersonKind) -> &'static str {
+    "Student"
+}
 """
 
     MODULE_CODE_UPDATED = (
-        MODULE_CODE
+        MODULE_CODE_INIT
         + """
+#[derive(SpacetimeType, Clone, Copy, PartialEq, Eq)]
+pub enum PersonKind {
+    Student,
+    Professor,
+}
+
+fn kind_from_string(kind: String) -> PersonKind {
+    match &*kind {
+        "Student" => Student,
+        "Professor" => Professor,
+        _ => panic!(),
+    }
+}
+
+fn kind_to_string(kind: PersonKind) -> &'static str {
+    match kind {
+        Student => "Student",
+        Professor => "Professor",
+    }
+}
+
 #[spacetimedb::table(name = book, public)]
 pub struct Book {
     isbn: String,
@@ -89,14 +129,14 @@ def test_add_table_auto_migration(self):
         logging.info("Initial publish complete")
         # initial module code is already published by test framework
 
-        self.call("add_person", "Robert")
-        self.call("add_person", "Julie")
-        self.call("add_person", "Samantha")
+        self.call("add_person", "Robert", "Student")
+        self.call("add_person", "Julie", "Student")
+        self.call("add_person", "Samantha", "Student")
         self.call("print_persons", "BEFORE")
         logs = self.logs(100)
-        self.assertIn("BEFORE: Samantha", logs)
-        self.assertIn("BEFORE: Julie", logs)
-        self.assertIn("BEFORE: Robert", logs)
+        self.assertIn("BEFORE: Samantha - Student", logs)
+        self.assertIn("BEFORE: Julie - Student", logs)
+        self.assertIn("BEFORE: Robert - Student", logs)
 
         logging.info(
             "Initial operations complete, updating module without clear",
@@ -120,16 +160,16 @@ def test_add_table_auto_migration(self):
 
         self.logs(100)
 
-        self.call("add_person", "Husserl")
+        self.call("add_person", "Husserl", "Professor")
         self.call("add_book", "1234567890")
         self.call("print_persons", "AFTER_PERSON")
         self.call("print_books", "AFTER_BOOK")
 
         logs = self.logs(100)
-        self.assertIn("AFTER_PERSON: Samantha", logs)
-        self.assertIn("AFTER_PERSON: Julie", logs)
-        self.assertIn("AFTER_PERSON: Robert", logs)
-        self.assertIn("AFTER_PERSON: Husserl", logs)
+        self.assertIn("AFTER_PERSON: Samantha - Student", logs)
+        self.assertIn("AFTER_PERSON: Julie - Student", logs)
+        self.assertIn("AFTER_PERSON: Robert - Student", logs)
+        self.assertIn("AFTER_PERSON: Husserl - Professor", logs)
         self.assertIn("AFTER_BOOK: 1234567890", logs)