chore: balance_as_of

cala-ledger/migrations/20231208110808_cala_ledger_setup.sql

@@ -190,6 +190,7 @@ CREATE TABLE cala_balance_history (
   FOREIGN KEY (data_source_id, journal_id, account_id, currency) REFERENCES cala_current_balances(data_source_id, journal_id, account_id, currency),
   FOREIGN KEY (data_source_id, latest_entry_id) REFERENCES cala_entries(data_source_id, id)
 );
+CREATE INDEX idx_cala_balance_history_recorded_at ON cala_balance_history (recorded_at);
 
 CREATE TABLE cala_velocity_limits (
   data_source_id UUID NOT NULL DEFAULT '00000000-0000-0000-0000-000000000000',

cala-ledger/src/balance/account_balance.rs

@@ -11,6 +11,25 @@ pub struct AccountBalance {
 }
 
 impl AccountBalance {
+    pub(super) fn derive_as_of(mut self, as_of: Self) -> Self {
+        self.details.settled = BalanceAmount {
+            dr_balance: self.details.settled.dr_balance - as_of.details.settled.dr_balance,
+            cr_balance: self.details.settled.cr_balance - as_of.details.settled.cr_balance,
+            ..self.details.settled
+        };
+        self.details.pending = BalanceAmount {
+            dr_balance: self.details.pending.dr_balance - as_of.details.pending.dr_balance,
+            cr_balance: self.details.pending.cr_balance - as_of.details.pending.cr_balance,
+            ..self.details.pending
+        };
+        self.details.encumbrance = BalanceAmount {
+            dr_balance: self.details.encumbrance.dr_balance - as_of.details.encumbrance.dr_balance,
+            cr_balance: self.details.encumbrance.cr_balance - as_of.details.encumbrance.cr_balance,
+            ..self.details.encumbrance
+        };
+        self
+    }
+
     pub fn pending(&self) -> Decimal {
         if self.balance_type == DebitOrCredit::Credit {
             self.details.pending.cr_balance - self.details.pending.dr_balance

cala-ledger/src/balance/mod.rs

@@ -64,6 +64,26 @@ impl Balances {
             .await
     }
 
+    #[instrument(name = "cala_ledger.balance.find_as_of", skip(self), err)]
+    pub async fn find_as_of(
+        &self,
+        journal_id: JournalId,
+        account_id: AccountId,
+        currency: Currency,
+        as_of: DateTime<Utc>,
+        up_until: Option<DateTime<Utc>>,
+    ) -> Result<AccountBalance, BalanceError> {
+        match self
+            .repo
+            .find_as_of(journal_id, account_id, currency, as_of, up_until)
+            .await?
+        {
+            (Some(last_before), Some(up_until)) => Ok(up_until.derive_as_of(last_before)),
+            (None, Some(up_until)) => Ok(up_until),
+            _ => Err(BalanceError::NotFound(journal_id, account_id, currency)),
+        }
+    }
+
     #[instrument(name = "cala_ledger.balance.find_all", skip(self), err)]
     pub async fn find_all(
         &self,

cala-ledger/src/balance/repo.rs

@@ -1,3 +1,4 @@
+use chrono::{DateTime, Utc};
 use sqlx::{Executor, PgPool, Postgres, QueryBuilder, Row, Transaction};
 use tracing::instrument;
 
@@ -86,6 +87,85 @@ impl BalanceRepo {
         }
     }
 
+    pub(super) async fn find_as_of(
+        &self,
+        journal_id: JournalId,
+        account_id: AccountId,
+        currency: Currency,
+        as_of: DateTime<Utc>,
+        up_until: Option<DateTime<Utc>>,
+    ) -> Result<(Option<AccountBalance>, Option<AccountBalance>), BalanceError> {
+        let rows = sqlx::query!(
+            r#"
+        WITH last_before_as_of AS (
+            SELECT
+              true AS last_before, false AS up_until, h.values,
+              a.normal_balance_type AS "normal_balance_type!: DebitOrCredit", h.recorded_at
+            FROM cala_balance_history h
+            JOIN cala_accounts a
+            ON h.data_source_id = a.data_source_id
+            AND h.account_id = a.id
+            WHERE h.data_source_id = '00000000-0000-0000-0000-000000000000'
+            AND h.journal_id = $1
+            AND h.account_id = $2
+            AND h.currency = $3
+            AND h.recorded_at < $4
+            ORDER BY h.recorded_at DESC
+            LIMIT 1
+        ),
+        last_before_or_equal_up_until AS (
+            SELECT 
+              false AS last_before, true AS up_until, h.values,
+              a.normal_balance_type AS "normal_balance_type!: DebitOrCredit", h.recorded_at
+            FROM cala_balance_history h
+            JOIN cala_accounts a
+            ON h.data_source_id = a.data_source_id
+            AND h.account_id = a.id
+            WHERE h.data_source_id = '00000000-0000-0000-0000-000000000000'
+            AND h.journal_id = $1
+            AND h.account_id = $2
+            AND h.currency = $3
+            AND h.recorded_at <= COALESCE($5, NOW())
+            ORDER BY h.recorded_at DESC
+            LIMIT 1
+        )
+        SELECT * FROM last_before_as_of
+        UNION ALL
+        SELECT * FROM last_before_or_equal_up_until
+        "#,
+            journal_id as JournalId,
+            account_id as AccountId,
+            currency.code(),
+            as_of,
+            up_until,
+        )
+        .fetch_all(&self.pool)
+        .await?;
+
+        let mut last_before = None;
+        let mut up_until = None;
+        for row in rows {
+            if row.last_before.expect("last_before is not null") {
+                let details: BalanceSnapshot =
+                    serde_json::from_value(row.values.expect("values is not null"))
+                        .expect("Failed to deserialize balance snapshot");
+                last_before = Some(AccountBalance {
+                    balance_type: row.normal_balance_type,
+                    details,
+                });
+            } else {
+                let details: BalanceSnapshot =
+                    serde_json::from_value(row.values.expect("values is not null"))
+                        .expect("Failed to deserialize balance snapshot");
+                up_until = Some(AccountBalance {
+                    balance_type: row.normal_balance_type,
+                    details,
+                });
+            }
+        }
+        Ok((last_before, up_until))
+    }
+
     pub(super) async fn find_all(
         &self,
         ids: &[BalanceId],

cala-server/schema.graphql

@@ -12,6 +12,7 @@ type Account {
 	createdAt: Timestamp!
 	modifiedAt: Timestamp!
 	balance(journalId: UUID!, currency: CurrencyCode!): Balance
+	balanceAsOf(journalId: UUID!, currency: CurrencyCode!, asOf: Timestamp!, upUntil: Timestamp): Balance
 	sets(first: Int!, after: String): AccountSetConnection!
 }
 
@@ -388,6 +389,7 @@ type Query {
 	accountSet(id: UUID!): AccountSet
 	journal(id: UUID!): Journal
 	balance(journalId: UUID!, accountId: UUID!, currency: CurrencyCode!): Balance
+	balanceAsOf(journalId: UUID!, accountId: UUID!, currency: CurrencyCode!, asOf: Timestamp!, upUntil: Timestamp): Balance
 	transaction(id: UUID!): Transaction
 	transactionByExternalId(externalId: String!): Transaction
 	txTemplate(id: UUID!): TxTemplate

cala-server/src/graphql/account.rs

@@ -61,6 +61,33 @@ impl Account {
         Ok(balance.map(Balance::from))
     }
 
+    async fn balance_as_of(
+        &self,
+        ctx: &Context<'_>,
+        journal_id: UUID,
+        currency: CurrencyCode,
+        as_of: Timestamp,
+        up_until: Option<Timestamp>,
+    ) -> async_graphql::Result<Option<Balance>> {
+        let app = ctx.data_unchecked::<CalaApp>();
+        match app
+            .ledger()
+            .balances()
+            .find_as_of(
+                JournalId::from(journal_id),
+                AccountId::from(self.account_id),
+                Currency::from(currency),
+                as_of.into_inner(),
+                up_until.map(|ts| ts.into_inner()),
+            )
+            .await
+        {
+            Ok(balance) => Ok(Some(balance.into())),
+            Err(cala_ledger::balance::error::BalanceError::NotFound(_, _, _)) => Ok(None),
+            Err(err) => Err(err.into()),
+        }
+    }
+
     async fn sets(
         &self,
         ctx: &Context<'_>,

cala-server/src/graphql/primitives.rs

@@ -160,6 +160,11 @@ impl From<NaiveDate> for Date {
         Self(value)
     }
 }
+impl From<Date> for NaiveDate {
+    fn from(value: Date) -> Self {
+        value.0
+    }
+}
 
 #[derive(Serialize, Deserialize)]
 #[serde(transparent)]

cala-server/src/graphql/schema.rs

@@ -124,6 +124,34 @@ impl<E: QueryExtensionMarker> CoreQuery<E> {
         Ok(balance.map(Balance::from))
     }
 
+    async fn balance_as_of(
+        &self,
+        ctx: &Context<'_>,
+        journal_id: UUID,
+        account_id: UUID,
+        currency: CurrencyCode,
+        as_of: Timestamp,
+        up_until: Option<Timestamp>,
+    ) -> async_graphql::Result<Option<Balance>> {
+        let app = ctx.data_unchecked::<CalaApp>();
+        match app
+            .ledger()
+            .balances()
+            .find_as_of(
+                JournalId::from(journal_id),
+                AccountId::from(account_id),
+                Currency::from(currency),
+                as_of.into_inner(),
+                up_until.map(|ts| ts.into_inner()),
+            )
+            .await
+        {
+            Ok(balance) => Ok(Some(balance.into())),
+            Err(cala_ledger::balance::error::BalanceError::NotFound(_, _, _)) => Ok(None),
+            Err(err) => Err(err.into()),
+        }
+    }
+
     async fn transaction(
         &self,
         ctx: &Context<'_>,