Migrating tables from redb v1 to v2

[Frando]

I wrote a migration for our database from redb v1 to v2, and I thought others might need this too so I'm sharing the code here. Not sure if there's a better place than an issue, but I guess people might look here.

License: MIT/Apache-2.0 https://github.com/n0-computer/iroh/pull/2120/files#diff-3e6c0d3203eaa07833f7afd7461bd486a127e2cdca7f98fc5f328c325bfd5fa0R1

redb = { version = "2.0.0" }
redb_v1  = { package = "redb", version = "1.5.1" }

mod db {
    use redb::{Database, MultimapTableDefinition, TableDefinition};

    use super::migrate_v1_v2;

    pub const FOO_TABLE: TableDefinition<u32, u64> = TableDefinition::new("foo-1");
    // more table definitions...

    struct Store {
        db: Database
    }

    impl Store {
        pub fn open(path: impl AsRef<Path>) -> Result<Self> {
            let db = match Database::create(&path) {
                Ok(db) => db,
                Err(DatabaseError::UpgradeRequired(1)) => migrate_v1_v2:::run(&path)?,
                Err(err) => return Err(err.into()),
            };
            Ok(Self { db }))
        }
    }
}

mod migrate_v1_v2 {
    use std::path::{Path, PathBuf};

    use anyhow::{Context, Result};
    use redb::{MultimapTableHandle, TableHandle};
    use redb_v1::{ReadableMultimapTable, ReadableTable};
    use tempfile::NamedTempFile;
    use tracing::info;

    macro_rules! migrate_table {
        ($rtx:expr, $wtx:expr, $old:expr, $new:expr) => {{
            let old_table = $rtx.open_table($old)?;
            let mut new_table = $wtx.open_table($new)?;
            let name = $new.name();
            let len = old_table.len()?;
            info!("migrate {name} ({len} rows)..");
            let ind = (len as usize / 1000) + 1;
            for (i, entry) in old_table.iter()?.enumerate() {
                let (key, value) = entry?;
                let key = key.value();
                let value = value.value();
                if i > 0 && i % 100 == 0 {
                    info!("    {name} {i:>ind$}/{len}");
                }
                new_table.insert(key, value)?;
            }
            info!("migrate {name} done");
        }};
    }

    macro_rules! migrate_multimap_table {
        ($rtx:expr, $wtx:expr, $old:expr, $new:expr) => {{
            let old_table = $rtx.open_multimap_table($old)?;
            let mut new_table = $wtx.open_multimap_table($new)?;
            let name = $new.name();
            let len = old_table.len()?;
            info!("migrate {name} ({len} rows)");
            let ind = (len as usize / 1000) + 1;
            for (i, entry) in old_table.iter()?.enumerate() {
                let (key, values) = entry?;
                let key = key.value();
                if i > 0 && i % 100 == 0 {
                    info!("    {name} {i:>ind$}/{len}");
                }
                for value in values {
                    let value = value?;
                    new_table.insert(key, value.value())?;
                }
            }
            info!("migrate {name} done");
        }};
    }

    pub fn run(source: impl AsRef<Path>) -> Result<redb::Database> {
        let source = source.as_ref();
        // create the new database in a tempfile
        let target = NamedTempFile::new()?.into_temp_path();
        info!("migrate {} to {}", source.display(), target.display());
        let old_db = redb_v1::Database::open(source)?;
        let new_db = redb::Database::create(&target)?;

        let rtx = old_db.begin_read()?;
        let wtx = new_db.begin_write()?;

        migrate_table!(rtx, wtx, old::FOO_TABLE, new::FOO_TABLE);
        // repeat for each table

        wtx.commit()?;
        drop(rtx);
        drop(old_db);
        drop(new_db);

        let backup_path: PathBuf = {
            let mut p = source.to_owned().into_os_string();
            p.push(".backup-redb-v1");
            p.into()
        };

        info!("rename {} to {}", source.display(), backup_path.display());
        std::fs::rename(source, &backup_path)?;
        info!("rename {} to {}", target.display(), source.display());
        target.persist_noclobber(source)?;
        info!("opening migrated database from {}", source.display());
        let db = redb::Database::open(source)?;
        Ok(db)
    }

    mod new {
        pub use super::db::*;
    }

    mod old {
        use redb_v1::{MultimapTableDefinition, TableDefinition};

        pub const FOO_TABLE: TableDefinition<u32, u64> = TableDefinition::new("foo-1");
       // .. more table definitions
    }
}