feat(Upgradable)!: enable batched fn call after deploy (#86)

* feat(Upgradable): enable batched fn call after deploy

* test: update existing tests

* test: add contract with state migration

* test: deploy and call migration function

* docs: describe new behavior

Author: mooori

near-plugins-derive/src/upgradable.rs

@@ -180,7 +180,7 @@ pub fn derive_upgradable(input: TokenStream) -> TokenStream {
             }
 
             #[#cratename::access_control_any(roles(#(#acl_roles_code_deployers),*))]
-            fn up_deploy_code(&mut self) -> near_sdk::Promise {
+            fn up_deploy_code(&mut self, function_call_args: Option<#cratename::upgradable::FunctionCallArgs>) -> near_sdk::Promise {
                 let staging_timestamp = self.up_get_timestamp(__UpgradableStorageKey::StagingTimestamp)
                     .unwrap_or_else(|| ::near_sdk::env::panic_str("Upgradable: staging timestamp isn't set"));
 
@@ -194,8 +194,18 @@ pub fn derive_upgradable(input: TokenStream) -> TokenStream {
                     );
                 }
 
-                near_sdk::Promise::new(near_sdk::env::current_account_id())
-                    .deploy_contract(self.up_staged_code().unwrap_or_else(|| ::near_sdk::env::panic_str("Upgradable: No staged code")))
+                let code = self.up_staged_code().unwrap_or_else(|| ::near_sdk::env::panic_str("Upgradable: No staged code"));
+                let promise = near_sdk::Promise::new(near_sdk::env::current_account_id())
+                    .deploy_contract(code);
+                match function_call_args {
+                    None => promise,
+                    Some(args) => {
+                        // Execute the `DeployContract` and `FunctionCall` actions in a batch
+                        // transaction to make a failure of the function call roll back the code
+                        // deployment.
+                        promise.function_call(args.function_name, args.arguments, args.amount, args.gas)
+                    },
+                }
             }
 
             #[#cratename::access_control_any(roles(#(#acl_roles_duration_initializers),*))]

near-plugins-derive/tests/common/upgradable_contract.rs

@@ -1,4 +1,4 @@
-use near_plugins::upgradable::UpgradableDurationStatus;
+use near_plugins::upgradable::{FunctionCallArgs, UpgradableDurationStatus};
 
 use near_sdk::serde_json::json;
 use near_sdk::CryptoHash;
@@ -71,9 +71,13 @@ impl UpgradableContract {
     pub async fn up_deploy_code(
         &self,
         caller: &Account,
+        function_call_args: Option<FunctionCallArgs>,
     ) -> workspaces::Result<ExecutionFinalResult> {
         caller
             .call(self.contract.id(), "up_deploy_code")
+            .args_json(json!({
+                "function_call_args": function_call_args,
+            }))
             .max_gas()
             .transact()
             .await

near-plugins-derive/tests/common/utils.rs

@@ -49,6 +49,11 @@ where
     assert_eq!(actual, expected);
 }
 
+/// Asserts transaction failure due `MethodNotFound` error.
+pub fn assert_method_not_found_failure(res: ExecutionFinalResult) {
+    assert_failure_with(res, "Action #0: MethodResolveError(MethodNotFound)");
+}
+
 /// Asserts transaction failure due to `method` being `#[private]`.
 pub fn assert_private_method_failure(res: ExecutionFinalResult, method: &str) {
     let err = res

near-plugins-derive/tests/contracts/upgradable_state_migration/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "upgradable_state_migration"
+version = "0.0.0"
+edition = "2018"
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+near-plugins = { path = "../../../../near-plugins" }
+near-sdk = "4.1.0"
+
+[profile.release]
+codegen-units = 1
+opt-level = "z"
+lto = true
+debug = false
+panic = "abort"
+overflow-checks = true
+
+[workspace]

near-plugins-derive/tests/contracts/upgradable_state_migration/Makefile

@@ -0,0 +1,8 @@
+build:
+	cargo build --target wasm32-unknown-unknown --release
+
+# Helpful for debugging. Requires `cargo-expand`.
+expand:
+	cargo expand > expanded.rs
+
+.PHONY: build expand

near-plugins-derive/tests/contracts/upgradable_state_migration/rust-toolchain

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "1.66.1"
+components = ["clippy", "rustfmt"]

near-plugins-derive/tests/contracts/upgradable_state_migration/src/lib.rs

@@ -0,0 +1,72 @@
+//! A simple contract to be deployed via `Upgradable`. It requires [state migration].
+//!
+//! [state migration]: https://docs.near.org/develop/upgrade#migrating-the-state
+
+use near_plugins::{access_control, AccessControlRole, AccessControllable, Upgradable};
+use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
+use near_sdk::serde::{Deserialize, Serialize};
+use near_sdk::{env, near_bindgen, PanicOnDefault};
+
+/// Roles correspond to those defined in the initial contract `../upgradable`, to make permissions
+/// granted before the upgrade remain valid.
+#[derive(AccessControlRole, Deserialize, Serialize, Copy, Clone)]
+#[serde(crate = "near_sdk::serde")]
+pub enum Role {
+    DAO,
+    CodeStager,
+    CodeDeployer,
+    DurationManager,
+}
+
+/// The struct differs from the one defined in the initial contract `../upgradable`, hence [state
+/// migration] is required.
+///
+/// [state migration]: https://docs.near.org/develop/upgrade#migrating-the-state
+#[access_control(role_type(Role))]
+#[near_bindgen]
+#[derive(Upgradable, PanicOnDefault, BorshDeserialize, BorshSerialize)]
+#[upgradable(access_control_roles(
+    code_stagers(Role::CodeStager, Role::DAO),
+    code_deployers(Role::CodeDeployer, Role::DAO),
+    duration_initializers(Role::DurationManager, Role::DAO),
+    duration_update_stagers(Role::DurationManager, Role::DAO),
+    duration_update_appliers(Role::DurationManager, Role::DAO),
+))]
+pub struct Contract {
+    is_migrated: bool,
+}
+
+#[near_bindgen]
+impl Contract {
+    /// Migrates state from [`OldContract`] to [`Contract`].
+    ///
+    /// It follows the state migration pattern described [here].
+    ///
+    /// [here]: https://docs.near.org/develop/upgrade#migrating-the-state
+    #[private]
+    #[init(ignore_state)]
+    pub fn migrate() -> Self {
+        // Ensure old state can be read and deserialized.
+        let _: OldContract = env::state_read().expect("Should be able to load old state");
+
+        Self { is_migrated: true }
+    }
+
+    /// A migration method that fails on purpose to test the rollback mechanism of
+    /// `Upgradable::up_deploy_code`.
+    #[private]
+    #[init(ignore_state)]
+    pub fn migrate_with_failure() -> Self {
+        env::panic_str("Failing migration on purpose");
+    }
+
+    /// This method is _not_ defined in the initial contract, so calling it successfully proves the
+    /// contract defined in this file was deployed and the old state was migrated.
+    pub fn is_migrated(&self) -> bool {
+        self.is_migrated
+    }
+}
+
+/// Corresponds to the state defined in the initial `../upgradable` contract.
+#[derive(BorshDeserialize)]
+pub struct OldContract;

near-plugins-derive/tests/upgradable.rs

@@ -6,19 +6,21 @@ use anyhow::Ok;
 use common::access_controllable_contract::AccessControllableContract;
 use common::upgradable_contract::UpgradableContract;
 use common::utils::{
-    assert_failure_with, assert_insufficient_acl_permissions, assert_success_with,
-    assert_success_with_unit_return, fast_forward_beyond, get_transaction_block,
-    sdk_duration_from_secs,
+    assert_failure_with, assert_insufficient_acl_permissions, assert_method_not_found_failure,
+    assert_success_with, assert_success_with_unit_return, fast_forward_beyond,
+    get_transaction_block, sdk_duration_from_secs,
 };
+use near_plugins::upgradable::FunctionCallArgs;
 use near_sdk::serde_json::json;
-use near_sdk::{CryptoHash, Duration, Timestamp};
+use near_sdk::{CryptoHash, Duration, Gas, Timestamp};
 use std::path::Path;
 use workspaces::network::Sandbox;
 use workspaces::result::ExecutionFinalResult;
 use workspaces::{Account, AccountId, Contract, Worker};
 
 const PROJECT_PATH: &str = "./tests/contracts/upgradable";
 const PROJECT_PATH_2: &str = "./tests/contracts/upgradable_2";
+const PROJECT_PATH_STATE_MIGRATION: &str = "./tests/contracts/upgradable_state_migration";
 
 const ERR_MSG_NO_STAGING_TS: &str = "Upgradable: staging timestamp isn't set";
 const ERR_MSG_DEPLOY_CODE_TOO_EARLY: &str = "Upgradable: Deploy code too early: staging ends on";
@@ -169,6 +171,18 @@ impl Setup {
             .await
     }
 
+    async fn call_is_migrated(&self, caller: &Account) -> workspaces::Result<ExecutionFinalResult> {
+        // `is_migrated` could be called via `view`, however here it is called via `transact` so we
+        // get an `ExecutionFinalResult` that can be passed to `assert_*` methods from
+        // `common::utils`. It is acceptable since all we care about is whether the method exists
+        // and can be called successfully.
+        caller
+            .call(self.contract.id(), "is_migrated")
+            .max_gas()
+            .transact()
+            .await
+    }
+
     /// Calls the contract's `is_set_up` method and asserts it returns `true`. Panics on failure.
     async fn assert_is_set_up(&self, caller: &Account) {
         let res = caller
@@ -419,14 +433,15 @@ async fn test_deploy_code_without_delay() -> anyhow::Result<()> {
     setup.assert_staged_code(Some(code)).await;
 
     // Deploy staged code.
-    let res = setup.upgradable_contract.up_deploy_code(&dao).await?;
+    let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?;
     assert_success_with_unit_return(res);
 
     Ok(())
 }
 
 /// Verifies the upgrade was successful by calling a method that's available only on the upgraded
-/// contract. Ensures the new contract can be deployed and state migration succeeds.
+/// contract. Ensures the new contract can be deployed and state remains valid without
+/// explicit state migration.
 #[tokio::test]
 async fn test_deploy_code_and_call_method() -> anyhow::Result<()> {
     let worker = workspaces::sandbox().await?;
@@ -435,7 +450,7 @@ async fn test_deploy_code_and_call_method() -> anyhow::Result<()> {
 
     // Verify function `is_upgraded` is not defined in the initial contract.
     let res = setup.call_is_upgraded(&setup.unauth_account).await?;
-    assert_failure_with(res, "Action #0: MethodResolveError(MethodNotFound)");
+    assert_method_not_found_failure(res);
 
     // Compile the other version of the contract and stage its code.
     let code = common::repo::compile_project(Path::new(PROJECT_PATH_2), "upgradable_2").await?;
@@ -447,7 +462,7 @@ async fn test_deploy_code_and_call_method() -> anyhow::Result<()> {
     setup.assert_staged_code(Some(code)).await;
 
     // Deploy staged code.
-    let res = setup.upgradable_contract.up_deploy_code(&dao).await?;
+    let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?;
     assert_success_with_unit_return(res);
 
     // The newly deployed contract defines the function `is_upgraded`. Calling it successfully
@@ -458,6 +473,94 @@ async fn test_deploy_code_and_call_method() -> anyhow::Result<()> {
     Ok(())
 }
 
+/// Deploys a new version of the contract that requires state migration and verifies the migration
+/// succeeded.
+#[tokio::test]
+async fn test_deploy_code_with_migration() -> anyhow::Result<()> {
+    let worker = workspaces::sandbox().await?;
+    let dao = worker.dev_create_account().await?;
+    let setup = Setup::new(worker.clone(), Some(dao.id().clone()), None).await?;
+
+    // Verify function `is_migrated` is not defined in the initial contract.
+    let res = setup.call_is_migrated(&setup.unauth_account).await?;
+    assert_method_not_found_failure(res);
+
+    // Compile the other version of the contract and stage its code.
+    let code = common::repo::compile_project(
+        Path::new(PROJECT_PATH_STATE_MIGRATION),
+        "upgradable_state_migration",
+    )
+    .await?;
+    let res = setup
+        .upgradable_contract
+        .up_stage_code(&dao, code.clone())
+        .await?;
+    assert_success_with_unit_return(res);
+    setup.assert_staged_code(Some(code)).await;
+
+    // Deploy staged code and call the new contract's `migrate` method.
+    let function_call_args = FunctionCallArgs {
+        function_name: "migrate".to_string(),
+        arguments: Vec::new(),
+        amount: 0,
+        gas: Gas::ONE_TERA,
+    };
+    let res = setup
+        .upgradable_contract
+        .up_deploy_code(&dao, Some(function_call_args))
+        .await?;
+    assert_success_with_unit_return(res);
+
+    // The newly deployed contract defines the function `is_migrated`. Calling it successfully
+    // verifies the staged contract is deployed and state migration succeeded.
+    let res = setup.call_is_migrated(&setup.unauth_account).await?;
+    assert_success_with(res, true);
+
+    Ok(())
+}
+
+/// Deploys a new version of the contract and, batched with the `DeployContractAction`, calls a
+/// migration method that fails. Verifies the failure rolls back the deployment, i.e. the initial
+/// code remains active.
+#[tokio::test]
+async fn test_deploy_code_with_migration_failure_rollback() -> anyhow::Result<()> {
+    let worker = workspaces::sandbox().await?;
+    let dao = worker.dev_create_account().await?;
+    let setup = Setup::new(worker.clone(), Some(dao.id().clone()), None).await?;
+
+    // Compile the other version of the contract and stage its code.
+    let code = common::repo::compile_project(
+        Path::new(PROJECT_PATH_STATE_MIGRATION),
+        "upgradable_state_migration",
+    )
+    .await?;
+    let res = setup
+        .upgradable_contract
+        .up_stage_code(&dao, code.clone())
+        .await?;
+    assert_success_with_unit_return(res);
+    setup.assert_staged_code(Some(code)).await;
+
+    // Deploy staged code and call the new contract's `migrate_with_failure` method.
+    let function_call_args = FunctionCallArgs {
+        function_name: "migrate_with_failure".to_string(),
+        arguments: Vec::new(),
+        amount: 0,
+        gas: Gas::ONE_TERA,
+    };
+    let res = setup
+        .upgradable_contract
+        .up_deploy_code(&dao, Some(function_call_args))
+        .await?;
+    assert_failure_with(res, "Failing migration on purpose");
+
+    // Verify `code` wasn't deployed by calling a function that is defined only in the initial
+    // contract but not in the contract contract corresponding to `code`.
+    setup.assert_is_set_up(&setup.unauth_account).await;
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_deploy_code_with_delay() -> anyhow::Result<()> {
     let worker = workspaces::sandbox().await?;
@@ -483,7 +586,7 @@ async fn test_deploy_code_with_delay() -> anyhow::Result<()> {
     fast_forward_beyond(&worker, staging_duration).await;
 
     // Deploy staged code.
-    let res = setup.upgradable_contract.up_deploy_code(&dao).await?;
+    let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?;
     assert_success_with_unit_return(res);
 
     Ok(())
@@ -513,9 +616,13 @@ async fn test_deploy_code_with_delay_failure_too_early() -> anyhow::Result<()> {
     fast_forward_beyond(&worker, sdk_duration_from_secs(1)).await;
 
     // Verify trying to deploy staged code fails.
-    let res = setup.upgradable_contract.up_deploy_code(&dao).await?;
+    let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?;
     assert_failure_with(res, ERR_MSG_DEPLOY_CODE_TOO_EARLY);
 
+    // Verify `code` wasn't deployed by calling a function that is defined only in the initial
+    // contract but not in the contract contract corresponding to `code`.
+    setup.assert_is_set_up(&setup.unauth_account).await;
+
     Ok(())
 }
 
@@ -538,7 +645,7 @@ async fn test_deploy_code_permission_failure() -> anyhow::Result<()> {
     // call this method.
     let res = setup
         .upgradable_contract
-        .up_deploy_code(&setup.unauth_account)
+        .up_deploy_code(&setup.unauth_account, None)
         .await?;
     assert_insufficient_acl_permissions(
         res,
@@ -577,7 +684,7 @@ async fn test_deploy_code_empty_failure() -> anyhow::Result<()> {
     // The staging timestamp is set when staging code and removed when unstaging code. So when there
     // is no code staged, there is no staging timestamp. Hence the error message regarding a missing
     // staging timestamp is expected.
-    let res = setup.upgradable_contract.up_deploy_code(&dao).await?;
+    let res = setup.upgradable_contract.up_deploy_code(&dao, None).await?;
     assert_failure_with(res, ERR_MSG_NO_STAGING_TS);
 
     Ok(())
@@ -825,7 +932,7 @@ async fn test_acl_permission_scope() -> anyhow::Result<()> {
     // deploy code.
     let res = setup
         .upgradable_contract
-        .up_deploy_code(&setup.unauth_account)
+        .up_deploy_code(&setup.unauth_account, None)
         .await?;
     assert_insufficient_acl_permissions(
         res,