[chain] Query balances at block_height

chain/src/main.rs

@@ -150,9 +150,10 @@ async fn crawling_fn(
         .into_rpc_error()?;
 
     let addresses = block.addresses_with_balance_change(native_token);
-    let balances = namada_service::query_balance(&client, &addresses)
-        .await
-        .into_rpc_error()?;
+    let balances =
+        namada_service::query_balance(&client, &addresses, Some(block_height))
+            .await
+            .into_rpc_error()?;
     tracing::info!("Updating balance for {} addresses...", addresses.len());
 
     let next_governance_proposal_id =
@@ -291,7 +292,9 @@ async fn initial_query(
         sleep(Duration::from_secs(initial_query_retry_time)).await;
     }
 
-    let balances = query_all_balances(client).await.into_rpc_error()?;
+    let balances = query_all_balances(client, Some(block_height))
+        .await
+        .into_rpc_error()?;
 
     tracing::info!("Querying bonds and unbonds...");
     let (bonds, unbonds) = query_all_bonds_and_unbonds(client, None, None)

chain/src/services/namada.rs

@@ -57,6 +57,7 @@ pub async fn get_epoch_at_block_height(
 pub async fn query_balance(
     client: &HttpClient,
     balance_changes: &HashSet<BalanceChange>,
+    block_height: Option<BlockHeight>,
 ) -> anyhow::Result<Balances> {
     Ok(futures::stream::iter(balance_changes)
         .filter_map(|balance_change| async move {
@@ -74,9 +75,14 @@ pub async fn query_balance(
                     .context("Failed to parse token address")
                     .ok()?;
 
-            let amount = rpc::get_token_balance(client, &token, &owner)
-                .await
-                .unwrap_or_default();
+            let amount = rpc::get_token_balance(
+                client,
+                &token,
+                &owner,
+                block_height.map(to_block_height),
+            )
+            .await
+            .unwrap_or_default();
 
             Some(Balance {
                 owner: Id::from(owner),
@@ -92,6 +98,7 @@ pub async fn query_balance(
 
 pub async fn query_all_balances(
     client: &HttpClient,
+    height: Option<BlockHeight>,
 ) -> anyhow::Result<Balances> {
     let token_addr = RPC
         .shell()
@@ -102,7 +109,7 @@ pub async fn query_all_balances(
     let balance_prefix = namada_token::storage_key::balance_prefix(&token_addr);
 
     let balances =
-        query_storage_prefix::<token::Amount>(client, &balance_prefix)
+        query_storage_prefix::<token::Amount>(client, &balance_prefix, height)
             .await
             .context("Failed to query all balances")?;
 
@@ -490,6 +497,8 @@ pub async fn query_all_votes(
     anyhow::Ok(votes.iter().flatten().cloned().collect())
 }
 
-fn to_block_height(block_height: u32) -> NamadaSdkBlockHeight {
+pub(super) fn to_block_height(
+    block_height: BlockHeight,
+) -> NamadaSdkBlockHeight {
     NamadaSdkBlockHeight::from(block_height as u64)
 }

chain/src/services/utils.rs

@@ -3,19 +3,28 @@ use namada_sdk::queries::RPC;
 use namada_sdk::storage::{self, PrefixValue};
 use tendermint_rpc::HttpClient;
 
+use shared::block::BlockHeight;
+
 /// Query a range of storage values with a matching prefix and decode them with
 /// [`BorshDeserialize`]. Returns an iterator of the storage keys paired with
 /// their associated values.
 pub async fn query_storage_prefix<T>(
     client: &HttpClient,
     key: &storage::Key,
+    height: Option<BlockHeight>,
 ) -> anyhow::Result<Option<impl Iterator<Item = (storage::Key, T)>>>
 where
     T: BorshDeserialize,
 {
     let values = RPC
         .shell()
-        .storage_prefix(client, None, None, false, key)
+        .storage_prefix(
+            client,
+            None,
+            height.map(super::namada::to_block_height),
+            false,
+            key,
+        )
         .await?;
 
     let decode = |PrefixValue { key, value }: PrefixValue| {