Adding get_balance() database function

docs/api/index.rst

@@ -11,4 +11,5 @@ API Documentation
     forms
     utilities_money
     utilities_currency
+    utilities_database
     exceptions

docs/api/models.rst

@@ -42,4 +42,3 @@ LegView (Database View)
 -----------------------
 
 .. autoclass:: hordak.models.LegView
-    :members:

docs/api/utilities_database.rst

@@ -0,0 +1,10 @@
+Database Utilities
+==================
+
+.. contents::
+
+GetBalance()
+------------
+
+.. autoclass:: hordak.utilities.db_functions.GetBalance
+    :members: __init__

hordak/migrations/0047_get_balance.mysql.sql

@@ -0,0 +1,55 @@
+-- ----
+CREATE FUNCTION get_balance(account_id BIGINT, as_of DATE)
+RETURNS JSON
+BEGIN
+    DECLARE account_lft INT;
+    DECLARE account_rght INT;
+    DECLARE account_tree_id INT;
+    DECLARE result_json JSON;
+
+    -- Fetch the account's hierarchical information
+    SELECT lft, rght, tree_id INTO account_lft, account_rght, account_tree_id
+    FROM hordak_account
+    WHERE id = account_id;
+
+    -- Prepare the result set with sums calculated in a derived table (subquery)
+    IF as_of IS NOT NULL THEN
+        SET result_json = (
+            SELECT JSON_ARRAYAGG(
+                JSON_OBJECT(
+                    'amount', sub.amount,
+                    'currency', sub.currency
+                )
+            )
+            FROM (
+                SELECT COALESCE(SUM(L.amount), 0.0) AS amount, L.amount_currency AS currency
+                FROM hordak_account A2
+                JOIN hordak_leg L ON L.account_id = A2.id
+                JOIN hordak_transaction T ON L.transaction_id = T.id
+                WHERE A2.lft >= account_lft AND A2.rght <= account_rght AND A2.tree_id = account_tree_id AND T.date <= as_of
+                GROUP BY L.amount_currency
+            ) AS sub
+        );
+    ELSE
+        SET result_json = (
+            SELECT JSON_ARRAYAGG(
+                JSON_OBJECT(
+                    'amount', sub.amount,
+                    'currency', sub.currency
+                )
+            )
+            FROM (
+                SELECT COALESCE(SUM(L.amount), 0.0) AS amount, L.amount_currency AS currency
+                FROM hordak_account A2
+                JOIN hordak_leg L ON L.account_id = A2.id
+                WHERE A2.lft >= account_lft AND A2.rght <= account_rght AND A2.tree_id = account_tree_id
+                GROUP BY L.amount_currency
+            ) AS sub
+        );
+    END IF;
+
+    -- Return the JSON result
+    RETURN result_json;
+END;
+-- - reverse:
+DROP FUNCTION get_balance;

hordak/migrations/0047_get_balance.pg.sql

@@ -0,0 +1,75 @@
+------
+CREATE FUNCTION get_balance_table(account_id BIGINT, as_of DATE = NULL)
+    RETURNS TABLE (amount DECIMAL, currency VARCHAR) AS
+$$
+DECLARE
+    account_lft int;
+    account_rght int;
+    account_tree_id int;
+BEGIN
+    -- Get the account's information
+    SELECT
+        lft,
+        rght,
+        tree_id
+    INTO
+        account_lft,
+        account_rght,
+        account_tree_id
+    FROM hordak_account
+    WHERE id = account_id;
+
+    IF as_of IS NOT NULL THEN
+        -- If `as_of` is specified then we need an extra join onto the
+        -- transactions table to get the transaction date
+        RETURN QUERY
+            SELECT
+                COALESCE(SUM(L.amount), 0.0) as amount,
+                L.amount_currency as currency
+            FROM hordak_account A2
+            INNER JOIN hordak_leg L on L.account_id = A2.id
+            INNER JOIN hordak_transaction T on L.transaction_id = T.id
+            WHERE
+                -- We want to include this account and all of its children
+                A2.lft >= account_lft AND
+                A2.rght <= account_rght AND
+                A2.tree_id = account_tree_id AND
+                -- Also respect the as_of parameter
+                T.date <= as_of
+            GROUP BY L.amount_currency;
+    ELSE
+        RETURN QUERY
+            SELECT
+                COALESCE(SUM(L.amount), 0.0) as amount,
+                L.amount_currency as currency
+            FROM hordak_account A2
+            INNER JOIN hordak_leg L on L.account_id = A2.id
+            WHERE
+                -- We want to include this account and all of its children
+                A2.lft >= account_lft AND
+                A2.rght <= account_rght AND
+                A2.tree_id = account_tree_id
+            GROUP BY L.amount_currency;
+    END IF;
+END;
+$$
+LANGUAGE plpgsql;
+--- reverse:
+DROP FUNCTION get_balance_table(BIGINT, DATE);
+
+
+------
+CREATE FUNCTION get_balance(account_id BIGINT, as_of DATE = NULL)
+    RETURNS JSONB AS
+$$
+BEGIN
+    -- Convert our balance table into JSONB in the form:
+    --     [{"amount": 100.00, "currency": "EUR"}]
+    RETURN
+        (SELECT jsonb_agg(jsonb_build_object('amount', amount, 'currency', currency)))
+        FROM get_balance_table(account_id, as_of);
+END;
+$$
+LANGUAGE plpgsql;
+--- reverse:
+DROP FUNCTION get_balance(BIGINT, DATE);

hordak/migrations/0047_get_balance.py

@@ -0,0 +1,22 @@
+# Generated by Django 4.2 on 2024-06-27 09:11
+from pathlib import Path
+
+from django.db import migrations
+
+from hordak.utilities.migrations import (
+    select_database_type,
+    migration_operations_from_sql,
+)
+
+PATH = Path(__file__).parent
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("hordak", "0046_alter_account_uuid_alter_leg_uuid_and_more"),
+    ]
+
+    operations = select_database_type(
+        postgresql=migration_operations_from_sql(PATH / "0047_get_balance.pg.sql"),
+        mysql=migration_operations_from_sql(PATH / "0047_get_balance.mysql.sql"),
+    )

hordak/models/core.py

@@ -20,10 +20,13 @@
 - ``StatementLine`` - Represents a statement line. ``StatementLine.create_transaction()`` may be called to
   create a transaction for the statement line.
 """
+
+from datetime import date
+
 from django.db import connection, models
 from django.db import transaction
 from django.db import transaction as db_transaction
-from django.db.models import JSONField
+from django.db.models import F, JSONField
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from djmoney.models.fields import MoneyField
@@ -40,6 +43,7 @@
     get_internal_currency,
 )
 from hordak.utilities.currency import Balance
+from hordak.utilities.db_functions import GetBalance
 from hordak.utilities.dreprecation import deprecated
 
 
@@ -61,6 +65,18 @@ class AccountQuerySet(models.QuerySet):
     def net_balance(self, raw=False):
         return sum((account.balance(raw) for account in self), Balance())
 
+    def with_balances(self, as_of: date = None):
+        """Annotate the account queryset with account balances
+
+        This is a much more performant way to calculate account balances,
+        especially when calculating balances for a lot of accounts.
+
+        Note that you will get better performance by setting the `as_of`
+        to `None`. This is because the underlying custom database function
+        can avoid a join.
+        """
+        return self.annotate(balance=GetBalance(F("id"), as_of=as_of))
+
 
 class AccountManager(TreeManager):
     def get_by_natural_key(self, uuid):

hordak/models/db_views.py

@@ -27,7 +27,27 @@ class LegView(models.Model):
     You can also improve query performance (in Postgresql) by deferring the
     `account_balance` field, assuming the value not required. For example:
 
-        HordakLegView.objects.defer('account_balance')
+        .. code-block:: python
+
+            HordakLegView.objects.defer('account_balance')
+
+    Attributes:
+
+        id (int): The leg ID
+        uuid (UUID): The leg UUID
+        transaction (Transaction): The transaction which contains this leg
+        account (Account): The account this leg is associated with
+        date (date): The date when the parent transaction actually occurred
+        amount (Balance): The balance of this leg (use ``amount.currency``
+                          to get the currency for the other ``Decimal`` fields on this view.
+        type (LegType): Either ``LegType.debit`` or ``LegType.credit``.
+        credit (Decimal): Amount of this credit, or NULL if not a credit
+        debit (Decimal): Amount of this debit, or NULL if not a debit
+        account_balance (Decimal): Account balance following this transaction.
+                                   For multiple-currency accounts this will
+                                   be the balance of the same currency as the leg amount.
+        leg_description (str): Description of the leg
+        transaction_description (str): Description of the transaction
 
     """
 

hordak/tests/models/test_core.py

@@ -547,6 +547,52 @@ def test_child_asset_account_can_be_bank_account(self):
         self.assertEqual(account2.type, AccountType.asset)
         self.assertEqual(account2.is_bank_account, True)
 
+    def test_with_balances_simple(self):
+        """Ensure with_balances() returns a valid value in the simplest case"""
+        src = self.account(type=AccountType.liability)
+        dst = self.account(type=AccountType.expense)
+        src.transfer_to(dst, Money(100, "EUR"))
+        src.transfer_to(dst, Money(10, "EUR"))
+
+        # Just some other transaction that should be ignored
+        self.account().transfer_to(self.account(), Money(50, "EUR"))
+
+        src = Account.objects.filter(type=AccountType.liability).with_balances().get()
+        self.assertEqual(src.balance, Balance([Money("-110", "EUR")]))
+
+        dst = Account.objects.filter(type=AccountType.expense).with_balances().get()
+        self.assertEqual(dst.balance, Balance([Money("110", "EUR")]))
+
+    def test_with_balances_child_accounts(self):
+        """Ensure with_balances() returns a valid value for child accounts"""
+        parent = self.account(type=AccountType.liability, name="Parent")
+
+        src = self.account(type=AccountType.liability, parent=parent)
+        dst = self.account(type=AccountType.expense)
+        src.transfer_to(dst, Money(100, "EUR"))
+
+        # Just some other transaction that should be ignored
+        self.account().transfer_to(self.account(), Money(50, "EUR"))
+
+        parent = Account.objects.filter(name="Parent").with_balances().get()
+        self.assertEqual(parent.balance, Balance([Money("-100", "EUR")]))
+
+    def test_with_balances_as_of(self):
+        src = self.account(type=AccountType.liability)
+        dst = self.account(type=AccountType.expense)
+        src.transfer_to(dst, Money(100, "EUR"), date="2000-01-15")
+        src.transfer_to(dst, Money(110, "EUR"), date="2000-01-16")
+
+        # Just some other transaction that should be ignored
+        self.account().transfer_to(self.account(), Money(50, "EUR"))
+
+        src = (
+            Account.objects.filter(type=AccountType.liability)
+            .with_balances(as_of="2000-01-15")
+            .get()
+        )
+        self.assertEqual(src.balance, Balance([Money("-100", "EUR")]))
+
 
 class LegTestCase(DataProvider, DbTransactionTestCase):
     def test_manager(self):

hordak/utilities/db_functions.py

@@ -0,0 +1,57 @@
+import json
+from datetime import date
+from functools import cached_property
+from typing import Union
+
+from django.db.models import Func
+from django.db.models.expressions import Combinable, Value
+from djmoney.models.fields import MoneyField
+from moneyed import Money
+
+from hordak.utilities.currency import Balance
+
+
+class GetBalance(Func):
+    """Django representation of the get_balance() custom database function provided by Hordak"""
+
+    function = "GET_BALANCE"
+
+    def __init__(
+        self,
+        account_id: Union[Combinable, int],
+        as_of: Union[Combinable, date, str] = None,
+        output_field=None,
+        **extra
+    ):
+        """Create a new GetBalance()
+
+        Examples:
+
+            .. code-block:: python
+
+                from hordak.utilities.db_functions import GetBalance
+
+                GetBalance(account_id=5)
+                GetBalance(account_id=5, as_of='2000-01-01')
+
+                Account.objects.all().annotate(
+                    balance=GetBalance(F("id"), as_of='2000-01-01')
+                )
+
+        """
+        if as_of is not None:
+            if not isinstance(as_of, Combinable):
+                as_of = Value(as_of)
+
+        output_field = output_field or MoneyField()
+        super().__init__(account_id, as_of, output_field=output_field, **extra)
+
+    @cached_property
+    def convert_value(self):
+        # Convert the JSON output into a Balance object. Example of a JSON response:
+        #    [{"amount": 100.00, "currency": "EUR"}]
+        def convertor(value, expression, connection):
+            value = json.loads(value)
+            return Balance([Money(v["amount"], v["currency"]) for v in value])
+
+        return convertor

hordak/utilities/migrations.py

@@ -7,6 +7,8 @@
 def migration_operations_from_sql(file_path: Path):
     operations = []
     sql: str = file_path.read_text(encoding="utf8").strip().strip("-")
+    # Mysql needs to have spaces after a '--' comment
+    sql = sql.replace("-- ----", "------").replace("-- -", "---")
     if not sql:
         return []
 

mypy.ini

@@ -5,7 +5,7 @@ exclude = /migrations/
 plugins =
     mypy_django_plugin.main
 disable_error_code = attr-defined, var-annotated, misc
-
+no_implicit_optional = False
 
 [mypy.plugins.django-stubs]
 django_settings_module = "example_project.settings"