feat: revamp migrations to support recreating cfs

Author: seriousben

server/state_store/src/lib.rs

@@ -28,6 +28,7 @@ use tracing::{debug, error, info, span, warn};
 pub mod in_memory_state;
 pub mod invocation_events;
 pub mod kv;
+pub mod migration_runner;
 pub mod migrations;
 pub mod requests;
 pub mod scanner;
@@ -108,7 +109,7 @@ impl IndexifyState {
         // Migrate the db before opening with all column families.
         // This is because the migration process may delete older column families.
         // If we open the db with all column families, it would fail to open.
-        let sm_meta = migrations::migrate(&path)?;
+        let sm_meta = migration_runner::run(&path)?;
 
         let sm_column_families = IndexifyObjectsColumns::iter()
             .map(|cf| ColumnFamilyDescriptor::new(cf.to_string(), Options::default()));
@@ -303,7 +304,7 @@ impl IndexifyState {
             &txn,
             &request.processed_state_changes,
         )?;
-        migrations::write_sm_meta(
+        migration_runner::write_sm_meta(
             &self.db,
             &txn,
             &StateMachineMetadata {

server/state_store/src/migration_runner.rs

@@ -0,0 +1,249 @@
+use std::path::Path;
+
+use anyhow::{Context, Result};
+use data_model::StateMachineMetadata;
+use rocksdb::{Transaction, TransactionDB};
+use tracing::info;
+
+use crate::{
+    migrations::{
+        contexts::{MigrationContext, PrepareContext},
+        registry::MigrationRegistry,
+    },
+    serializer::{JsonEncode, JsonEncoder},
+    state_machine::IndexifyObjectsColumns,
+};
+
+/// Main function to run all necessary migrations on a database at the given
+/// path
+pub fn run(path: &Path) -> Result<StateMachineMetadata> {
+    // Initialize prepare context
+    let prepare_ctx = PrepareContext::new(path.to_path_buf());
+
+    // Initialize registry
+    let registry = MigrationRegistry::new()?;
+    let latest_version = registry.latest_version();
+
+    // Check if DB exists
+    let db = match prepare_ctx.open_db() {
+        Ok(db) => db,
+        Err(e) if e.to_string().contains("No such file or directory") => {
+            // New DB, return default metadata
+            info!(
+                "No database found. Initializing at version {}",
+                latest_version
+            );
+            return Ok(StateMachineMetadata {
+                db_version: latest_version,
+                last_change_idx: 0,
+            });
+        }
+        Err(e) => return Err(anyhow::anyhow!("Error opening database: {:?}", e)),
+    };
+
+    // Read current metadata
+    let mut sm_meta = read_sm_meta(&db)?;
+    drop(db); // Close DB before migrations
+
+    // Find applicable migrations
+    let migrations = registry.find_migrations(sm_meta.db_version);
+
+    // No migrations needed
+    if migrations.is_empty() {
+        info!(
+            "Database already at version {}. No migrations needed.",
+            sm_meta.db_version
+        );
+        return Ok(sm_meta);
+    }
+
+    info!(
+        "Starting migrations from v{} to v{}",
+        sm_meta.db_version, latest_version
+    );
+
+    // Execute each migration in sequence
+    for migration in migrations {
+        let from_version = sm_meta.db_version;
+        let to_version = migration.version();
+
+        info!(
+            "Running migration {}: v{} → v{}",
+            migration.name(),
+            from_version,
+            to_version
+        );
+
+        // Each migration prepares the DB as needed
+        let db = migration
+            .prepare(&prepare_ctx)
+            .with_context(|| format!("Preparing DB for migration to v{}", to_version))?;
+
+        // Apply migration in a transaction
+        let txn = db.transaction();
+
+        // Create migration context
+        let migration_ctx = MigrationContext::new(&db, &txn);
+
+        // Apply the migration
+        migration
+            .apply(&migration_ctx)
+            .with_context(|| format!("Applying migration to v{}", to_version))?;
+
+        // Update metadata in the same transaction
+        sm_meta.db_version = to_version;
+        write_sm_meta(&db, &txn, &sm_meta)?;
+
+        info!("Committing migration to v{}", to_version);
+        txn.commit()
+            .with_context(|| format!("Committing migration to v{}", to_version))?;
+
+        // Close DB after each migration to ensure clean state
+        drop(db);
+    }
+
+    info!(
+        "Completed all migrations. DB now at version {}",
+        sm_meta.db_version
+    );
+    Ok(sm_meta)
+}
+
+/// Read state machine metadata from the database
+pub fn read_sm_meta(db: &TransactionDB) -> Result<StateMachineMetadata> {
+    let meta = db.get_cf(
+        &IndexifyObjectsColumns::StateMachineMetadata.cf_db(db),
+        b"sm_meta",
+    )?;
+    match meta {
+        Some(meta) => Ok(JsonEncoder::decode(&meta)?),
+        None => Ok(StateMachineMetadata {
+            db_version: 0,
+            last_change_idx: 0,
+        }),
+    }
+}
+
+/// Write state machine metadata to the database
+pub fn write_sm_meta(
+    db: &TransactionDB,
+    txn: &Transaction<TransactionDB>,
+    sm_meta: &StateMachineMetadata,
+) -> Result<()> {
+    let serialized_meta = JsonEncoder::encode(sm_meta)?;
+    txn.put_cf(
+        &IndexifyObjectsColumns::StateMachineMetadata.cf_db(db),
+        b"sm_meta",
+        &serialized_meta,
+    )?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use rocksdb::{ColumnFamilyDescriptor, Options, TransactionDBOptions};
+    use tempfile::TempDir;
+
+    use super::*;
+    use crate::migrations::migration_trait::Migration;
+
+    #[derive(Clone)]
+    struct MockMigration {
+        version: u64,
+        name: &'static str,
+    }
+
+    impl Migration for MockMigration {
+        fn version(&self) -> u64 {
+            self.version
+        }
+
+        fn name(&self) -> &'static str {
+            self.name
+        }
+
+        fn prepare(&self, ctx: &PrepareContext) -> Result<TransactionDB> {
+            // Simple mock - just open DB
+            ctx.open_db()
+        }
+
+        fn apply(&self, _ctx: &MigrationContext) -> Result<()> {
+            // No-op for test
+            Ok(())
+        }
+
+        fn box_clone(&self) -> Box<dyn Migration> {
+            Box::new(self.clone())
+        }
+    }
+
+    #[test]
+    fn test_migration_new_db() -> Result<()> {
+        let temp_dir = TempDir::new()?;
+        let path = temp_dir.path();
+
+        // Create a mock migration
+        let mock_migration = MockMigration {
+            version: 1,
+            name: "MockMigration",
+        };
+
+        // Use the mock migration (e.g., log its name)
+        info!("Testing with migration: {}", mock_migration.name());
+
+        // Run migrations on non-existent DB
+        let sm_meta = run(path)?;
+
+        // Check migration resulted in latest version
+        assert_eq!(
+            sm_meta.db_version,
+            MigrationRegistry::new()?.latest_version()
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_migration_existing_db() -> Result<()> {
+        let temp_dir = TempDir::new()?;
+        let path = temp_dir.path();
+
+        // Create DB with initial metadata
+        let cf_desc = ColumnFamilyDescriptor::new(
+            IndexifyObjectsColumns::StateMachineMetadata.as_ref(),
+            Options::default(),
+        );
+
+        let mut db_opts = Options::default();
+        db_opts.create_missing_column_families(true);
+        db_opts.create_if_missing(true);
+
+        let db = TransactionDB::open_cf_descriptors(
+            &db_opts,
+            &TransactionDBOptions::default(),
+            path,
+            vec![cf_desc],
+        )?;
+
+        // Set initial version to 1
+        let txn = db.transaction();
+        let initial_meta = StateMachineMetadata {
+            db_version: 1,
+            last_change_idx: 0,
+        };
+        write_sm_meta(&db, &txn, &initial_meta)?;
+        txn.commit()?;
+        drop(db);
+
+        // Run migrations
+        let sm_meta = run(path)?;
+
+        // Check migration resulted in latest version
+        assert_eq!(
+            sm_meta.db_version,
+            MigrationRegistry::new()?.latest_version()
+        );
+
+        Ok(())
+    }
+}