Bullet Proofing Django Models

We recently added a bank account like functionality into one of our products. During the development we encountered some textbook problems and I thought it can be a good opportunity to go over some of the patterns we use in our Django models.

A Bank Account

This article was written in the order in which we usually address new problems:

  1. Define the business requirements.
  2. Write down a naive implementation and model definition.
  3. Challenge the solution.
  4. Refine and repeat.

Business Requirements

  • Each user can have only one account but not every user must have one.
  • The user can deposit and withdraw from the account up to a certain amount.
  • The account balance cannot be negative.
  • There is a max limit to the user's account balance.
  • The total amount of all balances in the app cannot exceed a certain amount.
  • There must be a record for every action on the account.
  • Actions on the account can be executed by the user from either the mobile app or the web interface and by support personnel from the admin interface.

Now that we have the business requirements we can start with a model definition.

Account Model

# models.py

import uuid

from django.conf import settings
from django.db import models

class Account(models.Model):
    class Meta:
        verbose_name = 'Account'
        verbose_name_plural = 'Accounts'

    MAX_TOTAL_BALANCES = 10000000

    MAX_BALANCE = 10000
    MIN_BALANCE = 0

    MAX_DEPOSIT = 1000
    MIN_DEPOSIT = 1

    MAX_WITHDRAW = 1000
    MIN_WITHDRAW = 1

    id = models.AutoField(
        primary_key=True,
    )
    uid = models.UUIDField(
        unique=True,
        editable=False,
        default=uuid.uuid4,
        verbose_name='Public identifier',
    )
    user = models.OneToOneField(
        settings.AUTH_USER_MODEL,
        on_delete=models.PROTECT,
    )
    created = models.DateTimeField(
        blank=True,
    )
    modified = models.DateTimeField(
        blank=True,
    )
    balance = models.PositiveIntegerField(
        verbose_name='Current balance',
    )

Let's break it down:

  • We use two unique identifiers  -  A private identifier which is an auto generated number (id) and a public id which is a uuid (uid). It's a good idea to keep enumerators private  - They expose important information about our data such as how many accounts we have and we don't want that.
  • We use OneToOneField for the user - It's like a ForeignKey but with a unique constraint. This ensures a user cannot have more than one account.
  • We set on_delete=models.PROTECT - Starting with Django 2.0 this argument will be mandatory. The default is CASCADE - when the user is deleted the related account is deleted as well. In our case this doesn't make sense - imagine the bank "deleting your money" when you close an account. Setting on_delete=models.PROTECT will raise an IntegrityError when attempting to delete a user with an account.
  • You probably noticed that the code is very... "vertical". Wwe write like that because it makes git diffs look nicer.

Account Action Model

Now that we have an account model we can create a model to log actions made to the account:

# models.py

class Action(models.Model):
    class Meta:
        verbose_name = 'Account Action'
        verbose_name_plural = 'Account Actions'

    ACTION_TYPE_CREATED = 'CREATED'
    ACTION_TYPE_DEPOSITED = 'DEPOSITED'
    ACTION_TYPE_WITHDRAWN = 'WITHDRAWN'
    ACTION_TYPE_CHOICES = (
        (ACTION_TYPE_CREATED, 'Created'),
        (ACTION_TYPE_DEPOSITED, 'Deposited'),
        (ACTION_TYPE_WITHDRAWN, 'Withdrawn'),
    )

    REFERENCE_TYPE_BANK_TRANSFER = 'BANK_TRANSFER'
    REFERENCE_TYPE_CHECK = 'CHECK'
    REFERENCE_TYPE_CASH = 'CASH'
    REFERENCE_TYPE_NONE = 'NONE'
    REFERENCE_TYPE_CHOICES = (
        (REFERENCE_TYPE_BANK_TRANSFER, 'Bank Transfer'),
        (REFERENCE_TYPE_CHECK, 'Check'),
        (REFERENCE_TYPE_CASH, 'Cash'),
        (REFERENCE_TYPE_NONE, 'None'),
    )

    id = models.AutoField(
        primary_key=True,
    )
    user_friendly_id = models.CharField(
        unique=True,
        editable=False,
        max_length=30,
    )
    user = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        on_delete=models.PROTECT,
        help_text='User who performed the action.',
    )
    created = models.DateTimeField(
        blank=True,
    )
    account = models.ForeignKey(
        Account,
    )
    type = models.CharField(
        max_length=30,
        choices=ACTION_TYPE_CHOICES,
    )
    delta = models.IntegerField(
        help_text='Balance delta.',
    )
    reference = models.TextField(
        blank=True,
    )
    reference_type = models.CharField(
        max_length=30,
        choices=REFERENCE_TYPE_CHOICES,
        default=REFERENCE_TYPE_NONE,
    )
    comment = models.TextField(
        blank=True,
    )

    # Fields used solely for debugging purposes.

    debug_balance = models.IntegerField(
        help_text='Balance after the action.',
    )

What do we have here?

  • Each record will hold a reference to the associated balance and the delta amount. A deposit of 100$ will have a delta of 100$, and a withdrawal of 50$ will have a delta of -50$. This way we can sum the deltas of all actions made to an account and get the current balance. This is important for validating our calculated balance.
  • We follow the same pattern of adding two identifiers - a private and a public one. The difference here is that reference numbers for actions are often used by users and support personnel to identify a specific action over the phone or in emails. A uuid is not user friendly- it's very long and it's not something users are used to see. I found a nice implementation of user-friendly ID's in django-invoice.
  • Two of the fields are only relevant for one type of action, deposit - reference and reference type. There are a lot of ways to tackle this issue - table inheritance and down casting, JSON fields, table polymorphism and the list of overly complicated solutions goes on. In our case we are going to use a sparse table.

Note about the design: Maintaining calculated fields in the model is usually bad design. Calculated fields such as the account's balance should be avoided whenever possible.

However, in our "real life" implementation there are additional action types and thousands of actions on each account - we treat calculated attribute as an optimization. Maintaining state poses some interesting challenges and we thought it can serve the purpose of this post so we decided to present it as well.

Challenges

Multiple Platforms

We have three client applications we need to support:

  • Mobile app  - Uses an API interface to manage the account.
  • Web client  - Uses either an API interface (if we have some sort of SPA), or good old server side rendering with Django forms.
  • Admin interface  - Uses Django's admin module with Django forms.

Our motivation is to keep things DRY and self contained as possible.

Validation

We have two types of validations hiding in the business requirements:

Input validation such as "amount must be between X and Y", "balance cannot exceed Z", etc - these types of validation are well supported by Django and can usually be expressed as database constraints or django validations.

The second validation is a bit more complicated. We need to ensure the total amount of all balances in the entire system does not exceed a certain amount. This forces us to validate an instance against all other instances of the model.

Atomicity

Race conditions are a very common issue in distributed systems and ever more so in models that maintain state such as bank account (you can read more about race conditions in Wikipedia).

To illustrate the problem consider an account with a balance of 100$. The user connects from two different devices at the exact same time and issue a withdraw of 100$. Since the two actions were executed at the exact same time it is possible that both of them get a current balance of 100$. Given that both session see sufficient balance they will both get approved and update the new balance to 0$. The user withdrawn a total of 200$ and the current balance is now 0$ - we have a race condition and we are down 100$.

Logging / History

The log serves two purposes:

  • Log and Audit - Information about historical actions - dates, amounts, users etc.
  • Check Consistency - We maintain state in the model so we want to be able to validate the calculated balance by aggregating the action deltas.

The history records must be 100% immutable.


The Naive Implementation

Let's start with a naive implementation of deposit (this is not a good implementation):

class Account(models.Model):

    # ...

    def deposit(self, amount, deposited_by, asof):
        assert amount > 0

        if not self.MIN_DEPOSIT <= amount <= self.MAX_DEPOSIT:
            raise InvalidAmount(amount)

        if self.balance + amount > self.MAX_BALANCE:
            raise ExceedsLimit()

        total = Account.objects.aggregate(
            total=Sum('balance')
        )['total']
        if total + amount > self.MAX_TOTAL_BALANCES:
            raise ExceedsLimit()

        action = self.actions.create(
            user=deposited_by,
            type=Action.ACTION_TYPE_DEPOSITED,
            delta=amount,
            asof=asof,
        )

    self.balance += amount
    self.modified = asof

    self.save()

And let's add a simple endpoint for it using DRF @api_view:

# api.py

# ...
from django.db.models import transaction
# ...

@api_view('POST')
def deposit(request):
    try:
        amount = int(request.data\['amount'\])
    except (KeyError, ValueError):
        return Response(status=status.HTTP_400_BAD_REQUEST)

    with transaction.atomic():
        try:
            account = (
                Account.objects
                .select_for_update()
                .get(user=request.user)
            )
        except Account.DoesNotExist:
            return Response(status=status.HTTP_404_NOT_FOUND)

        try:
            account.deposit(
                amount=amount,
                deposited_by=request.user,
                asof=timezone.now(),
            )
        except (ExceedsLimit, InvalidAmount):
            return Response(status=status.HTTP_400_BAD_REQUEST)

        return Response(status=status.HTTP_200_OK)

So what is the problem?

Locking the account  -  An instance cannot lock itself because it had already been fetched. We gave up control over the locking and fetching so we have to trust the caller to properly obtain a lock - this is very bad design. Don't take my word for it, let's take a glimpse at Django's design philosophy:

Loose coupling A fundamental goal of Django's stack is loose coupling and tight cohesion. The various layers of the framework shouldn't "know" about each other unless absolutely necessary.

So is it really the business of our API, forms and django admin to fetch the account for us and obtain a proper lock? I think not.

Validation  -  The account has to validate itself against all other accounts - this just feels awkward.

A Better Approach

We need to hook into the process before the account is fetched (to obtain a lock) and in a place where it makes sense to validate and process more than one account.

Let's start with a function to create an Action instance and write it as a classmethod:

# models.py

from django.core.exceptions import ValidationError


class Action(models.Model):

    # ...

    @classmethod
    def create(
        cls,
        user,
        account,
        type,
        delta,
        asof,
        reference=None,
        reference_type=None,
        comment=None,
    ):
        """Create Action.

        user (User):
            User who executed the action.
        account (Account):
            Account the action executed on.
        type (str, one of Action.ACTION_TYPE_\*):
            Type of action.
        delta (int):
            Change in balance.
        asof (datetime.datetime):
            When was the action executed.
        reference (str or None):
            Reference number when appropriate.
        reference_type(str or None):
            Type of reference.
            Defaults to "NONE".
        comment (str or None):
            Optional comment on the action.

        Raises:
            ValidationError

        Returns (Action)
        """
        assert asof is not None

        if (type == cls.ACTION_TYPE_DEPOSITED and
            reference_type is None):
            raise errors.ValidationError({
                'reference_type': 'required for deposit.',
            })

        if reference_type is None:
            reference_type = cls.REFERENCE_TYPE_NONE

        # Don't store null in text field.

        if reference is None:
            reference = ''

        if comment is None:
            comment = ''

        user_friendly_id = generate_user_friendly_id()

        return cls.objects.create(
            user_friendly_id=user_friendly_id,
            created=asof,
            user=user,
            account=account,
            type=type,
            delta=delta,
            reference=reference,
            reference_type=reference_type,
            comment=comment,
            debug_balance=account.balance,
        )

What do we have here:

  • We used a classmethod that accepts all necessary data to validate and create the new instance. By *not* using the default manager's create function (Action.objects.create) we encapsulate all the business logic in the creation process.
  • We easily introduced custom validation and raise proper ValidationError.
  • We accept the creation time as an argument. That might seem a bit strange at first glance - why not use the built in auto_time_add? For starters It's much easier to test with predictable values. Second, as we are going to see in just a bit, we can make sure the modified time of the account is exactly the same as the action created time.

Before moving over to the implementation of the account methods let's define custom exceptions for our Account module:

# errors.py

class Error(Exception):
    pass

class ExceedsLimit(Error):
    pass

class InvalidAmount(Error):
    def __init__(self, amount):
        self.amount = amount

    def __str__(str):
        return 'Invalid Amount: {}'.format(amount)

class InsufficientFunds(Error):
    def __init__(self, balance, amount):
        self.balance = balance
        self.amount = amount

    def __str__(self):
        return 'amount: {}, current balance: {}'.format(
                self.amount, self.balance)

We define a base Error class that inherits from Exception. This is something we found very useful and we use it a lot. A base error class allows us to catch all errors coming from a certain module:

from account.errors import Error as AccountError

try:
   # action on account
except AccountError:
   # Handle all errors from account

A similar pattern can be found in the popular requests package.

Let's implement the method to create a new Account:

class Account(models.Model):

    # ...

    @classmethod
    def create(cls, user, created_by, asof):
        """Create account.

        user (User):
            Owner of the account.
        created_by (User):
            User that created the account.
        asof (datetime.datetime):
            Time of creation.

        Returns (tuple):
            [0] Account
            [1] Action
        """
        with transaction.atomic():
            account = cls.objects.create(
                user=user,
                created=asof,
                modified=asof,
                balance=0,
            )

            action = Action.create(
                user=created_by,
                account=account,
                type=Action.ACTION_TYPE_CREATED,
                delta=0,
                asof=asof,
            )

        return account, action

Pretty straight forward - create the instance, create the action and return them both.

Notice how we accept asof here as well - modified, created and the action creation time are all equal - you cant do that with auto_add and auto_add_now.

Now to the business logic:

# models.py

@classmethod
def deposit(
    cls,
    uid,
    deposited_by,
    amount,
    asof,
    comment=None,
):
    """Deposit to account.

    uid (uuid.UUID):
        Account public identifier.
    deposited_by (User):
        Deposited by.
    amount (positive int):
        Amount to deposit.
    asof (datetime.datetime):
        Time of deposit.
    comment(str or None):
       Optional comment.

    Raises
        Account.DoesNotExist
        InvalidAmount
        ExceedsLimit

    Returns (tuple):
        [0] (Account) Updated account instance.
        [1] (Action) Deposit action.
    """
    assert amount > 0

    with transaction.atomic():
        account = cls.objects.select_for_update().get(uid=uid)

        if not (cls.MIN_DEPOSIT <= amount <= cls.MAX_DEPOSIT):
            raise errors.InvalidAmount(amount)

        if account.balance + amount > cls.MAX_BALANCE:
            raise errors.ExceedsLimit()

        total = cls.objects.aggregate(total=Sum('balance'))\['total'\]
        if total + amount > cls.MAX_TOTAL_BALANCES:
            raise errors.ExceedsLimit()

        account.balance += amount
        account.modified = asof

        account.save(update_fields=[
            'balance',
            'modified',
        ])

        action = Action.create(
            user=deposited_by,
            account=account,
            type=Action.ACTION_TYPE_DEPOSITED,
            delta=amount,
            asof=asof,
        )

    return account, action

@classmethod
def withdraw(
    cls,
    uid,
    withdrawn_by,
    amount,
    asof
    comment=None,
):
    """Withdraw from account.

    uid (uuid.UUID):
        Account public identifier.
    withdrawn_by (User):
        The withdrawing user.
    amount (positive int):
        Amount to withdraw.
    asof (datetime.datetime):
        Time of withdraw.
    comment (str or None):
       Optional comment.

    Raises:
        Account.DoesNotExist
        InvalidAmount
        InsufficientFunds

    Returns (tuple):
        [0] (Account) Updated account instance.
        [1] (Action) Withdraw action.
    """
    assert amount > 0

    with transaction.atomic():
        account = cls.objects.select_for_update().get(uid=uid)

        if not (cls.MIN_WITHDRAW <= amount <= cls.MAX_WITHDRAW):
            raise InvalidAmount(amount)

        if account.balance - amount < cls.MIN_BALANCE:
            raise InsufficientFunds(amount, account.balance)

        account.balance -= amount
        account.modified = asof

        account.save(update_fields=[
            'balance',
            'modified',
        ])

        action = Action.create(
            user=withdrawn_by,
            account=account,
            type=Action.ACTION_TYPE_WITHDRAWN,
            delta=-amount,
            asof=asof,
        )

    return account, action

We can start to see the pattern here:

  1. Acquire a lock on the account using select_for_update. This will lock the account row in the database and make sure no one can update the account instance until the transaction is completed (either committed or rolled-back).
  2. Perform validation checks and raise proper exceptions - raising an exception will cause the transaction to rollback.
  3. If all the validations passed update the state (current balance), set the modification time, save the instance and create the log (action).

So how does the model hold up to our challenges?

  • Multiple Platforms and Validation- We encapsulated all of our business logic including input and system wide validation inside the model method so consumers such as API, admin action or forms only need to handle exceptions and serialization / UI.
  • Atomicity - Each method obtains its own lock so there is no risk of race condition.
  • Logging / History - We created an action model and made sure each function registers the proper action.

Profit!


Testing

Our app will be incomplete without proper tests. I previously wrote about class based testings - we are going to take a slightly different approach but still have a base class with utility functions:

# tests/common.py

class TestAccountBase:
    DEFAULT = object()

    @classmethod
    def default(value, default_value):
        return default_value if value is cls.DEFAULT else value

    @classmethod
    def setUpTestData(cls):
        super().setUpTestData()

        # Set up some default values
        cls.admin = User.objects.create_superuser(
            'Admin',
            'admin',
            'admin@testing.test',
        )

        cls.user_A = User.objects.create_user(
            'user_A',
            'user_A',
            'A@testing.test',
        )

    @classmethod
    def create(
        cls,
        user=DEFAULT,
        created_by=DEFAULT,
        asof=DEFAULT
    ):
        user = cls.default(user, cls.user_A)
        created_by = cls.default(created_by, cls.admin)
        asof = cls.default(asof, timezone.now())

        account, action = Account.create(user, created_by, asof)
        return cls.account, action

    def deposit(
        self,
        amount,
        account=DEFAULT,
        deposited_by=DEFAULT,
        asof=DEFAULT,
        comment=DEFAULT,
    ):
        account = self.default(account, self.account)
        deposited_by = self.default(deposited_by, self.admin)
        asof = self.default(asof, timezone.now())
        comment = self.default(comment, 'deposit comment')

        self.account, action = Account.deposit(
            uid=account.uid,
            deposited_by=deposited_by,
            amount=amount,
            asof=asof,
        )

        self.assertEqual(action.type, Action.ACTION_TYPE_DEPOSITED)
        self.assertIsNotNone(action.user_friendly_id)
        self.assertEqual(action.created, asof)
        self.assertEqual(action.delta, amount)
        self.assertEqual(action.user, deposited_by)

        return action

    def withdraw(
        self,
        amount,
        account=DEFAULT,
        withdrawn_by=DEFAULT,
        asof=DEFAULT,
        comment=DEFAULT,
    ):
        account = self.default(account, self.account)
        withdrawn_by = self.default(withdrawn_by, self.admin)
        asof = self.default(asof, timezone.now())
        comment = self.default(comment, 'withdraw comment')

        self.account, action = Account.withdraw(
            uid=account.uid,
            withdrawn_by=withdrawn_by,
            amount=amount,
            asof=asof,
        )

        self.assertEqual(action.type, Action.ACTION_TYPE_WITHDRAWN)
        self.assertIsNotNone(action.user_friendly_id)
        self.assertEqual(action.created, asof)
        self.assertEqual(action.delta, amount)
        self.assertEqual(action.user, withdrawn_by)

        return action

To make testing easier to write we use utility functions to reduce the boilerplate of specifying the user, the account etc each time by providing default values and operating on self.account.

Lets use our base class to write some tests:

# tests/test_account.py

from unittest import mock
from django.test import TestCase

from .common import TestAccoutBase
from ..models import Account, Action
from ..errors import (
    InvalidAmount,
    ExceedsLimit,
    InsuficientFunds,
)

class TestAccount(TestAccountBase, TestCase):

    def setUp(self):
 **self.account, _ = cls.create()**

    def test_should_start_with_zero_balance(self):
        self.assertEqual(self.account.balance, 0)

    def test_should_deposit(self):
        self.deposit(100)
        self.assertEqual(self.account.balance, 100)
        self.deposit(150)
        self.assertEqual(self.account.balance, 250)

    def test_should_fail_to_deposit_less_than_minimum(self):
        with self.assertRaises(InvalidAmount):
            self.deposit(Account.MIN_DEPOSIT - 1)
        self.assertEqual(self.account.balance, 0)

    def test_should_fail_to_deposit_more_than_maximum(self):
        with self.assertRaises(InvalidAmount):
            self.deposit(Account.MAX_DEPOSIT + 1)
        self.assertEqual(self.account.balance, 0)

    @mock.patch('account.models.Account.MAX_BALANCE', 500)
    @mock.patch('account.models.Account.MAX_DEPOSIT', 502)
    def test_should_fail_to_deposit_more_than_max_balance(self):
        with self.assertRaises(ExceedsLimit):
            self.deposit(501)
        self.assertEqual(self.account.balance, 0)

    @mock.patch('account.models.Account.MAX_BALANCE', 500)
    @mock.patch('account.models.Account.MAX_DEPOSIT', 500)
    @mock.patch('account.models.Account.MAX_TOTAL_BALANCES', 600)
    def test_should_fail_when_exceed_max_total_balances(self):

        # Exceed max total balances for the same account

        self.deposit(500)
        with self.assertRaises(ExceedsLimit):
            self.deposit(500)
        self.assertEqual(self.account.balance, 500)

        # Exceed max total balances in other account

        other_user = User.objects.create_user('foo', 'bar', 'baz')
        other_account = self.create(user=other_user)

        with self.assertRaises(ExceedsLimit):
            self.deposit(200, account=other_account)
        self.assertEqual(other_account.balance, 0)

    def test_should_withdraw(self):
        self.deposit(100)
        self.withdraw(50)
        self.assertEqual(self.account.balance, 50)
        self.withdraw(30)
        self.assertEqual(self.account.balance, 20)

    def test_should_fail_when_insufficient_funds(self):
        self.deposit(100)
        with self.assertRaises(InsufficientFunds):
            self.withdraw(101)
        self.assertEqual(self.account.balance, 100)

Final Words

The classmethod approach has proved itself in our development for quite some time now. We found that it provides the necessary flexibility, readability and testability with very little overhead.

In this article we presented two common issues we encounter frequently - validation and concurrency. This method can be extended to handle access control (permissions) and caching (we have total control over the fetch, remember?), performance optimization (use select_related and update_fields...), audit and monitoring and additional business logic.

We usually support several interfaces for each model - admin interface for support, API for mobile and SPA clients, and a dashboard. Encapsulating the business logic inside the model reduced the amount of code duplication and required tests which leads to overall quality code that is easy to maintain.

In a follow up post I (might) present the admin interface for this model with some neat tricks (such as custom actions, intermediate pages etc) and possibly an RPC implementation using DRF to interact with the account as an API.




Similar articles