datastore: add columns support (#3230)

# Description of Changes
- `add_columns_to_table` api to handle `AutoMigrateStep::AddColumns` --
Look at doc comment for more Info.

depends on: #3261

TODO: handle `AutoMigrateStep::DisconnectAllUsers`.

# API and ABI breaking changes
N/A.

# Expected complexity level and risk
2? Changes are not in hotpath but a bug in migration can corrupt tables.

# Testing
a test to exercise the API.

---------

Signed-off-by: Mazdak Farrokhzad <twingoow@gmail.com>
Signed-off-by: Shubham Mishra <shivam828787@gmail.com>
Co-authored-by: Mazdak Farrokhzad <twingoow@gmail.com>

Author: Shubham8287

crates/core/src/db/relational_db.rs

@@ -1088,6 +1088,18 @@ impl RelationalDB {
         Ok(self.inner.alter_table_row_type_mut_tx(tx, table_id, column_schemas)?)
     }
 
+    pub(crate) fn add_columns_to_table(
+        &self,
+        tx: &mut MutTx,
+        table_id: TableId,
+        column_schemas: Vec<ColumnSchema>,
+        default_values: Vec<AlgebraicValue>,
+    ) -> Result<TableId, DBError> {
+        Ok(self
+            .inner
+            .add_columns_to_table_mut_tx(tx, table_id, column_schemas, default_values)?)
+    }
+
     /// Reports the `TxMetrics`s passed.
     ///
     /// Should only be called after the tx lock has been fully released.

crates/core/src/db/update.rs

@@ -252,6 +252,18 @@ fn auto_migrate_database(
                 log!(logger, "Removing-row level security `{sql_rls}`");
                 stdb.drop_row_level_security(tx, sql_rls.clone())?;
             }
+            spacetimedb_schema::auto_migrate::AutoMigrateStep::AddColumns(table_name) => {
+                let table_def = plan.new.stored_in_table_def(table_name).expect("table must exist");
+                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);
+
+                let default_values: Vec<AlgebraicValue> = table_def
+                    .columns
+                    .iter()
+                    .filter_map(|col_def| col_def.default_value.clone())
+                    .collect();
+                stdb.add_columns_to_table(tx, table_id, column_schemas, default_values)?;
+            }
             _ => anyhow::bail!("migration step not implemented: {step:?}"),
         }
     }

crates/datastore/src/error.rs

@@ -81,6 +81,8 @@ pub enum TableError {
     #[error(transparent)]
     // Error here is `Box`ed to avoid triggering https://rust-lang.github.io/rust-clippy/master/index.html#result_large_err .
     ChangeColumnsError(#[from] Box<table::ChangeColumnsError>),
+    #[error(transparent)]
+    AddColumnsError(#[from] Box<table::AddColumnsError>),
 }
 
 #[derive(Error, Debug, PartialEq, Eq)]

crates/datastore/src/locking_tx_datastore/committed_state.rs

@@ -706,6 +706,7 @@ impl CommittedState {
         for change in pending_schema_changes {
             if let PendingSchemaChange::TableRemoved(table_id, mut table) = change {
                 let row_ptrs = table.scan_all_row_ptrs();
+                truncates.insert(table_id);
                 delete_rows(
                     tx_data,
                     table_id,

crates/datastore/src/locking_tx_datastore/datastore.rs

@@ -335,6 +335,16 @@ impl Locking {
     ) -> Result<()> {
         tx.alter_table_row_type(table_id, column_schemas)
     }
+
+    pub fn add_columns_to_table_mut_tx(
+        &self,
+        tx: &mut MutTxId,
+        table_id: TableId,
+        column_schemas: Vec<ColumnSchema>,
+        defaults: Vec<AlgebraicValue>,
+    ) -> Result<TableId> {
+        tx.add_columns_to_table(table_id, column_schemas, defaults)
+    }
 }
 
 impl DataRow for Locking {
@@ -1264,6 +1274,7 @@ mod tests {
     use core::{fmt, mem};
     use itertools::Itertools;
     use pretty_assertions::{assert_eq, assert_matches};
+    use spacetimedb_execution::dml::MutDatastore as _;
     use spacetimedb_execution::Datastore;
     use spacetimedb_lib::db::auth::{StAccess, StTableType};
     use spacetimedb_lib::error::ResultTest;
@@ -3228,11 +3239,12 @@ mod tests {
         // Now drop the table again and commit.
         assert!(datastore.drop_table_mut_tx(&mut tx, table_id).is_ok());
         let tx_data = commit(&datastore, tx)?;
-        let (_, deleted) = tx_data
+        let (_, deleted_rows) = tx_data
             .deletes()
             .find(|(id, _)| **id == table_id)
             .expect("should have deleted rows for `table_id`");
-        assert_eq!(&**deleted, [row]);
+        assert_eq!(&**deleted_rows, [row]);
+        assert!(tx_data.truncates().contains(&table_id), "table should be truncated");
 
         // In the next transaction, the table doesn't exist.
         assert!(
@@ -3453,8 +3465,148 @@ mod tests {
 
         Ok(())
     }
-
     // TODO: Add the following tests
     // - Create a tx that inserts 2000 rows with an auto_inc column
     // - Create a tx that inserts 2000 rows with an auto_inc column and then rolls back
+
+    #[test]
+    fn test_add_columns_to_table() -> ResultTest<()> {
+        let datastore = get_datastore()?;
+
+        let mut tx = begin_mut_tx(&datastore);
+
+        let initial_sum_type = AlgebraicType::sum([("ba", AlgebraicType::U16)]);
+        let initial_columns = [
+            ColumnSchema::for_test(0, "a", AlgebraicType::U64),
+            ColumnSchema::for_test(1, "b", initial_sum_type.clone()),
+        ];
+
+        let initial_indices = [
+            IndexSchema::for_test("index_a", BTreeAlgorithm::from(0)),
+            IndexSchema::for_test("index_b", BTreeAlgorithm::from(1)),
+        ];
+
+        let sequence = SequenceRow {
+            id: SequenceId::SENTINEL.into(),
+            table: 0,
+            col_pos: 0,
+            name: "Foo_id_seq",
+            start: 5,
+        };
+
+        let schema = user_public_table(
+            initial_columns,
+            initial_indices.clone(),
+            [],
+            map_array([sequence]),
+            None,
+            None,
+        );
+
+        let table_id = datastore.create_table_mut_tx(&mut tx, schema)?;
+
+        let mut columns_original = tx.get_schema(table_id).unwrap().columns.to_vec();
+
+        // Insert initial rows
+        let initial_row = product![0u64, AlgebraicValue::sum(0, AlgebraicValue::U16(1))];
+        insert(&datastore, &mut tx, table_id, &initial_row).unwrap();
+        insert(&datastore, &mut tx, table_id, &initial_row).unwrap();
+        commit(&datastore, tx)?;
+
+        // Alter Table: Add Variant and Column
+        //
+        // Change the `b` column, adding a variant.
+        let vars_ref = &mut columns_original[1].col_type.as_sum_mut().unwrap().variants;
+        let mut vars = Vec::from(mem::take(vars_ref));
+        vars.push(SumTypeVariant::new_named(AlgebraicType::U8, "bb"));
+        *vars_ref = vars.into();
+        // Add column `c`
+        let mut new_columns = columns_original.clone();
+        new_columns.push(ColumnSchema::for_test(2, "c", AlgebraicType::U8));
+        let defaults = vec![AlgebraicValue::U8(42)];
+
+        let mut tx = begin_mut_tx(&datastore);
+        // insert a row in tx_state before adding column
+        tx.insert_product_value(table_id, &initial_row).unwrap();
+        // add column and then rollback
+        let rollback_table_id =
+            datastore.add_columns_to_table_mut_tx(&mut tx, table_id, new_columns.clone(), defaults.clone())?;
+        let _ = tx.rollback();
+
+        let old_rows = [
+            product![5u64, AlgebraicValue::sum(0, 1u16.into())],
+            product![6u64, AlgebraicValue::sum(0, 1u16.into())],
+        ];
+
+        let mut tx = begin_mut_tx(&datastore);
+        // check rollback was successful
+        let rows = tx
+            .table_scan(table_id)
+            .unwrap()
+            .map(|row| row.to_product_value())
+            .collect::<Vec<_>>();
+        assert_eq!(rows, old_rows, "Rows shouldn't be changed if rolledback");
+        let table = tx.table(rollback_table_id);
+        assert!(table.is_none(), "new table shouldn't be created if rolledback");
+
+        // Add column and actually commit this time.
+        tx.insert_product_value(table_id, &initial_row).unwrap();
+        let new_table_id = datastore.add_columns_to_table_mut_tx(&mut tx, table_id, new_columns.clone(), defaults)?;
+
+        let tx_data = commit(&datastore, tx)?;
+
+        assert_ne!(
+            new_table_id, table_id,
+            "New table ID after migration should differ from old one"
+        );
+
+        //  Validate Commitlog Changes
+        let (_, deletes) = tx_data
+            .deletes()
+            .find(|(id, _)| **id == table_id)
+            .expect("Expected delete log for original table");
+
+        assert_eq!(
+            &**deletes, &old_rows,
+            "Unexpected delete entries after altering the table"
+        );
+
+        let inserted_rows = [
+            product![5u64, AlgebraicValue::sum(0, 1u16.into()), 42u8],
+            product![6u64, AlgebraicValue::sum(0, 1u16.into()), 42u8],
+            product![8u64, AlgebraicValue::sum(0, 1u16.into()), 42u8],
+        ];
+
+        let (_, inserts) = tx_data
+            .inserts()
+            .find(|(id, _)| **id == new_table_id)
+            .expect("Expected insert log for new table");
+
+        assert_eq!(
+            &**inserts, &inserted_rows,
+            "Unexpected insert entries after altering the table"
+        );
+
+        // Insert Rows into New Table
+        let mut tx = begin_mut_tx(&datastore);
+
+        let new_row = product![0u64, AlgebraicValue::sum(0, 1u16.into()), 0u8];
+        tx.insert_product_value(new_table_id, &new_row).unwrap();
+        commit(&datastore, tx)?;
+
+        // test for auto_inc feields
+        let tx = begin_mut_tx(&datastore);
+        let rows = tx.table_scan(new_table_id).unwrap().map(|row| row.to_product_value());
+
+        let mut last_row_auto_inc = 0;
+        for row in rows {
+            let auto_inc_col = row.get_field(0, None)?;
+            if let AlgebraicValue::U64(val) = auto_inc_col {
+                assert!(val > &last_row_auto_inc, "Auto-increment value did not increase");
+                last_row_auto_inc = *val;
+            }
+        }
+
+        Ok(())
+    }
 }

crates/datastore/src/locking_tx_datastore/mut_tx.rs

@@ -59,6 +59,7 @@ use spacetimedb_table::{
     table_index::TableIndex,
 };
 use std::{
+    collections::HashMap,
     sync::Arc,
     time::{Duration, Instant},
 };
@@ -494,6 +495,134 @@ impl MutTxId {
         Ok(())
     }
 
+    /// Change the row type of the table identified by `table_id`.
+    ///
+    /// This is an incompatible change that requires a new table to be created.
+    /// The existing rows are copied over to the new table,
+    /// with `default_values` appended to each row.
+    ///
+    /// `column_schemas` must contain all the old columns in same order as before,
+    /// followed by the new columns to be added.
+    /// All new columns must have a default value in `default_values`.
+    ///
+    /// After calling this method, Table referred by `table_id` is dropped,
+    /// and a new table is created with the same name but a new `table_id`.
+    /// The new `table_id` is returned.
+    pub(crate) fn add_columns_to_table(
+        &mut self,
+        table_id: TableId,
+        column_schemas: Vec<ColumnSchema>,
+        mut default_values: Vec<AlgebraicValue>,
+    ) -> Result<TableId> {
+        let original_table_schema: TableSchema = (*self.schema_for_table(table_id)?).clone();
+
+        // Only keep the default values of new columns.
+        let new_cols = column_schemas
+            .len()
+            .checked_sub(original_table_schema.columns.len())
+            .ok_or_else(|| {
+                anyhow::anyhow!(
+                    "new column schemas must be more than existing ones for table_id: {}",
+                    table_id
+                )
+            })?;
+        let older_defaults = default_values.len().checked_sub(new_cols).ok_or_else(|| {
+            anyhow::anyhow!(
+                "not enough default values provided for new columns for table_id: {}",
+                table_id
+            )
+        })?;
+        default_values.drain(..older_defaults);
+
+        let ((tx_table, ..), _) = self.get_or_create_insert_table_mut(table_id)?;
+
+        tx_table
+            .validate_add_columns_schema(&column_schemas, &default_values)
+            .map_err(TableError::from)?;
+
+        // Copy all rows as `ProductValue` before dropping the table.
+        // This approach isn't ideal, as allocating a large contiguous block of memory
+        // can lead to memory overflow errors.
+        // Ideally, we should use iterators to avoid this, but that's not sufficient on its own.
+        // `CommittedState::merge_apply_inserts` also allocates a similar `Vec` to hold the data.
+        // So, if we really find this problematic in practice, this should be fixed in both places.
+        let mut table_rows: Vec<ProductValue> = iter(&self.tx_state, &self.committed_state_write_lock, table_id)?
+            .map(|r| r.to_product_value())
+            .collect();
+
+        log::debug!(
+            "ADDING TABLE COLUMN (incompatible layout): {}, table_id: {}",
+            original_table_schema.table_name,
+            table_id
+        );
+
+        // Store sequence values to restore them later with new table.
+        // Using a map from name to value as the new sequence ids will be different.
+        // and I am not sure if we should rely on the order of sequences in the table schema.
+        let seq_values: HashMap<Box<str>, i128> = original_table_schema
+            .sequences
+            .iter()
+            .map(|s| {
+                (
+                    s.sequence_name.clone(),
+                    self.sequence_state_lock
+                        .get_sequence_mut(s.sequence_id)
+                        .expect("sequence exists in original schema and should in sequence state.")
+                        .get_value(),
+                )
+            })
+            .collect();
+
+        // Drop existing table first due to unique constraints on table name in `st_table`
+        self.drop_table(table_id)?;
+
+        // Update existing (dropped) table schema with provided columns, reset Ids.
+        let mut updated_table_schema = original_table_schema;
+        updated_table_schema.columns = column_schemas;
+        updated_table_schema.reset();
+        let new_table_id = self.create_table_and_update_seq(updated_table_schema, seq_values)?;
+
+        // Populate rows with default values for new columns
+        for product_value in table_rows.iter_mut() {
+            let mut row_elements = product_value.elements.to_vec();
+            row_elements.extend(default_values.iter().cloned());
+            product_value.elements = row_elements.into();
+        }
+
+        let (new_table, tx_blob_store) = self
+            .tx_state
+            .get_table_and_blob_store(new_table_id)
+            .ok_or(TableError::IdNotFoundState(new_table_id))?;
+
+        for row in table_rows {
+            new_table.insert(&self.committed_state_write_lock.page_pool, tx_blob_store, &row)?;
+        }
+
+        Ok(new_table_id)
+    }
+
+    fn create_table_and_update_seq(
+        &mut self,
+        table_schema: TableSchema,
+        seq_values: HashMap<Box<str>, i128>,
+    ) -> Result<TableId> {
+        let table_id = self.create_table(table_schema)?;
+        let table_schema = self.schema_for_table(table_id)?;
+
+        for seq in table_schema.sequences.iter() {
+            let new_seq = self
+                .sequence_state_lock
+                .get_sequence_mut(seq.sequence_id)
+                .expect("sequence just created");
+            let value = *seq_values
+                .get(&seq.sequence_name)
+                .ok_or_else(|| SequenceError::NotFound(seq.sequence_id))?;
+            new_seq.update_value(value);
+        }
+
+        Ok(table_id)
+    }
+
     /// Create an index.
     ///
     /// Requires: