Add NEAR to balances

src/main.rs

@@ -264,15 +264,17 @@ async fn get_balances(
 
     let client = reqwest::Client::new();
     let mut handles = vec![];
+
     for account in accounts {
         let client = client.clone();
         let ft_service = ft_service.clone();
         let start_block_id = start_block_id;
         let end_block_id = end_block_id;
 
         let handle = spawn(async move {
+            info!("Getting balances for {}", account);
             let mut rows: Vec<GetBalancesResultRow> = vec![];
-            let _a = account.clone();
+
             let likely_tokens = client
                 .get(format!(
                     "https://api.kitwallet.app/account/{account}/likelyTokens"
@@ -281,56 +283,107 @@ async fn get_balances(
                 .await?
                 .json::<Vec<String>>()
                 .await?;
-            info!("{}: {:?}", account, likely_tokens);
-            for token in likely_tokens {
-                let metadata = match ft_service.assert_ft_metadata(&token).await {
-                    Ok(v) => v,
-                    Err(e) => {
-                        debug!("{}: {}", account, e);
-                        continue;
-                    }
-                };
-                let start_balance = match ft_service
-                    .assert_ft_balance(&token, &account, start_block_id as u64)
-                    .await
-                {
-                    Ok(v) => v,
-                    Err(e) => {
-                        debug!("{}: {}", account, e);
-                        0.0
+            info!("Account {} likely tokens: {:?}", account, likely_tokens);
+
+            let token_handles: Vec<_> = likely_tokens
+                .iter()
+                .map(|token| {
+                    let token = token.clone();
+                    let account = account.clone();
+                    let ft_service = ft_service.clone();
+                    async move {
+                        let metadata = match ft_service.assert_ft_metadata(&token).await {
+                            Ok(v) => v,
+                            Err(e) => {
+                                debug!("{}: {}", account, e);
+                                return Err(e);
+                            }
+                        };
+                        let start_balance = match ft_service
+                            .assert_ft_balance(&token, &account, start_block_id as u64)
+                            .await
+                        {
+                            Ok(v) => v,
+                            Err(e) => {
+                                debug!("{}: {}", account, e);
+                                0.0
+                            }
+                        };
+                        let end_balance = match ft_service
+                            .assert_ft_balance(&token, &account, end_block_id as u64)
+                            .await
+                        {
+                            Ok(v) => v,
+                            Err(e) => {
+                                debug!("{}: {}", account, e);
+                                0.0
+                            }
+                        };
+                        let record = GetBalancesResultRow {
+                            account: account.clone(),
+                            start_date: start_date.to_rfc3339(),
+                            end_date: end_date.to_rfc3339(),
+                            start_block_id,
+                            end_block_id,
+                            start_balance,
+                            end_balance,
+                            token_id: token.clone(),
+                            symbol: metadata.symbol,
+                        };
+                        Ok(record)
                     }
-                };
-                let end_balance = match ft_service
-                    .assert_ft_balance(&token, &account, end_block_id as u64)
-                    .await
-                {
-                    Ok(v) => v,
+                })
+                .collect();
+
+            let token_results: Vec<_> = join_all(token_handles).await;
+            for result in token_results {
+                match result {
+                    Ok(record) => rows.push(record),
                     Err(e) => {
-                        debug!("{}: {}", account, e);
-                        0.0
+                        debug!("Token fetch error: {:?}", e);
                     }
-                };
-                let record = GetBalancesResultRow {
-                    account: account.clone(),
-                    start_date: start_date.to_rfc3339(),
-                    end_date: end_date.to_rfc3339(),
-                    start_block_id,
-                    end_block_id,
-                    start_balance,
-                    end_balance,
-                    token_id: token.clone(),
-                    symbol: metadata.symbol,
-                };
-                rows.push(record)
+                }
             }
-            println!("{:?}", rows);
+
+            let start_near_balance = match ft_service
+                .get_near_balance(&account, start_block_id as u64)
+                .await
+            {
+                Ok(v) => v,
+                Err(e) => {
+                    debug!("{}: {}", account, e);
+                    0.0
+                }
+            };
+            let end_near_balance = match ft_service
+                .get_near_balance(&account, end_block_id as u64)
+                .await
+            {
+                Ok(v) => v,
+                Err(e) => {
+                    debug!("{}: {}", account, e);
+                    0.0
+                }
+            };
+            let record = GetBalancesResultRow {
+                account: account.clone(),
+                start_date: start_date.to_rfc3339(),
+                end_date: end_date.to_rfc3339(),
+                start_block_id,
+                end_block_id,
+                start_balance: start_near_balance,
+                end_balance: end_near_balance,
+                token_id: "NEAR".to_string(),
+                symbol: "NEAR".to_string(),
+            };
+            rows.push(record);
+
             anyhow::Ok(rows)
         });
         handles.push(handle);
     }
 
     let mut rows = vec![];
-
     join_all(handles).await.iter().for_each(|row| match row {
         Ok(result) => match result {
             Ok(res) => rows.extend(res.iter().cloned()),
@@ -348,10 +401,8 @@ async fn get_balances(
         wtr.serialize(row).unwrap();
     }
     wtr.flush()?;
-    // Create a response with the CSV data
     let response = Response::builder()
         .header("Content-Type", "text/csv")
-        // .header("Content-Disposition", "attachment; filename=data.csv")
         .body(Body::from(wtr.into_inner().unwrap()))?;
     Ok(response)
 }

src/tta/ft_metadata.rs

@@ -10,7 +10,7 @@ use near_primitives::{
         Finality::{self},
         FunctionArgs,
     },
-    views::{CallResult, QueryRequest},
+    views::{AccountView, CallResult, QueryRequest},
 };
 use serde::{Deserialize, Serialize};
 use serde_json::json;
@@ -205,6 +205,34 @@ impl FtService {
 
         Ok(amount)
     }
+
+    pub async fn get_near_balance(&self, account_id: &str, block_id: u64) -> Result<f64> {
+        self.archival_rate_limiter.write().await.until_ready().await;
+        let RpcQueryResponse { kind, .. } = match self
+            .near_client
+            .call(RpcQueryRequest {
+                request: QueryRequest::ViewAccount {
+                    account_id: account_id.parse().unwrap(),
+                },
+                block_reference: BlockReference::BlockId(Height(block_id)),
+            })
+            .await
+        {
+            Ok(v) => v,
+            Err(e) => {
+                bail!("Error calling ViewAccount: {:?}", e);
+            }
+        };
+        let amount = match kind {
+            QueryResponseKind::ViewAccount(AccountView { amount, .. }) => amount,
+            _ => {
+                bail!("Received unexpected kind: {:?}", kind); // <-- Add this line
+            }
+        };
+
+        let amount = safe_divide_u128(amount, 24);
+        Ok(amount)
+    }
 }
 
 pub async fn view_function_call(
@@ -229,9 +257,11 @@ pub async fn view_function_call(
         }
     };
 
-    let QueryResponseKind::CallResult(CallResult{result, ..}) = kind else {
-      bail!("Unexpected response kind");
-    };
-
-    Ok(result)
+    match kind {
+        QueryResponseKind::CallResult(CallResult { result, .. }) => Ok(result),
+        _ => {
+            eprintln!("Received unexpected kind: {:?}", kind); // <-- Add this line
+            bail!("Unexpected response kind");
+        }
+    }
 }