Storage Runtime Migration (Writing & Testing)

Overview

Storage Runtime Upgrades are a dangerous beast but with the right tools it can be tamed and become harmless. There are two ways how we can achieve this:

1. To write foolproof Storage Migrations functions.
2. To test those function with real production chain data.
3. To write tests

Writing Storage Migration Functions

This can be achieved by writing the on_runtime_upgrade() function inside the Hooks part of a pallet. Example:

  #[pallet::hooks]
  impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
    #[cfg(feature = "try-runtime")]
    fn pre_upgrade() -> Result<(), &'static str> {
      // This function is called before a runtime upgrade is executed. Here we can
      // test if our remote data is in the desired state. It's important to say that
      // pre_upgrade won't be called when a real runtime upgrade is executed.

      log::info!("Pre Upgrade.");
      log::info!("{:?}", Marketplaces::<T>::iter().count());

      Ok(())
    }

    fn on_runtime_upgrade() -> frame_support::weights::Weight {
      // This function is called when a runtime upgrade is called. We need to make sure that
      // what ever we do here won't brick the chain or leave the data in a invalid state.
      let version = StorageVersion::get::<Pallet<T>>();
      if version == StorageVersion::new(0) {
        // Do magic
        // ...

        // Update the storage version.
        StorageVersion::put::<Pallet<T>>(&StorageVersion::new(1));
      }

      frame_support::weights::Weight::MAX
    }

    #[cfg(feature = "try-runtime")]
    fn post_upgrade() -> Result<(), &'static str> {
      // This function is called after a runtime upgrade is executed. Here we can
      // test if the new state of blockchain data is valid. It's important to say that
      // post_upgrade won't be called when a real runtime upgrade is executed.

      log::info!("Post Upgrade.");
      log::info!("{:?}", Marketplaces::<T>::iter().count());

      Ok(())
    }
  }

This can work for small scale testing but for production ready code you would want to move the pre_upgrade, on_runtime_upgrade and post_upgrade to a different struct. That struct would be located inside the migration.rs file. This is how the struct would look like:

mod v1 {
  use super::*;
  use frame_support::traits::OnRuntimeUpgrade;

  // Here you would copy and paste the old struct instead of typedefing it.
  pub type OldMarketplaceData<T> = MarketplaceData<
    <T as frame_system::Config>::AccountId,
    BalanceOf<T>,
    <T as Config>::AccountSizeLimit,
    <T as Config>::OffchainDataLimit,
  >;

  pub struct MigrationV1<T>(sp_std::marker::PhantomData<T>);
  impl<T: Config> OnRuntimeUpgrade for MigrationV1<T> {
    #[cfg(feature = "try-runtime")]
    fn pre_upgrade() -> Result<(), &'static str> {
      log::info!("Pre-upgrade inside MigrationV1");
      Ok(())
    }

    fn on_runtime_upgrade() -> frame_support::weights::Weight {
      // If you want to change the existing storage so that it looks differently (adding new
      // fields, deleting existing ones,....) then you would use translate.
      Marketplaces::<T>::translate(|_id, old: OldMarketplaceData<T>| {
        let mut new_val = old.clone();
        new_val.commission_fee = None;
        new_val.listing_fee = None;
        new_val.account_list = None;
        new_val.offchain_data = None;

        Some(new_val)
      });

      frame_support::weights::Weight::MAX
    }

    #[cfg(feature = "try-runtime")]
    fn post_upgrade() -> Result<(), &'static str> {
      log::info!("Post-upgrade inside MigrationV1");
      Ok(())
    }
  }
}

And the Hook would look like this:

  #[pallet::hooks]
  impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
    // This function is called before a runtime upgrade is executed. Here we can
    // test if our remote data is in the desired state. It's important to say that
    // pre_upgrade won't be called when a real runtime upgrade is executed.
    #[cfg(feature = "try-runtime")]
    fn pre_upgrade() -> Result<(), &'static str> {
      <v1::MigrationV1<T> as OnRuntimeUpgrade>::pre_upgrade()
    }

    // This function is called when a runtime upgrade is called. We need to make sure that
    // what ever we do here won't brick the chain or leave the data in a invalid state.
    fn on_runtime_upgrade() -> frame_support::weights::Weight {
      let mut weight = 0;

      let version = StorageVersion::get::<Pallet<T>>();
      if version == StorageVersion::new(0) {
        weight = <v1::MigrationV1<T> as OnRuntimeUpgrade>::on_runtime_upgrade();

        // Update the storage version.
        StorageVersion::put::<Pallet<T>>(&StorageVersion::new(1));
      }

      weight
    }

    // This function is called after a runtime upgrade is executed. Here we can
    // test if the new state of blockchain data is valid. It's important to say that
    // post_upgrade won't be called when a real runtime upgrade is executed.
    #[cfg(feature = "try-runtime")]
    fn post_upgrade() -> Result<(), &'static str> {
      <v1::MigrationV1<T> as OnRuntimeUpgrade>::post_upgrade()
    }
  }

Testing Storage Migration Functions with Real Data

Once the pre_upgrade, on_runtime_upgrade and post_upgrade functions have been written, it's time to test it with real data from either Alphanet or Mainnet. To do it we first need to build our Binary with the try-runtime feature flag:

cargo build --features try-runtime

After that we execute the try-runtime command:

./target/debug/ternoa try-runtime --chain alphanet-dev --execution Native on-runtime-upgrade live -p Marketplace -u ws://51.15.235.53:9944

The -u flag would need to point to an existing public node IP address.