Skip to content

Commit

Permalink
Add pagination for query_user_lockups, query_expired_lockups, and que…
Browse files Browse the repository at this point in the history
…ry_proposal_tributes
  • Loading branch information
p-offtermatt committed Jul 1, 2024
1 parent f7f1bc9 commit 24f4353
Show file tree
Hide file tree
Showing 7 changed files with 72 additions and 24 deletions.
55 changes: 41 additions & 14 deletions contracts/atom_wars/src/contract.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
))));
}

Expand Down Expand Up @@ -624,12 +623,18 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
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)?)
}
Expand Down Expand Up @@ -685,22 +690,41 @@ pub fn query_constants(deps: Deps) -> StdResult<Constants> {
CONSTANTS.load(deps.storage)
}

pub fn query_all_user_lockups(deps: Deps, address: String) -> StdResult<UserLockupsResponse> {
pub fn query_all_user_lockups(
deps: Deps,
address: String,
start_from: u32,
limit: u32,
) -> StdResult<UserLockupsResponse> {
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,
),
})
}

pub fn query_expired_user_lockups(
deps: Deps,
env: Env,
address: String,
start_from: u32,
limit: u32,
) -> StdResult<UserLockupsResponse> {
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,
),
})
}

Expand Down Expand Up @@ -838,13 +862,16 @@ fn query_user_lockups(
deps: Deps,
user_address: Addr,
predicate: impl FnMut(&LockEntry) -> bool,
start_from: u32,
limit: u32,
) -> Vec<LockEntry> {
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()
}

Expand Down
4 changes: 4 additions & 0 deletions contracts/atom_wars/src/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
16 changes: 8 additions & 8 deletions contracts/atom_wars/src/testing.rs
Original file line number Diff line number Diff line change
@@ -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::{
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -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());
Expand All @@ -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());
Expand All @@ -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());
Expand Down
8 changes: 7 additions & 1 deletion contracts/atom_wars/src/testing_queries.rs
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,13 @@ fn get_expired_user_lockups(
env: Env,
user_address: String,
) -> Vec<LockEntry> {
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();

Expand Down
9 changes: 8 additions & 1 deletion contracts/tribute/src/contract.rs
Original file line number Diff line number Diff line change
Expand Up @@ -298,11 +298,15 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
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,
)),
}
}
Expand All @@ -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> {
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()
}

Expand Down
2 changes: 2 additions & 0 deletions contracts/tribute/src/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,5 +9,7 @@ pub enum QueryMsg {
round_id: u64,
tranche_id: u64,
proposal_id: u64,
start_from: u32,
limit: u32,
},
}
2 changes: 2 additions & 0 deletions contracts/tribute/src/testing.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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());

Expand Down

0 comments on commit 24f4353

Please sign in to comment.