From 014bc9c1d0ba754394f7cf1d89fdba4887ebbc92 Mon Sep 17 00:00:00 2001
From: Joel Nordell <joel@unit410.com>
Date: Thu, 18 Jul 2024 15:13:17 -0500
Subject: [PATCH] [chain] Query balances at block_height

---
 chain/src/main.rs            | 11 +++++++----
 chain/src/services/namada.rs | 19 ++++++++++++++-----
 chain/src/services/utils.rs  | 11 ++++++++++-
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/chain/src/main.rs b/chain/src/main.rs
index c5973abe..cad4a120 100644
--- a/chain/src/main.rs
+++ b/chain/src/main.rs
@@ -157,9 +157,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 =
@@ -326,7 +327,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)
diff --git a/chain/src/services/namada.rs b/chain/src/services/namada.rs
index 5c57d59d..36862777 100644
--- a/chain/src/services/namada.rs
+++ b/chain/src/services/namada.rs
@@ -69,6 +69,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 {
@@ -86,9 +87,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),
@@ -104,6 +110,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()
@@ -114,7 +121,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")?;
 
@@ -522,6 +529,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)
 }
diff --git a/chain/src/services/utils.rs b/chain/src/services/utils.rs
index 4e131158..64dcfc57 100644
--- a/chain/src/services/utils.rs
+++ b/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| {