From 24f435357c46f5046b89e4a7058c85a420f0cd26 Mon Sep 17 00:00:00 2001 From: Philip Offtermatt Date: Mon, 1 Jul 2024 13:29:22 +0200 Subject: [PATCH] Add pagination for query_user_lockups, query_expired_lockups, and query_proposal_tributes --- contracts/atom_wars/src/contract.rs | 55 ++++++++++++++++------ contracts/atom_wars/src/query.rs | 4 ++ contracts/atom_wars/src/testing.rs | 16 +++---- contracts/atom_wars/src/testing_queries.rs | 8 +++- contracts/tribute/src/contract.rs | 9 +++- contracts/tribute/src/query.rs | 2 + contracts/tribute/src/testing.rs | 2 + 7 files changed, 72 insertions(+), 24 deletions(-) diff --git a/contracts/atom_wars/src/contract.rs b/contracts/atom_wars/src/contract.rs index fe9cbdf..001941d 100644 --- a/contracts/atom_wars/src/contract.rs +++ b/contracts/atom_wars/src/contract.rs @@ -26,7 +26,7 @@ const CONTRACT_NAME: &str = "atom_wars"; /// Contract version that is used for migration. const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION"); -pub const DEFAULT_MAX_ENTRIES: usize = 100; +pub const MAX_LOCK_ENTRIES: usize = 100; #[entry_point] pub fn instantiate( @@ -147,11 +147,10 @@ fn lock_tokens( } // validate that the user does not have too many locks - if get_lock_count(deps.as_ref(), info.sender.clone()) >= DEFAULT_MAX_ENTRIES.try_into().unwrap() - { + if get_lock_count(deps.as_ref(), info.sender.clone()) >= MAX_LOCK_ENTRIES.try_into().unwrap() { return Err(ContractError::Std(StdError::generic_err(format!( "User has too many locks, only {} locks allowed", - DEFAULT_MAX_ENTRIES + MAX_LOCK_ENTRIES )))); } @@ -624,12 +623,18 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult { match msg { QueryMsg::Constants {} => to_json_binary(&query_constants(deps)?), QueryMsg::Tranches {} => to_json_binary(&query_tranches(deps)?), - QueryMsg::AllUserLockups { address } => { - to_json_binary(&query_all_user_lockups(deps, address)?) - } - QueryMsg::ExpiredUserLockups { address } => { - to_json_binary(&query_expired_user_lockups(deps, env, address)?) - } + QueryMsg::AllUserLockups { + address, + start_from, + limit, + } => to_json_binary(&query_all_user_lockups(deps, address, start_from, limit)?), + QueryMsg::ExpiredUserLockups { + address, + start_from, + limit, + } => to_json_binary(&query_expired_user_lockups( + deps, env, address, start_from, limit, + )?), QueryMsg::UserVotingPower { address } => { to_json_binary(&query_user_voting_power(deps, env, address)?) } @@ -685,9 +690,20 @@ pub fn query_constants(deps: Deps) -> StdResult { CONSTANTS.load(deps.storage) } -pub fn query_all_user_lockups(deps: Deps, address: String) -> StdResult { +pub fn query_all_user_lockups( + deps: Deps, + address: String, + start_from: u32, + limit: u32, +) -> StdResult { Ok(UserLockupsResponse { - lockups: query_user_lockups(deps, deps.api.addr_validate(&address)?, |_| true), + lockups: query_user_lockups( + deps, + deps.api.addr_validate(&address)?, + |_| true, + start_from, + limit, + ), }) } @@ -695,12 +711,20 @@ pub fn query_expired_user_lockups( deps: Deps, env: Env, address: String, + start_from: u32, + limit: u32, ) -> StdResult { let user_address = deps.api.addr_validate(&address)?; let expired_lockup_predicate = |l: &LockEntry| l.lock_end < env.block.time; Ok(UserLockupsResponse { - lockups: query_user_lockups(deps, user_address, expired_lockup_predicate), + lockups: query_user_lockups( + deps, + user_address, + expired_lockup_predicate, + start_from, + limit, + ), }) } @@ -838,13 +862,16 @@ fn query_user_lockups( deps: Deps, user_address: Addr, predicate: impl FnMut(&LockEntry) -> bool, + start_from: u32, + limit: u32, ) -> Vec { LOCKS_MAP .prefix(user_address) .range(deps.storage, None, None, Order::Ascending) .map(|l| l.unwrap().1) .filter(predicate) - .take(DEFAULT_MAX_ENTRIES) + .skip(start_from as usize) + .take(limit as usize) .collect() } diff --git a/contracts/atom_wars/src/query.rs b/contracts/atom_wars/src/query.rs index 2446002..68539c4 100644 --- a/contracts/atom_wars/src/query.rs +++ b/contracts/atom_wars/src/query.rs @@ -9,9 +9,13 @@ pub enum QueryMsg { Tranches {}, AllUserLockups { address: String, + start_from: u32, + limit: u32, }, ExpiredUserLockups { address: String, + start_from: u32, + limit: u32, }, UserVotingPower { address: String, diff --git a/contracts/atom_wars/src/testing.rs b/contracts/atom_wars/src/testing.rs index 735889a..2f3f27e 100644 --- a/contracts/atom_wars/src/testing.rs +++ b/contracts/atom_wars/src/testing.rs @@ -1,4 +1,4 @@ -use crate::contract::DEFAULT_MAX_ENTRIES; +use crate::contract::MAX_LOCK_ENTRIES; use crate::query::QueryMsg; use crate::state::Tranche; use crate::{ @@ -86,7 +86,7 @@ fn lock_tokens_basic_test() { let res = execute(deps.as_mut(), env.clone(), info2.clone(), msg); assert!(res.is_ok()); - let res = query_all_user_lockups(deps.as_ref(), user_address.to_string()); + let res = query_all_user_lockups(deps.as_ref(), user_address.to_string(), 0, 2000); assert!(res.is_ok()); let res = res.unwrap(); assert_eq!(2, res.lockups.len()); @@ -775,9 +775,9 @@ fn test_too_many_locks() { let lock_msg = ExecuteMsg::LockTokens { lock_duration: ONE_MONTH_IN_NANO_SECONDS, }; - for i in 0..DEFAULT_MAX_ENTRIES + 10 { + for i in 0..MAX_LOCK_ENTRIES + 10 { let res = execute(deps.as_mut(), env.clone(), info.clone(), lock_msg.clone()); - if i < DEFAULT_MAX_ENTRIES { + if i < MAX_LOCK_ENTRIES { assert!(res.is_ok()); } else { assert!(res.is_err()); @@ -790,9 +790,9 @@ fn test_too_many_locks() { // now test that another user can still lock tokens let info2 = mock_info("addr0001", &[Coin::new(1000, STATOM.to_string())]); - for i in 0..DEFAULT_MAX_ENTRIES + 10 { + for i in 0..MAX_LOCK_ENTRIES + 10 { let res = execute(deps.as_mut(), env.clone(), info2.clone(), lock_msg.clone()); - if i < DEFAULT_MAX_ENTRIES { + if i < MAX_LOCK_ENTRIES { assert!(res.is_ok()); } else { assert!(res.is_err()); @@ -810,9 +810,9 @@ fn test_too_many_locks() { assert!(res.is_ok()); // now the first user can lock tokens again - for i in 0..DEFAULT_MAX_ENTRIES + 10 { + for i in 0..MAX_LOCK_ENTRIES + 10 { let res = execute(deps.as_mut(), env.clone(), info.clone(), lock_msg.clone()); - if i < DEFAULT_MAX_ENTRIES { + if i < MAX_LOCK_ENTRIES { assert!(res.is_ok()); } else { assert!(res.is_err()); diff --git a/contracts/atom_wars/src/testing_queries.rs b/contracts/atom_wars/src/testing_queries.rs index c941c3c..b2b2441 100644 --- a/contracts/atom_wars/src/testing_queries.rs +++ b/contracts/atom_wars/src/testing_queries.rs @@ -145,7 +145,13 @@ fn get_expired_user_lockups( env: Env, user_address: String, ) -> Vec { - let res = query_expired_user_lockups(deps.as_ref(), env.clone(), user_address.to_string()); + let res = query_expired_user_lockups( + deps.as_ref(), + env.clone(), + user_address.to_string(), + 0, + 2000, + ); assert!(res.is_ok()); let res = res.unwrap(); diff --git a/contracts/tribute/src/contract.rs b/contracts/tribute/src/contract.rs index c5cfc8c..af38b17 100644 --- a/contracts/tribute/src/contract.rs +++ b/contracts/tribute/src/contract.rs @@ -298,11 +298,15 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult { round_id, tranche_id, proposal_id, + start_from, + limit, } => to_json_binary(&query_proposal_tributes( deps, round_id, tranche_id, proposal_id, + start_from, + limit, )), } } @@ -312,12 +316,15 @@ pub fn query_proposal_tributes( round_id: u64, tranche_id: u64, proposal_id: u64, + start_from: u32, + limit: u32, ) -> Vec { TRIBUTE_MAP .prefix(((round_id, tranche_id), proposal_id)) .range(deps.storage, None, None, Order::Ascending) .map(|l| l.unwrap().1) - .take(DEFAULT_MAX_ENTRIES) + .skip(start_from as usize) + .take(limit as usize) .collect() } diff --git a/contracts/tribute/src/query.rs b/contracts/tribute/src/query.rs index 88bfd4d..4a790b6 100644 --- a/contracts/tribute/src/query.rs +++ b/contracts/tribute/src/query.rs @@ -9,5 +9,7 @@ pub enum QueryMsg { round_id: u64, tranche_id: u64, proposal_id: u64, + start_from: u32, + limit: u32, }, } diff --git a/contracts/tribute/src/testing.rs b/contracts/tribute/src/testing.rs index 22cfc66..1d4a019 100644 --- a/contracts/tribute/src/testing.rs +++ b/contracts/tribute/src/testing.rs @@ -275,6 +275,8 @@ fn add_tribute_test() { test.mock_data.0, test.proposal_info.0, test.proposal_info.1, + 0, + 3000, ); assert_eq!(test.tributes_to_add.len(), res.len());