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

use frame_support::{
    pallet_prelude::{PalletInfoAccess, StorageVersion, Weight},
    traits::OnRuntimeUpgrade,
};

/// To run both pre- and post-checks around every migration, we entangle methods of
/// `OnRuntimeUpgrade` into the desired flow and expose it with `migrate` method.
///
/// This way, `try-runtime` no longer triggers checks. We do it by hand.
pub trait StorageMigration: OnRuntimeUpgrade {
    #[allow(clippy::let_and_return)]
    fn migrate() -> Weight {
        #[cfg(feature = "try-runtime")]
        let state = match Self::pre_upgrade() {
            Ok(state) => state,
            Err(_err) => b"Pre-upgrade failed".to_vec(),
        };

        let weight = Self::on_runtime_upgrade();

        #[cfg(feature = "try-runtime")]
        match Self::post_upgrade(state) {
            Ok(()) => (),
            Err(err) => {
                log::error!(
                    "Post-upgrade failed for pallet. Error: {:?}",
                    err
                );
            }
        };

        weight
    }
}

impl<T: OnRuntimeUpgrade> StorageMigration for T {}

/// Ensure that the current pallet storage version matches `version`.
pub fn ensure_storage_version<P: PalletInfoAccess>(version: u16) -> Result<(), &'static str> {
    if StorageVersion::get::<P>() == StorageVersion::new(version) {
        Ok(())
    } else {
        Err("Bad storage version")
    }
}