Enable adding columns in auto migrations in schema

Author: kazimuth

crates/core/src/db/update.rs

@@ -252,6 +252,7 @@ fn auto_migrate_database(
                 log!(logger, "Removing-row level security `{sql_rls}`");
                 stdb.drop_row_level_security(tx, sql_rls.clone())?;
             }
+            _ => unimplemented!(),
         }
     }
 

crates/lib/src/db/raw_def/v9.rs

@@ -14,6 +14,7 @@ use spacetimedb_primitives::*;
 use spacetimedb_sats::typespace::TypespaceBuilder;
 use spacetimedb_sats::AlgebraicType;
 use spacetimedb_sats::AlgebraicTypeRef;
+use spacetimedb_sats::AlgebraicValue;
 use spacetimedb_sats::ProductType;
 use spacetimedb_sats::ProductTypeElement;
 use spacetimedb_sats::SpacetimeType;
@@ -361,7 +362,25 @@ pub struct RawRowLevelSecurityDefV9 {
 #[sats(crate = crate)]
 #[cfg_attr(feature = "test", derive(PartialEq, Eq, PartialOrd, Ord))]
 #[non_exhaustive]
-pub enum RawMiscModuleExportV9 {}
+pub enum RawMiscModuleExportV9 {
+    ColumnDefaultValue(RawColumnDefaultValueV9),
+}
+
+/// Marks a field as having a particularly default
+#[derive(Debug, Clone, SpacetimeType)]
+#[sats(crate = crate)]
+#[cfg_attr(feature = "test", derive(PartialEq, Eq, PartialOrd, Ord))]
+pub struct RawColumnDefaultValueV9 {
+    /// Must point to a `ProductType` or an error will be thrown during validation.
+    pub table: RawIdentifier,
+
+    /// Must be the index of a valid column within `ty`.
+    pub col_id: ColId,
+
+    /// A BSATN-encoded AlgebraicValue valid at `ty`.
+    /// (We can't use AlgebraicValue directly because it isn't serializable!)
+    pub value: Box<[u8]>,
+}
 
 /// A type declaration.
 ///
@@ -794,6 +813,20 @@ impl RawTableDefBuilder<'_> {
         self
     }
 
+    /// Adds a default value for the field.
+    /// Note: this sets the value for all other tables defined on the same
+    pub fn with_default_column_value(self, column: impl Into<ColId>, value: AlgebraicValue) -> Self {
+        // Added to misc_exports for backwards-compatibility reasons
+        self.module_def
+            .misc_exports
+            .push(RawMiscModuleExportV9::ColumnDefaultValue(RawColumnDefaultValueV9 {
+                table: self.table.name.clone(),
+                col_id: column.into(),
+                value: spacetimedb_sats::bsatn::to_vec(&value).unwrap().into(),
+            }));
+        self
+    }
+
     /// Build the table and add it to the module, returning the `product_type_ref` of the table.
     pub fn finish(self) -> AlgebraicTypeRef {
         self.table.product_type_ref

crates/sats/src/buffer.rs

@@ -10,7 +10,7 @@ use core::fmt;
 use core::str::Utf8Error;
 
 /// An error that occurred when decoding.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub enum DecodeError {
     /// Not enough data was provided in the input.
     BufferLength {

crates/schema/src/auto_migrate.rs

@@ -109,6 +109,15 @@ pub enum AutoMigrateStep<'def> {
     ///
     /// This should be done before any new indices are added.
     ChangeColumns(<TableDef as ModuleDefLookup>::Key<'def>),
+    /// Add columns to a table, in a layout-INCOMPATIBLE way.
+    ///
+    /// This is a destructive operation that requires first running a `DisconnectAllUsers`.
+    ///
+    /// The added columns are guaranteed to be contiguous and at the end of the table. They are also
+    /// guaranteed to have default values set.
+    ///
+    /// This suppresses any `ChangeColumns` steps for the same table.
+    AddColumns(<TableDef as ModuleDefLookup>::Key<'def>),
 
     /// Add a table, including all indexes, constraints, and sequences.
     /// There will NOT be separate steps in the plan for adding indexes, constraints, and sequences.
@@ -124,6 +133,9 @@ pub enum AutoMigrateStep<'def> {
 
     /// Change the access of a table.
     ChangeAccess(<TableDef as ModuleDefLookup>::Key<'def>),
+
+    /// Disconnect all users connected to the module.
+    DisconnectAllUsers,
 }
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
@@ -137,7 +149,7 @@ pub struct ChangeColumnTypeParts {
 /// Something that might prevent an automatic migration.
 #[derive(thiserror::Error, Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub enum AutoMigrateError {
-    #[error("Adding a column {column} to table {table} requires a manual migration")]
+    #[error("Adding a column {column} to table {table} requires a default value annotation")]
     AddColumn { table: Identifier, column: Identifier },
 
     #[error("Removing a column {column} from table {table} requires a manual migration")]
@@ -386,11 +398,18 @@ fn auto_migrate_table<'def>(plan: &mut AutoMigratePlan<'def>, old: &'def TableDe
     })
     .map(|col_diff| -> Result<_> {
         match col_diff {
-            Diff::Add { new } => Err(AutoMigrateError::AddColumn {
-                table: new.table_name.clone(),
-                column: new.name.clone(),
+            Diff::Add { new } => {
+                if new.default_value.is_some() {
+                    // row_type_changed, columns_added
+                    Ok(ProductMonoid(Any(false), Any(true)))
+                } else {
+                    Err(AutoMigrateError::AddColumn {
+                        table: new.table_name.clone(),
+                        column: new.name.clone(),
+                    }
+                    .into())
+                }
             }
-            .into()),
             Diff::Remove { old } => Err(AutoMigrateError::RemoveColumn {
                 table: old.table_name.clone(),
                 column: old.name.clone(),
@@ -408,6 +427,7 @@ fn auto_migrate_table<'def>(plan: &mut AutoMigratePlan<'def>, old: &'def TableDe
 
                 // Note that the diff algorithm relies on `ModuleDefLookup` for `ColumnDef`,
                 // which looks up columns by NAME, NOT position: precisely to allow this step to work!
+                // We reject changes to
                 let positions_ok = if old.col_id == new.col_id {
                     Ok(())
                 } else {
@@ -417,22 +437,34 @@ fn auto_migrate_table<'def>(plan: &mut AutoMigratePlan<'def>, old: &'def TableDe
                     .into())
                 };
 
-                (types_ok, positions_ok).combine_errors().map(|(x, _)| x)
+                (types_ok, positions_ok)
+                    .combine_errors()
+                    .map(|(x, _)| ProductMonoid(x, Any(false)))
             }
         }
     })
-    .collect_all_errors::<Any>();
+    .collect_all_errors::<ProductMonoid<Any, Any>>();
 
-    let ((), Any(row_type_changed)) = (type_ok, columns_ok).combine_errors()?;
+    let ((), ProductMonoid(Any(row_type_changed), Any(columns_added))) = (type_ok, columns_ok).combine_errors()?;
 
-    if row_type_changed {
+    if columns_added {
+        if !plan
+            .steps
+            .iter()
+            .any(|step| matches!(step, AutoMigrateStep::DisconnectAllUsers))
+        {
+            plan.steps.push(AutoMigrateStep::DisconnectAllUsers);
+        }
+        plan.steps.push(AutoMigrateStep::AddColumns(key));
+    } else if row_type_changed {
         plan.steps.push(AutoMigrateStep::ChangeColumns(key));
     }
 
     Ok(())
 }
 
 /// An "any" monoid with `false` as identity and `|` as the operator.
+#[derive(Default)]
 struct Any(bool);
 
 impl FromIterator<Any> for Any {
@@ -448,6 +480,26 @@ impl BitOr for Any {
     }
 }
 
+/// A monoid that allows running two `Any`s in parallel.
+#[derive(Default)]
+struct ProductMonoid<M1, M2>(M1, M2);
+
+impl<M1: BitOr<Output = M1>, M2: BitOr<Output = M2>> BitOr for ProductMonoid<M1, M2> {
+    type Output = Self;
+
+    fn bitor(self, rhs: Self) -> Self::Output {
+        Self(self.0 | rhs.0, self.1 | rhs.1)
+    }
+}
+
+impl<M1: BitOr<Output = M1> + Default, M2: BitOr<Output = M2> + Default> FromIterator<ProductMonoid<M1, M2>>
+    for ProductMonoid<M1, M2>
+{
+    fn from_iter<T: IntoIterator<Item = ProductMonoid<M1, M2>>>(iter: T) -> Self {
+        iter.into_iter().reduce(|p1, p2| p1 | p2).unwrap_or_default()
+    }
+}
+
 fn ensure_old_ty_upgradable_to_new(
     within: bool,
     old: &ColumnDef,
@@ -700,7 +752,7 @@ mod tests {
     use spacetimedb_data_structures::expect_error_matching;
     use spacetimedb_lib::{
         db::raw_def::{v9::btree, *},
-        AlgebraicType, ProductType, ScheduleAt,
+        AlgebraicType, AlgebraicValue, ProductType, ScheduleAt,
     };
     use spacetimedb_primitives::ColId;
     use v9::{RawModuleDefV9Builder, TableAccess};
@@ -813,11 +865,13 @@ mod tests {
                     ("id", AlgebraicType::U64),
                     ("name", AlgebraicType::String),
                     ("count", AlgebraicType::U16),
+                    ("freshness", AlgebraicType::U32), // added column!
                 ]),
                 true,
             )
             // add column sequence
             .with_column_sequence(0)
+            .with_default_column_value(3, AlgebraicValue::U32(5))
             // change access
             .with_access(TableAccess::Private)
             .finish();
@@ -963,6 +1017,9 @@ mod tests {
             steps.contains(&AutoMigrateStep::ChangeColumns(&deliveries)),
             "{steps:?}"
         );
+
+        assert!(steps.contains(&AutoMigrateStep::DisconnectAllUsers), "{steps:?}");
+        assert!(steps.contains(&AutoMigrateStep::AddColumns(&bananas)), "{steps:?}");
     }
 
     #[test]

crates/schema/src/def.rs

@@ -32,13 +32,14 @@ use spacetimedb_data_structures::error_stream::{CollectAllErrors, CombineErrors,
 use spacetimedb_data_structures::map::HashMap;
 use spacetimedb_lib::db::raw_def;
 use spacetimedb_lib::db::raw_def::v9::{
-    Lifecycle, RawConstraintDataV9, RawConstraintDefV9, RawIdentifier, RawIndexAlgorithm, RawIndexDefV9,
-    RawModuleDefV9, RawReducerDefV9, RawRowLevelSecurityDefV9, RawScheduleDefV9, RawScopedTypeNameV9, RawSequenceDefV9,
-    RawSql, RawTableDefV9, RawTypeDefV9, RawUniqueConstraintDataV9, TableAccess, TableType,
+    Lifecycle, RawColumnDefaultValueV9, RawConstraintDataV9, RawConstraintDefV9, RawIdentifier, RawIndexAlgorithm,
+    RawIndexDefV9, RawMiscModuleExportV9, RawModuleDefV9, RawReducerDefV9, RawRowLevelSecurityDefV9, RawScheduleDefV9,
+    RawScopedTypeNameV9, RawSequenceDefV9, RawSql, RawTableDefV9, RawTypeDefV9, RawUniqueConstraintDataV9, TableAccess,
+    TableType,
 };
 use spacetimedb_lib::{ProductType, RawModuleDef};
 use spacetimedb_primitives::{ColId, ColList, ColOrCols, ColSet, ReducerId, TableId};
-use spacetimedb_sats::AlgebraicType;
+use spacetimedb_sats::{AlgebraicType, AlgebraicValue};
 use spacetimedb_sats::{AlgebraicTypeRef, Typespace};
 
 pub mod deserialize;
@@ -659,6 +660,9 @@ pub struct ColumnDef {
 
     /// The table this `ColumnDef` is stored in.
     pub table_name: Identifier,
+
+    /// The default value of this column, if present.
+    pub default_value: Option<AlgebraicValue>,
 }
 
 /// A constraint definition attached to a table.