Add optional delete field to burner migrate msg

contracts/burner/schema/burner.json

@@ -19,6 +19,13 @@
       "payout"
     ],
     "properties": {
+      "delete": {
+        "description": "Optional amount of items to delete in this call. If it is not provided, nothing will be deleted. You can delete further items in a subsequent execute call.",
+        "default": 0,
+        "type": "integer",
+        "format": "uint32",
+        "minimum": 0.0
+      },
       "payout": {
         "description": "The address we send all remaining balance to",
         "type": "string"

contracts/burner/schema/raw/migrate.json

@@ -6,6 +6,13 @@
     "payout"
   ],
   "properties": {
+    "delete": {
+      "description": "Optional amount of items to delete in this call. If it is not provided, nothing will be deleted. You can delete further items in a subsequent execute call.",
+      "default": 0,
+      "type": "integer",
+      "format": "uint32",
+      "minimum": 0.0
+    },
     "payout": {
       "description": "The address we send all remaining balance to",
       "type": "string"

contracts/burner/src/contract.rs

@@ -1,5 +1,5 @@
 use cosmwasm_std::{
-    entry_point, BankMsg, DepsMut, Env, MessageInfo, Order, Response, StdError, StdResult,
+    entry_point, BankMsg, DepsMut, Env, MessageInfo, Order, Response, StdError, StdResult, Storage,
 };
 
 use crate::msg::{ExecuteMsg, InstantiateMsg, MigrateMsg};
@@ -24,10 +24,14 @@ pub fn migrate(deps: DepsMut, env: Env, msg: MigrateMsg) -> StdResult<Response>
         to_address: msg.payout.clone(),
         amount: balance,
     };
+
+    let deleted = cleanup(deps.storage, msg.delete as usize);
+
     Ok(Response::new()
         .add_message(send)
-        .add_attribute("action", "burn")
-        .add_attribute("payout", msg.payout))
+        .add_attribute("action", "migrate")
+        .add_attribute("payout", msg.payout)
+        .add_attribute("deleted_entries", deleted.to_string()))
 }
 
 #[entry_point]
@@ -43,32 +47,36 @@ pub fn execute_cleanup(
     _info: MessageInfo,
     limit: Option<u32>,
 ) -> StdResult<Response> {
-    // the number of elements we can still take (decreasing over time)
-    let mut limit = limit.unwrap_or(u32::MAX) as usize;
+    let limit = limit.unwrap_or(u32::MAX) as usize;
+    let deleted = cleanup(deps.storage, limit);
+
+    Ok(Response::new()
+        .add_attribute("action", "cleanup")
+        .add_attribute("deleted_entries", deleted.to_string()))
+}
 
+fn cleanup(storage: &mut dyn Storage, mut limit: usize) -> usize {
     let mut deleted = 0;
     const PER_SCAN: usize = 20;
     loop {
         let take_this_scan = std::cmp::min(PER_SCAN, limit);
-        let keys: Vec<_> = deps
-            .storage
+        let keys: Vec<_> = storage
             .range_keys(None, None, Order::Ascending)
             .take(take_this_scan)
             .collect();
         let deleted_this_scan = keys.len();
         for k in keys {
-            deps.storage.remove(&k);
+            storage.remove(&k);
         }
         deleted += deleted_this_scan;
+        // decrease the number of elements we can still take
         limit -= deleted_this_scan;
         if limit == 0 || deleted_this_scan < take_this_scan {
             break;
         }
     }
 
-    Ok(Response::new()
-        .add_attribute("action", "burn")
-        .add_attribute("deleted_entries", deleted.to_string()))
+    deleted
 }
 
 #[cfg(test)]
@@ -114,6 +122,7 @@ mod tests {
         let payout = String::from("someone else");
         let msg = MigrateMsg {
             payout: payout.clone(),
+            delete: 0,
         };
         let res = migrate(deps.as_mut(), mock_env(), msg).unwrap();
         // check payout
@@ -128,6 +137,29 @@ mod tests {
         );
     }
 
+    #[test]
+    fn migrate_with_delete() {
+        let mut deps = mock_dependencies_with_balance(&coins(123456, "gold"));
+
+        // store some sample data
+        deps.storage.set(b"foo", b"bar");
+        deps.storage.set(b"key2", b"data2");
+        deps.storage.set(b"key3", b"cool stuff");
+        let cnt = deps.storage.range(None, None, Order::Ascending).count();
+        assert_eq!(cnt, 3);
+
+        // migrate all of the data in one go
+        let msg = MigrateMsg {
+            payout: "user".to_string(),
+            delete: 100,
+        };
+        migrate(deps.as_mut(), mock_env(), msg).unwrap();
+
+        // no more data
+        let cnt = deps.storage.range(None, None, Order::Ascending).count();
+        assert_eq!(cnt, 0);
+    }
+
     #[test]
     fn execute_cleans_up_data() {
         let mut deps = mock_dependencies_with_balance(&coins(123456, "gold"));
@@ -141,7 +173,7 @@ mod tests {
 
         // change the verifier via migrate
         let payout = String::from("someone else");
-        let msg = MigrateMsg { payout };
+        let msg = MigrateMsg { payout, delete: 0 };
         let _res = migrate(deps.as_mut(), mock_env(), msg).unwrap();
 
         let res = execute(

contracts/burner/src/msg.rs

@@ -4,6 +4,11 @@ use cosmwasm_schema::cw_serde;
 pub struct MigrateMsg {
     /// The address we send all remaining balance to
     pub payout: String,
+    /// Optional amount of items to delete in this call.
+    /// If it is not provided, nothing will be deleted.
+    /// You can delete further items in a subsequent execute call.
+    #[serde(default)]
+    pub delete: u32,
 }
 
 /// A placeholder where we don't take any input

contracts/burner/tests/integration.rs

@@ -62,6 +62,7 @@ fn migrate_sends_funds() {
     let payout = String::from("someone else");
     let msg = MigrateMsg {
         payout: payout.clone(),
+        delete: 0,
     };
     let res: Response = migrate(&mut deps, mock_env(), msg).unwrap();
     // check payout
@@ -94,7 +95,7 @@ fn execute_cleans_up_data() {
 
     // change the verifier via migrate
     let payout = String::from("someone else");
-    let msg = MigrateMsg { payout };
+    let msg = MigrateMsg { payout, delete: 0 };
     let _res: Response = migrate(&mut deps, mock_env(), msg).unwrap();
 
     let res: Response = execute(