diff --git a/crates/datastore/src/locking_tx_datastore/replay.rs b/crates/datastore/src/locking_tx_datastore/replay.rs index d803360551d..d439553c01f 100644 --- a/crates/datastore/src/locking_tx_datastore/replay.rs +++ b/crates/datastore/src/locking_tx_datastore/replay.rs @@ -1085,7 +1085,7 @@ impl StateView for ReplayCommittedState<'_> { /// then falling back to [`CommittedState::iter_by_col_eq`]. fn find_st_table_row(&self, table_id: TableId) -> Result { if let Some(row_ptr) = self.replay_table_updated.get(&table_id) { - let (table, blob_store, _) = self.state.get_table_and_blob_store(table_id)?; + let (table, blob_store, _) = self.state.get_table_and_blob_store(ST_TABLE_ID)?; // SAFETY: `row_ptr` is stored in `self.replay_table_updated`, // meaning it was inserted into `st_table` by `replay_insert` // and has not yet been deleted by `replay_delete_by_rel`. diff --git a/crates/dst/Cargo.toml b/crates/dst/Cargo.toml index 7d47c9a1b47..7c0bceb8cd9 100644 --- a/crates/dst/Cargo.toml +++ b/crates/dst/Cargo.toml @@ -11,15 +11,18 @@ fallocate = ["spacetimedb-commitlog/fallocate"] [dependencies] anyhow.workspace = true clap.workspace = true +futures.workspace = true spacetimedb-datastore.workspace = true spacetimedb-commitlog.workspace = true spacetimedb-durability.workspace = true spacetimedb-engine.workspace = true +spacetimedb-fs-utils.workspace = true spacetimedb-lib.workspace = true spacetimedb-primitives.workspace = true spacetimedb-runtime = { workspace = true, features = ["simulation"] } spacetimedb-sats.workspace = true spacetimedb-schema.workspace = true +spacetimedb-snapshot.workspace = true spacetimedb-table.workspace = true tracing.workspace = true tracing-subscriber.workspace = true diff --git a/crates/dst/src/engine.rs b/crates/dst/src/engine.rs index 451b7417e41..5341de93641 100644 --- a/crates/dst/src/engine.rs +++ b/crates/dst/src/engine.rs @@ -6,28 +6,40 @@ use spacetimedb_datastore::traits::{IsolationLevel, TxData}; use spacetimedb_engine::error::{DBError, DatastoreError, IndexError}; use spacetimedb_engine::persistence::{DiskSizeFn, Durability as EngineDurability, Persistence}; use spacetimedb_engine::relational_db::{MutTx, RelationalDB}; -use spacetimedb_lib::{Identity, RawModuleDef}; +use spacetimedb_engine::snapshot::{Compression, SnapshotWorker}; +use spacetimedb_engine::update::{update_database, UpdateLogger}; +use spacetimedb_lib::identity::AuthCtx; +use spacetimedb_lib::Identity; use spacetimedb_primitives::TableId; -use spacetimedb_runtime::sim::{Rng, Runtime as SimRuntime}; +use spacetimedb_runtime::sim::{yield_now, Rng}; use spacetimedb_runtime::Handle; +use spacetimedb_schema::auto_migrate::{ponder_migrate, MigratePlan}; use spacetimedb_schema::def::ModuleDef; use spacetimedb_schema::schema::{Schema, TableSchema}; +use spacetimedb_snapshot::DynSnapshotRepo; use spacetimedb_table::page_pool::PagePool; +mod generation; +mod migrations; mod model; mod properties; +mod row; +mod state; mod workload; -use self::workload::{ - normalize_rows, row_to_bytes, CommitDelta, CountState, InsertOutcome, Interaction, Observation, TableDelta, - TableRowCount, +use self::migrations::{Migration, MigrationExpectation}; +use self::row::{normalize_rows, row_to_bytes}; +use self::state::{ + table_schema_state_for_schema, CommitDelta, CountState, SchemaState, TableDelta, TableRowCount, TableRows, }; +use self::workload::{InsertOutcome, Interaction, Observation}; use crate::engine::model::Model; use crate::engine::properties::EngineProperties; use crate::engine::workload::WorkloadGen; -use crate::schema::{default_schema, to_raw_def, SchemaPlan}; +use crate::schema::{canonical_schema, default_schema, to_module_def, SchemaPlan}; use crate::sim::commitlog::{InMemoryCommitlog, InMemoryCommitlogHandle}; +use crate::sim::snapshot::InMemorySnapshotRepo; use crate::traits::{TargetDriver, TestSuite}; pub struct EngineTarget { @@ -36,17 +48,19 @@ pub struct EngineTarget { active_mut_tx: Option, commitlog: InMemoryCommitlog, runtime_handle: Handle, + snapshot_repo: Arc, + schema: SchemaPlan, } impl EngineTarget { - pub fn init(schema: SchemaPlan, runtime_seed: u64) -> anyhow::Result { - let runtime = SimRuntime::new(runtime_seed); - let runtime_handle = Handle::simulation(runtime.handle()); + pub fn init(schema: SchemaPlan, runtime_handle: Handle) -> anyhow::Result { + let canonical = canonical_schema(&schema)?; let commitlog = InMemoryCommitlog::new(); - let db = Self::open_db(&commitlog, runtime_handle.clone())?; + let snapshot_repo = Arc::new(InMemorySnapshotRepo::new(Identity::ZERO, 0)); + let db = Self::open_db(&commitlog, snapshot_repo.clone(), runtime_handle.clone())?; Self::install_schema(&db, &schema)?; - let table_ids = Self::load_table_ids(&db, &schema)?; + let table_ids = Self::load_table_ids(&db, &canonical)?; Ok(Self { db: Some(db), @@ -54,12 +68,21 @@ impl EngineTarget { active_mut_tx: None, commitlog, runtime_handle, + snapshot_repo, + schema: canonical, }) } - fn open_db(commitlog: &InMemoryCommitlog, runtime_handle: Handle) -> anyhow::Result { + fn open_db( + commitlog: &InMemoryCommitlog, + snapshot_repo: Arc, + runtime_handle: Handle, + ) -> anyhow::Result { + let snapshot_repo: Arc = snapshot_repo; + let snapshot_worker = SnapshotWorker::new(snapshot_repo, Compression::Disabled, runtime_handle.clone()); + let snapshot_requester = snapshot_worker.clone(); let history = commitlog.open_handle()?; - let persistence = Self::persistence(history.clone(), runtime_handle); + let persistence = Self::persistence(history.clone(), Some(snapshot_worker), runtime_handle); let (db, connected_clients) = RelationalDB::open( Identity::ZERO, Identity::ZERO, @@ -69,10 +92,17 @@ impl EngineTarget { PagePool::new_for_test(), )?; anyhow::ensure!(connected_clients.is_empty(), "replay produced connected clients"); + commitlog.set_on_new_segment(Some(Arc::new(move || { + snapshot_requester.request_snapshot_ignore_closed(); + }))); Ok(db) } - fn persistence(handle: InMemoryCommitlogHandle, runtime_handle: Handle) -> Persistence { + fn persistence( + handle: InMemoryCommitlogHandle, + snapshots: Option, + runtime_handle: Handle, + ) -> Persistence { let durability: Arc = Arc::new(handle); let disk_size: DiskSizeFn = Arc::new(|| { io::Result::Ok(SizeOnDisk { @@ -83,16 +113,17 @@ impl EngineTarget { Persistence { durability, disk_size, - snapshots: None, + snapshots, runtime: runtime_handle, } } + fn module_def(schema: &SchemaPlan) -> anyhow::Result { + to_module_def(schema) + } + fn install_schema(db: &RelationalDB, schema: &SchemaPlan) -> anyhow::Result<()> { - let raw = to_raw_def(schema); - let raw_module_def = RawModuleDef::V10(raw); - let module_def = - ModuleDef::try_from(raw_module_def).map_err(|e| anyhow::anyhow!("schema validation failed: {e}"))?; + let module_def = Self::module_def(schema)?; db.with_auto_commit(Workload::Internal, |tx| -> Result<(), DBError> { for table_def in module_def.tables() { @@ -127,7 +158,11 @@ impl EngineTarget { drop(db); - self.db = Some(Self::open_db(&self.commitlog, self.runtime_handle.clone())?); + self.db = Some(Self::open_db( + &self.commitlog, + self.snapshot_repo.clone(), + self.runtime_handle.clone(), + )?); Ok(()) } @@ -138,27 +173,46 @@ impl EngineTarget { .ok_or_else(|| anyhow::anyhow!("database is not open"))?; let tx = db.begin_tx(Workload::Internal); let mut row_counts = Vec::with_capacity(self.table_ids.len()); + let mut table_rows = Vec::with_capacity(self.table_ids.len()); + let mut schema_tables = Vec::with_capacity(self.table_ids.len()); for (table, table_id) in self.table_ids.iter().enumerate() { - let count = match db.iter(&tx, *table_id) { - Ok(iter) => iter.count() as u64, + let rows = match db.iter(&tx, *table_id) { + Ok(iter) => normalize_rows(iter.map(|row| row.to_product_value()).collect()), Err(err) => { let _ = db.release_tx(tx); return Err(err.into()); } }; + let count = rows.len() as u64; row_counts.push(TableRowCount { table, count }); + table_rows.push(TableRows { table, rows }); + + let schema = match db.schema_for_table(&tx, *table_id) { + Ok(schema) => schema, + Err(err) => { + let _ = db.release_tx(tx); + return Err(err.into()); + } + }; + schema_tables.push(table_schema_state_for_schema(table, &schema)); } let _ = db.release_tx(tx); - Ok(CountState { row_counts }) + Ok(CountState { + row_counts, + table_rows, + schema: SchemaState { tables: schema_tables }, + }) } - fn is_unique_constraint_violation(error: &DBError) -> bool { - matches!( - error, - DBError::Datastore(DatastoreError::Index(IndexError::UniqueConstraintViolation(_))) - ) + fn unique_constraint_violation_details(error: &DBError) -> Option { + match error { + DBError::Datastore(DatastoreError::Index(IndexError::UniqueConstraintViolation(violation))) => { + Some(violation.to_string()) + } + _ => None, + } } fn commit_delta_from_tx_data(&self, tx_data: &TxData) -> CommitDelta { @@ -187,6 +241,79 @@ impl EngineTarget { CommitDelta { tables } } + fn migrate(&mut self, migration: &Migration) -> anyhow::Result { + anyhow::ensure!( + self.active_mut_tx.is_none(), + "migration while mutable transaction is active" + ); + anyhow::ensure!( + migration.target_schema() != &self.schema, + "engine DST generated a no-op migration" + ); + + let old_module_def = Self::module_def(&self.schema)?; + let new_module_def = match Self::module_def(migration.target_schema()) { + Ok(module_def) => module_def, + Err(error) => { + return match migration.expectation() { + MigrationExpectation::Rejected => Ok(Observation::MigrationRejected), + MigrationExpectation::Accepted => Err(error), + }; + } + }; + + let plan = match ponder_migrate(&old_module_def, &new_module_def) { + Ok(plan) => plan, + Err(error) => { + return match migration.expectation() { + MigrationExpectation::Rejected => Ok(Observation::MigrationRejected), + MigrationExpectation::Accepted => Err(error.into()), + }; + } + }; + + match migration.expectation() { + MigrationExpectation::Accepted => { + ensure_auto_plan(&plan)?; + self.apply_auto_migration(migration, plan)?; + Ok(Observation::Migrated) + } + MigrationExpectation::Rejected => self.expect_rejected_migration(plan), + } + } + + fn apply_auto_migration(&mut self, migration: &Migration, plan: MigratePlan<'_>) -> anyhow::Result<()> { + let db = self + .db + .as_ref() + .ok_or_else(|| anyhow::anyhow!("database is not open"))?; + let mut tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::Internal); + let _ = update_database(db, &mut tx, AuthCtx::for_testing(), plan, &DstUpdateLogger)?; + let Some((_tx_offset, _tx_data, _tx_metrics, _reducer)) = db.commit_tx(tx)? else { + anyhow::bail!("migration commit produced no transaction data"); + }; + + self.schema = migration.schema().clone(); + self.table_ids = Self::load_table_ids(db, &self.schema)?; + Ok(()) + } + + fn expect_rejected_migration(&self, plan: MigratePlan<'_>) -> anyhow::Result { + if matches!(&plan, MigratePlan::Manual(_)) { + return Ok(Observation::MigrationRejected); + } + + let db = self + .db + .as_ref() + .ok_or_else(|| anyhow::anyhow!("database is not open"))?; + let mut tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::Internal); + match update_database(db, &mut tx, AuthCtx::for_testing(), plan, &DstUpdateLogger) { + Ok(_) => anyhow::bail!("engine applied migration expected to be rejected"), + Err(_) => Ok(Observation::MigrationRejected), + } + } + pub fn execute(&mut self, interaction: &Interaction) -> anyhow::Result { tracing::debug!(?interaction, "executing interaction"); @@ -216,11 +343,14 @@ impl EngineTarget { .ok_or_else(|| anyhow::anyhow!("insert without active mutable transaction"))?; let outcome = match db.insert(tx, table_id, &bytes) { Ok((_generated_columns, row, _flags)) => InsertOutcome::Accepted(row.to_product_value()), - // Generated rows can intentionally hit unique constraints; the oracle validates that rejection. - Err(error) if Self::is_unique_constraint_violation(&error) => { - InsertOutcome::UniqueConstraintViolation + Err(error) => { + // Generated rows can intentionally hit unique constraints; the oracle validates that rejection. + if let Some(details) = Self::unique_constraint_violation_details(&error) { + InsertOutcome::UniqueConstraintViolation { details } + } else { + return Err(error.into()); + } } - Err(error) => return Err(error.into()), }; Ok(Observation::Inserted { outcome }) } @@ -249,10 +379,10 @@ impl EngineTarget { let Some((_tx_offset, tx_data, _tx_metrics, _reducer)) = db.commit_tx(tx)? else { anyhow::bail!("commit produced no transaction data"); }; - Ok(Observation::Committed { - delta: self.commit_delta_from_tx_data(&tx_data), - }) + let delta = self.commit_delta_from_tx_data(&tx_data); + Ok(Observation::Committed { delta }) } + Interaction::Migrate(migration) => self.migrate(migration), Interaction::Replay => { let _ = self.active_mut_tx.take(); self.reopen_from_commitlog()?; @@ -271,11 +401,35 @@ impl EngineTarget { } } +fn ensure_auto_plan(plan: &MigratePlan<'_>) -> anyhow::Result<()> { + let MigratePlan::Auto(_) = plan else { + anyhow::bail!("engine DST generated a manual migration plan"); + }; + Ok(()) +} + +struct DstUpdateLogger; + +impl UpdateLogger for DstUpdateLogger { + fn info(&self, msg: &str) { + tracing::debug!(%msg, "engine DST migration update"); + } +} + impl TargetDriver for EngineTarget { type Observation = Observation; + fn progress_status(&self) -> Option { + Some(self.snapshot_repo.stats().to_string()) + } + async fn execute<'a>(&'a mut self, interaction: &'a Interaction) -> Result { - EngineTarget::execute(self, interaction) + let observation = EngineTarget::execute(self, interaction)?; + if matches!(interaction, Interaction::CommitTx | Interaction::Migrate(_)) { + // DO not yeild when tx lock is held. + yield_now().await; + } + Ok(observation) } } @@ -290,10 +444,14 @@ impl TestSuite for EngineTest { type Properties = EngineProperties; - async fn build(&self, rng: Rng) -> Result<(Self::Interactions, Self::Target, Self::Properties), anyhow::Error> { + async fn build( + &self, + rng: Rng, + runtime: Handle, + ) -> Result<(Self::Interactions, Self::Target, Self::Properties), anyhow::Error> { let schema = default_schema(rng.clone()); - let runtime_seed = rng.next_u64(); - let target = EngineTarget::init(schema.clone(), runtime_seed)?; + let target = EngineTarget::init(schema, runtime)?; + let schema = target.schema.clone(); let properties = EngineProperties::new(schema.clone()); let model = Model::new(schema); @@ -302,3 +460,244 @@ impl TestSuite for EngineTest { Ok((interactions, target, properties)) } } + +#[cfg(test)] +mod tests { + use spacetimedb_lib::AlgebraicValue; + use spacetimedb_runtime::sim::Runtime as SimRuntime; + use spacetimedb_sats::product; + + use super::migrations::{Migration, MigrationExpectation, SchemaRewrite, TableMigrationOp}; + use super::*; + use crate::schema::{ColumnPlan, IndexAlgorithm, IndexPlan, TablePlan, Type, UniqueConstraintPlan}; + + fn migration_replay_schema() -> SchemaPlan { + SchemaPlan { + tables: vec![TablePlan { + name: "items".into(), + columns: vec![ + ColumnPlan { + name: "id".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "kind".into(), + ty: Type::Sum { variants: 1 }, + }, + ], + primary_key: Some(0), + indexes: vec![IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }], + unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], + sequences: vec![], + is_public: true, + is_event: false, + }], + } + } + + fn add_column_replay_schema() -> SchemaPlan { + SchemaPlan { + tables: vec![TablePlan { + name: "items".into(), + columns: vec![ + ColumnPlan { + name: "id".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "score".into(), + ty: Type::U64, + }, + ], + primary_key: Some(0), + indexes: vec![IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }], + unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], + sequences: vec![], + is_public: true, + is_event: false, + }], + } + } + + fn change_index_replay_schema() -> SchemaPlan { + SchemaPlan { + tables: vec![TablePlan { + name: "items".into(), + columns: vec![ + ColumnPlan { + name: "id".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "score".into(), + ty: Type::U64, + }, + ], + primary_key: Some(0), + indexes: vec![ + IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }, + IndexPlan { + columns: vec![1], + algorithm: IndexAlgorithm::BTree, + }, + ], + unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], + sequences: vec![], + is_public: true, + is_event: false, + }], + } + } + + fn insert_u64_rows(target: &mut EngineTarget) -> anyhow::Result<()> { + target.execute(&Interaction::BeginMutTx)?; + for id in 0..128u64 { + target.execute(&Interaction::Insert { + table: 0, + row: product![id, id * 10], + })?; + } + target.execute(&Interaction::CommitTx)?; + Ok(()) + } + + #[test] + fn engine_dst_smoke_runs_random_workload() -> anyhow::Result<()> { + let mut runtime = SimRuntime::new(0); + let runtime_handle = Handle::simulation(runtime.handle()); + runtime.block_on(EngineTest.run(Rng::new(0), runtime_handle, 1_000))?; + Ok(()) + } + + #[test] + fn init_installs_raw_schema_and_stores_canonical_schema() -> anyhow::Result<()> { + let raw_schema = SchemaPlan { + tables: vec![TablePlan { + name: "RawTable".into(), + columns: vec![ColumnPlan { + name: "SomeValue".into(), + ty: Type::U64, + }], + primary_key: None, + indexes: vec![], + unique_constraints: vec![], + sequences: vec![], + is_public: true, + is_event: false, + }], + }; + + let expected = canonical_schema(&raw_schema)?; + assert_ne!(expected, raw_schema); + + let target = EngineTarget::init(raw_schema, Handle::simulation(SimRuntime::new(0).handle()))?; + assert_eq!(target.schema, expected); + + Ok(()) + } + + #[test] + fn rejected_migration_is_expected_rejected_by_engine() -> anyhow::Result<()> { + let mut target = EngineTarget::init( + add_column_replay_schema(), + Handle::simulation(SimRuntime::new(0).handle()), + )?; + let model = Model::new(target.schema.clone()); + let migration = Migration::choose_rejected(&Rng::new(0), &target.schema, &model) + .expect("schema should produce a duplicate-object rejected migration"); + + assert_eq!(migration.expectation(), MigrationExpectation::Rejected); + assert_eq!( + target.execute(&Interaction::Migrate(migration))?, + Observation::MigrationRejected + ); + + Ok(()) + } + + #[test] + fn add_column_migration_replays_with_existing_rows() -> anyhow::Result<()> { + let mut target = EngineTarget::init( + add_column_replay_schema(), + Handle::simulation(SimRuntime::new(0).handle()), + )?; + insert_u64_rows(&mut target)?; + + target.execute(&Interaction::Migrate(Migration::from_rewrites( + &target.schema, + vec![SchemaRewrite::AlterTable { + table: "items".into(), + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::AddColumn { ty: Type::U64 }, + ], + }], + )?))?; + target.execute(&Interaction::Replay)?; + + Ok(()) + } + + #[test] + fn change_index_migration_replays_with_existing_rows() -> anyhow::Result<()> { + let mut target = EngineTarget::init( + change_index_replay_schema(), + Handle::simulation(SimRuntime::new(0).handle()), + )?; + insert_u64_rows(&mut target)?; + + target.execute(&Interaction::Migrate(Migration::from_rewrites( + &target.schema, + vec![SchemaRewrite::AlterTable { + table: "items".into(), + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeIndex { columns: vec![1] }, + ], + }], + )?))?; + target.execute(&Interaction::Replay)?; + + Ok(()) + } + + #[test] + fn migration_that_updates_st_table_and_st_column_replays() -> anyhow::Result<()> { + let mut target = EngineTarget::init( + migration_replay_schema(), + Handle::simulation(SimRuntime::new(0).handle()), + )?; + + target.execute(&Interaction::BeginMutTx)?; + for id in 0..128u64 { + target.execute(&Interaction::Insert { + table: 0, + row: product![id, AlgebraicValue::sum(0, AlgebraicValue::U8(1))], + })?; + } + target.execute(&Interaction::CommitTx)?; + + target.execute(&Interaction::Migrate(Migration::from_rewrites( + &target.schema, + vec![SchemaRewrite::AlterTable { + table: "items".into(), + ops: vec![ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeColumnType { column: 1 }, + ], + }], + )?))?; + target.execute(&Interaction::Replay)?; + + Ok(()) + } +} diff --git a/crates/dst/src/engine/generation.rs b/crates/dst/src/engine/generation.rs new file mode 100644 index 00000000000..8db69fe96ba --- /dev/null +++ b/crates/dst/src/engine/generation.rs @@ -0,0 +1,709 @@ +//! Workload value and migration generation for the engine DST driver. +//! +//! Read this file as policy first, mechanics second: +//! - `ValueGen::gen_insert_row` chooses the row-level insert shape: valid row, +//! whole-row duplicate, arbitrary candidate, or targeted uniqueness conflict. +//! - `ValueGen::gen_value_for_case` chooses one generated column value: fresh +//! random, sampled from the accumulated pool, or near an accumulated pool value. +//! - `MigrationGen` only chooses accepted vs. rejected migration work; concrete +//! schema rewrite rules live in `migrations.rs`. + +use std::collections::BTreeMap; + +use spacetimedb_lib::{AlgebraicValue, ProductValue}; +use spacetimedb_runtime::sim::Rng; +use spacetimedb_sats::ArrayValue; + +use super::migrations::Migration; +use super::model::Model; +use super::row::Row; +use crate::rng::{choice, Choice, WeightedChoice}; +use crate::schema::{SchemaPlan, TablePlan, Type}; + +// Bound the valid-insert search: random generation may collide with existing +// unique constraints, but a failed search must not stall the workload. +const INSERT_CANDIDATE_ATTEMPTS: usize = 32; +const VALUE_POOL_VALUES_PER_TYPE: usize = 4096; + +/// Generation context for one model state plus accumulated generator memory. +pub(crate) struct GenCtx<'a> { + rng: &'a Rng, + model: &'a Model, + generation: &'a mut GenerationState, +} + +impl<'a> GenCtx<'a> { + pub(crate) fn new(rng: &'a Rng, model: &'a Model, generation: &'a mut GenerationState) -> Self { + Self { rng, model, generation } + } + + pub(crate) fn gen_insert_row(&mut self, table: usize) -> Row { + ValueGen::new(self.rng, self.model, self.generation).gen_insert_row(table) + } + + pub(crate) fn gen_migration(&self) -> Option { + MigrationGen::new(self.rng, self.model).choose() + } +} + +// Row-level insert cases choose what kind of insert operation to try. +#[derive(Clone, Copy)] +enum InsertRowCase { + Valid, + AnyCandidate, + ExistingRow, + UniqueConflict, +} + +impl InsertRowCase { + const CHOICES: [Choice; 4] = [ + choice(80, Self::Valid), + choice(10, Self::AnyCandidate), + choice(5, Self::ExistingRow), + choice(5, Self::UniqueConflict), + ]; +} + +impl WeightedChoice for InsertRowCase {} + +// Column-level cases choose how to synthesize each non-sequence column value. +#[derive(Clone, Copy)] +enum ColumnValueCase { + Random, + Pooled, + Near, +} + +impl ColumnValueCase { + const CHOICES: [Choice; 3] = [ + choice(50, Self::Random), + choice(30, Self::Pooled), + choice(20, Self::Near), + ]; +} + +impl WeightedChoice for ColumnValueCase {} + +pub(crate) struct GenerationState { + generators: BTreeMap>, +} + +impl GenerationState { + pub(crate) fn seeded(schema: &SchemaPlan) -> Self { + let mut state = Self { + generators: BTreeMap::new(), + }; + state.seed_schema(schema); + state + } + + pub(crate) fn seed_schema(&mut self, schema: &SchemaPlan) { + for table in &schema.tables { + for column in &table.columns { + self.generator_mut(column.ty); + } + } + } + + pub(crate) fn observe_row(&mut self, table: &TablePlan, row: &Row) { + for (column, value) in table.columns.iter().zip(row.elements.iter()) { + self.generator_mut(column.ty).observe(value.clone()); + } + } + + fn generator_mut(&mut self, ty: Type) -> &mut dyn TypeValueGen { + self.generators + .entry(TypeKey::from(ty)) + .or_insert_with(|| new_generator_for(ty)) + .as_mut() + } +} + +#[derive(Default)] +struct ValueBucket { + values: Vec, + observed: usize, +} + +impl ValueBucket { + fn contains(&self, value: &AlgebraicValue) -> bool { + self.values.contains(value) + } + + fn observe(&mut self, value: AlgebraicValue) { + if self.values.len() < VALUE_POOL_VALUES_PER_TYPE { + self.values.push(value); + } else { + let slot = self.observed % VALUE_POOL_VALUES_PER_TYPE; + self.values[slot] = value; + } + self.observed = self.observed.wrapping_add(1); + } + + fn sample(&self, rng: &Rng) -> Option { + (!self.values.is_empty()).then(|| self.values[rng.index(self.values.len())].clone()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum TypeKey { + Bool, + I64, + U64, + String, + Bytes, + Sum { variants: u8 }, +} + +impl From for TypeKey { + fn from(ty: Type) -> Self { + match ty { + Type::Bool => Self::Bool, + Type::I64 => Self::I64, + Type::U64 => Self::U64, + Type::String => Self::String, + Type::Bytes => Self::Bytes, + Type::Sum { variants } => Self::Sum { variants }, + } + } +} + +/// Stateful value generator for one exact DST column type. +/// +/// Each implementation owns its accumulated value pool. `seeds` supplies the +/// stable starting corpus installed when the generator is first stored; +/// `observe` appends runtime values; `near` samples the pool itself and applies +/// the type-local perturbation. Keeping the pool behind this trait makes pooled +/// and near-value generation follow the same exact type key. +trait TypeValueGen { + fn pool(&self) -> &ValueBucket; + fn pool_mut(&mut self) -> &mut ValueBucket; + + /// Stable boundary values installed once when this exact type first appears. + fn seeds(&self) -> Vec; + + /// Produce a fresh valid value without consulting accumulated history. + fn random(&self, rng: &Rng) -> AlgebraicValue; + + /// Make a type-valid nearby value from an already selected value. + /// Returning `value` unchanged is allowed when no sensible nearby value + /// exists or the input does not match this generator's type. + fn near_from(&self, rng: &Rng, value: AlgebraicValue) -> AlgebraicValue; + + /// Check whether a materialized value is valid for this exact type. + fn matches(&self, value: &AlgebraicValue) -> bool; + + fn observe(&mut self, value: AlgebraicValue) { + if self.matches(&value) { + self.pool_mut().observe(value); + } + } + + fn pooled(&self, rng: &Rng) -> Option { + self.pool().sample(rng) + } + + fn near(&self, rng: &Rng) -> Option { + let value = self.pooled(rng)?; + let original = value.clone(); + let near = self.near_from(rng, value); + Some(if self.matches(&near) { near } else { original }) + } +} + +fn seed_generator(generator: &mut dyn TypeValueGen) { + for value in generator.seeds() { + let should_seed = generator.matches(&value) && !generator.pool().contains(&value); + if should_seed { + generator.pool_mut().observe(value); + } + } +} + +macro_rules! value_pool_methods { + () => { + fn pool(&self) -> &ValueBucket { + &self.pool + } + + fn pool_mut(&mut self) -> &mut ValueBucket { + &mut self.pool + } + }; +} + +#[derive(Default)] +struct BoolGen { + pool: ValueBucket, +} + +#[derive(Default)] +struct I64Gen { + pool: ValueBucket, +} + +#[derive(Default)] +struct U64Gen { + pool: ValueBucket, +} + +#[derive(Default)] +struct StringGen { + pool: ValueBucket, +} + +#[derive(Default)] +struct BytesGen { + pool: ValueBucket, +} + +struct SumGen { + variants: u8, + pool: ValueBucket, +} + +impl TypeValueGen for BoolGen { + value_pool_methods!(); + + fn seeds(&self) -> Vec { + vec![AlgebraicValue::Bool(false), AlgebraicValue::Bool(true)] + } + + fn random(&self, rng: &Rng) -> AlgebraicValue { + AlgebraicValue::Bool(rng.next_u64().is_multiple_of(2)) + } + + fn near_from(&self, _rng: &Rng, value: AlgebraicValue) -> AlgebraicValue { + match value { + AlgebraicValue::Bool(value) => AlgebraicValue::Bool(!value), + other => other, + } + } + + fn matches(&self, value: &AlgebraicValue) -> bool { + matches!(value, AlgebraicValue::Bool(_)) + } +} + +impl TypeValueGen for I64Gen { + value_pool_methods!(); + + fn seeds(&self) -> Vec { + [-3, -2, -1, 0, 1, 2, 3, i64::MIN, i64::MIN + 1, i64::MAX - 1, i64::MAX] + .into_iter() + .map(AlgebraicValue::I64) + .collect() + } + + fn random(&self, rng: &Rng) -> AlgebraicValue { + AlgebraicValue::I64(rng.next_u64() as i64) + } + + fn near_from(&self, rng: &Rng, value: AlgebraicValue) -> AlgebraicValue { + match value { + AlgebraicValue::I64(value) if rng.next_u64().is_multiple_of(2) => { + AlgebraicValue::I64(value.wrapping_add(1)) + } + AlgebraicValue::I64(value) => AlgebraicValue::I64(value.wrapping_sub(1)), + other => other, + } + } + + fn matches(&self, value: &AlgebraicValue) -> bool { + matches!(value, AlgebraicValue::I64(_)) + } +} + +impl TypeValueGen for U64Gen { + value_pool_methods!(); + + fn seeds(&self) -> Vec { + [0, 1, 2, 3, 4, 5, u64::MAX - 1, u64::MAX] + .into_iter() + .map(AlgebraicValue::U64) + .collect() + } + + fn random(&self, rng: &Rng) -> AlgebraicValue { + AlgebraicValue::U64(rng.next_u64()) + } + + fn near_from(&self, rng: &Rng, value: AlgebraicValue) -> AlgebraicValue { + match value { + AlgebraicValue::U64(value) if rng.next_u64().is_multiple_of(2) => { + AlgebraicValue::U64(value.wrapping_add(1)) + } + AlgebraicValue::U64(value) => AlgebraicValue::U64(value.wrapping_sub(1)), + other => other, + } + } + + fn matches(&self, value: &AlgebraicValue) -> bool { + matches!(value, AlgebraicValue::U64(_)) + } +} + +impl TypeValueGen for StringGen { + value_pool_methods!(); + + fn seeds(&self) -> Vec { + [ + "", + "a", + "aa", + "aaa", + "ab", + "aba", + "abb", + "b", + "z", + "v_0", + "v_1", + "quote'", + "double\"quote", + "back\\slash", + "line\nbreak", + "nul\0byte", + ] + .into_iter() + .map(|value| AlgebraicValue::String(value.into())) + .chain(std::iter::once(AlgebraicValue::String("x".repeat(128).into()))) + .collect() + } + + fn random(&self, rng: &Rng) -> AlgebraicValue { + let mut state = rng.next_u64(); + let len = match state % 100 { + 0..=9 => 0, + 10..=79 => (state as usize >> 8) % 32, + 80..=97 => 32 + ((state as usize >> 8) % 224), + _ => 256 + ((state as usize >> 8) % 768), + }; + let alphabet = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-'\"\\\n\t\0 %.,:;()[]{}"; + let value = (0..len) + .map(|_| { + state = state.wrapping_mul(6364136223846793005).wrapping_add(1); + alphabet[(state as usize >> 32) % alphabet.len()] as char + }) + .collect::(); + AlgebraicValue::String(value.into()) + } + + fn near_from(&self, rng: &Rng, value: AlgebraicValue) -> AlgebraicValue { + match value { + AlgebraicValue::String(value) => { + let mut value = value.into_string(); + if !value.is_empty() && rng.next_u64().is_multiple_of(4) { + value.pop(); + } else { + value.push((b'a' + rng.index(26) as u8) as char); + } + AlgebraicValue::String(value.into()) + } + other => other, + } + } + + fn matches(&self, value: &AlgebraicValue) -> bool { + matches!(value, AlgebraicValue::String(_)) + } +} + +impl TypeValueGen for BytesGen { + value_pool_methods!(); + + fn seeds(&self) -> Vec { + vec![ + AlgebraicValue::Array(ArrayValue::U8(Vec::new().into())), + AlgebraicValue::Array(ArrayValue::U8(vec![0].into())), + AlgebraicValue::Array(ArrayValue::U8(vec![1].into())), + AlgebraicValue::Array(ArrayValue::U8(vec![0, 255].into())), + AlgebraicValue::Array(ArrayValue::U8(vec![0; 32].into())), + AlgebraicValue::Array(ArrayValue::U8(vec![255; 32].into())), + AlgebraicValue::Array(ArrayValue::U8(vec![0, 255, 0, 255, 0, 255].into())), + ] + } + + fn random(&self, rng: &Rng) -> AlgebraicValue { + let len = (rng.next_u64() % 16) as usize; + let value = (0..len).map(|_| rng.next_u64() as u8).collect::>(); + AlgebraicValue::Array(ArrayValue::U8(value.into())) + } + + fn near_from(&self, rng: &Rng, value: AlgebraicValue) -> AlgebraicValue { + match value { + AlgebraicValue::Array(ArrayValue::U8(value)) => { + let mut value = value.to_vec(); + if value.is_empty() { + value.push(rng.next_u64() as u8); + } else { + let edits = 1 + rng.index(3); + for _ in 0..edits { + let index = rng.index(value.len()); + let delta = match rng.index(4) { + 0 => 1, + 1 => u8::MAX, + 2 => 0x80, + _ => 0x7f, + }; + value[index] = value[index].wrapping_add(delta); + } + } + AlgebraicValue::Array(ArrayValue::U8(value.into())) + } + other => other, + } + } + + fn matches(&self, value: &AlgebraicValue) -> bool { + matches!(value, AlgebraicValue::Array(ArrayValue::U8(_))) + } +} + +impl TypeValueGen for SumGen { + value_pool_methods!(); + + fn seeds(&self) -> Vec { + let mut values = vec![AlgebraicValue::sum(0, AlgebraicValue::U8(0))]; + if self.variants > 1 { + values.push(AlgebraicValue::sum(self.variants - 1, AlgebraicValue::U8(u8::MAX))); + } + values + } + + fn random(&self, rng: &Rng) -> AlgebraicValue { + debug_assert!(self.variants > 0); + let tag = rng.index(self.variants as usize) as u8; + AlgebraicValue::sum(tag, AlgebraicValue::U8(rng.next_u64() as u8)) + } + + fn near_from(&self, _rng: &Rng, value: AlgebraicValue) -> AlgebraicValue { + match value { + AlgebraicValue::Sum(sum) if self.variants > 0 => { + AlgebraicValue::sum(sum.tag.wrapping_add(1) % self.variants, *sum.value) + } + other => other, + } + } + + fn matches(&self, value: &AlgebraicValue) -> bool { + matches!(value, AlgebraicValue::Sum(sum) if sum.tag < self.variants) + } +} + +fn new_generator_for(ty: Type) -> Box { + let mut generator: Box = match ty { + Type::Bool => Box::::default(), + Type::I64 => Box::::default(), + Type::U64 => Box::::default(), + Type::String => Box::::default(), + Type::Bytes => Box::::default(), + Type::Sum { variants } => Box::new(SumGen { + variants, + pool: ValueBucket::default(), + }), + }; + seed_generator(generator.as_mut()); + generator +} + +// Row generation deliberately mixes normal traffic with rows the engine should +// reject. The model is the oracle for whether a candidate is valid; this module +// only decides which kind of pressure to apply. +struct ValueGen<'a> { + rng: &'a Rng, + model: &'a Model, + generation: &'a mut GenerationState, +} + +impl<'a> ValueGen<'a> { + fn new(rng: &'a Rng, model: &'a Model, generation: &'a mut GenerationState) -> Self { + Self { rng, model, generation } + } + + fn gen_insert_row(&mut self, table: usize) -> Row { + // Start here when tuning insert behavior; everything below materializes + // one arm of this match. + match InsertRowCase::pick(self.rng, &InsertRowCase::CHOICES) { + InsertRowCase::Valid => self.gen_valid_insert_row(table), + InsertRowCase::AnyCandidate => self.gen_row_candidate(table), + InsertRowCase::ExistingRow => self + .existing_row(table) + .unwrap_or_else(|| self.gen_valid_insert_row(table)), + InsertRowCase::UniqueConflict => self + .unique_conflict_row(table) + .unwrap_or_else(|| self.gen_valid_insert_row(table)), + } + } + + fn gen_valid_insert_row(&mut self, table: usize) -> Row { + // The valid path samples the same value domain as rejected paths. After + // bounded attempts, prefer an existing row as an accepted no-op over a + // hand-built per-type escape value. + for _ in 0..INSERT_CANDIDATE_ATTEMPTS { + let row = self.gen_row_candidate(table); + if !self.model.insert_would_violate_unique_constraint(table, &row) { + return row; + } + } + + self.existing_row(table) + .unwrap_or_else(|| self.gen_row_candidate(table)) + } + + fn gen_row_candidate(&mut self, table: usize) -> Row { + let column_types = self + .table(table) + .columns + .iter() + .map(|column| column.ty) + .collect::>(); + column_types + .into_iter() + .enumerate() + .map(|(column, ty)| self.gen_insert_value(table, column, ty)) + .collect::() + } + + fn gen_insert_value(&mut self, table: usize, column: usize, ty: Type) -> AlgebraicValue { + if self.is_sequence_column(table, column) { + return sequence_placeholder(ty); + } + + self.gen_value(ty) + } + + fn gen_value(&mut self, ty: Type) -> AlgebraicValue { + let case = ColumnValueCase::pick(self.rng, &ColumnValueCase::CHOICES); + if let Some(value) = self.gen_value_for_case(ty, case) { + return value; + } + + self.gen_value_for_case(ty, ColumnValueCase::Random) + .expect("random value generation cannot fail") + } + + fn gen_value_for_case(&mut self, ty: Type, case: ColumnValueCase) -> Option { + match case { + ColumnValueCase::Random => Some(self.generation.generator_mut(ty).random(self.rng)), + ColumnValueCase::Pooled => self.pooled_value(ty), + ColumnValueCase::Near => Some(self.near_value(ty)), + } + } + + fn existing_row(&self, table: usize) -> Option { + let count = self.model.row_count(table); + (count > 0).then(|| { + self.model + .row(table, self.rng.index(count)) + .expect("sampled row index is in bounds") + .clone() + }) + } + + fn unique_conflict_row(&mut self, table: usize) -> Option { + // Target a specific unique constraint without turning the operation into + // an exact-row duplicate: copy constrained columns from an existing row, + // then perturb an unconstrained column if necessary. + let constraints = self.table(table).unique_constraints.clone(); + if constraints.is_empty() || self.model.row_count(table) == 0 { + return None; + } + + let start = self.rng.index(constraints.len()); + for offset in 0..constraints.len() { + let constraint = &constraints[(start + offset) % constraints.len()]; + let base = self.existing_row(table)?; + let mut row = self.gen_row_candidate(table); + + for &column in &constraint.columns { + row.elements[column] = base.elements[column].clone(); + } + + if row == base { + let Some(column) = + (0..self.table(table).columns.len()).find(|column| !constraint.columns.contains(column)) + else { + continue; + }; + let ty = self.table(table).columns[column].ty; + row.elements[column] = self.near_value(ty); + } + + if row != base && self.model.insert_would_violate_unique_constraint(table, &row) { + return Some(row); + } + } + + None + } + + fn pooled_value(&mut self, ty: Type) -> Option { + self.generation.generator_mut(ty).pooled(self.rng) + } + + fn near_value(&mut self, ty: Type) -> AlgebraicValue { + let type_gen = self.generation.generator_mut(ty); + type_gen.near(self.rng).unwrap_or_else(|| type_gen.random(self.rng)) + } + + fn is_sequence_column(&self, table: usize, column: usize) -> bool { + self.table(table) + .sequences + .iter() + .any(|sequence| sequence.column == column) + } + + fn table(&self, table: usize) -> &TablePlan { + &self.model.schema().tables[table] + } +} + +// Sequence columns are engine-filled on insert. The command still needs a SATS +// value with the right shape, so this placeholder should be ignored downstream. +fn sequence_placeholder(ty: Type) -> AlgebraicValue { + match ty { + Type::I64 => AlgebraicValue::I64(0), + Type::U64 => AlgebraicValue::U64(0), + _ => unreachable!("sequence columns are integral"), + } +} + +#[derive(Debug, Clone, Copy)] +enum MigrationCase { + Accepted, + Rejected, +} + +impl MigrationCase { + const CHOICES: [Choice; 2] = [choice(95, Self::Accepted), choice(5, Self::Rejected)]; +} + +impl WeightedChoice for MigrationCase {} + +// Migration generation is intentionally shallow here: accepted migrations are +// short rewrite chains, while rejected migrations are one invalidated accepted +// rule composed with otherwise valid accepted cases. +struct MigrationGen<'a> { + rng: &'a Rng, + model: &'a Model, +} + +impl<'a> MigrationGen<'a> { + fn new(rng: &'a Rng, model: &'a Model) -> Self { + Self { rng, model } + } + + fn choose(&self) -> Option { + match MigrationCase::pick(self.rng, &MigrationCase::CHOICES) { + MigrationCase::Accepted => self.choose_accepted(), + MigrationCase::Rejected => { + Migration::choose_rejected(self.rng, self.model.schema(), self.model).or_else(|| self.choose_accepted()) + } + } + } + + fn choose_accepted(&self) -> Option { + Migration::choose_accepted(self.rng, self.model.schema(), self.model) + } +} diff --git a/crates/dst/src/engine/migrations.rs b/crates/dst/src/engine/migrations.rs new file mode 100644 index 00000000000..6058e395c94 --- /dev/null +++ b/crates/dst/src/engine/migrations.rs @@ -0,0 +1,1811 @@ +//! Schema migration candidates for the engine DST driver. +//! +//! This module generates accepted schema changes for engine auto-migration and +//! rejected schema changes by asking the same rule set for one intentionally +//! invalidated precondition. + +use super::model::Model; +use crate::rng::choose_index; +use crate::schema::{ + canonical_schema, ColumnPlan, IndexAlgorithm, IndexPlan, SchemaGenerator, SchemaNames, SchemaPlan, SchemaProfile, + SequencePlan, TablePlan, Type, UniqueConstraintPlan, +}; +use spacetimedb_runtime::sim::Rng; +use std::collections::HashSet; + +const MAX_SUM_VARIANTS: u8 = 32; +const MAX_EVENT_COLUMNS: usize = 32; +const MAX_TABLE_COLUMNS: usize = 32; +const MAX_TABLES: usize = 128; + +/// A fully materialized schema migration interaction. +/// +/// `target_schema` is the raw schema handed to engine validation. `schema` is +/// the canonical schema the model should observe after the engine applies an +/// accepted migration. Rejected migrations never advance model state. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Migration { + target_schema: SchemaPlan, + schema: SchemaPlan, + expectation: MigrationExpectation, +} + +/// Whether the target schema should be accepted or refused by the engine. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MigrationExpectation { + Accepted, + Rejected, +} + +/// A deterministic schema edit used while building a migration candidate. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SchemaRewrite { + AddTable { table: TablePlan }, + RemoveTable { table: String }, + AlterTable { table: String, ops: Vec }, +} + +/// Table-local migration operations generated by the DST harness. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TableMigrationOp { + ChangeAccess, + AddColumn { + ty: Type, + }, + AddIndex { + columns: Vec, + algorithm: IndexAlgorithm, + }, + RemoveIndex { + columns: Vec, + }, + AddSequence { + sequence: SequencePlan, + }, + RemoveSequence { + column: usize, + }, + AddUniqueConstraint { + columns: Vec, + }, + RemoveUniqueConstraint { + columns: Vec, + }, + ChangePrimaryKey { + column: Option, + }, + ChangeIndex { + columns: Vec, + }, + ChangeColumnType { + column: usize, + }, + ReschemaEventTable, +} + +/// Stable identity for migration rules. +/// +/// Keep this ID separate from weights and rule implementation. It is the +/// greppable handle for triage, swarm toggles, and future campaign profiles. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum RuleId { + AddTable, + RemovePristineNonEventTable, + RemoveEmptyEventTable, + AddColumn, + AddIndex, + RemoveIndex, + ChangeIndex, + AddSequenceOnPristineTable, + RemoveSequence, + AddUniqueConstraintOnPristineTable, + RemoveUniqueConstraint, + DropPrimaryKeyAndWidenSum, + WidenSumColumn, + ReschemaEmptyEventTable, +} + +#[derive(Clone, Copy)] +struct RuleEntry { + id: RuleId, + rule: &'static dyn MigrationRule, +} + +struct RuleWeights { + rules: &'static [(RuleId, u32)], +} + +static MIGRATION_RULES: &[RuleEntry] = &[ + RuleEntry { + id: RuleId::AddTable, + rule: &AddTableRule, + }, + RuleEntry { + id: RuleId::RemovePristineNonEventTable, + rule: &RemovePristineNonEventTableRule, + }, + RuleEntry { + id: RuleId::RemoveEmptyEventTable, + rule: &RemoveEmptyEventTableRule, + }, + RuleEntry { + id: RuleId::AddColumn, + rule: &AddColumnRule, + }, + RuleEntry { + id: RuleId::AddIndex, + rule: &AddIndexRule, + }, + RuleEntry { + id: RuleId::RemoveIndex, + rule: &RemoveIndexRule, + }, + RuleEntry { + id: RuleId::ChangeIndex, + rule: &ChangeIndexRule, + }, + RuleEntry { + id: RuleId::AddSequenceOnPristineTable, + rule: &AddSequenceOnPristineTableRule, + }, + RuleEntry { + id: RuleId::RemoveSequence, + rule: &RemoveSequenceRule, + }, + RuleEntry { + id: RuleId::AddUniqueConstraintOnPristineTable, + rule: &AddUniqueConstraintOnPristineTableRule, + }, + RuleEntry { + id: RuleId::RemoveUniqueConstraint, + rule: &RemoveUniqueConstraintRule, + }, + RuleEntry { + id: RuleId::DropPrimaryKeyAndWidenSum, + rule: &DropPrimaryKeyAndWidenSumRule, + }, + RuleEntry { + id: RuleId::WidenSumColumn, + rule: &WidenSumColumnRule, + }, + RuleEntry { + id: RuleId::ReschemaEmptyEventTable, + rule: &ReschemaEmptyEventTableRule, + }, +]; + +static DEFAULT_RULE_WEIGHTS: RuleWeights = RuleWeights { + rules: &[ + (RuleId::AddTable, 2), + (RuleId::RemovePristineNonEventTable, 1), + (RuleId::RemoveEmptyEventTable, 1), + (RuleId::AddColumn, 14), + (RuleId::AddIndex, 10), + (RuleId::RemoveIndex, 8), + (RuleId::ChangeIndex, 10), + (RuleId::AddSequenceOnPristineTable, 8), + (RuleId::RemoveSequence, 8), + (RuleId::AddUniqueConstraintOnPristineTable, 12), + (RuleId::RemoveUniqueConstraint, 8), + (RuleId::DropPrimaryKeyAndWidenSum, 6), + (RuleId::WidenSumColumn, 12), + (RuleId::ReschemaEmptyEventTable, 8), + ], +}; + +fn migration_rules() -> &'static [RuleEntry] { + MIGRATION_RULES +} + +#[derive(Debug, Clone, Copy)] +struct RuleCtx<'a> { + schema: &'a SchemaPlan, + model: &'a Model, +} + +/// A positive case emitted by a migration rule. +/// +/// `writes` are derived from the rewrite and used by the batch validator. +/// `protects` are rule-owned semantic dependencies such as "this table was +/// pristine." +#[derive(Debug, Clone, PartialEq, Eq)] +struct AcceptedCase { + rule: RuleId, + rewrite: SchemaRewrite, + writes: Vec, + protects: Vec, +} + +impl AcceptedCase { + fn new(rule: RuleId, rewrite: SchemaRewrite) -> Self { + let writes = write_resources(&rewrite); + Self { + rule, + rewrite, + writes, + protects: Vec::new(), + } + } + + fn protecting(mut self, protects: impl Into>) -> Self { + self.protects = protects.into(); + self + } +} + +/// A rejected case emitted by intentionally violating one accepted rule precondition. +#[derive(Debug, Clone, PartialEq, Eq)] +struct RejectedCase { + rule: RuleId, + rewrite: SchemaRewrite, + writes: Vec, + protects: Vec, +} + +impl RejectedCase { + fn new(rule: RuleId, rewrite: SchemaRewrite) -> Self { + let writes = write_resources(&rewrite); + Self { + rule, + rewrite, + writes, + protects: Vec::new(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +enum Resource { + Table(String), + Column { table: String, column: usize }, + Index { table: String, columns: Vec }, + Sequence { table: String, column: usize }, + UniqueConstraint { table: String, columns: Vec }, + PrimaryKey { table: String }, +} + +trait MigrationRule: Sync { + fn accepted_cases(&self, _rng: &Rng, _ctx: RuleCtx<'_>, _rule: RuleId) -> Vec { + Vec::new() + } + + fn invalid_cases(&self, _rng: &Rng, _ctx: RuleCtx<'_>, _rule: RuleId) -> Vec { + Vec::new() + } + + /// Rule-owned semantic validation for accepted cases. + /// + /// Structural validation is shared by applying the rewrite batch to a draft + /// schema. Data/model assumptions stay with the rule which made them. + fn validate_accepted( + &self, + _case: &AcceptedCase, + _original: RuleCtx<'_>, + _final_schema: &SchemaPlan, + ) -> anyhow::Result<()> { + Ok(()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AcceptedComposition { + Capped { max_rules: usize }, +} + +impl Migration { + fn accepted(target_schema: SchemaPlan) -> anyhow::Result { + let schema = canonical_schema(&target_schema)?; + Ok(Self { + target_schema, + schema, + expectation: MigrationExpectation::Accepted, + }) + } + + fn rejected(target_schema: SchemaPlan) -> Self { + Self { + schema: target_schema.clone(), + target_schema, + expectation: MigrationExpectation::Rejected, + } + } + + #[cfg(test)] + pub(crate) fn from_rewrites(base: &SchemaPlan, rewrites: Vec) -> anyhow::Result { + let mut schema = base.clone(); + for rewrite in &rewrites { + rewrite.apply_to(&mut schema)?; + } + Self::accepted(schema) + } + + /// Canonical schema visible to the DST model after an accepted migration. + pub(crate) fn schema(&self) -> &SchemaPlan { + &self.schema + } + + /// Raw target schema handed to engine validation and auto-migration. + pub(crate) fn target_schema(&self) -> &SchemaPlan { + &self.target_schema + } + + pub(crate) fn expectation(&self) -> MigrationExpectation { + self.expectation + } + + pub(crate) fn choose_accepted(rng: &Rng, schema: &SchemaPlan, model: &Model) -> Option { + let max_rules = 1 + rng.index(10); + build_accepted_migration( + rng, + RuleCtx { schema, model }, + AcceptedComposition::Capped { max_rules }, + ) + } + + pub(crate) fn choose_rejected(rng: &Rng, schema: &SchemaPlan, model: &Model) -> Option { + build_rejected_migration(rng, RuleCtx { schema, model }) + } +} + +fn build_accepted_migration(rng: &Rng, ctx: RuleCtx<'_>, mode: AcceptedComposition) -> Option { + let AcceptedComposition::Capped { max_rules } = mode; + let max_rules = max_rules.max(1); + + for _ in 0..16 { + let cases = pick_accepted_cases(rng, ctx, max_rules)?; + if let Ok(migration) = build_validated_accepted_batch(ctx, cases) { + return Some(migration); + } + } + + let case = pick_accepted_case(rng, ctx, &DEFAULT_RULE_WEIGHTS)?; + build_validated_accepted_batch(ctx, vec![case]).ok() +} + +fn build_validated_accepted_batch(ctx: RuleCtx<'_>, cases: Vec) -> anyhow::Result { + anyhow::ensure!(!cases.is_empty(), "empty accepted migration batch"); + validate_no_resource_conflicts(&cases)?; + + let mut draft = ctx.schema.clone(); + for case in &cases { + case.rewrite.apply_to(&mut draft)?; + } + + let final_schema = canonical_schema(&draft)?; + + validate_accepted_cases(ctx, &cases, &final_schema)?; + + anyhow::ensure!(final_schema != *ctx.schema, "accepted migration batch was no-op"); + Migration::accepted(draft) +} + +fn validate_accepted_cases(ctx: RuleCtx<'_>, cases: &[AcceptedCase], final_schema: &SchemaPlan) -> anyhow::Result<()> { + for case in cases { + let rule = rule_by_id(case.rule) + .ok_or_else(|| anyhow::anyhow!("accepted migration case references missing rule {:?}", case.rule))?; + rule.rule.validate_accepted(case, ctx, final_schema)?; + } + Ok(()) +} + +fn build_rejected_migration(rng: &Rng, ctx: RuleCtx<'_>) -> Option { + for _ in 0..16 { + let rejected = pick_rejected_case(rng, ctx, &DEFAULT_RULE_WEIGHTS)?; + let accepted = pick_accepted_cases_around_rejected(rng, ctx, &rejected); + if let Ok(migration) = build_validated_rejected_batch(ctx, accepted, rejected) { + return Some(migration); + } + } + + let rejected = pick_rejected_case(rng, ctx, &DEFAULT_RULE_WEIGHTS)?; + build_validated_rejected_batch(ctx, Vec::new(), rejected).ok() +} + +fn build_validated_rejected_batch( + ctx: RuleCtx<'_>, + accepted: Vec, + rejected: RejectedCase, +) -> anyhow::Result { + validate_no_resource_conflicts(&accepted)?; + validate_rejected_case_isolated(&accepted, &rejected)?; + + let mut draft = ctx.schema.clone(); + for case in &accepted { + case.rewrite.apply_to(&mut draft)?; + } + + validate_accepted_cases(ctx, &accepted, &draft)?; + rejected.rewrite.apply_unchecked_to(&mut draft)?; + anyhow::ensure!(draft != *ctx.schema, "rejected migration batch was no-op"); + Ok(Migration::rejected(draft)) +} + +fn rule_by_id(id: RuleId) -> Option<&'static RuleEntry> { + migration_rules().iter().find(|entry| entry.id == id) +} + +fn pick_accepted_cases(rng: &Rng, ctx: RuleCtx<'_>, max_rules: usize) -> Option> { + let mut cases = Vec::new(); + for _ in 0..max_rules { + let Some(case) = pick_accepted_case(rng, ctx, &DEFAULT_RULE_WEIGHTS) else { + break; + }; + cases.push(case); + } + (!cases.is_empty()).then_some(cases) +} + +fn pick_accepted_cases_around_rejected(rng: &Rng, ctx: RuleCtx<'_>, rejected: &RejectedCase) -> Vec { + let mut cases = Vec::new(); + + for &(id, weight) in DEFAULT_RULE_WEIGHTS.rules { + if weight == 0 || id == rejected.rule { + continue; + } + let Some(entry) = rule_by_id(id) else { + continue; + }; + let Some(case) = pick_candidate(rng, entry.rule.accepted_cases(rng, ctx, entry.id)) else { + continue; + }; + + let mut trial = cases.clone(); + trial.push(case.clone()); + if validate_no_resource_conflicts(&trial).is_ok() && validate_rejected_case_isolated(&trial, rejected).is_ok() { + cases.push(case); + } + } + + cases +} + +fn pick_accepted_case(rng: &Rng, ctx: RuleCtx<'_>, weights: &RuleWeights) -> Option { + for _ in 0..16 { + let entry = pick_weighted_rule(rng, weights.rules)?; + if let Some(case) = pick_candidate(rng, entry.rule.accepted_cases(rng, ctx, entry.id)) { + return Some(case); + } + } + + pick_candidate(rng, accepted_cases_for_weights(rng, ctx, weights.rules)) +} + +fn pick_rejected_case(rng: &Rng, ctx: RuleCtx<'_>, weights: &RuleWeights) -> Option { + // Enumerate first so rules with no invalidation hook do not consume weighted retries. + pick_candidate(rng, rejected_cases_for_weights(rng, ctx, weights.rules)) +} + +fn pick_weighted_rule(rng: &Rng, weights: &[(RuleId, u32)]) -> Option<&'static RuleEntry> { + let total = weights.iter().map(|(_, weight)| *weight).sum::(); + if total == 0 { + return None; + } + + let mut ticket = rng.index(total as usize) as u32; + for &(id, weight) in weights { + if ticket < weight { + return Some(rule_by_id(id).expect("rule weight references unregistered migration rule")); + } + ticket -= weight; + } + + None +} + +fn accepted_cases_for_weights(rng: &Rng, ctx: RuleCtx<'_>, weights: &[(RuleId, u32)]) -> Vec { + weights + .iter() + .filter(|(_, weight)| *weight > 0) + .filter_map(|(id, _)| rule_by_id(*id)) + .flat_map(|entry| entry.rule.accepted_cases(rng, ctx, entry.id)) + .collect() +} + +fn rejected_cases_for_weights(rng: &Rng, ctx: RuleCtx<'_>, weights: &[(RuleId, u32)]) -> Vec { + weights + .iter() + .filter(|(_, weight)| *weight > 0) + .filter_map(|(id, _)| rule_by_id(*id)) + .flat_map(|entry| entry.rule.invalid_cases(rng, ctx, entry.id)) + .collect() +} + +fn pick_candidate(rng: &Rng, mut candidates: Vec) -> Option { + let idx = choose_index(rng, candidates.len())?; + Some(candidates.swap_remove(idx)) +} + +fn validate_rejected_case_isolated(accepted: &[AcceptedCase], rejected: &RejectedCase) -> anyhow::Result<()> { + for case in accepted { + validate_resource_set_is_disjoint(&case.writes, &rejected.writes)?; + validate_resource_set_is_disjoint(&case.writes, &rejected.protects)?; + validate_resource_set_is_disjoint(&case.protects, &rejected.writes)?; + validate_resource_set_is_disjoint(&case.protects, &rejected.protects)?; + } + Ok(()) +} + +fn validate_no_resource_conflicts(cases: &[AcceptedCase]) -> anyhow::Result<()> { + for (left_idx, left) in cases.iter().enumerate() { + for right in &cases[left_idx + 1..] { + validate_resource_set_is_disjoint(&left.writes, &right.writes)?; + validate_resource_set_is_disjoint(&left.writes, &right.protects)?; + validate_resource_set_is_disjoint(&left.protects, &right.writes)?; + } + } + Ok(()) +} + +fn validate_resource_set_is_disjoint(left: &[Resource], right: &[Resource]) -> anyhow::Result<()> { + for left in left { + for right in right { + anyhow::ensure!( + !resources_conflict(left, right), + "accepted migration batch has conflicting resources {left:?} and {right:?}" + ); + } + } + Ok(()) +} + +fn write_resources(rewrite: &SchemaRewrite) -> Vec { + let mut resources = match rewrite { + SchemaRewrite::AddTable { table } => vec![Resource::Table(table.name.clone())], + SchemaRewrite::RemoveTable { table } => vec![Resource::Table(table.clone())], + SchemaRewrite::AlterTable { table, ops } => ops + .iter() + .map(|op| match op { + TableMigrationOp::ChangeAccess + | TableMigrationOp::AddColumn { .. } + | TableMigrationOp::ReschemaEventTable => Resource::Table(table.clone()), + TableMigrationOp::AddIndex { columns, .. } + | TableMigrationOp::RemoveIndex { columns } + | TableMigrationOp::ChangeIndex { columns } => Resource::Index { + table: table.clone(), + columns: columns.clone(), + }, + TableMigrationOp::AddSequence { sequence } => Resource::Sequence { + table: table.clone(), + column: sequence.column, + }, + TableMigrationOp::RemoveSequence { column } => Resource::Sequence { + table: table.clone(), + column: *column, + }, + TableMigrationOp::AddUniqueConstraint { columns } + | TableMigrationOp::RemoveUniqueConstraint { columns } => Resource::UniqueConstraint { + table: table.clone(), + columns: columns.clone(), + }, + TableMigrationOp::ChangePrimaryKey { .. } => Resource::PrimaryKey { table: table.clone() }, + TableMigrationOp::ChangeColumnType { column } => Resource::Column { + table: table.clone(), + column: *column, + }, + }) + .collect(), + }; + dedup_resources(&mut resources); + resources +} + +fn dedup_resources(resources: &mut Vec) { + let mut seen = HashSet::new(); + resources.retain(|resource| seen.insert(resource.clone())); +} + +fn resources_conflict(left: &Resource, right: &Resource) -> bool { + if left == right { + return true; + } + + match (left, right) { + (Resource::Table(left), right) => resource_table(right).is_some_and(|right| right == left), + (left, Resource::Table(right)) => resource_table(left).is_some_and(|left| left == right), + ( + Resource::Column { + table: left_table, + column: left_column, + }, + Resource::Sequence { + table: right_table, + column: right_column, + }, + ) + | ( + Resource::Sequence { + table: left_table, + column: left_column, + }, + Resource::Column { + table: right_table, + column: right_column, + }, + ) => left_table == right_table && left_column == right_column, + _ => false, + } +} + +fn resource_table(resource: &Resource) -> Option<&str> { + match resource { + Resource::Table(table) + | Resource::Column { table, .. } + | Resource::Index { table, .. } + | Resource::Sequence { table, .. } + | Resource::UniqueConstraint { table, .. } + | Resource::PrimaryKey { table } => Some(table), + } +} + +fn accepted_cases_from_rewrites(rule: RuleId, rewrites: Vec) -> Vec { + rewrites + .into_iter() + .map(|rewrite| AcceptedCase::new(rule, rewrite)) + .collect() +} + +fn rejected_cases_from_rewrites(rule: RuleId, rewrites: Vec) -> Vec { + rewrites + .into_iter() + .map(|rewrite| RejectedCase::new(rule, rewrite)) + .collect() +} + +fn table_is_pristine(model: &Model, table: &TablePlan) -> bool { + model.row_count_by_table_name(&table.name) == 0 && !model.ever_inserted_by_table_name(&table.name) +} + +fn original_table<'a>(ctx: RuleCtx<'a>, table: &str) -> anyhow::Result<&'a TablePlan> { + ctx.schema + .tables + .iter() + .find(|table_plan| table_plan.name == table) + .ok_or_else(|| anyhow::anyhow!("migration rule references missing original table {table}")) +} + +fn alter_table_sequence(case: &AcceptedCase) -> anyhow::Result<(&str, &SequencePlan)> { + let SchemaRewrite::AlterTable { table, ops } = &case.rewrite else { + anyhow::bail!("add-sequence rule emitted non-table rewrite"); + }; + let Some(sequence) = ops.iter().find_map(|op| match op { + TableMigrationOp::AddSequence { sequence } => Some(sequence), + _ => None, + }) else { + anyhow::bail!("add-sequence rule emitted rewrite without add-sequence op"); + }; + Ok((table, sequence)) +} + +struct AddTableRule; + +impl MigrationRule for AddTableRule { + fn accepted_cases(&self, rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, add_table_rewrites(rng, ctx.schema)) + } +} + +struct RemovePristineNonEventTableRule; + +impl MigrationRule for RemovePristineNonEventTableRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + let non_event_tables = ctx.schema.tables.iter().filter(|table| !table.is_event).count(); + ctx.schema + .tables + .iter() + .filter(|table| !table.is_event && non_event_tables > 1 && table_is_pristine(ctx.model, table)) + .map(|table| { + AcceptedCase::new( + rule, + SchemaRewrite::RemoveTable { + table: table.name.clone(), + }, + ) + .protecting(vec![Resource::Table(table.name.clone())]) + }) + .collect() + } + + fn validate_accepted( + &self, + case: &AcceptedCase, + original: RuleCtx<'_>, + _final_schema: &SchemaPlan, + ) -> anyhow::Result<()> { + let SchemaRewrite::RemoveTable { table } = &case.rewrite else { + anyhow::bail!("remove-pristine-table rule emitted non-remove rewrite"); + }; + let table_plan = original_table(original, table)?; + let non_event_tables = original.schema.tables.iter().filter(|table| !table.is_event).count(); + anyhow::ensure!(!table_plan.is_event, "remove-pristine-table rule selected event table"); + anyhow::ensure!( + non_event_tables > 1, + "remove-pristine-table rule selected last non-event table" + ); + anyhow::ensure!( + table_is_pristine(original.model, table_plan), + "remove-pristine-table rule selected non-pristine table" + ); + Ok(()) + } +} + +struct RemoveEmptyEventTableRule; + +impl MigrationRule for RemoveEmptyEventTableRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + ctx.schema + .tables + .iter() + .filter(|table| table.is_event && ctx.model.row_count_by_table_name(&table.name) == 0) + .map(|table| { + AcceptedCase::new( + rule, + SchemaRewrite::RemoveTable { + table: table.name.clone(), + }, + ) + .protecting(vec![Resource::Table(table.name.clone())]) + }) + .collect() + } + + fn validate_accepted( + &self, + case: &AcceptedCase, + original: RuleCtx<'_>, + _final_schema: &SchemaPlan, + ) -> anyhow::Result<()> { + let SchemaRewrite::RemoveTable { table } = &case.rewrite else { + anyhow::bail!("remove-empty-event-table rule emitted non-remove rewrite"); + }; + let table_plan = original_table(original, table)?; + anyhow::ensure!( + table_plan.is_event, + "remove-empty-event-table rule selected non-event table" + ); + anyhow::ensure!( + original.model.row_count_by_table_name(table) == 0, + "remove-empty-event-table rule selected non-empty table" + ); + Ok(()) + } +} + +struct AddColumnRule; + +impl MigrationRule for AddColumnRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, add_column_rewrites(ctx.schema)) + } +} + +struct AddIndexRule; + +impl MigrationRule for AddIndexRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, add_index_rewrites(ctx.schema)) + } + + fn invalid_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + rejected_cases_from_rewrites(rule, duplicate_index_rewrites(ctx.schema)) + } +} + +struct RemoveIndexRule; + +impl MigrationRule for RemoveIndexRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, remove_index_rewrites(ctx.schema)) + } +} + +struct ChangeIndexRule; + +impl MigrationRule for ChangeIndexRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, change_index_rewrites(ctx.schema)) + } +} + +struct AddSequenceOnPristineTableRule; + +impl MigrationRule for AddSequenceOnPristineTableRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + ctx.schema + .tables + .iter() + .filter(|table| !table.is_event && table_is_pristine(ctx.model, table)) + .flat_map(|table| { + addable_sequences(table).into_iter().map(move |sequence| { + AcceptedCase::new( + rule, + SchemaRewrite::alter_table(table, [TableMigrationOp::AddSequence { sequence }]), + ) + .protecting(vec![Resource::Table(table.name.clone())]) + }) + }) + .collect() + } + + fn invalid_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + duplicate_sequence_rewrites(ctx.schema, ctx.model) + .into_iter() + .map(|rewrite| RejectedCase::new(rule, rewrite)) + .collect() + } + + fn validate_accepted( + &self, + case: &AcceptedCase, + original: RuleCtx<'_>, + _final_schema: &SchemaPlan, + ) -> anyhow::Result<()> { + let (table, sequence) = alter_table_sequence(case)?; + let table_plan = original_table(original, table)?; + anyhow::ensure!(!table_plan.is_event, "add-sequence-pristine rule selected event table"); + anyhow::ensure!( + table_is_pristine(original.model, table_plan), + "add-sequence-pristine rule selected non-pristine table" + ); + let column_plan = table_plan + .columns + .get(sequence.column) + .ok_or_else(|| anyhow::anyhow!("add-sequence-pristine rule references missing column"))?; + anyhow::ensure!( + column_plan.ty.is_integral(), + "add-sequence-pristine rule selected non-integral column" + ); + anyhow::ensure!( + table_plan + .sequences + .iter() + .all(|existing| existing.column != sequence.column), + "add-sequence-pristine rule selected already sequenced column" + ); + Ok(()) + } +} + +struct RemoveSequenceRule; + +impl MigrationRule for RemoveSequenceRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, remove_sequence_rewrites(ctx.schema)) + } +} + +struct AddUniqueConstraintOnPristineTableRule; + +impl MigrationRule for AddUniqueConstraintOnPristineTableRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + ctx.schema + .tables + .iter() + .filter(|table| !table.is_event && table_is_pristine(ctx.model, table)) + .flat_map(|table| { + addable_unique_constraint_columns(table) + .into_iter() + .map(move |columns| { + let mut ops = Vec::new(); + if !has_index(table, &columns) { + ops.push(TableMigrationOp::AddIndex { + columns: columns.clone(), + algorithm: IndexAlgorithm::BTree, + }); + } + ops.push(TableMigrationOp::AddUniqueConstraint { columns }); + AcceptedCase::new(rule, SchemaRewrite::alter_table(table, ops)) + .protecting(vec![Resource::Table(table.name.clone())]) + }) + }) + .collect() + } + + fn invalid_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + duplicate_unique_constraint_rewrites(ctx.schema, ctx.model) + .into_iter() + .map(|rewrite| RejectedCase::new(rule, rewrite)) + .collect() + } + + fn validate_accepted( + &self, + case: &AcceptedCase, + original: RuleCtx<'_>, + _final_schema: &SchemaPlan, + ) -> anyhow::Result<()> { + let SchemaRewrite::AlterTable { table, ops } = &case.rewrite else { + anyhow::bail!("add-unique-constraint rule emitted non-table rewrite"); + }; + let table_plan = original_table(original, table)?; + anyhow::ensure!(!table_plan.is_event, "add-unique-constraint rule selected event table"); + anyhow::ensure!( + table_is_pristine(original.model, table_plan), + "add-unique-constraint rule selected non-pristine table" + ); + anyhow::ensure!( + ops.iter() + .any(|op| matches!(op, TableMigrationOp::AddUniqueConstraint { .. })), + "add-unique-constraint rule emitted rewrite without unique-constraint op" + ); + Ok(()) + } +} + +struct RemoveUniqueConstraintRule; + +impl MigrationRule for RemoveUniqueConstraintRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, remove_unique_constraint_rewrites(ctx.schema)) + } +} + +struct DropPrimaryKeyAndWidenSumRule; + +impl MigrationRule for DropPrimaryKeyAndWidenSumRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, drop_primary_key_and_widen_sum_rewrites(ctx.schema)) + } +} + +struct WidenSumColumnRule; + +impl MigrationRule for WidenSumColumnRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, widen_sum_column_rewrites(ctx.schema)) + } +} + +struct ReschemaEmptyEventTableRule; + +impl MigrationRule for ReschemaEmptyEventTableRule { + fn accepted_cases(&self, _rng: &Rng, ctx: RuleCtx<'_>, rule: RuleId) -> Vec { + accepted_cases_from_rewrites(rule, reschema_event_table_rewrites(ctx.schema, ctx.model)) + .into_iter() + .map(|case| { + let protected_table = match &case.rewrite { + SchemaRewrite::AlterTable { table, .. } => Some(table.clone()), + _ => None, + }; + if let Some(table) = protected_table { + case.protecting(vec![Resource::Table(table)]) + } else { + case + } + }) + .collect() + } + + fn validate_accepted( + &self, + case: &AcceptedCase, + original: RuleCtx<'_>, + _final_schema: &SchemaPlan, + ) -> anyhow::Result<()> { + let SchemaRewrite::AlterTable { table, .. } = &case.rewrite else { + anyhow::bail!("reschema-empty-event-table rule emitted non-table rewrite"); + }; + let table_plan = original_table(original, table)?; + anyhow::ensure!( + table_plan.is_event, + "reschema-empty-event-table rule selected non-event table" + ); + anyhow::ensure!( + original.model.row_count_by_table_name(table) == 0, + "reschema-empty-event-table rule selected non-empty table" + ); + anyhow::ensure!( + table_plan.columns.len() < MAX_EVENT_COLUMNS, + "reschema-empty-event-table rule selected full event table" + ); + Ok(()) + } +} + +fn add_table_rewrites(rng: &Rng, schema: &SchemaPlan) -> Vec { + if schema.tables.len() >= MAX_TABLES { + return Vec::new(); + } + + let generator = SchemaGenerator::new(rng.clone(), SchemaProfile::engine_dst()); + [false, true] + .into_iter() + .map(|is_event| SchemaRewrite::AddTable { + table: generator.gen_table_for_schema(schema, is_event), + }) + .collect() +} + +fn add_column_rewrites(schema: &SchemaPlan) -> Vec { + let types = addable_column_types(schema); + schema + .tables + .iter() + .filter(|table| !table.is_event && table.columns.len() < MAX_TABLE_COLUMNS) + .flat_map(|table| { + types.iter().copied().map(|ty| { + SchemaRewrite::alter_table( + table, + [TableMigrationOp::ChangeAccess, TableMigrationOp::AddColumn { ty }], + ) + }) + }) + .collect() +} + +fn addable_column_types(schema: &SchemaPlan) -> Vec { + Type::ALL + .iter() + .copied() + .filter(|ty| !matches!(ty, Type::Sum { .. }) || !schema_has_sum_column(schema)) + .collect() +} + +fn schema_has_sum_column(schema: &SchemaPlan) -> bool { + schema.tables.iter().any(table_has_sum_column) +} + +fn table_has_sum_column(table: &TablePlan) -> bool { + table.columns.iter().any(|column| matches!(column.ty, Type::Sum { .. })) +} + +fn add_index_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + addable_indexes(table).into_iter().map(|(columns, algorithm)| { + SchemaRewrite::alter_table(table, [TableMigrationOp::AddIndex { columns, algorithm }]) + }) + }) + .collect() +} + +fn duplicate_index_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + table.indexes.iter().map(|index| { + SchemaRewrite::alter_table( + table, + [TableMigrationOp::AddIndex { + columns: index.columns.clone(), + algorithm: index.algorithm, + }], + ) + }) + }) + .collect() +} + +fn remove_index_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + changeable_index_columns(table) + .into_iter() + .map(|columns| SchemaRewrite::alter_table(table, [TableMigrationOp::RemoveIndex { columns }])) + }) + .collect() +} + +fn change_index_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + changeable_index_columns(table).into_iter().map(|columns| { + SchemaRewrite::alter_table( + table, + [ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeIndex { columns }, + ], + ) + }) + }) + .collect() +} + +fn remove_sequence_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + removable_sequence_columns(table) + .into_iter() + .map(|column| SchemaRewrite::alter_table(table, [TableMigrationOp::RemoveSequence { column }])) + }) + .collect() +} + +fn duplicate_sequence_rewrites(schema: &SchemaPlan, model: &Model) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event && table_is_pristine(model, table)) + .flat_map(|table| { + table + .sequences + .iter() + .cloned() + .map(|sequence| SchemaRewrite::alter_table(table, [TableMigrationOp::AddSequence { sequence }])) + }) + .collect() +} + +fn remove_unique_constraint_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + removable_unique_constraint_columns(table).into_iter().map(|columns| { + SchemaRewrite::alter_table(table, [TableMigrationOp::RemoveUniqueConstraint { columns }]) + }) + }) + .collect() +} + +fn duplicate_unique_constraint_rewrites(schema: &SchemaPlan, model: &Model) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event && table_is_pristine(model, table)) + .flat_map(|table| { + table.unique_constraints.iter().map(|constraint| { + SchemaRewrite::alter_table( + table, + [TableMigrationOp::AddUniqueConstraint { + columns: constraint.columns.clone(), + }], + ) + }) + }) + .collect() +} + +fn drop_primary_key_and_widen_sum_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event && table.primary_key.is_some() && table.sequences.is_empty()) + .flat_map(|table| { + widenable_sum_columns(table).into_iter().map(|column| { + SchemaRewrite::alter_table( + table, + [ + TableMigrationOp::ChangePrimaryKey { column: None }, + TableMigrationOp::ChangeColumnType { column }, + ], + ) + }) + }) + .collect() +} + +fn widen_sum_column_rewrites(schema: &SchemaPlan) -> Vec { + schema + .tables + .iter() + .filter(|table| !table.is_event) + .flat_map(|table| { + widenable_sum_columns(table).into_iter().map(|column| { + SchemaRewrite::alter_table( + table, + [ + TableMigrationOp::ChangeAccess, + TableMigrationOp::ChangeColumnType { column }, + ], + ) + }) + }) + .collect() +} + +fn reschema_event_table_rewrites(schema: &SchemaPlan, model: &Model) -> Vec { + schema + .tables + .iter() + .filter(|table| table.is_event && model.row_count_by_table_name(&table.name) == 0) + .filter(|table| table.columns.len() < MAX_EVENT_COLUMNS) + .map(|table| { + SchemaRewrite::alter_table( + table, + [TableMigrationOp::ChangeAccess, TableMigrationOp::ReschemaEventTable], + ) + }) + .collect() +} + +impl SchemaRewrite { + fn alter_table(table: &TablePlan, ops: impl Into>) -> Self { + Self::AlterTable { + table: table.name.clone(), + ops: ops.into(), + } + } + + pub(crate) fn apply_to(&self, schema: &mut SchemaPlan) -> anyhow::Result<()> { + match self { + Self::AddTable { table } => { + anyhow::ensure!( + schema.tables.len() < MAX_TABLES, + "migration would exceed the configured maximum number of tables" + ); + anyhow::ensure!( + schema.tables.iter().all(|existing| existing.name != table.name), + "add-table migration selected existing table name" + ); + anyhow::ensure!( + !table_has_sum_column(table) || !schema_has_sum_column(schema), + "migration would add a second sum column" + ); + schema.tables.push(table.clone()); + } + Self::RemoveTable { table } => { + let table = table_position(schema, table)?; + if !schema.tables[table].is_event { + let non_event_tables = schema.tables.iter().filter(|table| !table.is_event).count(); + anyhow::ensure!(non_event_tables > 1, "migration would remove the last non-event table"); + } + schema.tables.remove(table); + } + Self::AlterTable { table, ops } => { + let table = table_position(schema, table)?; + for op in ops { + if matches!(op, TableMigrationOp::AddColumn { ty: Type::Sum { .. } }) { + anyhow::ensure!( + !schema_has_sum_column(schema), + "migration would add a second sum column" + ); + } + let table_plan = &mut schema.tables[table]; + apply_table_op(table_plan, op.clone())?; + } + } + } + Ok(()) + } + + fn apply_unchecked_to(&self, schema: &mut SchemaPlan) -> anyhow::Result<()> { + match self { + Self::AddTable { table } => schema.tables.push(table.clone()), + Self::RemoveTable { table } => { + let table = table_position(schema, table)?; + schema.tables.remove(table); + } + Self::AlterTable { table, ops } => { + let table = table_position(schema, table)?; + for op in ops { + let table_plan = &mut schema.tables[table]; + apply_table_op_unchecked(table_plan, op.clone())?; + } + } + } + Ok(()) + } +} + +fn table_position(schema: &SchemaPlan, table: &str) -> anyhow::Result { + schema + .tables + .iter() + .position(|table_plan| table_plan.name == table) + .ok_or_else(|| anyhow::anyhow!("migration references missing table {table}")) +} + +fn apply_table_op(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result<()> { + match op { + TableMigrationOp::ChangeAccess => { + table.is_public = !table.is_public; + } + TableMigrationOp::AddColumn { ty } => { + anyhow::ensure!(!table.is_event, "add-column migration selected event table"); + anyhow::ensure!( + table.columns.len() < MAX_TABLE_COLUMNS, + "table already has the configured maximum number of columns" + ); + table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "added_col"), + ty, + }); + } + TableMigrationOp::AddIndex { columns, algorithm } => { + ensure_columns_exist(table, &columns)?; + anyhow::ensure!( + !has_index(table, &columns), + "add-index migration selected existing index columns" + ); + table.indexes.push(IndexPlan { columns, algorithm }); + } + TableMigrationOp::RemoveIndex { columns } => { + ensure_columns_exist(table, &columns)?; + let Some(index) = table.indexes.iter().position(|index| index.columns == columns) else { + anyhow::bail!("remove-index migration references missing index on columns {columns:?}"); + }; + anyhow::ensure!( + !is_required_index(table, &columns), + "remove-index migration selected a required index" + ); + table.indexes.remove(index); + } + TableMigrationOp::AddSequence { sequence } => { + let column = sequence.column; + let Some(column_plan) = table.columns.get(column) else { + anyhow::bail!("add-sequence migration references missing column {column}"); + }; + anyhow::ensure!( + column_plan.ty.is_integral(), + "add-sequence migration selected non-integral column" + ); + anyhow::ensure!( + table.sequences.iter().all(|existing| existing.column != column), + "add-sequence migration selected an already sequenced column" + ); + table.sequences.push(sequence); + } + TableMigrationOp::RemoveSequence { column } => { + anyhow::ensure!( + column < table.columns.len(), + "remove-sequence migration references missing column" + ); + let Some(sequence) = table.sequences.iter().position(|sequence| sequence.column == column) else { + anyhow::bail!("remove-sequence migration references missing sequence on column {column}"); + }; + table.sequences.remove(sequence); + } + TableMigrationOp::AddUniqueConstraint { columns } => { + ensure_columns_exist(table, &columns)?; + anyhow::ensure!( + !table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == columns), + "add-constraint migration selected existing constraint columns" + ); + anyhow::ensure!( + has_index(table, &columns), + "add-constraint migration requires a matching index" + ); + table.unique_constraints.push(UniqueConstraintPlan { columns }); + } + TableMigrationOp::RemoveUniqueConstraint { columns } => { + ensure_columns_exist(table, &columns)?; + let Some(constraint) = table + .unique_constraints + .iter() + .position(|constraint| constraint.columns == columns) + else { + anyhow::bail!("remove-constraint migration references missing constraint on columns {columns:?}"); + }; + anyhow::ensure!( + !table.primary_key.is_some_and(|primary_key| columns == [primary_key]), + "remove-constraint migration selected primary-key constraint" + ); + table.unique_constraints.remove(constraint); + } + TableMigrationOp::ChangePrimaryKey { column } => { + if let Some(column) = column { + anyhow::ensure!( + column < table.columns.len(), + "primary-key migration references missing column" + ); + } + table.primary_key = column; + } + TableMigrationOp::ChangeIndex { columns } => { + ensure_columns_exist(table, &columns)?; + anyhow::ensure!( + !is_required_index(table, &columns), + "index migration selected a required index" + ); + let Some(index_plan) = table.indexes.iter_mut().find(|index| index.columns == columns) else { + anyhow::bail!("index migration references missing index on columns {columns:?}"); + }; + index_plan.algorithm = match index_plan.algorithm { + IndexAlgorithm::BTree => IndexAlgorithm::Hash, + IndexAlgorithm::Hash => IndexAlgorithm::BTree, + }; + } + TableMigrationOp::ChangeColumnType { column } => { + widen_sum_column(table, Some(column))?; + } + TableMigrationOp::ReschemaEventTable => { + anyhow::ensure!(table.is_event, "event-table migration selected non-event table"); + anyhow::ensure!( + table.columns.len() < MAX_EVENT_COLUMNS, + "event table already has the configured maximum number of columns" + ); + table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "reschema_payload"), + ty: Type::U64, + }); + } + } + + Ok(()) +} + +fn apply_table_op_unchecked(table: &mut TablePlan, op: TableMigrationOp) -> anyhow::Result<()> { + match op { + TableMigrationOp::ChangeAccess => table.is_public = !table.is_public, + TableMigrationOp::AddColumn { ty } => table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "added_col"), + ty, + }), + TableMigrationOp::AddIndex { columns, algorithm } => table.indexes.push(IndexPlan { columns, algorithm }), + TableMigrationOp::RemoveIndex { columns } => { + let index = table + .indexes + .iter() + .position(|index| index.columns == columns) + .ok_or_else(|| anyhow::anyhow!("unchecked remove-index references missing columns {columns:?}"))?; + table.indexes.remove(index); + } + TableMigrationOp::AddSequence { sequence } => table.sequences.push(sequence), + TableMigrationOp::RemoveSequence { column } => { + let sequence = table + .sequences + .iter() + .position(|sequence| sequence.column == column) + .ok_or_else(|| anyhow::anyhow!("unchecked remove-sequence references missing column {column}"))?; + table.sequences.remove(sequence); + } + TableMigrationOp::AddUniqueConstraint { columns } => { + table.unique_constraints.push(UniqueConstraintPlan { columns }); + } + TableMigrationOp::RemoveUniqueConstraint { columns } => { + let constraint = table + .unique_constraints + .iter() + .position(|constraint| constraint.columns == columns) + .ok_or_else(|| anyhow::anyhow!("unchecked remove-constraint references missing columns {columns:?}"))?; + table.unique_constraints.remove(constraint); + } + TableMigrationOp::ChangePrimaryKey { column } => table.primary_key = column, + TableMigrationOp::ChangeIndex { columns } => { + let index_plan = table + .indexes + .iter_mut() + .find(|index| index.columns == columns) + .ok_or_else(|| anyhow::anyhow!("unchecked change-index references missing columns {columns:?}"))?; + index_plan.algorithm = match index_plan.algorithm { + IndexAlgorithm::BTree => IndexAlgorithm::Hash, + IndexAlgorithm::Hash => IndexAlgorithm::BTree, + }; + } + TableMigrationOp::ChangeColumnType { column } => widen_sum_column(table, Some(column))?, + TableMigrationOp::ReschemaEventTable => table.columns.push(ColumnPlan { + name: SchemaNames::fresh_column_name(table, "reschema_payload"), + ty: Type::U64, + }), + } + Ok(()) +} + +fn widenable_sum_columns(table: &TablePlan) -> Vec { + table + .columns + .iter() + .enumerate() + .filter_map(|(column, column_plan)| match column_plan.ty { + Type::Sum { variants } if variants < MAX_SUM_VARIANTS => Some(column), + _ => None, + }) + .collect() +} + +fn addable_indexes(table: &TablePlan) -> Vec<(Vec, IndexAlgorithm)> { + (0..table.columns.len()) + .filter_map(|column| { + let columns = vec![column]; + (!has_index(table, &columns)).then_some((columns, IndexAlgorithm::BTree)) + }) + .collect() +} + +fn changeable_index_columns(table: &TablePlan) -> Vec> { + table + .indexes + .iter() + .filter_map(|index| (!is_required_index(table, &index.columns)).then_some(index.columns.clone())) + .collect() +} + +fn addable_sequences(table: &TablePlan) -> Vec { + table + .columns + .iter() + .enumerate() + .filter_map(|(column, column_plan)| { + (column_plan.ty.is_integral() && table.sequences.iter().all(|sequence| sequence.column != column)) + .then(|| SequencePlan::new(column, column_plan.ty).expect("column type checked above")) + }) + .collect() +} + +fn removable_sequence_columns(table: &TablePlan) -> Vec { + table.sequences.iter().map(|sequence| sequence.column).collect() +} + +fn addable_unique_constraint_columns(table: &TablePlan) -> Vec> { + (0..table.columns.len()) + .filter_map(|column| { + let columns = vec![column]; + (!table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == columns)) + .then_some(columns) + }) + .collect() +} + +fn removable_unique_constraint_columns(table: &TablePlan) -> Vec> { + table + .unique_constraints + .iter() + .filter_map(|constraint| { + (!table + .primary_key + .is_some_and(|primary_key| constraint.columns == [primary_key])) + .then_some(constraint.columns.clone()) + }) + .collect() +} + +fn has_index(table: &TablePlan, columns: &[usize]) -> bool { + table.indexes.iter().any(|index| index.columns == columns) +} + +fn is_required_index(table: &TablePlan, columns: &[usize]) -> bool { + table.primary_key.is_some_and(|primary_key| columns == [primary_key]) + || table + .unique_constraints + .iter() + .any(|constraint| constraint.columns == columns) +} + +fn ensure_columns_exist(table: &TablePlan, columns: &[usize]) -> anyhow::Result<()> { + anyhow::ensure!(!columns.is_empty(), "migration selected empty column list"); + anyhow::ensure!( + columns.iter().all(|&column| column < table.columns.len()), + "migration references missing column" + ); + Ok(()) +} + +fn widen_sum_column(table: &mut TablePlan, column: Option) -> anyhow::Result<()> { + let column = column.ok_or_else(|| anyhow::anyhow!("sum-widening migration missing column"))?; + let Some(column_plan) = table.columns.get_mut(column) else { + anyhow::bail!("sum-widening migration references missing column {column}"); + }; + + let Type::Sum { variants } = &mut column_plan.ty else { + anyhow::bail!("sum-widening migration selected a non-sum column"); + }; + anyhow::ensure!( + *variants < MAX_SUM_VARIANTS, + "sum-widening migration selected maxed-out sum column" + ); + *variants += 1; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn i64_table(name: impl Into, is_event: bool) -> TablePlan { + TablePlan { + name: name.into(), + columns: vec![ColumnPlan { + name: "id".into(), + ty: Type::I64, + }], + primary_key: None, + indexes: vec![], + unique_constraints: vec![], + sequences: vec![], + is_public: true, + is_event, + } + } + + fn indexed_unique_i64_schema() -> SchemaPlan { + let mut table = i64_table("items", false); + table.indexes.push(IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }); + table.unique_constraints.push(UniqueConstraintPlan { columns: vec![0] }); + SchemaPlan { tables: vec![table] } + } + + #[test] + fn migration_rule_registry_has_unique_ids_and_default_weights_resolve() { + let mut seen = HashSet::new(); + for entry in migration_rules() { + assert!(seen.insert(entry.id), "duplicate registered RuleId: {:?}", entry.id); + } + + for &(id, weight) in DEFAULT_RULE_WEIGHTS.rules { + assert!(weight > 0, "default rule weights should not contain zero entries"); + assert!( + rule_by_id(id).is_some(), + "default weights reference unregistered rule {id:?}" + ); + } + } + + #[test] + fn rejected_case_pool_uses_registered_rule_invalidations() { + let schema = indexed_unique_i64_schema(); + let model = Model::new(schema.clone()); + let cases = rejected_cases_for_weights( + &Rng::new(0), + RuleCtx { + schema: &schema, + model: &model, + }, + DEFAULT_RULE_WEIGHTS.rules, + ); + + assert!(!cases.is_empty(), "expected at least one rejected migration candidate"); + assert!(cases.iter().all(|case| rule_by_id(case.rule).is_some())); + } + + #[test] + fn rejected_batch_contains_one_invalid_case() { + let schema = indexed_unique_i64_schema(); + let model = Model::new(schema.clone()); + let ctx = RuleCtx { + schema: &schema, + model: &model, + }; + let rejected = pick_rejected_case(&Rng::new(0), ctx, &DEFAULT_RULE_WEIGHTS) + .expect("schema should have duplicate-index or duplicate-constraint invalidation"); + + let migration = build_validated_rejected_batch(ctx, Vec::new(), rejected).expect("rejected batch should build"); + + assert_eq!(migration.expectation(), MigrationExpectation::Rejected); + assert_ne!(migration.target_schema(), &schema); + } + + #[test] + fn accepted_batch_rejects_write_to_protected_table() { + let schema = indexed_unique_i64_schema(); + let model = Model::new(schema.clone()); + let table = &schema.tables[0]; + let sequence = SequencePlan::new(0, Type::I64).expect("i64 sequence"); + + let add_sequence = AcceptedCase::new( + RuleId::AddSequenceOnPristineTable, + SchemaRewrite::alter_table(table, [TableMigrationOp::AddSequence { sequence }]), + ) + .protecting(vec![Resource::Table(table.name.clone())]); + let add_column = AcceptedCase::new( + RuleId::AddColumn, + SchemaRewrite::alter_table(table, [TableMigrationOp::AddColumn { ty: Type::U64 }]), + ); + + let err = build_validated_accepted_batch( + RuleCtx { + schema: &schema, + model: &model, + }, + vec![add_sequence, add_column], + ) + .expect_err("protected table write should be rejected before emission"); + + assert!(err.to_string().contains("conflicting resources")); + } + + #[test] + fn accepted_batch_rejects_removing_every_non_event_table() { + let schema = SchemaPlan { + tables: vec![i64_table("left", false), i64_table("right", false)], + }; + let model = Model::new(schema.clone()); + + let remove_left = AcceptedCase::new( + RuleId::RemovePristineNonEventTable, + SchemaRewrite::RemoveTable { table: "left".into() }, + ); + let remove_right = AcceptedCase::new( + RuleId::RemovePristineNonEventTable, + SchemaRewrite::RemoveTable { table: "right".into() }, + ); + + let err = build_validated_accepted_batch( + RuleCtx { + schema: &schema, + model: &model, + }, + vec![remove_left, remove_right], + ) + .expect_err("batch should not remove all non-event tables"); + + assert!(err.to_string().contains("last non-event table")); + } + + #[test] + fn accepted_batch_rejects_exceeding_max_tables() { + let schema = SchemaPlan { + tables: (0..MAX_TABLES - 1) + .map(|idx| i64_table(format!("table_{idx}"), false)) + .collect(), + }; + let model = Model::new(schema.clone()); + + let add_first = AcceptedCase::new( + RuleId::AddTable, + SchemaRewrite::AddTable { + table: i64_table("new_a", false), + }, + ); + let add_second = AcceptedCase::new( + RuleId::AddTable, + SchemaRewrite::AddTable { + table: i64_table("new_b", false), + }, + ); + + let err = build_validated_accepted_batch( + RuleCtx { + schema: &schema, + model: &model, + }, + vec![add_first, add_second], + ) + .expect_err("batch should not exceed table cap"); + + assert!(err.to_string().contains("maximum number of tables")); + } + + #[test] + fn accepted_batch_rejects_adding_two_sum_columns() { + let schema = SchemaPlan { + tables: vec![i64_table("left", false), i64_table("right", false)], + }; + let model = Model::new(schema.clone()); + + let add_left_sum = AcceptedCase::new( + RuleId::AddColumn, + SchemaRewrite::alter_table( + &schema.tables[0], + [TableMigrationOp::AddColumn { + ty: Type::Sum { variants: 1 }, + }], + ), + ); + let add_right_sum = AcceptedCase::new( + RuleId::AddColumn, + SchemaRewrite::alter_table( + &schema.tables[1], + [TableMigrationOp::AddColumn { + ty: Type::Sum { variants: 1 }, + }], + ), + ); + + let err = build_validated_accepted_batch( + RuleCtx { + schema: &schema, + model: &model, + }, + vec![add_left_sum, add_right_sum], + ) + .expect_err("batch should not add two sum columns"); + + assert!(err.to_string().contains("second sum column")); + } +} diff --git a/crates/dst/src/engine/model.rs b/crates/dst/src/engine/model.rs index 5d6913212e0..70fb14c1bb1 100644 --- a/crates/dst/src/engine/model.rs +++ b/crates/dst/src/engine/model.rs @@ -1,7 +1,10 @@ -use super::workload::{ - normalize_rows, CommitDelta, CountState, InsertOutcome, Interaction, Observation, Row, TableDelta, TableRowCount, -}; -use crate::schema::SchemaPlan; +use spacetimedb_lib::{AlgebraicValue, ProductValue}; + +use super::migrations::MigrationExpectation; +use super::row::{normalize_rows, Row}; +use super::state::{schema_state_for_plan, CommitDelta, CountState, TableDelta, TableRowCount, TableRows}; +use super::workload::{InsertOutcome, Interaction, Observation}; +use crate::schema::{ColumnPlan, SchemaPlan, TablePlan}; #[derive(Debug)] pub struct Model { @@ -13,6 +16,7 @@ pub struct Model { #[derive(Debug)] struct TableState { rows: Vec, + ever_inserted: bool, } #[derive(Debug)] @@ -44,7 +48,14 @@ impl PendingTx { impl Model { pub fn new(schema: SchemaPlan) -> Self { - let committed_tables = schema.tables.iter().map(|_| TableState { rows: vec![] }).collect(); + let committed_tables = schema + .tables + .iter() + .map(|_| TableState { + rows: vec![], + ever_inserted: false, + }) + .collect(); Self { schema, committed_tables, @@ -90,9 +101,18 @@ impl Model { self.visible_rows(table).any(matches) } - fn violates_unique_constraint(&self, table: usize, row: &Row) -> bool { + fn violates_unique_constraint(&self, table: usize, row: &Row, ignore_sequence_constraints: bool) -> bool { let table_plan = &self.schema.tables[table]; for constraint in &table_plan.unique_constraints { + if ignore_sequence_constraints + && constraint + .columns + .iter() + .any(|&column| column_has_sequence(table_plan, column)) + { + continue; + } + if self.any_visible_row(table, |visible_row| { constraint .columns @@ -105,35 +125,51 @@ impl Model { false } - pub fn apply(&mut self, interaction: &Interaction) -> Observation { + pub(crate) fn insert_would_violate_unique_constraint(&self, table: usize, row: &Row) -> bool { + !self.any_visible_row(table, |visible_row| visible_row == row) + && self.violates_unique_constraint(table, row, true) + } + + pub fn apply(&mut self, interaction: &Interaction, observation: &Observation) -> anyhow::Result { match interaction { Interaction::BeginMutTx => { + anyhow::ensure!( + matches!(observation, Observation::BeganMutTx), + "begin-mut-tx produced unexpected observation" + ); debug_assert!(self.pending_tx.is_none()); - self.pending_tx = Some(PendingTx::new(self.committed_tables.len())); - Observation::BeganMutTx + self.pending_tx = Some(PendingTx::new(self.schema.tables.len())); + Ok(Observation::BeganMutTx) } - Interaction::Insert { table, row } => { + Interaction::Insert { table, .. } => { debug_assert!(self.pending_tx.is_some()); - // Properties feed the target-returned row here, so sequence-generated - // values become part of the oracle before commit/replay checks run. - if self.any_visible_row(*table, |visible_row| visible_row == row) { - return Observation::Inserted { - outcome: InsertOutcome::Accepted(row.clone()), - }; - } - - if self.violates_unique_constraint(*table, row) { - return Observation::Inserted { - outcome: InsertOutcome::UniqueConstraintViolation, - }; - } - - self.pending_table_mut(*table).inserts.push(row.clone()); - Observation::Inserted { - outcome: InsertOutcome::Accepted(row.clone()), + let Observation::Inserted { outcome } = observation else { + anyhow::bail!("insert produced unexpected observation"); + }; + + self.committed_tables[*table].ever_inserted = true; + match outcome { + InsertOutcome::Accepted(row) => { + let already_visible = self.any_visible_row(*table, |visible_row| visible_row == row); + anyhow::ensure!( + already_visible || !self.violates_unique_constraint(*table, row, false), + "target accepted row that violates a visible unique constraint" + ); + if !already_visible { + self.pending_table_mut(*table).inserts.push(row.clone()); + } + Ok(Observation::Inserted { + outcome: InsertOutcome::Accepted(row.clone()), + }) + } + InsertOutcome::UniqueConstraintViolation { .. } => Ok(observation.clone()), } } Interaction::Delete { table, row } => { + anyhow::ensure!( + matches!(observation, Observation::Deleted), + "delete produced unexpected observation" + ); debug_assert!(self.pending_tx.is_some()); if self.any_visible_row(*table, |visible_row| visible_row == row) { let committed_has_row = self.visible_committed_rows(*table).any(|committed| committed == row); @@ -143,27 +179,48 @@ impl Model { pending_table.deletes.push(row.clone()); } } - Observation::Deleted + Ok(Observation::Deleted) } Interaction::CommitTx => { + anyhow::ensure!( + matches!(observation, Observation::Committed { .. }), + "commit produced unexpected observation" + ); debug_assert!(self.pending_tx.is_some()); let pending_tx = self.pending_tx.take().expect("active transaction"); let delta = self.commit_pending(pending_tx); - Observation::Committed { delta } + Ok(Observation::Committed { delta }) + } + Interaction::Migrate(migration) => { + debug_assert!(self.pending_tx.is_none()); + match migration.expectation() { + MigrationExpectation::Accepted => { + let old_schema = std::mem::replace(&mut self.schema, migration.schema().clone()); + let old_tables = std::mem::take(&mut self.committed_tables); + self.committed_tables = remap_table_states(&old_schema, &self.schema, old_tables); + Ok(Observation::Migrated) + } + MigrationExpectation::Rejected => Ok(Observation::MigrationRejected), + } } Interaction::Replay => { + anyhow::ensure!( + matches!(observation, Observation::Replayed { .. }), + "replay produced unexpected observation" + ); self.pending_tx = None; - Observation::Replayed { + Ok(Observation::Replayed { state: self.count_state(), - } + }) } } } fn commit_pending(&mut self, pending_tx: PendingTx) -> CommitDelta { + let PendingTx { tables: pending_tables } = pending_tx; let mut tables = Vec::new(); - for (table, pending_table) in pending_tx.tables.into_iter().enumerate() { + for (table, pending_table) in pending_tables.into_iter().enumerate() { if pending_table.inserts.is_empty() && pending_table.deletes.is_empty() { continue; } @@ -217,6 +274,26 @@ impl Model { self.visible_count(table) as usize } + pub fn ever_inserted(&self, table: usize) -> bool { + self.committed_tables[table].ever_inserted + } + + pub(crate) fn row_count_by_table_name(&self, table: &str) -> usize { + self.table_index(table).map_or(0, |table| self.row_count(table)) + } + + pub(crate) fn ever_inserted_by_table_name(&self, table: &str) -> bool { + self.table_index(table) + .is_some_and(|table| self.committed_tables[table].ever_inserted) + } + + fn table_index(&self, table: &str) -> Option { + self.schema + .tables + .iter() + .position(|table_plan| table_plan.name == table) + } + pub fn row(&self, table: usize, row: usize) -> Option<&Row> { self.visible_rows(table).nth(row) } @@ -233,16 +310,92 @@ impl Model { count: self.visible_count(table), }) .collect(); - CountState { row_counts } + let table_rows = (0..self.schema.tables.len()) + .map(|table| TableRows { + table, + rows: normalize_rows(self.visible_rows(table).cloned().collect()), + }) + .collect(); + + CountState { + row_counts, + table_rows, + schema: schema_state_for_plan(&self.schema), + } + } +} + +fn column_has_sequence(table: &TablePlan, column: usize) -> bool { + table.sequences.iter().any(|sequence| sequence.column == column) +} + +fn remap_table_states( + old_schema: &SchemaPlan, + new_schema: &SchemaPlan, + old_tables: Vec, +) -> Vec { + let mut old_tables = old_tables.into_iter().map(Some).collect::>(); + new_schema + .tables + .iter() + .map(|new_table| { + let Some(old_table_idx) = old_schema + .tables + .iter() + .position(|old_table| old_table.name == new_table.name) + else { + return TableState { + rows: vec![], + ever_inserted: false, + }; + }; + + let old_table = &old_schema.tables[old_table_idx]; + let old_state = old_tables[old_table_idx] + .take() + .expect("old table state is consumed once"); + remap_table_state(old_table, new_table, old_state) + }) + .collect() +} + +fn remap_table_state(old_table: &TablePlan, new_table: &TablePlan, state: TableState) -> TableState { + TableState { + rows: state + .rows + .into_iter() + .map(|row| remap_row(old_table, new_table, row)) + .collect(), + ever_inserted: state.ever_inserted, + } +} + +fn remap_row(old_table: &TablePlan, new_table: &TablePlan, row: Row) -> Row { + let elements = new_table + .columns + .iter() + .map(|new_column| remap_value(old_table, new_column, &row)) + .collect::>(); + ProductValue { + elements: elements.into_boxed_slice(), } } +fn remap_value(old_table: &TablePlan, new_column: &ColumnPlan, row: &Row) -> AlgebraicValue { + old_table + .columns + .iter() + .position(|old_column| old_column.name == new_column.name) + .map(|old_column| row.elements[old_column].clone()) + .unwrap_or_else(|| new_column.ty.default_value()) +} + #[cfg(test)] mod tests { use spacetimedb_lib::AlgebraicValue; use super::*; - use crate::schema::{ColumnPlan, IndexAlgorithm, IndexPlan, TablePlan, Type, UniqueConstraintPlan}; + use crate::schema::{ColumnPlan, IndexAlgorithm, IndexPlan, SequencePlan, TablePlan, Type, UniqueConstraintPlan}; fn schema() -> SchemaPlan { SchemaPlan { @@ -260,6 +413,7 @@ mod tests { unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], sequences: vec![], is_public: true, + is_event: false, }], } } @@ -270,12 +424,90 @@ mod tests { } } + fn observe(model: &mut Model, interaction: Interaction, observation: Observation) -> Observation { + model + .apply(&interaction, &observation) + .expect("model should accept target observation") + } + + fn accepted(row: Row) -> Observation { + Observation::Inserted { + outcome: InsertOutcome::Accepted(row), + } + } + + fn committed() -> Observation { + Observation::Committed { + delta: CommitDelta { tables: Vec::new() }, + } + } + + fn replayed(model: &Model) -> Observation { + Observation::Replayed { + state: model.count_state(), + } + } + + fn keyed_payload_schema(sequence: bool) -> SchemaPlan { + SchemaPlan { + tables: vec![TablePlan { + name: "items".into(), + columns: vec![ + ColumnPlan { + name: "id".into(), + ty: Type::U64, + }, + ColumnPlan { + name: "payload".into(), + ty: Type::String, + }, + ], + primary_key: Some(0), + indexes: vec![IndexPlan { + columns: vec![0], + algorithm: IndexAlgorithm::BTree, + }], + unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], + sequences: sequence + .then(|| SequencePlan::new(0, Type::U64).expect("u64 sequence")) + .into_iter() + .collect(), + is_public: true, + is_event: false, + }], + } + } + + fn payload_row(id: u64, payload: &str) -> Row { + Row { + elements: vec![AlgebraicValue::U64(id), AlgebraicValue::String(payload.into())].into(), + } + } + + #[test] + fn insert_would_violate_unique_constraint_distinguishes_duplicates_from_conflicts() { + let mut model = Model::new(keyed_payload_schema(false)); + model.committed_tables[0].rows.push(payload_row(1, "a")); + + assert!(!model.insert_would_violate_unique_constraint(0, &payload_row(1, "a"))); + assert!(model.insert_would_violate_unique_constraint(0, &payload_row(1, "b"))); + assert!(!model.insert_would_violate_unique_constraint(0, &payload_row(2, "b"))); + } + + #[test] + fn insert_would_violate_unique_constraint_treats_sequence_constraints_as_ambiguous() { + let mut model = Model::new(keyed_payload_schema(true)); + model.committed_tables[0].rows.push(payload_row(1, "a")); + + assert!(!model.insert_would_violate_unique_constraint(0, &payload_row(0, "b"))); + } + #[test] fn begin_mut_tx_does_not_clone_committed_tables() { let mut model = Model::new(schema()); model.committed_tables[0].rows.push(row(1)); - model.apply(&Interaction::BeginMutTx); + observe(&mut model, Interaction::BeginMutTx, Observation::BeganMutTx); let pending_tx = model.pending_tx.as_ref().expect("active transaction"); assert!(pending_tx @@ -290,8 +522,12 @@ mod tests { let mut model = Model::new(schema()); model.committed_tables[0].rows.push(row(1)); - model.apply(&Interaction::BeginMutTx); - model.apply(&Interaction::Insert { table: 0, row: row(2) }); + observe(&mut model, Interaction::BeginMutTx, Observation::BeganMutTx); + observe( + &mut model, + Interaction::Insert { table: 0, row: row(2) }, + accepted(row(2)), + ); let pending_table = &model.pending_tx.as_ref().expect("active transaction").tables[0]; assert_eq!(pending_table.inserts, vec![row(2)]); @@ -306,8 +542,12 @@ mod tests { model.committed_tables[0].rows.push(row(1)); model.committed_tables[0].rows.push(row(2)); - model.apply(&Interaction::BeginMutTx); - model.apply(&Interaction::Delete { table: 0, row: row(1) }); + observe(&mut model, Interaction::BeginMutTx, Observation::BeganMutTx); + observe( + &mut model, + Interaction::Delete { table: 0, row: row(1) }, + Observation::Deleted, + ); let pending_table = &model.pending_tx.as_ref().expect("active transaction").tables[0]; assert!(pending_table.inserts.is_empty()); @@ -320,12 +560,17 @@ mod tests { fn insert_is_visible_before_commit_and_replay_rolls_back() { let mut model = Model::new(schema()); - model.apply(&Interaction::BeginMutTx); - model.apply(&Interaction::Insert { table: 0, row: row(1) }); + observe(&mut model, Interaction::BeginMutTx, Observation::BeganMutTx); + observe( + &mut model, + Interaction::Insert { table: 0, row: row(1) }, + accepted(row(1)), + ); assert_eq!(model.row_count(0), 1); - model.apply(&Interaction::Replay); - model.apply(&Interaction::BeginMutTx); + let replay = replayed(&model); + observe(&mut model, Interaction::Replay, replay); + observe(&mut model, Interaction::BeginMutTx, Observation::BeganMutTx); assert_eq!(model.row_count(0), 0); } @@ -333,9 +578,13 @@ mod tests { fn commit_applies_only_pending_overlay() { let mut model = Model::new(schema()); - model.apply(&Interaction::BeginMutTx); - model.apply(&Interaction::Insert { table: 0, row: row(1) }); - let observation = model.apply(&Interaction::CommitTx); + observe(&mut model, Interaction::BeginMutTx, Observation::BeganMutTx); + observe( + &mut model, + Interaction::Insert { table: 0, row: row(1) }, + accepted(row(1)), + ); + let observation = observe(&mut model, Interaction::CommitTx, committed()); let Observation::Committed { delta, .. } = observation else { panic!("expected commit observation"); @@ -350,8 +599,12 @@ mod tests { let mut model = Model::new(schema()); model.committed_tables[0].rows.push(row(1)); - model.apply(&Interaction::BeginMutTx); - model.apply(&Interaction::Delete { table: 0, row: row(1) }); + observe(&mut model, Interaction::BeginMutTx, Observation::BeganMutTx); + observe( + &mut model, + Interaction::Delete { table: 0, row: row(1) }, + Observation::Deleted, + ); assert_eq!(model.row_count(0), 0); } diff --git a/crates/dst/src/engine/properties.rs b/crates/dst/src/engine/properties.rs index 667eec09510..0f41398c10a 100644 --- a/crates/dst/src/engine/properties.rs +++ b/crates/dst/src/engine/properties.rs @@ -1,5 +1,6 @@ use super::model::Model; -use super::workload::{InsertOutcome, Interaction, Observation, Row}; +use super::state::CountState; +use super::workload::{Interaction, Observation}; use crate::schema::SchemaPlan; use crate::traits::Properties; @@ -13,8 +14,8 @@ impl EngineProperties { Self { oracle: EngineOracle::new(schema), properties: vec![ - Box::new(InsertMatches), Box::new(CommitMatches), + Box::new(MigrateMatches), Box::new(ReplayMatchesModel), ], } @@ -54,39 +55,15 @@ impl EngineOracle { } fn apply(&mut self, interaction: &Interaction, observation: &Observation) -> anyhow::Result { - let observation = match (interaction, observation) { - ( - Interaction::Insert { table, .. }, - Observation::Inserted { - outcome: InsertOutcome::Accepted(row), - }, - ) => self.apply_insert(*table, row), - ( - Interaction::Insert { .. }, - Observation::Inserted { - outcome: InsertOutcome::UniqueConstraintViolation, - }, - ) => self.model.apply(interaction), - (Interaction::Insert { .. }, _) => anyhow::bail!("insert produced unexpected observation"), - _ => self.model.apply(interaction), - }; - - Ok(observation) - } - - fn apply_insert(&mut self, table: usize, row: &Row) -> Observation { - self.model.apply(&Interaction::Insert { - table, - row: row.clone(), - }) + self.model.apply(interaction, observation) } } -struct InsertMatches; +struct CommitMatches; -impl EngineProperty for InsertMatches { +impl EngineProperty for CommitMatches { fn observes(&self, interaction: &Interaction) -> bool { - matches!(interaction, Interaction::Insert { .. }) + matches!(interaction, Interaction::CommitTx) } fn check( @@ -95,51 +72,38 @@ impl EngineProperty for InsertMatches { observation: &Observation, expected: &Observation, ) -> anyhow::Result<()> { - let Observation::Inserted { outcome } = observation else { - anyhow::bail!("insert_matches: insert produced unexpected observation"); + let Observation::Committed { delta, .. } = observation else { + anyhow::bail!("commit_matches: commit produced unexpected observation"); }; - let Observation::Inserted { outcome: expected } = expected else { - unreachable!("InsertMatches only subscribes to insert interactions"); + let Observation::Committed { delta: expected, .. } = expected else { + unreachable!("CommitMatches only subscribes to commit interactions"); }; - match (outcome, expected) { - (InsertOutcome::Accepted(row), InsertOutcome::Accepted(expected)) => { - anyhow::ensure!(row == expected, "insert_matches: accepted row diverged from model"); - } - (InsertOutcome::UniqueConstraintViolation, InsertOutcome::UniqueConstraintViolation) => {} - (InsertOutcome::Accepted(_), InsertOutcome::UniqueConstraintViolation) => { - anyhow::bail!("insert_matches: target accepted row rejected by model"); - } - (InsertOutcome::UniqueConstraintViolation, InsertOutcome::Accepted(_)) => { - anyhow::bail!("insert_matches: target rejected row accepted by model"); - } - } - + anyhow::ensure!(delta == expected, "commit_matches: committed delta diverged from model"); Ok(()) } } -struct CommitMatches; +struct MigrateMatches; -impl EngineProperty for CommitMatches { +impl EngineProperty for MigrateMatches { fn observes(&self, interaction: &Interaction) -> bool { - matches!(interaction, Interaction::CommitTx) + matches!(interaction, Interaction::Migrate(_)) } fn check( &self, - _interaction: &Interaction, + interaction: &Interaction, observation: &Observation, expected: &Observation, ) -> anyhow::Result<()> { - let Observation::Committed { delta, .. } = observation else { - anyhow::bail!("commit_matches: commit produced unexpected observation"); - }; - let Observation::Committed { delta: expected, .. } = expected else { - unreachable!("CommitMatches only subscribes to commit interactions"); - }; - - anyhow::ensure!(delta == expected, "commit_matches: committed delta diverged from model"); + anyhow::ensure!( + observation == expected, + "migrate_matches: migration outcome diverged from model +interaction: {interaction:#?} +target: {observation:#?} +model: {expected:#?}" + ); Ok(()) } } @@ -166,8 +130,107 @@ impl EngineProperty for ReplayMatchesModel { anyhow::ensure!( state == expected, - "replay_matches_model: replayed state diverged from model" + "replay_matches_model: replayed state diverged from model: {}", + count_state_mismatch(state, expected) ); Ok(()) } } + +fn count_state_mismatch(target: &CountState, model: &CountState) -> String { + if target.schema != model.schema { + if target.schema.tables.len() != model.schema.tables.len() { + return format!( + "schema table count target={} model={}", + target.schema.tables.len(), + model.schema.tables.len() + ); + } + + for (table_idx, (target_table, model_table)) in + target.schema.tables.iter().zip(&model.schema.tables).enumerate() + { + if target_table == model_table { + continue; + } + + if target_table.name != model_table.name { + return format!( + "schema table {table_idx} name target={:?} model={:?}", + target_table.name, model_table.name + ); + } + if target_table.columns != model_table.columns { + return format!( + "schema table {table_idx} columns target={:?} model={:?}", + target_table.columns, model_table.columns + ); + } + if target_table.indexes != model_table.indexes { + return format!( + "schema table {table_idx} indexes target={:?} model={:?}", + target_table.indexes, model_table.indexes + ); + } + if target_table.unique_constraints != model_table.unique_constraints { + return format!( + "schema table {table_idx} unique constraints target={:?} model={:?}", + target_table.unique_constraints, model_table.unique_constraints + ); + } + if target_table.sequences != model_table.sequences { + return format!( + "schema table {table_idx} sequences target={:?} model={:?}", + target_table.sequences, model_table.sequences + ); + } + return format!("schema table {table_idx} target={target_table:?} model={model_table:?}"); + } + } + + if target.row_counts != model.row_counts { + return format!("row counts target={:?} model={:?}", target.row_counts, model.row_counts); + } + + if target.table_rows.len() != model.table_rows.len() { + return format!( + "table row set count target={} model={}", + target.table_rows.len(), + model.table_rows.len() + ); + } + + for (table_idx, (target_rows, model_rows)) in target.table_rows.iter().zip(&model.table_rows).enumerate() { + if target_rows == model_rows { + continue; + } + if target_rows.table != model_rows.table { + return format!( + "table row entry {table_idx} id target={} model={}", + target_rows.table, model_rows.table + ); + } + if target_rows.rows.len() != model_rows.rows.len() { + return format!( + "table {} row count target={} model={}", + target_rows.table, + target_rows.rows.len(), + model_rows.rows.len() + ); + } + if let Some(row_idx) = target_rows + .rows + .iter() + .zip(&model_rows.rows) + .position(|(target_row, model_row)| target_row != model_row) + { + return format!( + "table {} row {row_idx} target={:?} model={:?}", + target_rows.table, target_rows.rows[row_idx], model_rows.rows[row_idx] + ); + } + return format!("table {} rows differ", target_rows.table); + } + + "unknown mismatch".into() +} diff --git a/crates/dst/src/engine/row.rs b/crates/dst/src/engine/row.rs new file mode 100644 index 00000000000..a3008d48d42 --- /dev/null +++ b/crates/dst/src/engine/row.rs @@ -0,0 +1,13 @@ +use spacetimedb_lib::bsatn::to_vec; +use spacetimedb_lib::ProductValue; + +pub type Row = ProductValue; + +pub fn row_to_bytes(row: &Row) -> Vec { + to_vec(row).expect("row serialization must not fail") +} + +pub fn normalize_rows(mut rows: Vec) -> Vec { + rows.sort_by_key(row_to_bytes); + rows +} diff --git a/crates/dst/src/engine/state.rs b/crates/dst/src/engine/state.rs new file mode 100644 index 00000000000..73f10a9d120 --- /dev/null +++ b/crates/dst/src/engine/state.rs @@ -0,0 +1,194 @@ +use spacetimedb_lib::db::auth::StAccess; +use spacetimedb_lib::AlgebraicType; +use spacetimedb_primitives::{ColId, TableId}; +use spacetimedb_schema::def::IndexAlgorithm as SchemaIndexAlgorithm; +use spacetimedb_schema::schema::{Schema, TableSchema}; + +use super::row::Row; +use crate::schema::{to_module_def, SchemaPlan}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CountState { + pub row_counts: Vec, + pub table_rows: Vec, + pub schema: SchemaState, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TableRowCount { + pub table: usize, + pub count: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableRows { + pub table: usize, + pub rows: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SchemaState { + pub tables: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableSchemaState { + pub table: usize, + pub name: String, + pub is_public: bool, + pub is_event: bool, + pub primary_key: Option, + pub columns: Vec, + pub indexes: Vec, + pub unique_constraints: Vec, + pub sequences: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ColumnState { + pub name: String, + pub ty: AlgebraicType, +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct IndexState { + pub columns: Vec, + pub algorithm: IndexAlgorithmState, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum IndexAlgorithmState { + BTree, + Hash, + Direct, + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct UniqueConstraintState { + pub columns: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct SequenceState { + pub column: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CommitDelta { + pub tables: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableDelta { + pub table: usize, + pub inserts: Vec, + pub deletes: Vec, + pub truncated: bool, +} + +impl From<&SchemaIndexAlgorithm> for IndexAlgorithmState { + fn from(algorithm: &SchemaIndexAlgorithm) -> Self { + match algorithm { + SchemaIndexAlgorithm::BTree(_) => Self::BTree, + SchemaIndexAlgorithm::Hash(_) => Self::Hash, + SchemaIndexAlgorithm::Direct(_) => Self::Direct, + _ => Self::Unknown, + } + } +} + +impl IndexState { + fn from_schema(algorithm: &SchemaIndexAlgorithm) -> Self { + Self { + columns: schema_index_columns(algorithm), + algorithm: algorithm.into(), + } + } +} + +impl UniqueConstraintState { + fn from_schema_columns(columns: impl IntoIterator) -> Self { + Self { + columns: columns.into_iter().map(|col| col.0 as usize).collect(), + } + } +} + +impl SequenceState { + fn from_schema_column(column: ColId) -> Self { + Self { + column: column.0 as usize, + } + } +} + +pub fn schema_state_for_plan(schema: &SchemaPlan) -> SchemaState { + schema_state_for_valid_plan(schema).expect("model schema must lower to a valid module definition") +} + +fn schema_state_for_valid_plan(schema: &SchemaPlan) -> anyhow::Result { + let module_def = to_module_def(schema)?; + + let tables = schema + .tables + .iter() + .enumerate() + .map(|(table, table_plan)| { + let table_def = module_def + .tables() + .find(|table_def| &*table_def.accessor_name == table_plan.name) + .ok_or_else(|| anyhow::anyhow!("validated schema is missing table accessor {:?}", table_plan.name))?; + let table_schema = TableSchema::from_module_def(&module_def, table_def, (), TableId::SENTINEL); + Ok(table_schema_state_for_schema(table, &table_schema)) + }) + .collect::>>()?; + + Ok(SchemaState { tables }) +} + +pub fn table_schema_state_for_schema(table: usize, schema: &TableSchema) -> TableSchemaState { + TableSchemaState { + table, + name: schema.table_name.to_string(), + is_public: schema.table_access == StAccess::Public, + is_event: schema.is_event, + primary_key: schema.primary_key.map(|col| col.0 as usize), + columns: schema + .columns + .iter() + .map(|column| ColumnState { + name: column.col_name.to_string(), + ty: column.col_type.clone(), + }) + .collect(), + indexes: sorted( + schema + .indexes + .iter() + .map(|index| IndexState::from_schema(&index.index_algorithm)), + ), + unique_constraints: sorted(schema.constraints.iter().filter_map(|constraint| { + constraint + .data + .unique_columns() + .map(|columns| UniqueConstraintState::from_schema_columns(columns.iter())) + })), + sequences: sorted( + schema + .sequences + .iter() + .map(|sequence| SequenceState::from_schema_column(sequence.col_pos)), + ), + } +} + +fn schema_index_columns(algorithm: &SchemaIndexAlgorithm) -> Vec { + algorithm.columns().iter().map(|col| col.0 as usize).collect() +} + +fn sorted(values: impl IntoIterator) -> Vec { + let mut values = values.into_iter().collect::>(); + values.sort(); + values +} diff --git a/crates/dst/src/engine/workload.rs b/crates/dst/src/engine/workload.rs index bc9e4fcc36e..a6f27cae48a 100644 --- a/crates/dst/src/engine/workload.rs +++ b/crates/dst/src/engine/workload.rs @@ -1,24 +1,30 @@ -use std::fmt::{Debug, Error, Formatter}; +//! Workload interaction generation for the engine DST driver. -use spacetimedb_lib::bsatn::to_vec; -use spacetimedb_lib::{AlgebraicValue, ProductValue}; -use spacetimedb_runtime::sim::Rng; -use spacetimedb_sats::ArrayValue; +use std::fmt::{Debug, Error as FmtError, Formatter}; +use super::generation::{GenCtx, GenerationState}; +use super::migrations::Migration; use super::model::Model; -use crate::schema::{SchemaPlan, TablePlan, Type}; - -pub type Row = ProductValue; +use super::row::Row; +use super::state::{CommitDelta, CountState}; +use crate::rng::{choice, pick_choice, Choice}; +use crate::schema::SchemaPlan; +use crate::traits::InteractionGen; +use anyhow::Error; +use spacetimedb_runtime::sim::Rng; +/// One generated action for the engine target to execute. #[derive(Debug, Clone)] pub enum Interaction { BeginMutTx, Insert { table: usize, row: Row }, Delete { table: usize, row: Row }, CommitTx, + Migrate(Migration), Replay, } +/// Counts of emitted workload interactions, reported at the end of each run. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct InteractionCounts { pub total: usize, @@ -26,6 +32,7 @@ pub struct InteractionCounts { pub insert: usize, pub delete: usize, pub commit_tx: usize, + pub migrate: usize, pub replay: usize, } @@ -38,67 +45,74 @@ impl InteractionCounts { Interaction::Insert { .. } => self.insert += 1, Interaction::Delete { .. } => self.delete += 1, Interaction::CommitTx => self.commit_tx += 1, + Interaction::Migrate(_) => self.migrate += 1, Interaction::Replay => self.replay += 1, } } } +/// Observable result of executing an interaction against the engine. #[derive(Debug, Clone, PartialEq, Eq)] pub enum Observation { BeganMutTx, Inserted { outcome: InsertOutcome }, Deleted, Committed { delta: CommitDelta }, + Migrated, + MigrationRejected, Replayed { state: CountState }, } #[derive(Debug, Clone, PartialEq, Eq)] pub enum InsertOutcome { Accepted(Row), - UniqueConstraintViolation, + UniqueConstraintViolation { details: String }, } +/// Runtime-tunable weights for top-level workload actions. #[derive(Debug, Clone, Copy)] pub struct InteractionWeights { pub insert: u64, pub delete: u64, pub commit_tx: u64, + pub migrate: u64, pub replay: u64, } impl Default for InteractionWeights { fn default() -> Self { Self { - insert: 50, - delete: 20, - commit_tx: 29, + insert: 5_000, + delete: 2_000, + commit_tx: 2_800, + migrate: 100, replay: 1, } } } -#[derive(Debug, Clone, Copy)] -enum InteractionChoice { - Insert, - Delete, - CommitTx, - Replay, -} - +/// Stateful interaction source that mirrors target observations into the model. pub struct WorkloadGen { rng: Rng, model: Model, + generation: GenerationState, stats: InteractionCounts, weights: InteractionWeights, } impl WorkloadGen { pub fn new(rng: Rng, model: Model) -> Self { + Self::with_weights(rng, model, InteractionWeights::default()) + } + + pub fn with_weights(rng: Rng, model: Model, weights: InteractionWeights) -> Self { + let generation = GenerationState::seeded(model.schema()); Self { rng, model, + generation, stats: InteractionCounts::default(), - weights: InteractionWeights::default(), + weights, } } @@ -106,70 +120,91 @@ impl WorkloadGen { self.stats } - fn schema(&self) -> &SchemaPlan { - self.model.schema() - } + pub fn next_interaction(&mut self) -> Interaction { + let choice = self.pick_interaction_choice(); + let interaction = self.interaction_from_choice(choice); - fn gen_value(&self, ty: Type) -> AlgebraicValue { - match ty { - Type::Bool => AlgebraicValue::Bool(self.rng.next_u64().is_multiple_of(2)), - Type::I64 => AlgebraicValue::I64(self.rng.next_u64() as i64), - Type::U64 => AlgebraicValue::U64(self.rng.next_u64()), - Type::String => AlgebraicValue::String(format!("v_{}", self.rng.next_u64()).into()), - Type::Bytes => { - let len = (self.rng.next_u64() % 16) as usize; - let bytes: Vec = (0..len).map(|_| self.rng.next_u64() as u8).collect(); - AlgebraicValue::Array(ArrayValue::U8(bytes.into())) - } - } - } + self.stats.record(&interaction); - fn gen_row(&self, table: &TablePlan) -> Row { - table - .columns - .iter() - .map(|c| self.gen_value(c.ty)) - .collect::() + interaction } - fn gen_insert_row(&self, table_idx: usize) -> Row { - let table = &self.schema().tables[table_idx]; - let mut row = self.gen_row(table); + fn observe_interaction(&mut self, interaction: &Interaction, observation: &Observation) -> Result<(), Error> { + let model_observation = self.model.apply(interaction, observation)?; - if let Some(sequence) = table.sequences.first() { - row.elements[sequence.column] = match table.columns[sequence.column].ty { - Type::I64 => AlgebraicValue::I64(0), - Type::U64 => AlgebraicValue::U64(0), - _ => unreachable!("sequence columns are integral"), - }; + self.observe_generation(interaction, &model_observation); + Ok(()) + } + + fn observe_generation(&mut self, interaction: &Interaction, model_observation: &Observation) { + match (interaction, model_observation) { + (Interaction::Insert { table, row }, Observation::Inserted { outcome }) => { + let table_plan = self.model.schema().tables[*table].clone(); + self.generation.observe_row(&table_plan, row); + if let InsertOutcome::Accepted(row) = outcome { + self.generation.observe_row(&table_plan, row); + } + } + (_, Observation::Migrated) => self.observe_schema_values(), + _ => {} } + } - row + fn observe_schema_values(&mut self) { + let schema = self.model.schema().clone(); + self.generation.seed_schema(&schema); + for (table, table_plan) in schema.tables.iter().enumerate() { + let rows = (0..self.model.row_count(table)) + .map(|row| self.model.row(table, row).expect("row index is in bounds").clone()) + .collect::>(); + for row in rows { + self.generation.observe_row(table_plan, &row); + } + } } +} - fn non_auto_inc_table_idx(&self) -> Option { - let auto_inc_table = self - .schema() - .auto_inc_table_and_column() - .map(|(table_idx, _)| table_idx); +#[derive(Debug, Clone, Copy)] +enum InteractionChoice { + Insert, + Delete, + CommitTx, + Migrate, + Replay, +} - (0..self.schema().tables.len()).find(|&table_idx| Some(table_idx) != auto_inc_table) +impl InteractionWeights { + fn choices(self) -> [Choice; 5] { + [ + choice(self.insert, InteractionChoice::Insert), + choice(self.delete, InteractionChoice::Delete), + choice(self.commit_tx, InteractionChoice::CommitTx), + choice(self.migrate, InteractionChoice::Migrate), + choice(self.replay, InteractionChoice::Replay), + ] } +} - pub fn next_interaction(&mut self) -> Interaction { - let choice = self.pick_interaction_choice(); - let interaction = self.interaction_from_choice(choice); - - self.model.apply(&interaction); - self.stats.record(&interaction); +impl WorkloadGen { + fn schema(&self) -> &SchemaPlan { + self.model.schema() + } - interaction + fn non_sequenced_table_idx(&self) -> Option { + (0..self.schema().tables.len()).find(|&table_idx| { + let table = &self.schema().tables[table_idx]; + !table.is_event && table.sequences.is_empty() + }) } fn interaction_from_choice(&mut self, choice: InteractionChoice) -> Interaction { if !self.model.in_mut_tx() { return match choice { InteractionChoice::Replay => Interaction::Replay, + InteractionChoice::Migrate => self + .gen_migration() + .map(Interaction::Migrate) + .unwrap_or(Interaction::Replay), // Insert/Delete/CommitTx are not legal outside a mutable tx. // Treat those weighted choices as pressure to start one. @@ -182,12 +217,15 @@ impl WorkloadGen { match choice { InteractionChoice::Replay => Interaction::Replay, + InteractionChoice::Migrate => Interaction::CommitTx, + InteractionChoice::Insert => { let table = self.insert_table_idx(); + let mut ctx = GenCtx::new(&self.rng, &self.model, &mut self.generation); Interaction::Insert { table, - row: self.gen_insert_row(table), + row: ctx.gen_insert_row(table), } } @@ -213,96 +251,70 @@ impl WorkloadGen { } fn pick_interaction_choice(&mut self) -> InteractionChoice { - let weights = self.weights; - - match self.pick_weighted(&[weights.insert, weights.delete, weights.commit_tx, weights.replay]) { - 0 => InteractionChoice::Insert, - 1 => InteractionChoice::Delete, - 2 => InteractionChoice::CommitTx, - 3 => InteractionChoice::Replay, - _ => unreachable!(), - } + let choices = self.weights.choices(); + pick_choice(&self.rng, &choices) } - fn pick_weighted(&mut self, weights: &[u64]) -> usize { - let total: u64 = weights.iter().sum(); - - assert!(total > 0, "at least one interaction weight must be non-zero"); - - let mut selected = self.rng.next_u64() % total; - - for (idx, weight) in weights.iter().copied().enumerate() { - if selected < weight { - return idx; - } + fn insert_table_idx(&self) -> usize { + let sequenced_tables = self.sequenced_table_indices(); + let data_tables = self.data_table_indices(); - selected -= weight; + if !sequenced_tables.is_empty() && !self.rng.next_u64().is_multiple_of(3) { + sequenced_tables[self.rng.index(sequenced_tables.len())] + } else { + data_tables[self.rng.index(data_tables.len())] } + } - unreachable!("selected value is always inside total weight") + fn sequenced_table_indices(&self) -> Vec { + self.schema() + .tables + .iter() + .enumerate() + .filter_map(|(table_idx, table)| (!table.is_event && !table.sequences.is_empty()).then_some(table_idx)) + .collect() } - fn insert_table_idx(&self) -> usize { - let auto_inc_table_idx = self + fn data_table_indices(&self) -> Vec { + let data_tables: Vec<_> = self .schema() - .auto_inc_table_and_column() - .map(|(table_idx, _)| table_idx); + .tables + .iter() + .enumerate() + .filter_map(|(table_idx, table)| (!table.is_event).then_some(table_idx)) + .collect(); + assert!( + !data_tables.is_empty(), + "engine DST schema must include a non-event table" + ); + data_tables + } - match auto_inc_table_idx { - Some(table_idx) if !self.rng.next_u64().is_multiple_of(3) => table_idx, - _ => self.rng.index(self.schema().tables.len()), - } + fn gen_migration(&mut self) -> Option { + GenCtx::new(&self.rng, &self.model, &mut self.generation).gen_migration() } fn deletable_table_idx(&self) -> Option { - self.non_auto_inc_table_idx() + self.non_sequenced_table_idx() .filter(|&table_idx| self.model.row_count(table_idx) > 0) } } impl Debug for WorkloadGen { - fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> { + fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), FmtError> { write!(f, "{:?}", self.stats()) } } -impl Iterator for WorkloadGen { - type Item = Interaction; +impl InteractionGen for WorkloadGen { + type Interaction = Interaction; + type Observation = Observation; - fn next(&mut self) -> Option { - Some(self.next_interaction()) + fn next_interaction(&mut self) -> Option { + Some(WorkloadGen::next_interaction(self)) } -} - -pub fn row_to_bytes(row: &Row) -> Vec { - to_vec(row).expect("row serialization must not fail") -} - -pub fn normalize_rows(mut rows: Vec) -> Vec { - rows.sort_by_key(row_to_bytes); - rows -} -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CountState { - pub row_counts: Vec, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct TableRowCount { - pub table: usize, - pub count: u64, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CommitDelta { - pub tables: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct TableDelta { - pub table: usize, - pub inserts: Vec, - pub deletes: Vec, - pub truncated: bool, + fn observe(&mut self, interaction: &Self::Interaction, observation: &Self::Observation) -> Result<(), Error> { + self.observe_interaction(interaction, observation) + } } diff --git a/crates/dst/src/lib.rs b/crates/dst/src/lib.rs index 8d12c575e4c..ef8e754f6de 100644 --- a/crates/dst/src/lib.rs +++ b/crates/dst/src/lib.rs @@ -1,4 +1,5 @@ pub mod engine; +mod rng; pub mod schema; pub mod sim; pub mod traits; diff --git a/crates/dst/src/rng.rs b/crates/dst/src/rng.rs new file mode 100644 index 00000000000..8c1360e5b4f --- /dev/null +++ b/crates/dst/src/rng.rs @@ -0,0 +1,81 @@ +//! Deterministic random-selection helpers shared by DST generators. + +use spacetimedb_runtime::sim::Rng; + +/// A weighted value in a deterministic choice table. +#[derive(Clone, Copy)] +pub(crate) struct Choice { + weight: u64, + value: T, +} + +/// Construct a weighted choice entry. +pub(crate) const fn choice(weight: u64, value: T) -> Choice { + Choice { weight, value } +} + +/// Pick one value from `choices`, using each entry's relative weight. +pub(crate) fn pick_choice(rng: &Rng, choices: &[Choice]) -> T { + let total: u64 = choices.iter().map(|choice| choice.weight).sum(); + + assert!(total > 0, "at least one choice weight must be non-zero"); + + let mut selected = rng.next_u64() % total; + + for choice in choices.iter().copied() { + if selected < choice.weight { + return choice.value; + } + + selected -= choice.weight; + } + + unreachable!("selected value is always inside total weight") +} + +/// Weighted choice helper for enum-like generator cases. +pub(crate) trait WeightedChoice: Copy { + fn pick(rng: &Rng, choices: &[Choice]) -> Self { + pick_choice(rng, choices) + } +} + +/// Pick an index from `0..len`, or return `None` for an empty collection. +pub(crate) fn choose_index(rng: &Rng, len: usize) -> Option { + (len > 0).then(|| rng.index(len)) +} + +/// Pick a value in the inclusive range `lo..=hi`. +pub(crate) fn range_inclusive(rng: &Rng, lo: usize, hi: usize) -> usize { + if lo >= hi { + return lo; + } + lo + (rng.next_u64() as usize % (hi - lo + 1)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum Case { + Default, + Mutated, + } + + impl Case { + const CHOICES: [Choice; 2] = [choice(1, Self::Default), choice(0, Self::Mutated)]; + } + + impl WeightedChoice for Case {} + + #[test] + fn fixed_choice_arrays_can_be_locally_mutated() { + let mut choices = Case::CHOICES; + choices[0] = choice(0, Case::Default); + choices[1] = choice(1, Case::Mutated); + + assert_eq!(Case::pick(&Rng::new(0), &choices), Case::Mutated); + assert_eq!(Case::pick(&Rng::new(0), &Case::CHOICES), Case::Default); + } +} diff --git a/crates/dst/src/schema.rs b/crates/dst/src/schema.rs index 641281db3c3..a817f9397bb 100644 --- a/crates/dst/src/schema.rs +++ b/crates/dst/src/schema.rs @@ -1,14 +1,84 @@ -use spacetimedb_lib::db::raw_def::v10::*; +//! Schema plans and raw module lowering for the engine DST harness. + +use crate::rng; use spacetimedb_lib::db::raw_def::v9::{RawIndexAlgorithm, TableAccess, TableType}; +use spacetimedb_lib::{db::raw_def::v10::*, RawModuleDef}; use spacetimedb_primitives::{ColId, ColList}; use spacetimedb_runtime::sim::Rng; -use spacetimedb_sats::{AlgebraicType, ArrayType, ProductType, ProductTypeElement}; - +use spacetimedb_sats::{ + AlgebraicType, AlgebraicValue, ArrayType, ArrayValue, ProductType, ProductTypeElement, SumType, SumTypeVariant, +}; +use spacetimedb_schema::def::ModuleDef; + +/// Generate the default engine DST schema. +/// +/// The initial schema is intentionally random and valid, not repaired into a +/// fixed coverage fixture. Long runs are expected to discover more surfaces via +/// migrations. pub fn default_schema(rng: Rng) -> SchemaPlan { - let profile = SchemaProfile::default(); - let mut plan = SchemaGenerator::new(rng, profile).gen_schema(); - plan.ensure_auto_inc_table(); - plan + SchemaGenerator::new(rng, SchemaProfile::engine_dst()).gen_schema() +} + +/// Lower a generated schema plan into the raw module format used by the engine. +pub fn to_raw_def(schema: &SchemaPlan) -> RawModuleDefV10 { + let mut builder = RawModuleDefV10Builder::new(); + builder.set_case_conversion_policy(CaseConversionPolicy::SnakeCase); + + for table in &schema.tables { + to_raw_def_table(&mut builder, table); + } + + let mut raw = builder.finish(); + apply_sequence_bounds(schema, &mut raw); + raw +} + +/// Lower and validate a schema plan exactly as the engine update path will see it. +pub fn to_module_def(schema: &SchemaPlan) -> anyhow::Result { + ModuleDef::try_from(RawModuleDef::V10(to_raw_def(schema))) + .map_err(|error| anyhow::anyhow!("schema validation failed: {error}")) +} + +/// Return the schema as the engine will see it after raw definition validation. +/// +/// The DST generator intentionally goes through the same `ModuleDef` validation +/// path as the engine instead of duplicating snake-case conversion rules in the +/// model. This keeps migration generation deterministic while making the model +/// use the canonical names that replay and auto-migration observe. +pub fn canonical_schema(schema: &SchemaPlan) -> anyhow::Result { + let module_def = to_module_def(schema)?; + + let mut tables = Vec::with_capacity(schema.tables.len()); + for table_plan in &schema.tables { + let table_def = module_def + .tables() + .find(|table_def| &*table_def.accessor_name == table_plan.name) + .ok_or_else(|| anyhow::anyhow!("validated schema is missing table accessor {:?}", table_plan.name))?; + + anyhow::ensure!( + table_def.columns.len() == table_plan.columns.len(), + "validated table {:?} has {} columns, plan has {}", + table_plan.name, + table_def.columns.len(), + table_plan.columns.len() + ); + + let mut table = table_plan.clone(); + table.name = table_def.name.to_string(); + for (column, column_def) in table.columns.iter_mut().zip(&table_def.columns) { + anyhow::ensure!( + &*column_def.accessor_name == column.name, + "validated table {:?} column accessor mismatch: expected {:?}, got {:?}", + table_plan.name, + column.name, + column_def.accessor_name.to_string() + ); + column.name = column_def.name.to_string(); + } + tables.push(table); + } + + Ok(SchemaPlan { tables }) } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -18,10 +88,19 @@ pub enum Type { U64, String, Bytes, + Sum { variants: u8 }, } impl Type { - pub const ALL: &'static [Type] = &[Type::Bool, Type::I64, Type::U64, Type::String, Type::Bytes]; + /// Representative column types used by migration candidate generation. + pub const ALL: &'static [Type] = &[ + Type::Bool, + Type::I64, + Type::U64, + Type::String, + Type::Bytes, + Type::Sum { variants: 1 }, + ]; pub fn to_algebraic(self) -> AlgebraicType { match self { @@ -32,6 +111,26 @@ impl Type { Type::Bytes => AlgebraicType::Array(ArrayType { elem_ty: Box::new(AlgebraicType::U8), }), + Type::Sum { variants } => { + debug_assert!(variants > 0); + AlgebraicType::Sum(SumType::new( + (0..variants) + .map(|variant| SumTypeVariant::new_named(AlgebraicType::U8, format!("variant_{variant}"))) + .collect::>() + .into_boxed_slice(), + )) + } + } + } + + pub fn default_value(self) -> AlgebraicValue { + match self { + Type::Bool => AlgebraicValue::Bool(false), + Type::I64 => AlgebraicValue::I64(0), + Type::U64 => AlgebraicValue::U64(0), + Type::String => AlgebraicValue::String("".into()), + Type::Bytes => AlgebraicValue::Array(ArrayValue::U8(Vec::new().into())), + Type::Sum { .. } => AlgebraicValue::sum(0, AlgebraicValue::U8(0)), } } @@ -40,55 +139,13 @@ impl Type { } } -// Schema plan — the canonical source of truth. -// This Schema should be able to translate to valid `RawModuleDefV10`. -#[derive(Debug, Clone)] +/// Canonical schema representation for the DST model and engine target. +#[derive(Debug, Clone, PartialEq, Eq)] pub struct SchemaPlan { pub tables: Vec, } -impl SchemaPlan { - pub fn auto_inc_table_and_column(&self) -> Option<(usize, usize)> { - self.tables - .iter() - .enumerate() - .find_map(|(table_idx, table)| table.sequences.first().map(|sequence| (table_idx, sequence.column))) - } - - pub fn ensure_auto_inc_table(&mut self) { - if self.auto_inc_table_and_column().is_some() { - return; - } - - let table = self.tables.first_mut().expect("schema must contain at least one table"); - if table.columns.is_empty() { - table.columns.push(ColumnPlan { - name: "id".into(), - ty: Type::U64, - }); - } else { - table.columns[0].ty = Type::U64; - } - - table.primary_key = Some(0); - if !table - .unique_constraints - .iter() - .any(|constraint| constraint.columns == [0]) - { - table.unique_constraints.push(UniqueConstraintPlan { columns: vec![0] }); - } - if !table.indexes.iter().any(|index| index.columns == [0]) { - table.indexes.push(IndexPlan { - columns: vec![0], - algorithm: IndexAlgorithm::BTree, - }); - } - table.sequences = vec![SequencePlan::new(0, Type::U64).expect("u64 is integral")]; - } -} - -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct TablePlan { pub name: String, pub columns: Vec, @@ -97,15 +154,16 @@ pub struct TablePlan { pub unique_constraints: Vec, pub sequences: Vec, pub is_public: bool, + pub is_event: bool, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct ColumnPlan { pub name: String, pub ty: Type, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct IndexPlan { /// Indices into `TablePlan.columns`. pub columns: Vec, @@ -118,17 +176,21 @@ pub enum IndexAlgorithm { Hash, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct UniqueConstraintPlan { /// Indices into `TablePlan.columns`. Non-empty. pub columns: Vec, } /// A sequence on a specific integral column. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct SequencePlan { /// Index into `TablePlan.columns`. pub column: usize, + pub start: Option, + pub min_value: Option, + pub max_value: Option, + pub increment: i128, } impl SequencePlan { @@ -137,82 +199,50 @@ impl SequencePlan { if !ty.is_integral() { return None; } - Some(Self { column }) + Some(Self { + column, + start: None, + min_value: None, + max_value: None, + increment: 1, + }) } -} - -// Lowering into RawModuleDefV10. -pub fn to_raw_def(schema: &SchemaPlan) -> RawModuleDefV10 { - let mut builder = RawModuleDefV10Builder::new(); - builder.set_case_conversion_policy(CaseConversionPolicy::None); - for table in &schema.tables { - to_raw_def_table(&mut builder, table); - } - - builder.finish() -} - -fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { - let product_type = ProductType { - elements: table - .columns - .iter() - .map(|col| ProductTypeElement { - name: Some(col.name.clone().into()), - algebraic_type: col.ty.to_algebraic(), - }) - .collect(), - }; - - let mut tbl = builder.build_table_with_new_type(table.name.clone(), product_type, true); - - tbl = tbl.with_type(TableType::User); - tbl = tbl.with_access(if table.is_public { - TableAccess::Public - } else { - TableAccess::Private - }); - // Primary key. - if let Some(pk) = table.primary_key { - tbl = tbl.with_primary_key(ColId(pk as u16)); + pub fn with_bounds( + column: usize, + ty: Type, + start: i128, + min_value: i128, + max_value: i128, + increment: i128, + ) -> Option { + if !ty.is_integral() || increment == 0 || min_value >= max_value || start < min_value || start > max_value { + return None; + } + Some(Self { + column, + start: Some(start), + min_value: Some(min_value), + max_value: Some(max_value), + increment, + }) } - // Unique constraints — all of them, including PK-matching. - for constraint in &table.unique_constraints { - let col_list: ColList = constraint.columns.iter().map(|&c| ColId(c as u16)).collect(); - tbl = tbl.with_unique_constraint(col_list); - } + /// Build a small bounded sequence whose max is an already-observed value. + /// + /// This intentionally creates a high-risk migration surface: if migration + /// prechecks accept an unsafe sequence, later inserts can collide with + /// existing unique values. + pub fn with_existing_value_as_max(column: usize, ty: Type, existing_value: i128) -> Option { + const DOMAIN_SIZE: i128 = 3; - // Indexes. - for index in &table.indexes { - let col_list: ColList = index.columns.iter().map(|&c| ColId(c as u16)).collect(); - - let algorithm = match index.algorithm { - IndexAlgorithm::BTree => RawIndexAlgorithm::BTree { columns: col_list }, - IndexAlgorithm::Hash => RawIndexAlgorithm::Hash { columns: col_list }, - }; - - let source_name = format!( - "{}_{}_idx", - table.name, - index - .columns - .iter() - .map(|&c| table.columns[c].name.as_str()) - .collect::>() - .join("_") - ); - - tbl = tbl.with_index_no_accessor_name(algorithm, source_name); - } + if existing_value < DOMAIN_SIZE { + return None; + } - // Sequences — all of them. - for seq in &table.sequences { - tbl = tbl.with_column_sequence(ColId(seq.column as u16)); + let min_value = existing_value - (DOMAIN_SIZE - 1); + Self::with_bounds(column, ty, min_value, min_value, existing_value, 1) } - - tbl.finish(); } /// Controls the shape of generated schemas. @@ -220,6 +250,9 @@ fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { pub struct SchemaProfile { pub table_count: (usize, usize), pub columns: (usize, usize), + pub table_kind_weights: TableKindWeights, + pub type_weights: TypeWeights, + pub sum_variants: (usize, usize), pub pk_prob: f64, pub auto_inc_prob: f64, pub indexes: (usize, usize), @@ -228,21 +261,65 @@ pub struct SchemaProfile { pub private_prob: f64, } +impl SchemaProfile { + pub fn engine_dst() -> Self { + Self { + table_count: (3, 10), + columns: (1, 20), + table_kind_weights: TableKindWeights::default(), + type_weights: TypeWeights::default(), + sum_variants: (1, 4), + pk_prob: 0.65, + auto_inc_prob: 0.20, + indexes: (0, 5), + unique_constraints: (0, 3), + btree_prob: 0.65, + private_prob: 0.10, + } + } +} + impl Default for SchemaProfile { + fn default() -> Self { + Self::engine_dst() + } +} + +#[derive(Debug, Clone, Copy)] +pub struct TableKindWeights { + pub data: u64, + pub event: u64, +} + +impl Default for TableKindWeights { + fn default() -> Self { + Self { data: 9, event: 1 } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct TypeWeights { + pub bool_: u64, + pub i64_: u64, + pub u64_: u64, + pub string: u64, + pub bytes: u64, + pub sum: u64, +} +impl Default for TypeWeights { fn default() -> Self { Self { - table_count: (1, 100), - columns: (1, 10), - pk_prob: 0.7, - auto_inc_prob: 0.3, - indexes: (0, 3), - unique_constraints: (0, 2), - btree_prob: 0.7, - private_prob: 0.1, + bool_: 12, + i64_: 24, + u64_: 28, + string: 16, + bytes: 12, + sum: 8, } } } +/// Random schema generator used by initial schema creation and add-table migrations. pub struct SchemaGenerator { rng: Rng, profile: SchemaProfile, @@ -253,60 +330,206 @@ impl SchemaGenerator { Self { rng, profile } } - fn range(&self, (lo, hi): (usize, usize)) -> usize { - if lo >= hi { - return lo; + /// Generate one table compatible with an existing schema. + /// + /// This is used by migration generation so add-table migrations use the same + /// shape policy as initial schema generation. Randomness happens before the + /// rewrite is applied; the rewrite itself remains deterministic. + pub fn gen_table_for_schema(&self, schema: &SchemaPlan, is_event: bool) -> TablePlan { + let mut sum_available = !schema_has_sum_column(schema); + self.gen_table(&schema.tables, is_event, &mut sum_available) + } + + /// Generate a complete schema from this generator's profile. + pub fn gen_schema(&self) -> SchemaPlan { + let table_count = SchemaDecisions::range(&self.rng, self.profile.table_count); + let mut tables: Vec = Vec::with_capacity(table_count); + let mut sum_available = true; + for table_idx in 0..table_count { + let must_be_data = table_idx + 1 == table_count && !tables.iter().any(|table| !table.is_event); + let is_event = !must_be_data && matches!(self.gen_table_kind(), TableKind::Event); + tables.push(self.gen_table(&tables, is_event, &mut sum_available)); + } + SchemaPlan { tables } + } +} + +/// Stable naming helpers for generated migration artifacts. +pub struct SchemaNames; + +impl SchemaNames { + pub fn fresh_column_name(table: &TablePlan, base: &str) -> String { + if table.columns.iter().all(|column| column.name != base) { + return base.into(); + } + + for suffix in 0.. { + let candidate = format!("{base}_{suffix}"); + if table.columns.iter().all(|column| column.name != candidate) { + return candidate; + } + } + + unreachable!("unbounded suffix search must find a unique column name") + } + + pub fn index_name(table: &TablePlan, index: &IndexPlan) -> String { + format!( + "{}_{}_idx", + table.name, + index + .columns + .iter() + .map(|&c| table.columns[c].name.as_str()) + .collect::>() + .join("_") + ) + } +} + +#[derive(Debug, Clone, Copy)] +enum TableKind { + Data, + Event, +} + +impl TableKindWeights { + fn choices(self) -> [rng::Choice; 2] { + [ + rng::choice(self.data, TableKind::Data), + rng::choice(self.event, TableKind::Event), + ] + } +} + +#[derive(Debug, Clone, Copy)] +enum TypeKind { + Bool, + I64, + U64, + String, + Bytes, + Sum, +} + +impl TypeWeights { + fn choices(self) -> [rng::Choice; 6] { + [ + rng::choice(self.bool_, TypeKind::Bool), + rng::choice(self.i64_, TypeKind::I64), + rng::choice(self.u64_, TypeKind::U64), + rng::choice(self.string, TypeKind::String), + rng::choice(self.bytes, TypeKind::Bytes), + rng::choice(self.sum, TypeKind::Sum), + ] + } + + fn non_sum_choices(self) -> [rng::Choice; 5] { + [ + rng::choice(self.bool_, TypeKind::Bool), + rng::choice(self.i64_, TypeKind::I64), + rng::choice(self.u64_, TypeKind::U64), + rng::choice(self.string, TypeKind::String), + rng::choice(self.bytes, TypeKind::Bytes), + ] + } +} + +struct SchemaDecisions; + +impl SchemaDecisions { + fn range(rng: &Rng, (lo, hi): (usize, usize)) -> usize { + rng::range_inclusive(rng, lo, hi) + } + + fn index(rng: &Rng, len: usize) -> usize { + rng::choose_index(rng, len).expect("len must be non-zero") + } + + fn sample_probability(rng: &Rng, probability: f64) -> bool { + rng.sample_probability(probability) + } + + fn gen_table_name(rng: &Rng, tables: &[TablePlan]) -> String { + loop { + let name = format!("tbl_{}", Self::gen_ident(rng)); + if tables.iter().all(|table| table.name != name) { + return name; + } } - lo + (self.rng.next_u64() as usize % (hi - lo + 1)) } - fn gen_type(&self) -> Type { - Type::ALL[self.rng.index(Type::ALL.len())] + fn gen_column_name(rng: &Rng, seen: &[String]) -> String { + loop { + let name = Self::gen_ident(rng); + if !seen.contains(&name) { + return name; + } + } } - fn gen_columns(&self) -> Vec { - let n = self.range(self.profile.columns); + fn gen_ident(rng: &Rng) -> String { + const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789_"; + const FIRST: &[u8] = b"abcdefghijklmnopqrstuvwxyz"; + let len = 4 + (rng.next_u64() as usize % 12); + let mut s = String::with_capacity(len); + s.push(FIRST[Self::index(rng, FIRST.len())] as char); + for _ in 1..len { + s.push(CHARS[Self::index(rng, CHARS.len())] as char); + } + s + } +} + +impl SchemaGenerator { + fn gen_columns(&self, sum_available: &mut bool) -> Vec { + let n = SchemaDecisions::range(&self.rng, self.profile.columns); let mut names = Vec::with_capacity(n); let mut seen = Vec::with_capacity(n); for _ in 0..n { - let name = self.gen_column_name(&seen); + let name = SchemaDecisions::gen_column_name(&self.rng, &seen); seen.push(name.clone()); names.push(ColumnPlan { name, - ty: self.gen_type(), + ty: self.gen_type(sum_available), }); } names } - fn gen_ident(&self) -> String { - const CHARS: &[u8] = b"abcdefghijklmnopqrstuvwxyz0123456789_"; - const FIRST: &[u8] = b"abcdefghijklmnopqrstuvwxyz_"; - let len = 4 + (self.rng.next_u64() as usize % 12); - let mut s = String::with_capacity(len); - s.push(FIRST[self.rng.index(FIRST.len())] as char); - for _ in 1..len { - s.push(CHARS[self.rng.index(CHARS.len())] as char); - } - s - } + fn gen_type(&self, sum_available: &mut bool) -> Type { + let kind = if *sum_available { + let choices = self.profile.type_weights.choices(); + rng::pick_choice(&self.rng, &choices) + } else { + let choices = self.profile.type_weights.non_sum_choices(); + rng::pick_choice(&self.rng, &choices) + }; - fn gen_column_name(&self, seen: &[String]) -> String { - loop { - let name = self.gen_ident(); - if !seen.contains(&name) { - return name; + match kind { + TypeKind::Bool => Type::Bool, + TypeKind::I64 => Type::I64, + TypeKind::U64 => Type::U64, + TypeKind::String => Type::String, + TypeKind::Bytes => Type::Bytes, + TypeKind::Sum => { + *sum_available = false; + Type::Sum { + variants: SchemaDecisions::range(&self.rng, self.profile.sum_variants) as u8, + } } } } fn gen_unique_constraints(&self, columns: &[ColumnPlan], pk: &Option) -> Vec { - let n = self.range(self.profile.unique_constraints); + let n = SchemaDecisions::range(&self.rng, self.profile.unique_constraints); let mut seen: Vec> = Vec::new(); let mut result = Vec::new(); for _ in 0..n { - let num_cols = 1 + self.rng.index(columns.len().min(3)); - let mut cols: Vec = (0..num_cols).map(|_| self.rng.index(columns.len())).collect(); + let num_cols = 1 + SchemaDecisions::index(&self.rng, columns.len().min(3)); + let mut cols: Vec = (0..num_cols) + .map(|_| SchemaDecisions::index(&self.rng, columns.len())) + .collect(); cols.sort(); cols.dedup(); if !cols.is_empty() && !seen.contains(&cols) { @@ -314,11 +537,11 @@ impl SchemaGenerator { result.push(UniqueConstraintPlan { columns: cols }); } } - // Ensure PK has a matching unique constraint. - if let Some(pk) = pk - && !seen.iter().any(|cols| cols.len() == 1 && cols[0] == *pk) - { - result.push(UniqueConstraintPlan { columns: vec![*pk] }); + // A primary key always has a matching unique constraint. + if let Some(pk) = pk { + if !seen.iter().any(|cols| cols.len() == 1 && cols[0] == *pk) { + result.push(UniqueConstraintPlan { columns: vec![*pk] }); + } } result } @@ -329,11 +552,9 @@ impl SchemaGenerator { unique_constraints: &[UniqueConstraintPlan], pk: &Option, ) -> Vec { - // Every unique constraint and PK needs a matching index. let mut seen_cols: Vec> = Vec::new(); let mut indexes: Vec = Vec::new(); - // Index for PK. if let Some(pk) = pk { seen_cols.push(vec![*pk]); indexes.push(IndexPlan { @@ -342,7 +563,6 @@ impl SchemaGenerator { }); } - // Indexes for unique constraints. for constraint in unique_constraints { if seen_cols.contains(&constraint.columns) { continue; @@ -354,18 +574,19 @@ impl SchemaGenerator { }); } - // Additional random indexes. - let n = self.range(self.profile.indexes); + let n = SchemaDecisions::range(&self.rng, self.profile.indexes); for _ in 0..n { - let num_cols = 1 + self.rng.index(columns.len().min(3)); - let mut cols: Vec = (0..num_cols).map(|_| self.rng.index(columns.len())).collect(); + let num_cols = 1 + SchemaDecisions::index(&self.rng, columns.len().min(3)); + let mut cols: Vec = (0..num_cols) + .map(|_| SchemaDecisions::index(&self.rng, columns.len())) + .collect(); cols.sort(); cols.dedup(); if cols.is_empty() || seen_cols.contains(&cols) { continue; } seen_cols.push(cols.clone()); - let algorithm = if self.rng.sample_probability(self.profile.btree_prob) { + let algorithm = if SchemaDecisions::sample_probability(&self.rng, self.profile.btree_prob) { IndexAlgorithm::BTree } else { IndexAlgorithm::Hash @@ -379,11 +600,27 @@ impl SchemaGenerator { indexes } - fn gen_table(&self, _table_index: usize) -> TablePlan { - let columns = self.gen_columns(); + fn gen_table(&self, existing_tables: &[TablePlan], is_event: bool, sum_available: &mut bool) -> TablePlan { + let columns = self.gen_columns(sum_available); + let name = SchemaDecisions::gen_table_name(&self.rng, existing_tables); + let is_public = !SchemaDecisions::sample_probability(&self.rng, self.profile.private_prob); - let primary_key = if self.rng.sample_probability(self.profile.pk_prob) && !columns.is_empty() { - Some(self.rng.index(columns.len())) + if is_event { + return TablePlan { + name, + columns, + primary_key: None, + indexes: vec![], + unique_constraints: vec![], + sequences: vec![], + is_public, + is_event: true, + }; + } + + let primary_key = if SchemaDecisions::sample_probability(&self.rng, self.profile.pk_prob) && !columns.is_empty() + { + Some(SchemaDecisions::index(&self.rng, columns.len())) } else { None }; @@ -391,7 +628,9 @@ impl SchemaGenerator { let unique_constraints = self.gen_unique_constraints(&columns, &primary_key); let sequences = if let Some(pk) = primary_key { - if columns[pk].ty.is_integral() && self.rng.sample_probability(self.profile.auto_inc_prob) { + if columns[pk].ty.is_integral() + && SchemaDecisions::sample_probability(&self.rng, self.profile.auto_inc_prob) + { SequencePlan::new(pk, columns[pk].ty).into_iter().collect() } else { vec![] @@ -402,8 +641,6 @@ impl SchemaGenerator { let indexes = self.gen_indexes(&columns, &unique_constraints, &primary_key); - let name = format!("tbl_{}", self.gen_ident()); - TablePlan { name, columns, @@ -411,20 +648,111 @@ impl SchemaGenerator { indexes, unique_constraints, sequences, - is_public: !self.rng.sample_probability(self.profile.private_prob), + is_public, + is_event: false, } } - pub fn gen_schema(&self) -> SchemaPlan { - let table_count = self.range(self.profile.table_count); - let tables = (0..table_count).map(|i| self.gen_table(i)).collect(); - SchemaPlan { tables } + fn gen_table_kind(&self) -> TableKind { + let choices = self.profile.table_kind_weights.choices(); + rng::pick_choice(&self.rng, &choices) + } +} + +fn apply_sequence_bounds(schema: &SchemaPlan, raw: &mut RawModuleDefV10) { + for (table_plan, raw_table) in schema.tables.iter().zip(raw.tables_mut_for_tests().iter_mut()) { + for (sequence_plan, raw_sequence) in table_plan.sequences.iter().zip(raw_table.sequences.iter_mut()) { + raw_sequence.start = sequence_plan.start; + raw_sequence.min_value = sequence_plan.min_value; + raw_sequence.max_value = sequence_plan.max_value; + raw_sequence.increment = sequence_plan.increment; + } } } +fn to_raw_def_table(builder: &mut RawModuleDefV10Builder, table: &TablePlan) { + let product_type = ProductType { + elements: table + .columns + .iter() + .map(|col| ProductTypeElement { + name: Some(col.name.clone().into()), + algebraic_type: col.ty.to_algebraic(), + }) + .collect(), + }; + + let mut tbl = builder.build_table_with_new_type_for_tests(table.name.clone(), product_type, true); + + tbl = tbl.with_type(TableType::User); + tbl = tbl.with_event(table.is_event); + tbl = tbl.with_access(if table.is_public { + TableAccess::Public + } else { + TableAccess::Private + }); + + if let Some(pk) = table.primary_key { + tbl = tbl.with_primary_key(ColId(pk as u16)); + } + + for constraint in &table.unique_constraints { + let col_list: ColList = constraint.columns.iter().map(|&c| ColId(c as u16)).collect(); + tbl = tbl.with_unique_constraint(col_list); + } + + for index in &table.indexes { + let col_list: ColList = index.columns.iter().map(|&c| ColId(c as u16)).collect(); + + let algorithm = match index.algorithm { + IndexAlgorithm::BTree => RawIndexAlgorithm::BTree { columns: col_list }, + IndexAlgorithm::Hash => RawIndexAlgorithm::Hash { columns: col_list }, + }; + + tbl = tbl.with_index_no_accessor_name(algorithm, SchemaNames::index_name(table, index)); + } + + for seq in &table.sequences { + tbl = tbl.with_column_sequence(ColId(seq.column as u16)); + } + + // AddColumns needs defaults when existing rows are present. Supplying stable + // defaults for all columns lets the engine keep only the newly-added tail. + for (col_id, column) in table.columns.iter().enumerate() { + tbl = tbl.with_default_column_value(ColId(col_id as u16), column.ty.default_value()); + } + + tbl.finish(); +} + +fn schema_has_sum_column(schema: &SchemaPlan) -> bool { + schema + .tables + .iter() + .flat_map(|table| table.columns.iter()) + .any(|column| matches!(column.ty, Type::Sum { .. })) +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn seed_3000_schema_canonicalizes_to_valid_snake_case_model_schema() { + let raw_schema = default_schema(Rng::new(3000)); + let schema = canonical_schema(&raw_schema).unwrap(); + let module_def = to_module_def(&schema).unwrap(); + + assert_eq!(canonical_schema(&schema).unwrap(), schema); + + for table_plan in &schema.tables { + let table_def = module_def.table(table_plan.name.as_str()).unwrap(); + assert_eq!(table_def.name.to_string(), table_plan.name); + + for (column_def, column_plan) in table_def.columns.iter().zip(&table_plan.columns) { + assert_eq!(column_def.name.to_string(), column_plan.name); + } + } + } #[test] fn lower_single_table() { @@ -453,6 +781,7 @@ mod tests { unique_constraints: vec![UniqueConstraintPlan { columns: vec![0] }], sequences: vec![SequencePlan::new(0, Type::U64).unwrap()], is_public: true, + is_event: false, }], }; diff --git a/crates/dst/src/sim/commitlog.rs b/crates/dst/src/sim/commitlog.rs index f2644d8767a..f87971d3e71 100644 --- a/crates/dst/src/sim/commitlog.rs +++ b/crates/dst/src/sim/commitlog.rs @@ -16,7 +16,12 @@ use spacetimedb_durability::{Close, Durability, DurableOffset, History, Prepared use spacetimedb_engine::relational_db::Txdata; use spacetimedb_runtime::sync::watch; -#[derive(Clone, Debug)] +pub const DST_COMMITLOG_SEGMENT_SIZE_MB: u64 = 32; +const BYTES_PER_MB: u64 = 1024 * 1024; + +pub type OnNewSegment = Arc; + +#[derive(Clone)] pub struct InMemoryCommitlog { repo: Memory, options: Options, @@ -30,15 +35,22 @@ impl Default for InMemoryCommitlog { impl InMemoryCommitlog { pub fn new() -> Self { + let mut options = Options::default(); + options.max_segment_size = DST_COMMITLOG_SEGMENT_SIZE_MB * BYTES_PER_MB; + Self { repo: Memory::unlimited(), - options: Options::default(), + options, } } pub fn open_handle(&self) -> io::Result { InMemoryCommitlogHandle::open(self.repo.clone(), self.options) } + + pub fn set_on_new_segment(&self, callback: Option) { + self.repo.set_on_new_segment(callback); + } } #[derive(Clone)] @@ -134,10 +146,11 @@ const PAGE_SIZE: usize = 4096; type SharedLock = Arc>; type SpaceOnDevice = Arc>; -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct Memory { space: SpaceOnDevice, segments: SharedLock>>, + on_new_segment: SharedLock>, } impl Memory { @@ -145,12 +158,17 @@ impl Memory { Self { space: Arc::new(Mutex::new(total_space)), segments: <_>::default(), + on_new_segment: <_>::default(), } } pub fn unlimited() -> Self { Self::new(u64::MAX) } + + fn set_on_new_segment(&self, callback: Option) { + *self.on_new_segment.write().unwrap() = callback; + } } impl fmt::Display for Memory { @@ -164,26 +182,32 @@ impl Repo for Memory { type SegmentReader = ReadOnlySegment; fn create_segment(&self, offset: u64, header: Header) -> io::Result { - let mut inner = self.segments.write().unwrap(); - let mut segment = match inner.entry(offset) { - btree_map::Entry::Occupied(entry) => { - let entry = entry.get(); - if entry.read().unwrap().is_empty() { - Segment::from_shared(self.space.clone(), entry.clone()) - } else { - return Err(io::Error::new( - io::ErrorKind::AlreadyExists, - format!("segment {offset} already exists"), - )); + let mut segment = { + let mut inner = self.segments.write().unwrap(); + match inner.entry(offset) { + btree_map::Entry::Occupied(entry) => { + let entry = entry.get(); + if entry.read().unwrap().is_empty() { + Segment::from_shared(self.space.clone(), entry.clone()) + } else { + return Err(io::Error::new( + io::ErrorKind::AlreadyExists, + format!("segment {offset} already exists"), + )); + } + } + btree_map::Entry::Vacant(entry) => { + let storage = entry.insert(Arc::new(RwLock::new(Storage::new()))); + Segment::from_shared(self.space.clone(), storage.clone()) } - } - btree_map::Entry::Vacant(entry) => { - let storage = entry.insert(Arc::new(RwLock::new(Storage::new()))); - Segment::from_shared(self.space.clone(), storage.clone()) } }; header.write(&mut segment)?; + if let Some(callback) = self.on_new_segment.read().unwrap().clone() { + callback(); + } + Ok(segment) } diff --git a/crates/dst/src/sim/mod.rs b/crates/dst/src/sim/mod.rs index 3a448644d48..f1acdc6ed18 100644 --- a/crates/dst/src/sim/mod.rs +++ b/crates/dst/src/sim/mod.rs @@ -1 +1,2 @@ pub mod commitlog; +pub mod snapshot; diff --git a/crates/dst/src/sim/snapshot.rs b/crates/dst/src/sim/snapshot.rs new file mode 100644 index 00000000000..0aefe5881f7 --- /dev/null +++ b/crates/dst/src/sim/snapshot.rs @@ -0,0 +1,335 @@ +use std::{ + collections::BTreeMap, + fmt, io, + ops::Range, + path::PathBuf, + sync::{Arc, Mutex}, +}; + +use spacetimedb_commitlog::repo::TxOffset; +use spacetimedb_fs_utils::compression::CompressType; +use spacetimedb_lib::Identity; +use spacetimedb_primitives::TableId; +use spacetimedb_sats::bsatn; +use spacetimedb_snapshot::{ + BoxedPendingSnapshot, CompressionStats, ObjectType, PendingSnapshot, ReconstructedSnapshot, SnapshotError, + SnapshotReadMetrics, SnapshotRepo, CURRENT_MODULE_ABI_VERSION, +}; +use spacetimedb_table::{ + blob_store::{BlobHash, BlobStore, HashMapBlobStore}, + page_pool::PagePool, + table::Table, +}; + +pub const DST_SNAPSHOT_RETENTION: usize = 2; + +#[derive(Clone)] +pub struct InMemorySnapshotRepo { + inner: Arc>, +} + +#[derive(Debug)] +struct Inner { + database_identity: Identity, + replica_id: u64, + retention: usize, + created_snapshots: u64, + deleted_snapshots: u64, + last_created_tx_offset: Option, + snapshots: BTreeMap, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SnapshotStats { + pub created: u64, + pub deleted: u64, + pub last_created_tx_offset: Option, + pub live_offsets: Vec, +} + +impl fmt::Display for SnapshotStats { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "snapshots: created={}, deleted={}, last_created=", + self.created, self.deleted + )?; + match self.last_created_tx_offset { + Some(offset) => write!(f, "{offset}")?, + None => f.write_str("none")?, + } + write!(f, ", live={:?}", self.live_offsets) + } +} + +#[derive(Clone, Debug)] +struct StoredSnapshot { + database_identity: Identity, + replica_id: u64, + tx_offset: TxOffset, + module_abi_version: [u16; 2], + blobs: Vec, + tables: Vec, +} + +#[derive(Clone, Debug)] +struct StoredBlob { + hash: BlobHash, + uses: usize, + bytes: Vec, +} + +#[derive(Clone, Debug)] +struct StoredTable { + table_id: TableId, + pages: Vec, +} + +#[derive(Clone, Debug)] +struct StoredPage { + hash: [u8; 32], + bytes: Vec, +} + +pub struct InMemoryPendingSnapshot { + repo: InMemorySnapshotRepo, + snapshot: StoredSnapshot, +} + +impl InMemorySnapshotRepo { + pub fn new(database_identity: Identity, replica_id: u64) -> Self { + Self::with_retention(database_identity, replica_id, DST_SNAPSHOT_RETENTION) + } + + pub fn with_retention(database_identity: Identity, replica_id: u64, retention: usize) -> Self { + Self { + inner: Arc::new(Mutex::new(Inner { + database_identity, + replica_id, + retention: retention.max(1), + created_snapshots: 0, + deleted_snapshots: 0, + last_created_tx_offset: None, + snapshots: BTreeMap::new(), + })), + } + } + + pub fn stats(&self) -> SnapshotStats { + let inner = self.inner.lock().unwrap(); + SnapshotStats { + created: inner.created_snapshots, + deleted: inner.deleted_snapshots, + last_created_tx_offset: inner.last_created_tx_offset, + live_offsets: inner.snapshots.keys().copied().collect(), + } + } + + fn missing_snapshot(tx_offset: TxOffset) -> SnapshotError { + SnapshotError::Io(io::Error::new( + io::ErrorKind::NotFound, + format!("snapshot {tx_offset} does not exist in DST memory repo"), + )) + } + + fn source_repo() -> PathBuf { + PathBuf::from("") + } +} + +impl SnapshotRepo for InMemorySnapshotRepo { + type Pending = BoxedPendingSnapshot; + + fn database_identity(&self) -> Identity { + self.inner.lock().unwrap().database_identity + } + + fn create_snapshot<'db>( + &self, + tables: &mut dyn Iterator, + blobs: &'db dyn BlobStore, + tx_offset: TxOffset, + ) -> Result { + let inner = self.inner.lock().unwrap(); + let database_identity = inner.database_identity; + let replica_id = inner.replica_id; + drop(inner); + + let blobs = blobs + .iter_blobs() + .map(|(hash, uses, bytes)| StoredBlob { + hash: *hash, + uses, + bytes: bytes.to_vec(), + }) + .collect(); + + let mut stored_tables = Vec::new(); + for table in tables { + let pages = table + .iter_pages_with_hashes() + .map(|(hash, page)| { + let bytes = bsatn::to_vec(page).map_err(|cause| SnapshotError::Serialize { + ty: ObjectType::Page(hash), + cause, + })?; + Ok(StoredPage { + hash: *hash.as_bytes(), + bytes, + }) + }) + .collect::, SnapshotError>>()?; + stored_tables.push(StoredTable { + table_id: table.schema.table_id, + pages, + }); + } + + Ok(Box::new(InMemoryPendingSnapshot { + repo: self.clone(), + snapshot: StoredSnapshot { + database_identity, + replica_id, + tx_offset, + module_abi_version: CURRENT_MODULE_ABI_VERSION, + blobs, + tables: stored_tables, + }, + })) + } + + fn read_snapshot(&self, tx_offset: TxOffset, page_pool: &PagePool) -> Result { + let snapshot = self + .inner + .lock() + .unwrap() + .snapshots + .get(&tx_offset) + .cloned() + .ok_or_else(|| Self::missing_snapshot(tx_offset))?; + + let mut read_metrics = SnapshotReadMetrics::default(); + read_metrics.blob.files = snapshot.blobs.len() as u64; + read_metrics.page.files = snapshot.tables.iter().map(|table| table.pages.len() as u64).sum(); + + let mut blob_store = HashMapBlobStore::default(); + for blob in &snapshot.blobs { + let computed = BlobHash::hash_from_bytes(&blob.bytes); + if computed != blob.hash { + return Err(SnapshotError::HashMismatch { + ty: ObjectType::Blob(blob.hash), + expected: blob.hash.data, + computed: computed.data, + source_repo: Self::source_repo(), + }); + } + read_metrics.blob.disk_bytes += blob.bytes.len() as u64; + blob_store.insert_with_uses(&blob.hash, blob.uses, blob.bytes.clone().into_boxed_slice()); + } + + let mut tables = BTreeMap::new(); + for table in &snapshot.tables { + let mut pages = Vec::with_capacity(table.pages.len()); + for stored_page in &table.pages { + read_metrics.page.disk_bytes += stored_page.bytes.len() as u64; + let page = page_pool.take_deserialize_from(&stored_page.bytes).map_err(|cause| { + SnapshotError::Deserialize { + ty: ObjectType::Snapshot, + source_repo: Self::source_repo(), + cause, + } + })?; + let computed = page.content_hash(); + if computed.as_bytes() != &stored_page.hash { + return Err(SnapshotError::HashMismatch { + ty: ObjectType::Page(computed), + expected: stored_page.hash, + computed: *computed.as_bytes(), + source_repo: Self::source_repo(), + }); + } + pages.push(page); + } + tables.insert(table.table_id, pages); + } + + Ok(ReconstructedSnapshot { + database_identity: snapshot.database_identity, + replica_id: snapshot.replica_id, + tx_offset: snapshot.tx_offset, + module_abi_version: snapshot.module_abi_version, + blob_store, + tables, + compress_type: CompressType::None, + read_metrics, + }) + } + + fn latest_snapshot_older_than(&self, upper_bound: TxOffset) -> Result, SnapshotError> { + Ok(self + .inner + .lock() + .unwrap() + .snapshots + .keys() + .rev() + .copied() + .find(|offset| *offset <= upper_bound)) + } + + fn compress_snapshots(&self, stats: &mut CompressionStats, range: Range) -> Result<(), SnapshotError> { + for offset in self.inner.lock().unwrap().snapshots.keys().copied() { + if range.contains(&offset) { + stats.skipped += 1; + stats.last_compressed = Some(offset); + } + } + Ok(()) + } + + fn invalidate_newer_snapshots(&self, upper_bound: TxOffset) -> Result<(), SnapshotError> { + let mut inner = self.inner.lock().unwrap(); + let newer = inner + .snapshots + .range((upper_bound.saturating_add(1))..) + .map(|(offset, _)| *offset) + .collect::>(); + for offset in newer { + if inner.snapshots.remove(&offset).is_some() { + inner.deleted_snapshots += 1; + } + } + Ok(()) + } + + fn invalidate_snapshot(&self, tx_offset: TxOffset) -> Result<(), SnapshotError> { + let mut inner = self.inner.lock().unwrap(); + inner + .snapshots + .remove(&tx_offset) + .map(|_| { + inner.deleted_snapshots += 1; + }) + .ok_or_else(|| Self::missing_snapshot(tx_offset)) + } +} + +impl PendingSnapshot for InMemoryPendingSnapshot { + fn sync_all(self: Box) -> Result { + let InMemoryPendingSnapshot { repo, snapshot } = *self; + let tx_offset = snapshot.tx_offset; + let mut inner = repo.inner.lock().unwrap(); + inner.snapshots.insert(tx_offset, snapshot); + inner.created_snapshots += 1; + inner.last_created_tx_offset = Some(tx_offset); + while inner.snapshots.len() > inner.retention { + let Some(oldest) = inner.snapshots.keys().next().copied() else { + break; + }; + if inner.snapshots.remove(&oldest).is_some() { + inner.deleted_snapshots += 1; + } + } + Ok(tx_offset) + } +} diff --git a/crates/dst/src/traits.rs b/crates/dst/src/traits.rs index 2185e0ec918..979bc3e3897 100644 --- a/crates/dst/src/traits.rs +++ b/crates/dst/src/traits.rs @@ -1,10 +1,22 @@ -use anyhow::Error; -use spacetimedb_runtime::sim::Rng; +use std::{ + fmt::Debug, + panic::{resume_unwind, AssertUnwindSafe}, +}; + +use anyhow::{Context, Error}; +use futures::FutureExt; +use spacetimedb_runtime::{sim::Rng, Handle}; + +const PROGRESS_INTERVAL: usize = 10_000; /// This should be implemented by System under test. pub trait TargetDriver { type Observation; + fn progress_status(&self) -> Option { + None + } + fn execute<'a>( &'a mut self, interaction: &'a I, @@ -16,6 +28,16 @@ pub trait Properties { fn observe(&mut self, interaction: &I, observation: &O) -> Result<(), Error>; } +/// Stateful source of interactions that can incorporate target feedback. +pub trait InteractionGen { + type Interaction: Debug; + type Observation; + + fn next_interaction(&mut self) -> Option; + + fn observe(&mut self, interaction: &Self::Interaction, observation: &Self::Observation) -> Result<(), Error>; +} + pub type TestSuiteParts = ( ::Interactions, ::Target, @@ -23,35 +45,82 @@ pub type TestSuiteParts = ( ); pub trait TestSuite { - type Interaction: std::fmt::Debug; - type Interactions: Iterator + std::fmt::Debug; + type Interaction: Debug; + type Interactions: InteractionGen< + Interaction = Self::Interaction, + Observation = >::Observation, + > + Debug; type Target: TargetDriver; type Properties: Properties>::Observation>; - fn build(&self, rng: Rng) -> impl std::future::Future, Error>> + '_ + fn build( + &self, + rng: Rng, + runtime: Handle, + ) -> impl std::future::Future, Error>> + '_ where Self: Sized; - fn run(&self, rng: Rng, max_interactions: usize) -> impl std::future::Future> + '_ + fn run( + &self, + rng: Rng, + runtime: Handle, + max_interactions: usize, + ) -> impl std::future::Future> + '_ where Self: Sized, { async move { - let (mut interactions, mut target, mut properties) = self.build(rng).await?; + let (mut interactions, mut target, mut properties) = self.build(rng, runtime).await?; + + let result = AssertUnwindSafe(async { + for step in 0..max_interactions { + let Some(interaction) = interactions.next_interaction() else { + break; + }; + + let observation = target + .execute(&interaction) + .await + .with_context(|| format!("DST target failed at interaction #{step}: {interaction:?}"))?; + + properties + .observe(&interaction, &observation) + .with_context(|| format!("DST property failed at interaction #{step}: {interaction:?}"))?; - let result = async { - for interaction in interactions.by_ref().take(max_interactions) { - let observation = target.execute(&interaction).await?; - properties.observe(&interaction, &observation)?; + interactions.observe(&interaction, &observation).with_context(|| { + format!("DST generator feedback failed at interaction #{step}: {interaction:?}") + })?; + + let completed = step + 1; + if completed % PROGRESS_INTERVAL == 0 { + if let Some(status) = target.progress_status() { + eprintln!( + "interaction progress: completed {completed} interactions: {interactions:?}; {status}" + ); + } else { + eprintln!("interaction progress: completed {completed} interactions: {interactions:?}"); + } + } } Ok(()) - } + }) + .catch_unwind() .await; - tracing::info!(interaction_counts = ?interactions, "final interaction counts"); + let final_status = target.progress_status(); + if let Some(status) = &final_status { + eprintln!("final interaction counts: {interactions:?}; {status}"); + } else { + eprintln!("final interaction counts: {interactions:?}"); + } + tracing::info!(interaction_counts = ?interactions, progress_status = ?final_status, "final interaction counts"); - result + match result { + Ok(result) => result, + Err(payload) => resume_unwind(payload), + } } } }