Add migrate to queue

Author: webmaster128

contracts/queue/examples/schema.rs

@@ -3,9 +3,8 @@ use std::fs::create_dir_all;
 
 use cosmwasm_schema::{export_schema, remove_schemas, schema_for};
 
-use queue::contract::{
-    CountResponse, HandleMsg, InitMsg, Item, ListResponse, QueryMsg, SumResponse,
-};
+use queue::contract::{CountResponse, HandleMsg, Item, ListResponse, QueryMsg, SumResponse};
+use queue::msg::{InitMsg, MigrateMsg};
 
 fn main() {
     let mut out_dir = current_dir().unwrap();
@@ -15,6 +14,7 @@ fn main() {
 
     export_schema(&schema_for!(InitMsg), &out_dir);
     export_schema(&schema_for!(HandleMsg), &out_dir);
+    export_schema(&schema_for!(MigrateMsg), &out_dir);
     export_schema(&schema_for!(QueryMsg), &out_dir);
     export_schema(&schema_for!(Item), &out_dir);
     export_schema(&schema_for!(CountResponse), &out_dir);

contracts/queue/schema/migrate_msg.json

@@ -0,0 +1,5 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "MigrateMsg",
+  "type": "object"
+}

contracts/queue/src/contract.rs

@@ -5,11 +5,10 @@ use serde::{Deserialize, Serialize};
 
 use cosmwasm_std::{
     from_slice, to_binary, to_vec, Binary, Deps, DepsMut, Env, HandleResponse, InitResponse,
-    MessageInfo, Order, QueryResponse, StdResult,
+    MessageInfo, MigrateResponse, Order, QueryResponse, StdResult, Storage,
 };
 
-#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
-pub struct InitMsg {}
+use crate::msg::{InitMsg, MigrateMsg};
 
 // we store one entry for each item in the queue
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
@@ -82,16 +81,21 @@ pub fn handle(
     msg: HandleMsg,
 ) -> StdResult<HandleResponse> {
     match msg {
-        HandleMsg::Enqueue { value } => enqueue(deps, value),
-        HandleMsg::Dequeue {} => dequeue(deps),
+        HandleMsg::Enqueue { value } => handle_enqueue(deps, value),
+        HandleMsg::Dequeue {} => handle_dequeue(deps),
     }
 }
 
 const FIRST_KEY: u8 = 0;
 
-fn enqueue(deps: DepsMut, value: i32) -> StdResult<HandleResponse> {
+fn handle_enqueue(deps: DepsMut, value: i32) -> StdResult<HandleResponse> {
+    enqueue(deps.storage, value)?;
+    Ok(HandleResponse::default())
+}
+
+fn enqueue(storage: &mut dyn Storage, value: i32) -> StdResult<()> {
     // find the last element in the queue and extract key
-    let last_item = deps.storage.range(None, None, Order::Descending).next();
+    let last_item = storage.range(None, None, Order::Descending).next();
 
     let new_key = match last_item {
         None => FIRST_KEY,
@@ -101,11 +105,11 @@ fn enqueue(deps: DepsMut, value: i32) -> StdResult<HandleResponse> {
     };
     let new_value = to_vec(&Item { value })?;
 
-    deps.storage.set(&[new_key], &new_value);
-    Ok(HandleResponse::default())
+    storage.set(&[new_key], &new_value);
+    Ok(())
 }
 
-fn dequeue(deps: DepsMut) -> StdResult<HandleResponse> {
+fn handle_dequeue(deps: DepsMut) -> StdResult<HandleResponse> {
     // find the first element in the queue and extract value
     let first = deps.storage.range(None, None, Order::Ascending).next();
 
@@ -120,6 +124,29 @@ fn dequeue(deps: DepsMut) -> StdResult<HandleResponse> {
     }
 }
 
+pub fn migrate(
+    deps: DepsMut,
+    _env: Env,
+    _info: MessageInfo,
+    _msg: MigrateMsg,
+) -> StdResult<MigrateResponse> {
+    // clear all
+    let keys: Vec<_> = deps
+        .storage
+        .range(None, None, Order::Ascending)
+        .map(|(key, _)| key)
+        .collect();
+    for key in keys {
+        deps.storage.remove(&key);
+    }
+
+    // Write new values
+    enqueue(deps.storage, 100)?;
+    enqueue(deps.storage, 101)?;
+    enqueue(deps.storage, 102)?;
+    Ok(MigrateResponse::default())
+}
+
 pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<QueryResponse> {
     match msg {
         QueryMsg::Count {} => to_binary(&query_count(deps)?),

contracts/queue/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod contract;
+pub mod msg;
 
 #[cfg(target_arch = "wasm32")]
-cosmwasm_std::create_entry_points!(contract);
+cosmwasm_std::create_entry_points_with_migration!(contract);

contracts/queue/src/msg.rs

@@ -0,0 +1,10 @@
+#![allow(clippy::field_reassign_with_default)] // see https://github.com/CosmWasm/cosmwasm/issues/685
+
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+pub struct InitMsg {}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+pub struct MigrateMsg {}

contracts/queue/tests/integration.rs

@@ -17,18 +17,21 @@
 //!      });
 //! 4. Anywhere you see query(&deps, ...) you must replace it with query(&mut deps, ...)
 
-use cosmwasm_std::{from_binary, from_slice, HandleResponse, HumanAddr, InitResponse, MessageInfo};
+use cosmwasm_std::{
+    from_binary, from_slice, HandleResponse, HumanAddr, InitResponse, MessageInfo, MigrateResponse,
+};
 use cosmwasm_vm::{
     testing::{
-        handle, init, mock_env, mock_info, mock_instance_with_gas_limit, query, MockApi,
+        handle, init, migrate, mock_env, mock_info, mock_instance_with_gas_limit, query, MockApi,
         MockQuerier, MockStorage,
     },
     Instance,
 };
 
 use queue::contract::{
-    CountResponse, HandleMsg, InitMsg, Item, ListResponse, QueryMsg, ReducerResponse, SumResponse,
+    CountResponse, HandleMsg, Item, ListResponse, QueryMsg, ReducerResponse, SumResponse,
 };
+use queue::msg::{InitMsg, MigrateMsg};
 
 static WASM: &[u8] = include_bytes!("../target/wasm32-unknown-unknown/release/queue.wasm");
 
@@ -170,6 +173,36 @@ fn push_and_reduce() {
     assert_eq!(counters, vec![(40, 85), (15, 125), (85, 0), (-10, 140)]);
 }
 
+#[test]
+fn migrate_works() {
+    let (mut deps, info) = create_contract();
+
+    let _: HandleResponse = handle(
+        &mut deps,
+        mock_env(),
+        info.clone(),
+        HandleMsg::Enqueue { value: 25 },
+    )
+    .unwrap();
+    let _: HandleResponse = handle(
+        &mut deps,
+        mock_env(),
+        info.clone(),
+        HandleMsg::Enqueue { value: 17 },
+    )
+    .unwrap();
+    assert_eq!(get_count(&mut deps), 2);
+    assert_eq!(get_sum(&mut deps), 25 + 17);
+
+    let msg = MigrateMsg {};
+    let info = mock_info("admin", &[]);
+    let res: MigrateResponse = migrate(&mut deps, mock_env(), info, msg).unwrap();
+    assert_eq!(res.messages.len(), 0);
+
+    assert_eq!(get_count(&mut deps), 3);
+    assert_eq!(get_sum(&mut deps), 100 + 101 + 102);
+}
+
 #[test]
 fn query_list() {
     let (mut deps, info) = create_contract();