1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

// Copyright 2019-2024 ChainSafe Systems
// SPDX-License-Identifier: Apache-2.0, MIT

use std::path::PathBuf;

use tracing::info;

use crate::{
    cli_shared::chain_path,
    db::{
        db_mode::{get_latest_versioned_database, DbMode},
        migration::migration_map::create_migration_chain,
    },
    utils::version::FOREST_VERSION,
    Config,
};

/// Governs the database migration process. This is the entry point for the migration process.
pub struct DbMigration {
    /// Forest configuration used.
    config: Config,
}

impl DbMigration {
    pub fn new(config: &Config) -> Self {
        Self {
            config: config.clone(),
        }
    }

    pub fn chain_data_path(&self) -> PathBuf {
        chain_path(&self.config)
    }

    /// Verifies if a migration is needed for the current Forest process.
    /// Note that migration can possibly happen only if the DB is in `current` mode.
    pub fn is_migration_required(&self) -> anyhow::Result<bool> {
        // No chain data means that this is a fresh instance. No migration required.
        if !self.chain_data_path().exists() {
            return Ok(false);
        }

        // migration is required only if the DB is in `current` mode and the current db version is
        // smaller than the current binary version
        if let DbMode::Current = DbMode::read() {
            let current_db = get_latest_versioned_database(&self.chain_data_path())?
                .unwrap_or_else(|| FOREST_VERSION.clone());
            Ok(current_db < *FOREST_VERSION)
        } else {
            Ok(false)
        }
    }

    /// Performs a database migration if required. Note that this may take a long time to complete
    /// and may need a lot of disk space (at least twice the size of the current database).
    /// On a successful migration, the current database will be removed and the new database will
    /// be used.
    /// This method is tested via integration tests.
    pub fn migrate(&self) -> anyhow::Result<()> {
        if !self.is_migration_required()? {
            info!("No database migration required");
            return Ok(());
        }

        let latest_db_version = get_latest_versioned_database(&self.chain_data_path())?
            .unwrap_or_else(|| FOREST_VERSION.clone());

        info!(
            "Migrating database from version {} to {}",
            latest_db_version, *FOREST_VERSION
        );

        let target_db_version = &FOREST_VERSION;

        let migrations = create_migration_chain(&latest_db_version, target_db_version)?;

        for migration in migrations {
            migration.migrate(&self.chain_data_path(), &self.config)?;
        }

        info!(
            "Migration to version {} complete",
            target_db_version.to_string()
        );

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::db::db_mode::FOREST_DB_DEV_MODE;

    use super::*;

    #[test]
    fn test_migration_not_required_no_chain_path() {
        let temp_dir = tempfile::tempdir().unwrap();
        let mut config = Config::default();
        config.client.data_dir = temp_dir.path().join("azathoth");
        let db_migration = DbMigration::new(&config);
        assert!(!db_migration.is_migration_required().unwrap());
    }

    #[test]
    fn test_migration_not_required_no_databases() {
        let temp_dir = tempfile::tempdir().unwrap();
        let mut config = Config::default();
        temp_dir.path().clone_into(&mut config.client.data_dir);
        let db_migration = DbMigration::new(&config);
        assert!(!db_migration.is_migration_required().unwrap());
    }

    #[test]
    fn test_migration_not_required_under_non_current_mode() {
        let temp_dir = tempfile::tempdir().unwrap();
        let mut config = Config::default();
        temp_dir.path().clone_into(&mut config.client.data_dir);

        let db_dir = temp_dir.path().join("mainnet/0.1.0");
        std::fs::create_dir_all(&db_dir).unwrap();
        let db_migration = DbMigration::new(&config);

        std::env::set_var(FOREST_DB_DEV_MODE, "latest");
        assert!(!db_migration.is_migration_required().unwrap());

        std::fs::remove_dir(db_dir).unwrap();
        std::fs::create_dir_all(temp_dir.path().join("mainnet/cthulhu")).unwrap();

        std::env::set_var(FOREST_DB_DEV_MODE, "cthulhu");
        assert!(!db_migration.is_migration_required().unwrap());
    }

    #[test]
    fn test_migration_required_current_mode() {
        let temp_dir = tempfile::tempdir().unwrap();
        let mut config = Config::default();
        temp_dir.path().clone_into(&mut config.client.data_dir);

        let db_dir = temp_dir.path().join("mainnet/0.1.0");
        std::fs::create_dir_all(db_dir).unwrap();
        let db_migration = DbMigration::new(&config);

        std::env::set_var(FOREST_DB_DEV_MODE, "current");
        assert!(db_migration.is_migration_required().unwrap());
        std::env::remove_var(FOREST_DB_DEV_MODE);
        assert!(db_migration.is_migration_required().unwrap());
    }
}