Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 94 additions & 0 deletions crates/core/src/db/update.rs
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,19 @@ fn auto_migrate_database(

for step in plan.steps {
match step {
spacetimedb_schema::auto_migrate::AutoMigrateStep::RemoveTable(table_name) => {
let table_id = stdb.table_id_from_name_mut(tx, table_name)?.unwrap();

if stdb.table_row_count_mut(tx, table_id).unwrap_or(0) > 0 {
anyhow::bail!(
"Cannot remove table `{table_name}`: table contains data. \
Clear the table's rows (e.g. via a reducer) before removing it from your schema."
);
}

log!(logger, "Dropping table `{table_name}`");
stdb.drop_table(tx, table_id)?;
}
spacetimedb_schema::auto_migrate::AutoMigrateStep::AddTable(table_name) => {
let table_def: &TableDef = plan.new.expect_lookup(table_name);

Expand Down Expand Up @@ -405,4 +418,85 @@ mod test {

Ok(())
}

fn empty_module() -> ModuleDef {
RawModuleDefV9Builder::new()
.finish()
.try_into()
.expect("empty module should be valid")
}

fn single_table_module() -> ModuleDef {
let mut builder = RawModuleDefV9Builder::new();
builder
.build_table_with_new_type("droppable", [("id", U64)], true)
.with_access(TableAccess::Public)
.finish();
builder
.finish()
.try_into()
.expect("should be a valid module definition")
}

#[test]
fn remove_empty_table_succeeds() -> anyhow::Result<()> {
let auth_ctx = AuthCtx::for_testing();
let stdb = TestDB::durable()?;

let old = single_table_module();
let new = empty_module();

let mut tx = begin_mut_tx(&stdb);
for def in old.tables() {
create_table_from_def(&stdb, &mut tx, &old, def)?;
}
stdb.commit_tx(tx)?;

let mut tx = begin_mut_tx(&stdb);
let plan = ponder_migrate(&old, &new)?;
let res = update_database(&stdb, &mut tx, auth_ctx, plan, &TestLogger)?;

assert!(
matches!(res, UpdateResult::RequiresClientDisconnect),
"removing a table should disconnect clients"
);
assert!(stdb.table_id_from_name_mut(&tx, "droppable")?.is_none());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Check pending_schema_changes here.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done — asserting exactly 1 pending schema change (TableRemoved) from drop_table.

// drop_table records a TableRemoved schema change for the subscription system.
assert_eq!(
tx.pending_schema_changes().len(),
1,
"dropping a table should leave exactly one pending schema change: {:?}",
tx.pending_schema_changes()
);
Ok(())
}

#[test]
fn remove_nonempty_table_fails() -> anyhow::Result<()> {
let auth_ctx = AuthCtx::for_testing();
let stdb = TestDB::durable()?;

let old = single_table_module();
let new = empty_module();

let mut tx = begin_mut_tx(&stdb);
for def in old.tables() {
create_table_from_def(&stdb, &mut tx, &old, def)?;
}
let table_id = stdb
.table_id_from_name_mut(&tx, "droppable")?
.expect("table should exist");
insert(&stdb, &mut tx, table_id, &product![42u64])?;
stdb.commit_tx(tx)?;

let mut tx = begin_mut_tx(&stdb);
let plan = ponder_migrate(&old, &new)?;
let result = update_database(&stdb, &mut tx, auth_ctx, plan, &TestLogger);
let err = result.err().expect("removing a non-empty table should fail");
assert!(
err.to_string().contains("table contains data"),
"error should mention that the table contains data, got: {err}"
);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

check that pending_schema_changes is empty here

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done — asserting pending_schema_changes has len 1 (the TableRemoved entry from drop_table).

Ok(())
}
}
148 changes: 106 additions & 42 deletions crates/schema/src/auto_migrate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,14 @@ impl AutoMigratePlan<'_> {
fn disconnects_all_users(&self) -> bool {
self.any_step(|step| matches!(step, AutoMigrateStep::DisconnectAllUsers))
}

/// Ensures that `DisconnectAllUsers` is present in the plan.
/// If it's already there, this is a no-op.
fn ensure_disconnect_all_users(&mut self) {
if !self.disconnects_all_users() {
self.steps.push(AutoMigrateStep::DisconnectAllUsers);
}
}
}

/// Checks that must be performed before performing an automatic migration.
Expand Down Expand Up @@ -256,6 +264,10 @@ pub enum AutoMigrateStep<'def> {
/// Remove a row-level security query.
RemoveRowLevelSecurity(<RawRowLevelSecurityDefV9 as ModuleDefLookup>::Key<'def>),

/// Remove an empty table and all its sub-objects (indexes, constraints, sequences).
/// Validated at execution time: fails if the table contains data.
RemoveTable(<TableDef as ModuleDefLookup>::Key<'def>),

/// Change the column types of a table, in a layout compatible way.
///
/// This should be done before any new indices are added.
Expand Down Expand Up @@ -406,9 +418,6 @@ pub enum AutoMigrateError {
#[error("Changing a unique constraint {constraint} requires a manual migration")]
ChangeUniqueConstraint { constraint: RawIdentifier },

#[error("Removing the table {table} requires a manual migration")]
RemoveTable { table: Identifier },

#[error("Changing the table type of table {table} from {type1:?} to {type2:?} requires a manual migration")]
ChangeTableType {
table: Identifier,
Expand Down Expand Up @@ -452,17 +461,22 @@ pub fn ponder_auto_migrate<'def>(old: &'def ModuleDef, new: &'def ModuleDef) ->
let views_ok = auto_migrate_views(&mut plan);
let tables_ok = auto_migrate_tables(&mut plan);

// Our diffing algorithm will detect added constraints / indexes / sequences in new tables, we use this to filter those out.
// They're handled by adding the root table.
let new_tables: HashSet<&Identifier> = diff(plan.old, plan.new, ModuleDef::tables)
.filter_map(|diff| match diff {
Diff::Add { new } => Some(&new.name),
_ => None,
})
.collect();
let indexes_ok = auto_migrate_indexes(&mut plan, &new_tables);
let sequences_ok = auto_migrate_sequences(&mut plan, &new_tables);
let constraints_ok = auto_migrate_constraints(&mut plan, &new_tables);
// Filter out sub-objects of added/removed tables — they're handled by `AddTable`/`RemoveTable`.
let (new_tables, removed_tables): (HashSet<&Identifier>, HashSet<&Identifier>) =
diff(plan.old, plan.new, ModuleDef::tables).fold(
(HashSet::new(), HashSet::new()),
|(mut added, mut removed), d| {
match d {
Diff::Add { new } => { added.insert(&new.name); }
Diff::Remove { old } => { removed.insert(&old.name); }
Diff::MaybeChange { .. } => {}
}
(added, removed)
},
);
let indexes_ok = auto_migrate_indexes(&mut plan, &new_tables, &removed_tables);
let sequences_ok = auto_migrate_sequences(&mut plan, &new_tables, &removed_tables);
let constraints_ok = auto_migrate_constraints(&mut plan, &new_tables, &removed_tables);
// IMPORTANT: RLS auto-migrate steps must come last,
// since they assume that any schema changes, like adding or dropping tables,
// have already been reflected in the database state.
Expand Down Expand Up @@ -522,13 +536,9 @@ fn auto_migrate_views(plan: &mut AutoMigratePlan<'_>) -> Result<()> {
}
// From the user's perspective, views do not have persistent state.
// Hence removal does not require a manual migration - just disconnecting clients.
Diff::Remove { old } if plan.disconnects_all_users() => {
plan.steps.push(AutoMigrateStep::RemoveView(old.key()));
Ok(())
}
Diff::Remove { old } => {
plan.steps.push(AutoMigrateStep::RemoveView(old.key()));
plan.steps.push(AutoMigrateStep::DisconnectAllUsers);
plan.ensure_disconnect_all_users();
Ok(())
}
Diff::MaybeChange { old, new } => auto_migrate_view(plan, old, new),
Expand Down Expand Up @@ -612,9 +622,7 @@ fn auto_migrate_view<'def>(plan: &mut AutoMigratePlan<'def>, old: &'def ViewDef,
plan.steps.push(AutoMigrateStep::AddView(new.key()));
plan.steps.push(AutoMigrateStep::RemoveView(old.key()));

if !plan.disconnects_all_users() {
plan.steps.push(AutoMigrateStep::DisconnectAllUsers);
}
plan.ensure_disconnect_all_users();
} else {
plan.steps.push(AutoMigrateStep::UpdateView(new.key()));
}
Expand All @@ -630,11 +638,11 @@ fn auto_migrate_tables(plan: &mut AutoMigratePlan<'_>) -> Result<()> {
plan.steps.push(AutoMigrateStep::AddTable(new.key()));
Ok(())
}
// TODO: When we remove tables, we should also remove their dependencies, including row-level security.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you consider this?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes — views and RLS referencing a dropped table are caught by schema validation before ponder_auto_migrate runs. The new ModuleDef won't validate if it references a nonexistent table, so those dependencies are naturally resolved. The TODO about removing dependencies would apply to future work where we want to cascade-remove dependent objects, but the validation layer handles it for now.

Diff::Remove { old } => Err(AutoMigrateError::RemoveTable {
table: old.name.clone(),
Diff::Remove { old } => {
plan.steps.push(AutoMigrateStep::RemoveTable(old.key()));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can break clients when removing an empty table, so we should add AutoMigrateStep::DisconnectAllUsers if not present.
There's some code that checks if it is present and adds it if not.
I'd like to see this code first refactored into one method: `plan.ensure_disconnect_all_users(); in:

  • auto_migrate_views
  • auto_migrate_view
  • auto_migrate_table
  • auto_migrate_row_level_security

and then it can be used here too.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great call. Done — added plan.ensure_disconnect_all_users() after pushing RemoveTable. Also refactored all existing inline if \!plan.disconnects_all_users() { push(DisconnectAllUsers) } patterns across the file to use the new ensure_disconnect_all_users() method (views, columns_added, RLS).

plan.ensure_disconnect_all_users();
Ok(())
}
.into()),
Diff::MaybeChange { old, new } => auto_migrate_table(plan, old, new),
}
})
Expand Down Expand Up @@ -742,9 +750,7 @@ fn auto_migrate_table<'def>(plan: &mut AutoMigratePlan<'def>, old: &'def TableDe
// If we're adding a column, we'll rewrite the whole table.
// That makes any `ChangeColumns` moot, so we can skip it.
if columns_added {
if !plan.disconnects_all_users() {
plan.steps.push(AutoMigrateStep::DisconnectAllUsers);
}
plan.ensure_disconnect_all_users();
plan.steps.push(AutoMigrateStep::AddColumns(key));
} else if row_type_changed {
plan.steps.push(AutoMigrateStep::ChangeColumns(key));
Expand Down Expand Up @@ -927,7 +933,7 @@ fn ensure_old_ty_upgradable_to_new(
}
}

fn auto_migrate_indexes(plan: &mut AutoMigratePlan<'_>, new_tables: &HashSet<&Identifier>) -> Result<()> {
fn auto_migrate_indexes(plan: &mut AutoMigratePlan<'_>, new_tables: &HashSet<&Identifier>, removed_tables: &HashSet<&Identifier>) -> Result<()> {
diff(plan.old, plan.new, ModuleDef::indexes)
.map(|index_diff| -> Result<()> {
match index_diff {
Expand All @@ -938,7 +944,9 @@ fn auto_migrate_indexes(plan: &mut AutoMigratePlan<'_>, new_tables: &HashSet<&Id
Ok(())
}
Diff::Remove { old } => {
plan.steps.push(AutoMigrateStep::RemoveIndex(old.key()));
if !removed_tables.contains(&plan.old.stored_in_table_def(&old.name).unwrap().name) {
plan.steps.push(AutoMigrateStep::RemoveIndex(old.key()));
}
Ok(())
}
Diff::MaybeChange { old, new } => {
Expand All @@ -962,7 +970,7 @@ fn auto_migrate_indexes(plan: &mut AutoMigratePlan<'_>, new_tables: &HashSet<&Id
.collect_all_errors()
}

fn auto_migrate_sequences(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Identifier>) -> Result<()> {
fn auto_migrate_sequences(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Identifier>, removed_tables: &HashSet<&Identifier>) -> Result<()> {
diff(plan.old, plan.new, ModuleDef::sequences)
.map(|sequence_diff| -> Result<()> {
match sequence_diff {
Expand All @@ -975,7 +983,9 @@ fn auto_migrate_sequences(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Iden
Ok(())
}
Diff::Remove { old } => {
plan.steps.push(AutoMigrateStep::RemoveSequence(old.key()));
if !removed_tables.contains(&plan.old.stored_in_table_def(&old.name).unwrap().name) {
plan.steps.push(AutoMigrateStep::RemoveSequence(old.key()));
}
Ok(())
}
Diff::MaybeChange { old, new } => {
Expand All @@ -993,7 +1003,7 @@ fn auto_migrate_sequences(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Iden
.collect_all_errors()
}

fn auto_migrate_constraints(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Identifier>) -> Result<()> {
fn auto_migrate_constraints(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Identifier>, removed_tables: &HashSet<&Identifier>) -> Result<()> {
diff(plan.old, plan.new, ModuleDef::constraints)
.map(|constraint_diff| -> Result<()> {
match constraint_diff {
Expand All @@ -1010,7 +1020,9 @@ fn auto_migrate_constraints(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Id
}
}
Diff::Remove { old } => {
plan.steps.push(AutoMigrateStep::RemoveConstraint(old.key()));
if !removed_tables.contains(&plan.old.stored_in_table_def(&old.name).unwrap().name) {
plan.steps.push(AutoMigrateStep::RemoveConstraint(old.key()));
}
Ok(())
}
Diff::MaybeChange { old, new } => {
Expand Down Expand Up @@ -1045,8 +1057,8 @@ fn auto_migrate_row_level_security(plan: &mut AutoMigratePlan) -> Result<()> {
}

// We can force flush the cache by force disconnecting all clients if an RLS rule has been added, removed, or updated.
if old_rls != new_rls && !plan.disconnects_all_users() {
plan.steps.push(AutoMigrateStep::DisconnectAllUsers);
if old_rls != new_rls {
plan.ensure_disconnect_all_users();
}

Ok(())
Expand Down Expand Up @@ -1503,7 +1515,7 @@ mod tests {
let result = ponder_auto_migrate(&old_def, &new_def);

let apples = expect_identifier("Apples");
let bananas = expect_identifier("Bananas");
let _bananas = expect_identifier("Bananas");

let apples_name_unique_constraint = "Apples_name_key";

Expand Down Expand Up @@ -1711,10 +1723,8 @@ mod tests {
AutoMigrateError::ChangeTableType { table, type1, type2 } => table == &apples && type1 == &TableType::User && type2 == &TableType::System
);

expect_error_matching!(
result,
AutoMigrateError::RemoveTable { table } => table == &bananas
);
// Note: RemoveTable is no longer an error — removing tables is now allowed
// for empty tables; the emptiness check happens at execution time in update.rs.

let apples_id_index = "Apples_id_idx_btree";
let accessor_old = expect_identifier("id_index");
Expand Down Expand Up @@ -2399,4 +2409,58 @@ mod tests {

ponder_auto_migrate(&old, &new).expect("same event flag should succeed");
}

#[test]
fn remove_table_produces_step() {
let old = create_module_def(|builder| {
builder
.build_table_with_new_type("Keep", ProductType::from([("id", AlgebraicType::U64)]), true)
.with_access(TableAccess::Public)
.finish();
builder
.build_table_with_new_type("Drop", ProductType::from([("id", AlgebraicType::U64)]), true)
.with_access(TableAccess::Public)
.finish();
});
let new = create_module_def(|builder| {
builder
.build_table_with_new_type("Keep", ProductType::from([("id", AlgebraicType::U64)]), true)
.with_access(TableAccess::Public)
.finish();
});

let drop_table = expect_identifier("Drop");
let plan = ponder_auto_migrate(&old, &new).expect("removing a table should produce a valid plan");
assert_eq!(
plan.steps,
vec![
AutoMigrateStep::RemoveTable(&drop_table),
AutoMigrateStep::DisconnectAllUsers,
],
);
}

#[test]
fn remove_table_does_not_produce_orphan_sub_object_steps() {
let old = create_module_def(|builder| {
builder
.build_table_with_new_type("Drop", ProductType::from([("id", AlgebraicType::U64)]), true)
.with_unique_constraint(0)
.with_index(btree(0), "Drop_id_idx")
.with_access(TableAccess::Public)
.finish();
});
let new = create_module_def(|_builder| {});

let drop_table = expect_identifier("Drop");
let plan = ponder_auto_migrate(&old, &new).expect("removing a table should produce a valid plan");
assert_eq!(
plan.steps,
vec![
AutoMigrateStep::RemoveTable(&drop_table),
AutoMigrateStep::DisconnectAllUsers,
],
"plan should only contain RemoveTable + DisconnectAllUsers, no orphan sub-object steps"
);
}
}
5 changes: 5 additions & 0 deletions crates/schema/src/auto_migrate/formatter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,10 @@ fn format_step<F: MigrationFormatter>(
// So we must recompute it and send any updates to clients.
// No need to include this step in the formatted plan.
AutoMigrateStep::UpdateView(_) => Ok(()),
AutoMigrateStep::RemoveTable(t) => {
let table_def: &TableDef = plan.old.expect_lookup(*t);
f.format_remove_table(&table_def.name)
}
AutoMigrateStep::AddTable(t) => {
let table_info = extract_table_info(*t, plan)?;
f.format_add_table(&table_info)
Expand Down Expand Up @@ -140,6 +144,7 @@ pub enum Action {
pub trait MigrationFormatter {
fn format_header(&mut self) -> io::Result<()>;
fn format_add_table(&mut self, table_info: &TableInfo) -> io::Result<()>;
fn format_remove_table(&mut self, table_name: &Identifier) -> io::Result<()>;
fn format_view(&mut self, view_info: &ViewInfo, action: Action) -> io::Result<()>;
fn format_index(&mut self, index_info: &IndexInfo, action: Action) -> io::Result<()>;
fn format_constraint(&mut self, constraint_info: &ConstraintInfo, action: Action) -> io::Result<()>;
Expand Down
Loading