fix migration which adds column and index together.

crates/core/src/db/update.rs

@@ -1,7 +1,6 @@
 use super::relational_db::RelationalDB;
 use crate::database_logger::SystemLogger;
 use crate::sql::parser::RowLevelExpr;
-use spacetimedb_data_structures::map::HashMap;
 use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_lib::db::auth::StTableType;
 use spacetimedb_lib::identity::AuthCtx;
@@ -10,7 +9,6 @@ use spacetimedb_primitives::{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};
-use std::sync::Arc;
 
 /// The logger used for by [`update_database`] and friends.
 pub trait UpdateLogger {
@@ -64,8 +62,8 @@ pub fn update_database(
     }
 
     match plan {
-        MigratePlan::Manual(plan) => manual_migrate_database(stdb, tx, plan, logger, existing_tables),
-        MigratePlan::Auto(plan) => auto_migrate_database(stdb, tx, auth_ctx, plan, logger, existing_tables),
+        MigratePlan::Manual(plan) => manual_migrate_database(stdb, tx, plan, logger),
+        MigratePlan::Auto(plan) => auto_migrate_database(stdb, tx, auth_ctx, plan, logger),
     }
 }
 
@@ -75,7 +73,6 @@ fn manual_migrate_database(
     _tx: &mut MutTxId,
     _plan: ManualMigratePlan,
     _logger: &dyn UpdateLogger,
-    _existing_tables: Vec<Arc<TableSchema>>,
 ) -> anyhow::Result<UpdateResult> {
     unimplemented!("Manual database migrations are not yet implemented")
 }
@@ -95,33 +92,23 @@ fn auto_migrate_database(
     auth_ctx: AuthCtx,
     plan: AutoMigratePlan,
     logger: &dyn UpdateLogger,
-    existing_tables: Vec<Arc<TableSchema>>,
 ) -> anyhow::Result<UpdateResult> {
-    // We have already checked in `migrate_database` that `existing_tables` are compatible with the `old` definition in `plan`.
-    // So we can look up tables in there using unwrap.
-
-    let table_schemas_by_name = existing_tables
-        .into_iter()
-        .map(|table| (table.table_name.clone(), table))
-        .collect::<HashMap<_, _>>();
-
     log::info!("Running database update prechecks: {}", stdb.database_identity());
 
     for precheck in plan.prechecks {
         match precheck {
             spacetimedb_schema::auto_migrate::AutoMigratePrecheck::CheckAddSequenceRangeValid(sequence_name) => {
                 let table_def = plan.new.stored_in_table_def(sequence_name).unwrap();
                 let sequence_def = &table_def.sequences[sequence_name];
-
-                let table_schema = &table_schemas_by_name[&table_def.name[..]];
+                let table_id = stdb.table_id_from_name_mut(tx, &table_def.name)?.unwrap();
 
                 let min: AlgebraicValue = sequence_def.min_value.unwrap_or(1).into();
                 let max: AlgebraicValue = sequence_def.max_value.unwrap_or(i128::MAX).into();
 
                 let range = min..max;
 
                 if stdb
-                    .iter_by_col_range_mut(tx, table_schema.table_id, sequence_def.column, range)?
+                    .iter_by_col_range_mut(tx, table_id, sequence_def.column, range)?
                     .next()
                     .is_some()
                 {
@@ -150,7 +137,7 @@ fn auto_migrate_database(
             spacetimedb_schema::auto_migrate::AutoMigrateStep::AddIndex(index_name) => {
                 let table_def = plan.new.stored_in_table_def(index_name).unwrap();
                 let index_def = table_def.indexes.get(index_name).unwrap();
-                let table_id = table_schemas_by_name[&table_def.name[..]].table_id;
+                let table_id = stdb.table_id_from_name_mut(tx, &table_def.name)?.unwrap();
 
                 let index_cols = ColSet::from(index_def.algorithm.columns());
 
@@ -169,7 +156,9 @@ fn auto_migrate_database(
             spacetimedb_schema::auto_migrate::AutoMigrateStep::RemoveIndex(index_name) => {
                 let table_def = plan.old.stored_in_table_def(index_name).unwrap();
 
-                let table_schema = &table_schemas_by_name[&table_def.name[..]];
+                let table_id = stdb.table_id_from_name_mut(tx, &table_def.name)?.unwrap();
+                let table_schema = stdb.schema_for_table_mut(tx, table_id)?;
+
                 let index_schema = table_schema
                     .indexes
                     .iter()
@@ -181,7 +170,9 @@ fn auto_migrate_database(
             }
             spacetimedb_schema::auto_migrate::AutoMigrateStep::RemoveConstraint(constraint_name) => {
                 let table_def = plan.old.stored_in_table_def(constraint_name).unwrap();
-                let table_schema = &table_schemas_by_name[&table_def.name[..]];
+
+                let table_id = stdb.table_id_from_name_mut(tx, &table_def.name)?.unwrap();
+                let table_schema = stdb.schema_for_table_mut(tx, table_id)?;
                 let constraint_schema = table_schema
                     .constraints
                     .iter()
@@ -199,7 +190,9 @@ fn auto_migrate_database(
             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();
-                let table_schema = &table_schemas_by_name[&table_def.name[..]];
+
+                let table_id = stdb.table_id_from_name_mut(tx, &table_def.name)?.unwrap();
+                let table_schema = stdb.schema_for_table_mut(tx, table_id)?;
 
                 log!(
                     logger,
@@ -213,7 +206,9 @@ fn auto_migrate_database(
             }
             spacetimedb_schema::auto_migrate::AutoMigrateStep::RemoveSequence(sequence_name) => {
                 let table_def = plan.old.stored_in_table_def(sequence_name).unwrap();
-                let table_schema = &table_schemas_by_name[&table_def.name[..]];
+
+                let table_id = stdb.table_id_from_name_mut(tx, &table_def.name)?.unwrap();
+                let table_schema = stdb.schema_for_table_mut(tx, table_id)?;
                 let sequence_schema = table_schema
                     .sequences
                     .iter()

smoketests/tests/auto_migration.py

@@ -265,6 +265,7 @@ class AddTableColumns(Smoketest):
 #[derive(Debug)]
 #[spacetimedb::table(name = person)]
 pub struct Person {
+    #[index(btree)]
     name: String,
     #[default(0)]
     age: u16,