Added xion code id for migration from xion account

AbstractSDK · Jan 8, 2025 · aec8b3d · aec8b3d
1 parent 1ae8172
commit aec8b3d
Show file tree

Hide file tree

Showing 7 changed files with 35 additions and 18 deletions.
diff --git a/framework/contracts/account/src/error.rs b/framework/contracts/account/src/error.rs
@@ -127,6 +127,9 @@ pub enum AccountError {
     #[error("Abstract Account don't have Authentication")]
     AbstractAccountNoAuth {},
 
+    #[error("The new_code_id field needs to be filled when migrating from a xion account")]
+    MissingCodeIdToMigrate {},
+
     #[cfg(feature = "xion")]
     #[error(transparent)]
     AbstractXion(#[from] abstract_xion::error::ContractError),

diff --git a/framework/contracts/account/src/migrate.rs b/framework/contracts/account/src/migrate.rs
@@ -14,7 +14,7 @@ use crate::contract::{AccountResponse, AccountResult, CONTRACT_VERSION};
 /// - XION Account, to allow upgrading their account to a more feature rich account (second part of the function)
 /// All other contracts cannot be migrated to this version
 #[cfg_attr(feature = "export", cosmwasm_std::entry_point)]
-pub fn migrate(deps: DepsMut, _env: Env, _msg: MigrateMsg) -> AccountResult {
+pub fn migrate(deps: DepsMut, _env: Env, msg: MigrateMsg) -> AccountResult {
     let version: Version = CONTRACT_VERSION.parse().unwrap();
 
     let current_contract_version = get_contract_version(deps.storage)?;
@@ -27,7 +27,13 @@ pub fn migrate(deps: DepsMut, _env: Env, _msg: MigrateMsg) -> AccountResult {
         #[cfg(feature = "xion")]
         {
             // We might want to migrate from a XION account
-            migrate_from_xion_account(deps, _env, current_contract_version)
+            migrate_from_xion_account(
+                deps,
+                _env,
+                current_contract_version,
+                msg.new_code_id
+                    .ok_or(crate::error::AccountError::MissingCodeIdToMigrate {})?,
+            )
         }
         #[cfg(not(feature = "xion"))]
         {
@@ -44,6 +50,7 @@ pub fn migrate_from_xion_account(
     mut deps: DepsMut,
     env: Env,
     current_contract_version: cw2::ContractVersion,
+    new_code_id: u64,
 ) -> AccountResult {
     use crate::modules::{_install_modules, MIGRATE_CONTEXT};
     use ::{
@@ -52,6 +59,7 @@ pub fn migrate_from_xion_account(
         abstract_std::account::ModuleInstallConfig,
         abstract_std::objects::module::ModuleInfo,
         abstract_std::objects::AccountId,
+        abstract_std::IBC_CLIENT,
         abstract_std::{
             account::state::{WhitelistedModules, SUSPENSION_STATUS, WHITELISTED_MODULES},
             objects::{
@@ -60,7 +68,6 @@ pub fn migrate_from_xion_account(
             },
             registry::state::LOCAL_ACCOUNT_SEQUENCE,
         },
-        abstract_std::{native_addrs, IBC_CLIENT},
         cosmwasm_std::wasm_execute,
     };
 
@@ -73,9 +80,7 @@ pub fn migrate_from_xion_account(
     // Use CW2 to set the contract version, this is needed for migrations
     set_contract_version(deps.storage, ACCOUNT, CONTRACT_VERSION)?;
 
-    let abstract_code_id =
-        native_addrs::abstract_code_id(&deps.querier, env.contract.address.clone())?;
-    let registry = RegistryContract::new(deps.as_ref(), abstract_code_id)?;
+    let registry = RegistryContract::new(deps.as_ref(), new_code_id)?;
 
     let account_id =
         AccountId::local(LOCAL_ACCOUNT_SEQUENCE.query(&deps.querier, registry.address.clone())?);
@@ -117,7 +122,7 @@ pub fn migrate_from_xion_account(
 
     // Install IBC Client module
     let (install_msgs, install_attribute) =
-        _install_modules(deps, install_modules, vec![], abstract_code_id)?;
+        _install_modules(deps, install_modules, vec![], new_code_id)?;
     response = response
         .add_submessages(install_msgs)
         .add_attribute(install_attribute.key, install_attribute.value);
@@ -147,7 +152,7 @@ mod tests {
 
         let version: Version = CONTRACT_VERSION.parse().unwrap();
 
-        let res = super::migrate(deps.as_mut(), env, MigrateMsg {});
+        let res = super::migrate(deps.as_mut(), env, MigrateMsg { new_code_id: None });
 
         assert_eq!(
             res,
@@ -175,7 +180,7 @@ mod tests {
 
         let version: Version = CONTRACT_VERSION.parse().unwrap();
 
-        let res = super::migrate(deps.as_mut(), env, MigrateMsg {});
+        let res = super::migrate(deps.as_mut(), env, MigrateMsg { new_code_id: None });
 
         assert_eq!(
             res,
@@ -202,7 +207,7 @@ mod tests {
         let old_name = "old:contract";
         set_contract_version(deps.as_mut().storage, old_name, old_version)?;
 
-        let res = super::migrate(deps.as_mut(), env, MigrateMsg {});
+        let res = super::migrate(deps.as_mut(), env, MigrateMsg { new_code_id: None });
 
         assert_eq!(
             res,
@@ -234,7 +239,7 @@ mod tests {
 
         set_contract_version(deps.as_mut().storage, ACCOUNT, small_version)?;
 
-        let res = super::migrate(deps.as_mut(), env, MigrateMsg {})?;
+        let res = super::migrate(deps.as_mut(), env, MigrateMsg { new_code_id: None })?;
         assert!(res.messages.is_empty());
 
         assert_eq!(

diff --git a/framework/contracts/account/tests/upgrades.rs b/framework/contracts/account/tests/upgrades.rs
@@ -915,7 +915,9 @@ mod upgrade_account {
             ),
             (
                 ModuleInfo::from_id_latest("abstract:account")?,
-                Some(to_json_binary(&abstract_std::account::MigrateMsg {})?),
+                Some(to_json_binary(&abstract_std::account::MigrateMsg {
+                    new_code_id: None,
+                })?),
             ),
             (
                 ModuleInfo::from_id_latest(adapter_1::MOCK_ADAPTER_ID)?,

diff --git a/framework/packages/abstract-client/src/account.rs b/framework/packages/abstract-client/src/account.rs
@@ -529,7 +529,7 @@ impl<Chain: CwEnv> Account<Chain> {
             .upgrade(vec![(
                 ModuleInfo::from_id(abstract_std::constants::ACCOUNT, version.clone())?,
                 Some(
-                    to_json_binary(&abstract_std::account::MigrateMsg {})
+                    to_json_binary(&abstract_std::account::MigrateMsg { new_code_id: None })
                         .map_err(Into::<CwOrchError>::into)?,
                 ),
             )])

diff --git a/framework/packages/abstract-client/src/interchain/remote_account.rs b/framework/packages/abstract-client/src/interchain/remote_account.rs
@@ -348,7 +348,7 @@ impl<Chain: IbcQueryHandler, IBC: InterchainEnv<Chain>> RemoteAccount<Chain, IBC
         let modules = vec![(
             ModuleInfo::from_id(abstract_std::constants::ACCOUNT, version.clone())?,
             Some(
-                to_json_binary(&abstract_std::account::MigrateMsg {})
+                to_json_binary(&abstract_std::account::MigrateMsg { new_code_id: None })
                     .map_err(Into::<CwOrchError>::into)?,
             ),
         )];

diff --git a/framework/packages/abstract-std/src/account.rs b/framework/packages/abstract-std/src/account.rs
@@ -104,7 +104,11 @@ pub mod state {
 }
 
 #[cosmwasm_schema::cw_serde]
-pub struct MigrateMsg {}
+pub struct MigrateMsg {
+    /// This field provides the new code id of the contract
+    /// This is only necessary for migrations from XION accounts.
+    pub new_code_id: Option<u64>,
+}
 
 /// Account Instantiate Msg
 /// https://github.com/burnt-labs/contracts/blob/main/contracts/account/src/msg.rs

diff --git a/interchain/framework-clone-testing/tests/migrate.rs b/interchain/framework-clone-testing/tests/migrate.rs
@@ -198,9 +198,12 @@ mod from_xion {
         let account = AccountI::new("account-xion", chain);
         account.set_address(&addr_contract);
 
-        account
-            .call_as(&addr_contract)
-            .migrate(&MigrateMsg {}, deployment.account_code_id()?)?;
+        account.call_as(&addr_contract).migrate(
+            &MigrateMsg {
+                new_code_id: Some(deployment.account_code_id()?),
+            },
+            deployment.account_code_id()?,
+        )?;
 
         account
             .update_info(None, None, Some("brand new abstract account".to_string()))