Pessimistic Locking in Django with select_for_update()

# Pessimistic Locking in Django with select_for_update()
#
# PURPOSE
# Pessimistic locking acquires a database lock before reading the state used for a write. It is
# useful when a decision spans several checks or related changes and concurrent writers must
# take turns.
#
# Consider two requests reserving the last item. Reading available=1 and later saving
# available=0 in both requests can report two successful reservations. The example protects the
# read, stock check, and decrement as one transaction.
#
# METHOD
# Enter transaction.atomic(), evaluate select_for_update() with get(), inspect the locked row,
# update it, and leave the transaction. On PostgreSQL at READ COMMITTED, another locking reader
# waits and then works with the committed state.
#
# The queryset is lazy: constructing it does not acquire a lock. Evaluation must happen inside
# the transaction. The lock lasts until the outermost transaction commits or rolls back,
# including an outer ATOMIC_REQUESTS transaction.
#
# A normal PostgreSQL SELECT can still read a committed snapshot while this row is locked. The
# lock principally coordinates conflicting writes and locking reads; it does not make all
# readers wait.
#
# COMPLEXITY
# Keep the transaction short: do not call payment providers, send email, or wait for user input
# while holding a row lock. Contention adds waiting time even when the query itself is fast.
#
# When locking several rows, acquire them in a consistent key order to reduce deadlocks. For a
# simple counter update, a conditional UPDATE with an F expression may be cheaper than fetching
# and locking the row first.
#
# Use an outbox or transaction.on_commit() for work that should happen only after a successful
# commit. External side effects still require their own idempotency design.
#
# CONTRACT
# Scope: Django 5.2/6.0-style APIs and PostgreSQL at READ COMMITTED. SQLite does not implement
# select_for_update row locking and cannot validate this behavior. Other backends support
# different option combinations.
#
# quantity must be a positive integer. The model supplies a nonnegative-stock database
# constraint through PositiveIntegerField. A missing SKU raises DoesNotExist; insufficient stock
# raises the domain exception and rolls back.
#
# The PostgreSQL-specific SET LOCAL lock_timeout bounds lock acquisition waits at three seconds.
# nowait=True instead fails immediately on a conflicting lock. Catch database errors outside the
# atomic block so rollback completes before handling them.
#
# skip_locked is useful for claiming queue work, but is a poor default for a specific customer
# SKU: skipping a busy row can look like missing stock. nowait and skip_locked cannot be
# combined.
#
# Classify transient failures narrowly. PostgreSQL 55P03 means a lock is unavailable or timed
# out; 40P01 is a deadlock; 40001 is a serialization failure. Any retry must begin a fresh
# transaction and re-evaluate the business decision.
#
# LANGUAGE DETAILS
# of=("self",) limits the locking target to this model. This matters when a query later gains
# select_related joins. Do not assume related rows are safe unless your transaction explicitly
# protects their invariants.
#
# The function uses the default database connection. In a multi-database application, use the
# same alias for the queryset, transaction, and timeout cursor. Returning from an inner atomic
# block does not release an outer transaction's lock.
#
# EXAMPLE
# Put the model in an installed Django app's models.py and the reservation function in a service
# module. Keep the model import consistent with your app. Generate and apply migrations before
# using the service.
#
# Usage: remaining = reserve_stock("book-1", 1). From a row with available=2, the result is 1. A
# second concurrent request sees the newly committed quantity before deciding whether it can
# reserve.
#
# The accompanying test project creates a separate temporary PostgreSQL database. Tests use
# distinct connections, verify an actual database lock wait, exercise nowait, confirm rollback
# and nested-transaction behavior, and check that concurrent requests cannot oversell.
#
# In a Django test suite use TransactionTestCase for transaction-boundary tests. TestCase wraps
# tests in transactions and can hide an accidental select_for_update call outside atomic().
#
# Download the complete examples, dependency lock file, and integration tests:
# https://www.dimkoug.com/static/downloads/locking-examples-tested.zip
#
# SOURCES
# Django QuerySet API:
# https://docs.djangoproject.com/en/6.0/ref/models/querysets/#select-for-update
#
# Django transactions: https://docs.djangoproject.com/en/6.0/topics/db/transactions/
#
# PostgreSQL explicit locking: https://www.postgresql.org/docs/current/explicit-locking.html
#
from django.db import connection, models, transaction


class InsufficientStock(Exception):
    pass


class PessimisticStock(models.Model):
    sku = models.CharField(max_length=80, primary_key=True)
    available = models.PositiveIntegerField(default=0)

    class Meta:
        db_table = "locking_pessimistic_stock"


def reserve_stock(sku: str, quantity: int, *, nowait: bool = False) -> int:
    """Reserve stock on PostgreSQL and return the remaining quantity."""
    if type(quantity) is not int or quantity <= 0:
        raise ValueError("quantity must be a positive integer")

    with transaction.atomic():
        # Static SQL: no user input is interpolated into the timeout setting.
        with connection.cursor() as cursor:
            cursor.execute("SET LOCAL lock_timeout = '3s'")

        # .get() evaluates the locking query INSIDE the transaction.
        stock = PessimisticStock.objects.select_for_update(
            nowait=nowait, of=("self",)
        ).get(pk=sku)

        if stock.available < quantity:
            raise InsufficientStock("not enough stock")

        stock.available -= quantity
        stock.save(update_fields=["available"])
        remaining = stock.available

    # Locks are released when the outermost enclosing transaction ends.
    return remaining