From f4613851d137cf25ee02fbf156e61e301c483137 Mon Sep 17 00:00:00 2001
From: Vaibhav <thevaibhavdixit@gmail.com>
Date: Mon, 17 Jun 2024 17:06:54 +0530
Subject: [PATCH] chore: read balance in op for account

---
 cala-server/src/graphql/account.rs | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/cala-server/src/graphql/account.rs b/cala-server/src/graphql/account.rs
index f9713dd9..f5bdaca8 100644
--- a/cala-server/src/graphql/account.rs
+++ b/cala-server/src/graphql/account.rs
@@ -10,6 +10,7 @@ use crate::app::CalaApp;
 
 use super::{
     account_set::*, balance::Balance, convert::ToGlobalId, loader::LedgerDataLoader, primitives::*,
+    schema::DbOp,
 };
 
 #[derive(Clone, SimpleObject)]
@@ -37,12 +38,26 @@ impl Account {
         journal_id: UUID,
         currency: CurrencyCode,
     ) -> async_graphql::Result<Option<Balance>> {
-        let loader = ctx.data_unchecked::<DataLoader<LedgerDataLoader>>();
         let journal_id = JournalId::from(journal_id);
         let account_id = AccountId::from(self.account_id);
         let currency = Currency::from(currency);
-        let balance: Option<AccountBalance> =
-            loader.load_one((journal_id, account_id, currency)).await?;
+
+        let balance: Option<AccountBalance> = match ctx.data_opt::<DbOp>() {
+            Some(op) => {
+                let app = ctx.data_unchecked::<CalaApp>();
+                let mut op = op.try_lock().expect("Lock held concurrently");
+                Some(
+                    app.ledger()
+                        .balances()
+                        .find_in_op(&mut op, journal_id, account_id, currency)
+                        .await?,
+                )
+            }
+            None => {
+                let loader = ctx.data_unchecked::<DataLoader<LedgerDataLoader>>();
+                loader.load_one((journal_id, account_id, currency)).await?
+            }
+        };
         Ok(balance.map(Balance::from))
     }