Merge pull request #130 from neutron-org/feat/improve-balances-query

feat: improve balances query

contracts/neutron_interchain_queries/schema/execute_msg.json

@@ -5,15 +5,15 @@
     {
       "type": "object",
       "required": [
-        "register_balance_query"
+        "register_balances_query"
       ],
       "properties": {
-        "register_balance_query": {
+        "register_balances_query": {
           "type": "object",
           "required": [
             "addr",
             "connection_id",
-            "denom",
+            "denoms",
             "update_period"
           ],
           "properties": {
@@ -23,8 +23,11 @@
             "connection_id": {
               "type": "string"
             },
-            "denom": {
-              "type": "string"
+            "denoms": {
+              "type": "array",
+              "items": {
+                "type": "string"
+              }
             },
             "update_period": {
               "type": "integer",

contracts/neutron_interchain_queries/src/contract.rs

@@ -24,7 +24,7 @@ use neutron_sdk::interchain_queries::{
     check_query_type, get_registered_query, query_kv_result,
     v047::{
         register_queries::{
-            new_register_balance_query_msg, new_register_bank_total_supply_query_msg,
+            new_register_balances_query_msg, new_register_bank_total_supply_query_msg,
             new_register_delegator_delegations_query_msg,
             new_register_delegator_unbonding_delegations_query_msg,
             new_register_distribution_fee_pool_query_msg, new_register_gov_proposals_query_msg,
@@ -69,12 +69,12 @@ pub fn execute(
     msg: ExecuteMsg,
 ) -> NeutronResult<Response<NeutronMsg>> {
     match msg {
-        ExecuteMsg::RegisterBalanceQuery {
+        ExecuteMsg::RegisterBalancesQuery {
             connection_id,
             addr,
-            denom,
+            denoms,
             update_period,
-        } => register_balance_query(connection_id, addr, denom, update_period),
+        } => register_balances_query(connection_id, addr, denoms, update_period),
         ExecuteMsg::RegisterBankTotalSupplyQuery {
             connection_id,
             denoms,
@@ -143,13 +143,13 @@ pub fn execute(
     }
 }
 
-pub fn register_balance_query(
+pub fn register_balances_query(
     connection_id: String,
     addr: String,
-    denom: String,
+    denoms: Vec<String>,
     update_period: u64,
 ) -> NeutronResult<Response<NeutronMsg>> {
-    let msg = new_register_balance_query_msg(connection_id, addr, denom, update_period)?;
+    let msg = new_register_balances_query_msg(connection_id, addr, denoms, update_period)?;
 
     Ok(Response::new().add_message(msg))
 }

contracts/neutron_interchain_queries/src/msg.rs

@@ -10,11 +10,11 @@ pub struct InstantiateMsg {}
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
 #[serde(rename_all = "snake_case")]
 pub enum ExecuteMsg {
-    RegisterBalanceQuery {
+    RegisterBalancesQuery {
         connection_id: String,
         update_period: u64,
         addr: String,
-        denom: String,
+        denoms: Vec<String>,
     },
     RegisterBankTotalSupplyQuery {
         connection_id: String,

contracts/neutron_interchain_queries/src/testing/tests.rs

@@ -212,20 +212,27 @@ fn build_interchain_query_gov_proposal_value(proposal_id: u64) -> StorageValue {
     }
 }
 
-fn build_interchain_query_balance_response(addr: Addr, denom: String, amount: String) -> Binary {
+fn build_interchain_query_balances_response(addr: Addr, balances: Vec<Coin>) -> Binary {
     let converted_addr_bytes = decode_and_convert(addr.as_str()).unwrap();
 
-    let balance_key = create_account_denom_balance_key(converted_addr_bytes, denom).unwrap();
+    let s: Vec<StorageValue> = balances
+        .iter()
+        .map(|c| {
+            let balance_key =
+                create_account_denom_balance_key(converted_addr_bytes.clone(), c.denom.clone())
+                    .unwrap();
+            StorageValue {
+                storage_prefix: "".to_string(),
+                key: Binary(balance_key),
+                value: Binary(c.amount.to_string().into_bytes()),
+            }
+        })
+        .collect();
 
-    let s = StorageValue {
-        storage_prefix: "".to_string(),
-        key: Binary(balance_key),
-        value: Binary(amount.into_bytes()),
-    };
     Binary::from(
         to_string(&QueryRegisteredQueryResultResponse {
             result: InterchainQueryResult {
-                kv_results: vec![s],
+                kv_results: s,
                 height: 123456,
                 revision: 2,
             },
@@ -256,11 +263,11 @@ fn register_query(
 fn test_query_balance() {
     let mut deps = dependencies(&[]);
 
-    let msg = ExecuteMsg::RegisterBalanceQuery {
+    let msg = ExecuteMsg::RegisterBalancesQuery {
         connection_id: "connection".to_string(),
         update_period: 10,
         addr: "osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs".to_string(),
-        denom: "uosmo".to_string(),
+        denoms: vec!["uosmo".to_string()],
     };
 
     let keys = register_query(&mut deps, mock_env(), mock_info("", &[]), msg);
@@ -271,10 +278,9 @@ fn test_query_balance() {
     deps.querier.add_registered_queries(1, registered_query);
     deps.querier.add_query_response(
         1,
-        build_interchain_query_balance_response(
+        build_interchain_query_balances_response(
             Addr::unchecked("osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs"),
-            "uosmo".to_string(),
-            "8278104".to_string(),
+            vec![Coin::new(8278104u128, "uosmo")],
         ),
     );
     let query_balance = QueryMsg::Balance { query_id: 1 };
@@ -291,6 +297,50 @@ fn test_query_balance() {
     )
 }
 
+#[test]
+fn test_query_balances() {
+    let mut deps = dependencies(&[]);
+
+    let msg = ExecuteMsg::RegisterBalancesQuery {
+        connection_id: "connection".to_string(),
+        update_period: 10,
+        addr: "osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs".to_string(),
+        denoms: vec!["uosmo".to_string(), "uatom".to_string()],
+    };
+
+    let keys = register_query(&mut deps, mock_env(), mock_info("", &[]), msg);
+
+    let registered_query =
+        build_registered_query_response(1, QueryParam::Keys(keys.0), QueryType::KV, 987);
+
+    deps.querier.add_registered_queries(1, registered_query);
+    deps.querier.add_query_response(
+        1,
+        build_interchain_query_balances_response(
+            Addr::unchecked("osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs"),
+            vec![
+                Coin::new(8278104u128, "uosmo"),
+                Coin::new(1234567u128, "uatom"),
+            ],
+        ),
+    );
+    let query_balance = QueryMsg::Balance { query_id: 1 };
+    let resp: BalanceResponse =
+        from_json(query(deps.as_ref(), mock_env(), query_balance).unwrap()).unwrap();
+    assert_eq!(
+        resp,
+        BalanceResponse {
+            last_submitted_local_height: 987,
+            balances: Balances {
+                coins: vec![
+                    Coin::new(8278104u128, "uosmo"),
+                    Coin::new(1234567u128, "uatom")
+                ]
+            },
+        }
+    )
+}
+
 #[test]
 fn test_bank_total_supply_query() {
     let mut deps = dependencies(&[]);

packages/neutron-sdk/src/interchain_queries/v045/helpers.rs

@@ -3,9 +3,9 @@ use crate::errors::error::NeutronResult;
 use crate::interchain_queries::helpers::{decode_and_convert, length_prefix};
 use crate::interchain_queries::types::AddressBytes;
 use crate::interchain_queries::v045::types::{
-    BALANCES_PREFIX, DELEGATION_KEY, FEE_POOL_KEY, PARAMS_STORE_DELIMITER, PROPOSALS_KEY_PREFIX,
-    SUPPLY_PREFIX, UNBONDING_DELEGATION_KEY, VALIDATORS_KEY, VALIDATOR_SIGNING_INFO_KEY,
-    WASM_CONTRACT_STORE_PREFIX,
+    BALANCES_PREFIX, BANK_STORE_KEY, DELEGATION_KEY, FEE_POOL_KEY, PARAMS_STORE_DELIMITER,
+    PROPOSALS_KEY_PREFIX, SUPPLY_PREFIX, UNBONDING_DELEGATION_KEY, VALIDATORS_KEY,
+    VALIDATOR_SIGNING_INFO_KEY, WASM_CONTRACT_STORE_PREFIX,
 };
 use crate::NeutronError;
 use cosmos_sdk_proto::cosmos::staking::v1beta1::Commission as ValidatorCommission;
@@ -46,6 +46,27 @@ pub fn create_account_denom_balance_key<AddrBytes: AsRef<[u8]>, S: AsRef<str>>(
     Ok(account_balance_key)
 }
 
+/// Creates keys for an Interchain Query to get balance of account on remote chain for list of denoms
+///
+/// * **addr** address of an account on remote chain for which you want to get balances;
+/// * **denoms** denominations of the coins for which you want to get balance;
+pub fn create_balances_query_keys(addr: String, denoms: Vec<String>) -> NeutronResult<Vec<KVKey>> {
+    let converted_addr_bytes = decode_and_convert(addr.as_str())?;
+    let mut kv_keys: Vec<KVKey> = Vec::with_capacity(denoms.len());
+
+    for denom in denoms {
+        let balance_key = create_account_denom_balance_key(&converted_addr_bytes, denom)?;
+
+        let kv_key = KVKey {
+            path: BANK_STORE_KEY.to_string(),
+            key: Binary(balance_key),
+        };
+
+        kv_keys.push(kv_key)
+    }
+    Ok(kv_keys)
+}
+
 /// Deconstructs a storage key for an **account** balance of a particular **denom**.
 /// Returns two values: **address** of an account and **denom**
 pub fn deconstruct_account_denom_balance_key<Key: IntoIterator<Item = u8>>(

packages/neutron-sdk/src/interchain_queries/v045/mod.rs

@@ -3,9 +3,10 @@ pub mod queries;
 pub mod register_queries;
 pub mod types;
 
+#[allow(deprecated)]
 pub use register_queries::{
-    new_register_balance_query_msg, new_register_bank_total_supply_query_msg,
-    new_register_delegator_delegations_query_msg,
+    new_register_balance_query_msg, new_register_balances_query_msg,
+    new_register_bank_total_supply_query_msg, new_register_delegator_delegations_query_msg,
     new_register_delegator_unbonding_delegations_query_msg,
     new_register_distribution_fee_pool_query_msg, new_register_gov_proposals_query_msg,
     new_register_staking_validators_query_msg, new_register_transfers_query_msg,

packages/neutron-sdk/src/interchain_queries/v045/register_queries.rs

@@ -10,40 +10,44 @@ use crate::{
     errors::error::NeutronResult,
     interchain_queries::helpers::decode_and_convert,
     interchain_queries::v045::helpers::{
-        create_account_denom_balance_key, create_delegation_key, create_fee_pool_key,
+        create_balances_query_keys, create_delegation_key, create_fee_pool_key,
         create_gov_proposal_keys, create_gov_proposals_voters_votes_keys, create_params_store_key,
         create_total_denom_key, create_unbonding_delegation_key, create_validator_key,
         create_validator_signing_info_key, create_wasm_contract_store_key,
     },
 };
 use cosmwasm_std::Binary;
 
-/// Creates a message to register an Interchain Query to get balance of account on remote chain for particular denom
+/// Creates a message to register an Interchain Query to get balance of account on remote chain for list of denoms
+///
+/// * **connection_id** is an IBC connection identifier between Neutron and remote chain;
+/// * **addr** address of an account on remote chain for which you want to get balances;
+/// * **denoms** denominations of the coins for which you want to get balance;
+/// * **update_period** is used to say how often the query must be updated.
+pub fn new_register_balances_query_msg(
+    connection_id: String,
+    addr: String,
+    denoms: Vec<String>,
+    update_period: u64,
+) -> NeutronResult<NeutronMsg> {
+    let kv_keys = create_balances_query_keys(addr, denoms)?;
+    NeutronMsg::register_interchain_query(QueryPayload::KV(kv_keys), connection_id, update_period)
+}
+
+/// Creates a message to register an Interchain Query to get balance of account on remote chain for a particular denom
 ///
 /// * **connection_id** is an IBC connection identifier between Neutron and remote chain;
 /// * **addr** address of an account on remote chain for which you want to get balances;
 /// * **denom** denomination of the coin for which you want to get balance;
 /// * **update_period** is used to say how often the query must be updated.
+#[deprecated(note = "Please use new_register_balances_query_msg instead")]
 pub fn new_register_balance_query_msg(
     connection_id: String,
     addr: String,
     denom: String,
     update_period: u64,
 ) -> NeutronResult<NeutronMsg> {
-    let converted_addr_bytes = decode_and_convert(addr.as_str())?;
-
-    let balance_key = create_account_denom_balance_key(converted_addr_bytes, denom)?;
-
-    let kv_key = KVKey {
-        path: BANK_STORE_KEY.to_string(),
-        key: Binary(balance_key),
-    };
-
-    NeutronMsg::register_interchain_query(
-        QueryPayload::KV(vec![kv_key]),
-        connection_id,
-        update_period,
-    )
+    new_register_balances_query_msg(connection_id, addr, vec![denom], update_period)
 }
 
 /// Creates a message to register an Interchain Query to get total supply on remote chain for particular denom

packages/neutron-sdk/src/interchain_queries/v045/testing.rs

@@ -1005,24 +1005,50 @@ fn test_delegations_reconstruct() {
 fn test_balance_reconstruct_from_hex() {
     let bytes = hex::decode(BALANCES_HEX_RESPONSE).unwrap(); // decode hex string to bytes
     let base64_input = BASE64_STANDARD.encode(bytes); // encode bytes to base64 string
+    let balance_key = create_account_denom_balance_key(
+        decode_and_convert("osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs").unwrap(),
+        "uosmo",
+    )
+    .unwrap();
 
     let s = StorageValue {
         storage_prefix: String::default(), // not used in reconstruct
-        key: Binary::default(),            // not used in reconstruct
+        key: Binary::from(balance_key),
         value: Binary::from_base64(base64_input.as_str()).unwrap(),
     };
     let bank_balances = Balances::reconstruct(&[s]).unwrap();
     assert_eq!(
         bank_balances,
         Balances {
             coins: vec![StdCoin {
-                denom: String::from("stake"),
+                denom: String::from("uosmo"),
                 amount: Uint128::from(99999000u64),
             }]
         }
     );
 }
 
+#[test]
+fn test_balance_reconstruct_from_empty_value() {
+    let balance_key = create_account_denom_balance_key(
+        decode_and_convert("osmo1yz54ncxj9csp7un3xled03q6thrrhy9cztkfzs").unwrap(),
+        "uosmo",
+    )
+    .unwrap();
+    let s = StorageValue {
+        storage_prefix: String::default(), // not used in reconstruct
+        key: Binary::from(balance_key),
+        value: Binary::from(vec![]),
+    };
+    let bank_balances = Balances::reconstruct(&[s]).unwrap();
+    assert_eq!(
+        bank_balances,
+        Balances {
+            coins: vec![StdCoin::new(0u128, "uosmo")]
+        }
+    );
+}
+
 #[test]
 fn test_bank_total_supply_reconstruct_from_hex() {
     let bytes = hex::decode(TOTAL_SUPPLY_HEX_RESPONSE).unwrap(); // decode hex string to bytes

packages/neutron-sdk/src/interchain_queries/v045/types.rs

@@ -1,5 +1,6 @@
 use crate::interchain_queries::helpers::uint256_to_u128;
 use crate::interchain_queries::types::KVReconstruct;
+use crate::interchain_queries::v045::helpers::deconstruct_account_denom_balance_key;
 use crate::{
     bindings::types::StorageValue,
     errors::error::{NeutronError, NeutronResult},
@@ -112,9 +113,15 @@ impl KVReconstruct for Balances {
         let mut coins: Vec<Coin> = Vec::with_capacity(storage_values.len());
 
         for kv in storage_values {
-            let balance: CosmosCoin = CosmosCoin::decode(kv.value.as_slice())?;
-            let amount = Uint128::from_str(balance.amount.as_str())?;
-            coins.push(Coin::new(amount.u128(), balance.denom));
+            let (_, denom) = deconstruct_account_denom_balance_key(kv.key.to_vec())?;
+            let amount = if kv.value.len() > 0 {
+                let balance: CosmosCoin = CosmosCoin::decode(kv.value.as_slice())?;
+                Uint128::from_str(balance.amount.as_str())?.u128()
+            } else {
+                0u128
+            };
+
+            coins.push(Coin::new(amount, denom))
         }
 
         Ok(Balances { coins })