Add columns to table and test

Author: Shubham8287

crates/core/src/db/relational_db.rs

@@ -10,6 +10,7 @@ use futures::channel::mpsc;
 use futures::StreamExt;
 use parking_lot::RwLock;
 use spacetimedb_commitlog as commitlog;
+use spacetimedb_data_structures::map::IntMap;
 use spacetimedb_datastore::db_metrics::DB_METRICS;
 use spacetimedb_datastore::error::{DatastoreError, TableError};
 use spacetimedb_datastore::execution_context::{ReducerContext, Workload, WorkloadType};
@@ -1052,6 +1053,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: IntMap<ColId, 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

@@ -1,12 +1,12 @@
 use super::relational_db::RelationalDB;
 use crate::database_logger::SystemLogger;
 use crate::sql::parser::RowLevelExpr;
-use spacetimedb_data_structures::map::HashMap;
+use spacetimedb_data_structures::map::{HashMap, IntMap};
 use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_lib::db::auth::StTableType;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::AlgebraicValue;
-use spacetimedb_primitives::{ColSet, TableId};
+use spacetimedb_primitives::{ColId, ColSet, TableId};
 use spacetimedb_schema::auto_migrate::{AutoMigratePlan, ManualMigratePlan, MigratePlan};
 use spacetimedb_schema::def::TableDef;
 use spacetimedb_schema::schema::{column_schemas_from_defs, IndexSchema, Schema, SequenceSchema, TableSchema};
@@ -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: IntMap<ColId, AlgebraicValue> = table_def
+                    .columns
+                    .iter()
+                    .filter_map(|col_def| col_def.default_value.as_ref().map(|v| (col_def.col_id, v.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/datastore.rs

@@ -27,12 +27,12 @@ use anyhow::{anyhow, Context};
 use core::{cell::RefCell, ops::RangeBounds};
 use parking_lot::{Mutex, RwLock};
 use spacetimedb_commitlog::payload::{txdata, Txdata};
-use spacetimedb_data_structures::map::{HashCollectionExt, HashMap};
+use spacetimedb_data_structures::map::{HashCollectionExt, HashMap, IntMap};
 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};
+use spacetimedb_primitives::{ColId, ColList, ConstraintId, IndexId, SequenceId, TableId};
 use spacetimedb_sats::memory_usage::MemoryUsage;
 use spacetimedb_sats::{bsatn, buffer::BufReader, AlgebraicValue, ProductValue};
 use spacetimedb_schema::schema::{ColumnSchema, IndexSchema, SequenceSchema, TableSchema};
@@ -313,6 +313,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: IntMap<ColId, AlgebraicValue>,
+    ) -> Result<TableId> {
+        tx.add_columns_to_table(table_id, column_schemas, defaults)
+    }
 }
 
 impl DataRow for Locking {
@@ -1159,6 +1169,8 @@ mod tests {
     use core::{fmt, mem};
     use itertools::Itertools;
     use pretty_assertions::{assert_eq, assert_matches};
+    use spacetimedb_data_structures::map::IntMap;
+    use spacetimedb_execution::dml::MutDatastore as _;
     use spacetimedb_execution::Datastore;
     use spacetimedb_lib::db::auth::{StAccess, StTableType};
     use spacetimedb_lib::error::ResultTest;
@@ -3320,8 +3332,138 @@ 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 columns_original = tx.get_schema(table_id).unwrap().columns.to_vec();
+
+        // Insert initial rows
+        let initial_row = ProductValue {
+            elements: [AlgebraicValue::U64(0), AlgebraicValue::sum(0, AlgebraicValue::U16(1))].into(),
+        };
+        tx.insert_product_value(table_id, &initial_row).unwrap();
+        tx.insert_product_value(table_id, &initial_row).unwrap();
+
+        commit(&datastore, tx)?;
+
+        // Alter Table: Add Variant and Column
+        let mut new_columns = columns_original.clone();
+        new_columns.push(ColumnSchema::for_test(2, "c", AlgebraicType::U8));
+        let defaults = IntMap::from_iter([(2.into(), AlgebraicValue::U8(42))]);
+
+        let mut tx = begin_mut_tx(&datastore);
+        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");
+
+        let deleted_rows = [
+            product![AlgebraicValue::U64(5), AlgebraicValue::sum(0, AlgebraicValue::U16(1))],
+            product![AlgebraicValue::U64(6), AlgebraicValue::sum(0, AlgebraicValue::U16(1))],
+        ];
+
+        assert_eq!(
+            &**deletes, &deleted_rows,
+            "Unexpected delete entries after altering the table"
+        );
+
+        let inserted_rows = [
+            product![
+                AlgebraicValue::U64(5),
+                AlgebraicValue::sum(0, AlgebraicValue::U16(1)),
+                AlgebraicValue::U8(42)
+            ],
+            product![
+                AlgebraicValue::U64(6),
+                AlgebraicValue::sum(0, AlgebraicValue::U16(1)),
+                AlgebraicValue::U8(42)
+            ],
+        ];
+
+        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 Altered Table
+        let mut tx = begin_mut_tx(&datastore);
+
+        let new_row = ProductValue {
+            elements: [
+                AlgebraicValue::U64(0),
+                AlgebraicValue::sum(0, AlgebraicValue::U16(1)),
+                AlgebraicValue::U8(0),
+            ]
+            .into(),
+        };
+
+        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

@@ -8,9 +8,9 @@ use super::{
     tx_state::{IndexIdMap, PendingSchemaChange, TxState, TxTableForInsertion},
     SharedMutexGuard, SharedWriteGuard,
 };
-use crate::execution_context::ExecutionContext;
 use crate::execution_context::Workload;
 use crate::traits::{InsertFlags, RowTypeForTable, TxData, UpdateFlags};
+use crate::{error::DatastoreError, execution_context::ExecutionContext};
 use crate::{
     error::{IndexError, SequenceError, TableError},
     system_tables::{
@@ -24,7 +24,9 @@ use crate::{
 use core::ops::RangeBounds;
 use core::{cell::RefCell, mem};
 use core::{iter, ops::Bound};
+use itertools::Itertools;
 use smallvec::SmallVec;
+use spacetimedb_data_structures::map::IntMap;
 use spacetimedb_execution::{dml::MutDatastore, Datastore, DeltaStore, Row};
 use spacetimedb_lib::{db::raw_def::v9::RawSql, metrics::ExecutionMetrics};
 use spacetimedb_lib::{
@@ -493,6 +495,115 @@ 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>,
+        default_values: IntMap<ColId, AlgebraicValue>,
+    ) -> Result<TableId> {
+        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)?;
+
+        // Intentionally using `schema_for_table_raw` instead of `schema_for_table`.
+        //
+        // Rationale:
+        // - `schema_for_table_raw` queries system tables directly, ensuring the
+        //   most up-to-date schema.
+        // - `schema_for_table` relies on an in-memory cache, which may not reflect
+        //   the latest sequence updates and can therefore return stale schemas.
+        let original_table_schema = self.schema_for_table_raw(table_id)?;
+
+        // Default values must be for a sequence of contiguous column IDs,
+        // beginning immediately after the last column ID in the existing schema.
+        // This has been checked inside `Table::validate_add_columns_schema`
+        let default_values: Vec<AlgebraicValue> = default_values
+            .into_iter()
+            .sorted_by_key(|(col_id, _)| *col_id)
+            .map(|(_, val)| val)
+            .collect();
+
+        log::debug!(
+            "TABLE ALTERING (incompatible layout): {}, table_id: {}",
+            original_table_schema.table_name,
+            table_id
+        );
+
+        // Drop existing table first due to unique constraints on table name in `st_table`
+        self.drop_table(table_id)?;
+        let original_commit_table = match self.tx_state.pending_schema_changes.pop() {
+            Some(PendingSchemaChange::TableRemoved(tid, table)) if tid == table_id => table,
+            _ => {
+                return Err(DatastoreError::Other(anyhow::anyhow!(
+                    "Expected last pending schema change to be dropping 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_bump_seq(updated_table_schema)?;
+
+        // Populate rows with default values for new columns
+        let committed_blob_store = &self.committed_state_write_lock.blob_store;
+        let migrated_rows = original_commit_table
+            .scan_rows(committed_blob_store)
+            .map(|table_row| table_row.to_product_value())
+            .map(|mut product_value| {
+                let mut row_elements = product_value.elements.to_vec();
+                row_elements.extend(default_values.iter().cloned());
+                product_value.elements = row_elements.into();
+                product_value
+            });
+
+        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 migrated_row in migrated_rows {
+            new_table.insert(&self.committed_state_write_lock.page_pool, tx_blob_store, &migrated_row)?;
+        }
+
+        self.push_schema_change(PendingSchemaChange::TableRemoved(table_id, original_commit_table));
+        self.push_schema_change(PendingSchemaChange::TableAdded(new_table_id));
+
+        Ok(new_table_id)
+    }
+
+    fn create_table_and_bump_seq(&mut self, table_schema: TableSchema) -> Result<TableId> {
+        let table_id = self.create_table(table_schema.clone())?;
+        let table_schema = self.schema_for_table(table_id)?;
+
+        // Bump each sequence value so that the next generated value requires
+        // a fresh allocation.
+        for seq_id in table_schema.sequences.iter().map(|s| s.sequence_id) {
+            let seq = self
+                .sequence_state_lock
+                .get_sequence_mut(seq_id)
+                .expect("sequence just created");
+            seq.value = seq.allocated() + 1;
+        }
+
+        Ok(table_id)
+    }
+
     /// Create an index.
     ///
     /// Requires: