PostgreSQL Isolation Levels Explained with Python

import os
import time

import psycopg
from psycopg import IsolationLevel

DSN = os.environ.get("DATABASE_URL", "postgresql://app:app@localhost:5432/spend")

LEVELS = {
    "read-committed": IsolationLevel.READ_COMMITTED,
    "repeatable-read": IsolationLevel.REPEATABLE_READ,
    "serializable": IsolationLevel.SERIALIZABLE,
}


def connect(name: str, level: IsolationLevel = IsolationLevel.READ_COMMITTED) -> psycopg.Connection:
    """One manager's session. The first statement opens a transaction at
    `level`, and it stays open until commit() or rollback()."""
    conn = psycopg.connect(DSN, application_name=name)
    conn.isolation_level = level
    return conn


def reset(budget: int = 5_000) -> None:
    """Back to the start: design team, nothing spent."""
    with psycopg.connect(DSN, autocommit=True) as conn:
        conn.execute("TRUNCATE expenses RESTART IDENTITY")
        conn.execute("UPDATE teams SET budget = %s, spent = 0 WHERE id = 'design'", (budget,))


def wait_until_blocked(name: str, timeout: float = 5.0) -> None:
    """Wait until session `name` is stuck waiting for a lock."""
    with psycopg.connect(DSN, autocommit=True) as conn:
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            row = conn.execute(
                "SELECT wait_event_type FROM pg_stat_activity WHERE application_name = %s",
                (name,),
            ).fetchone()
            if row and row[0] == "Lock":
                return
            time.sleep(0.01)
    raise TimeoutError(f"{name} never blocked")


def log(who: str, message: str) -> None:
    print(f"[{who}] {message}", flush=True)


def usd(amount: int) -> str:
    return f"${amount:,}"

"""Lost update: two managers add to the team's running total at the same time.

Usage: python lost_update.py read-committed | repeatable-read | atomic
"""
import sys
import threading

from psycopg.errors import SerializationFailure

from common import LEVELS, connect, log, reset, usd, wait_until_blocked

ALICE_AMOUNT = 1_200
BOB_AMOUNT = 800


def read_modify_write(level_name: str) -> None:
    level = LEVELS[level_name]
    alice = connect("alice", level)
    bob = connect("bob", level)

    # Both managers read the total before either writes.
    a_spent = alice.execute("SELECT spent FROM teams WHERE id = 'design'").fetchone()[0]
    log("alice", f"read spent = {usd(a_spent)}")
    b_spent = bob.execute("SELECT spent FROM teams WHERE id = 'design'").fetchone()[0]
    log("bob", f"read spent = {usd(b_spent)}")

    # Each adds in Python and writes the result back.
    alice.execute("UPDATE teams SET spent = %s WHERE id = 'design'", (a_spent + ALICE_AMOUNT,))
    log("alice", f"set spent = {usd(a_spent + ALICE_AMOUNT)}")
    alice.commit()
    log("alice", "commit")

    try:
        bob.execute("UPDATE teams SET spent = %s WHERE id = 'design'", (b_spent + BOB_AMOUNT,))
        log("bob", f"set spent = {usd(b_spent + BOB_AMOUNT)}")
        bob.commit()
        log("bob", "commit")
    except SerializationFailure as e:
        log("bob", f"UPDATE failed: SQLSTATE {e.sqlstate}: {e.diag.message_primary}")
        bob.rollback()


def atomic() -> None:
    alice = connect("alice")
    bob = connect("bob")

    # Let the database do the adding, on whatever the row holds right now.
    new = alice.execute(
        "UPDATE teams SET spent = spent + %s WHERE id = 'design' RETURNING spent",
        (ALICE_AMOUNT,),
    ).fetchone()[0]
    log("alice", f"add {usd(ALICE_AMOUNT)} -> spent = {usd(new)} (not committed yet)")

    def bob_adds() -> None:
        new = bob.execute(
            "UPDATE teams SET spent = spent + %s WHERE id = 'design' RETURNING spent",
            (BOB_AMOUNT,),
        ).fetchone()[0]
        log("bob", f"add {usd(BOB_AMOUNT)} -> spent = {usd(new)}")
        bob.commit()
        log("bob", "commit")

    t = threading.Thread(target=bob_adds)
    t.start()
    wait_until_blocked("bob")
    log("bob", "waiting for alice's row lock")
    log("alice", "commit")
    alice.commit()
    t.join()


def main() -> None:
    variant = sys.argv[1] if len(sys.argv) > 1 else "read-committed"
    reset()
    print(f"--- lost update, {variant} ---")
    if variant == "atomic":
        atomic()
    else:
        read_modify_write(variant)

    with connect("check") as conn:
        spent = conn.execute("SELECT spent FROM teams WHERE id = 'design'").fetchone()[0]
    print(f"final spent = {usd(spent)} (expected {usd(ALICE_AMOUNT + BOB_AMOUNT)})")


if __name__ == "__main__":
    main()

"""Read skew: a month-end report reads two numbers that should always agree.

Usage: python read_skew.py read-committed | repeatable-read
"""
import sys

import psycopg

from common import LEVELS, connect, log, reset, usd


def approve(conn: psycopg.Connection, amount: int, approver: str) -> None:
    """Record an expense and add it to the team's total, in ONE transaction."""
    conn.execute(
        "INSERT INTO expenses (team_id, amount, approved_by) VALUES ('design', %s, %s)",
        (amount, approver),
    )
    conn.execute("UPDATE teams SET spent = spent + %s WHERE id = 'design'", (amount,))
    conn.commit()


def main() -> None:
    level_name = sys.argv[1] if len(sys.argv) > 1 else "read-committed"
    reset()
    print(f"--- read skew, report at {level_name} ---")

    carol = connect("carol")
    approve(carol, 1_000, "carol")
    log("carol", "approved $1,000 (earlier this month)")

    report = connect("report", LEVELS[level_name])
    spent = report.execute("SELECT spent FROM teams WHERE id = 'design'").fetchone()[0]
    log("report", f"teams.spent          = {usd(spent)}")

    approve(carol, 250, "carol")
    log("carol", "approved $250 and committed")

    total = report.execute(
        "SELECT coalesce(sum(amount), 0) FROM expenses WHERE team_id = 'design'"
    ).fetchone()[0]
    log("report", f"sum(expenses.amount) = {usd(total)}")
    report.commit()

    verdict = "agree" if spent == total else "DISAGREE"
    print(f"report numbers {verdict}: {usd(spent)} vs {usd(total)}")


if __name__ == "__main__":
    main()

"""Write skew: two managers each check the budget, then each insert an expense.

Usage: python write_skew.py read-committed | repeatable-read | serializable
"""
import sys

import psycopg
from psycopg.errors import SerializationFailure

from common import LEVELS, connect, log, reset, usd

AMOUNT = 3_000  # each, against a $5,000 budget


def check(conn: psycopg.Connection, who: str) -> bool:
    budget = conn.execute("SELECT budget FROM teams WHERE id = 'design'").fetchone()[0]
    approved = conn.execute(
        "SELECT coalesce(sum(amount), 0) FROM expenses WHERE team_id = 'design'"
    ).fetchone()[0]
    ok = approved + AMOUNT <= budget
    log(who, f"check: approved {usd(approved)} + {usd(AMOUNT)} <= {usd(budget)}? {ok}")
    return ok


def insert(conn: psycopg.Connection, who: str) -> None:
    conn.execute(
        "INSERT INTO expenses (team_id, amount, approved_by) VALUES ('design', %s, %s)",
        (AMOUNT, who),
    )
    log(who, f"insert {usd(AMOUNT)}")


def main() -> None:
    level_name = sys.argv[1] if len(sys.argv) > 1 else "read-committed"
    reset()
    print(f"--- write skew, {level_name} ---")
    alice = connect("alice", LEVELS[level_name])
    bob = connect("bob", LEVELS[level_name])

    # Both checks happen before either insert.
    check(alice, "alice")
    check(bob, "bob")
    insert(alice, "alice")
    insert(bob, "bob")
    alice.commit()
    log("alice", "commit")
    try:
        bob.commit()
        log("bob", "commit")
    except SerializationFailure as e:
        log("bob", f"COMMIT failed: SQLSTATE {e.sqlstate}: {e.diag.message_primary}")
        bob.rollback()
        log("bob", "retry the whole transaction")
        if check(bob, "bob"):
            insert(bob, "bob")
        else:
            log("bob", "rejected: over budget")
        bob.commit()

    with connect("check") as conn:
        total, count = conn.execute(
            "SELECT coalesce(sum(amount), 0), count(*) FROM expenses WHERE team_id = 'design'"
        ).fetchone()
    print(f"approved {count} expense(s) totalling {usd(total)} on a $5,000 budget")


if __name__ == "__main__":
    main()

"""Three safe ways to approve an expense against a team budget, and a demo of each.

Usage: python approve.py serializable | for-update | conditional
"""
import random
import sys
import threading
import time
from collections.abc import Callable

import psycopg
from psycopg import IsolationLevel
from psycopg.errors import DeadlockDetected, SerializationFailure

from common import connect, log, reset, usd, wait_until_blocked

MAX_ATTEMPTS = 10
RETRYABLE = (SerializationFailure, DeadlockDetected)  # SQLSTATE 40001, 40P01
VERBOSE = True


def say(who: str, message: str) -> None:
    if VERBOSE:
        log(who, message)


def nothing() -> None:
    pass


def check_and_insert(
    conn: psycopg.Connection, team: str, amount: int, who: str, pause: Callable[[], None]
) -> bool:
    """The business rule: approve only if the expense fits what is left of the budget."""
    budget = conn.execute("SELECT budget FROM teams WHERE id = %s", (team,)).fetchone()[0]
    approved = conn.execute(
        "SELECT coalesce(sum(amount), 0) FROM expenses WHERE team_id = %s", (team,)
    ).fetchone()[0]
    if approved + amount > budget:
        say(who, f"check: {usd(approved)} + {usd(amount)} > {usd(budget)}, rejected")
        return False
    say(who, f"check: {usd(approved)} + {usd(amount)} <= {usd(budget)}, ok")
    pause()
    conn.execute(
        "INSERT INTO expenses (team_id, amount, approved_by) VALUES (%s, %s, %s)",
        (team, amount, who),
    )
    say(who, f"insert {usd(amount)}")
    return True


def approve_naive(conn: psycopg.Connection, team: str, amount: int, who: str) -> bool:
    """Check, then insert, at whatever level `conn` uses. Not safe."""
    ok = check_and_insert(conn, team, amount, who, nothing)
    conn.commit()
    return ok


def approve_serializable(
    conn: psycopg.Connection,
    team: str,
    amount: int,
    who: str,
    pause: Callable[[], None] = nothing,
) -> tuple[bool, int]:
    """Fix 1: run at SERIALIZABLE and retry the whole transaction when PostgreSQL
    aborts it. Returns (approved, attempts)."""
    for attempt in range(1, MAX_ATTEMPTS + 1):
        try:
            ok = check_and_insert(conn, team, amount, who, pause)
            conn.commit()
            say(who, "commit")
            return ok, attempt
        except RETRYABLE as e:
            conn.rollback()
            say(who, f"attempt {attempt} failed: {e.diag.message_primary}")
            if attempt == MAX_ATTEMPTS:
                raise
            time.sleep(random.uniform(0, min(0.1, 0.002 * 2**attempt)))
    raise AssertionError("unreachable")


def approve_for_update(
    conn: psycopg.Connection,
    team: str,
    amount: int,
    who: str,
    pause: Callable[[], None] = nothing,
) -> bool:
    """Fix 2: READ COMMITTED, but lock the team's row first, so approvals
    for the same team take turns."""
    conn.execute("SELECT 1 FROM teams WHERE id = %s FOR UPDATE", (team,))
    say(who, f"locked team {team}")
    ok = check_and_insert(conn, team, amount, who, pause)
    say(who, "commit")
    conn.commit()
    return ok


def approve_conditional(
    conn: psycopg.Connection,
    team: str,
    amount: int,
    who: str,
    pause: Callable[[], None] = nothing,
) -> bool:
    """Fix 3: READ COMMITTED, check and reserve the money in one UPDATE."""
    row = conn.execute(
        "UPDATE teams SET spent = spent + %s"
        " WHERE id = %s AND spent + %s <= budget"
        " RETURNING spent",
        (amount, team, amount),
    ).fetchone()
    if row is None:
        say(who, f"UPDATE matched 0 rows: {usd(amount)} does not fit, rejected")
        conn.rollback()
        return False
    say(who, f"UPDATE reserved {usd(amount)}, spent now {usd(row[0])}")
    conn.execute(
        "INSERT INTO expenses (team_id, amount, approved_by) VALUES (%s, %s, %s)",
        (team, amount, who),
    )
    pause()
    say(who, "commit")
    conn.commit()
    return True


AMOUNT = 3_000  # each, against a $5,000 budget


def demo_serializable() -> None:
    alice = connect("alice", IsolationLevel.SERIALIZABLE)
    bob = connect("bob", IsolationLevel.SERIALIZABLE)
    alice_checked, bob_checked, alice_done = threading.Event(), threading.Event(), threading.Event()

    def alice_pause() -> None:
        alice_checked.set()
        bob_checked.wait()

    def bob_pause() -> None:
        bob_checked.set()
        alice_done.wait()

    def run_alice() -> None:
        approve_serializable(alice, "design", AMOUNT, "alice", alice_pause)
        alice_done.set()

    a = threading.Thread(target=run_alice)
    a.start()
    alice_checked.wait()
    ok, attempts = approve_serializable(bob, "design", AMOUNT, "bob", bob_pause)
    a.join()
    say("bob", f"approved={ok} after {attempts} attempt(s)")


def demo_locking(approve: Callable[..., bool], waiting: str) -> None:
    alice = connect("alice")
    bob = connect("bob")
    bob_thread = threading.Thread(target=approve, args=(bob, "design", AMOUNT, "bob"))

    def alice_pause() -> None:
        bob_thread.start()
        wait_until_blocked("bob")
        say("bob", waiting)

    approve(alice, "design", AMOUNT, "alice", alice_pause)
    bob_thread.join()


def main() -> None:
    variant = sys.argv[1] if len(sys.argv) > 1 else "serializable"
    reset()
    print(f"--- approve, {variant} ---")
    if variant == "serializable":
        demo_serializable()
    elif variant == "for-update":
        demo_locking(approve_for_update, "SELECT ... FOR UPDATE is waiting for alice")
    else:
        demo_locking(approve_conditional, "UPDATE is waiting for alice's row lock")

    with connect("check") as conn:
        total, count = conn.execute(
            "SELECT coalesce(sum(amount), 0), count(*) FROM expenses WHERE team_id = 'design'"
        ).fetchone()
    print(f"approved {count} expense(s) totalling {usd(total)} on a $5,000 budget")


if __name__ == "__main__":
    main()

"""Many managers at once: 16 workers share 400 random approvals for one team.

Usage: python load_test.py naive-rc | naive-rr | serializable | for-update | conditional | all
"""
import queue
import random
import sys
import threading

from psycopg import IsolationLevel

import approve
from common import connect, reset, usd

WORKERS = 16
APPROVALS = 400
BUDGET = 50_000  # the 400 requests ask for roughly twice this

STRATEGIES = {
    "naive-rc": IsolationLevel.READ_COMMITTED,
    "naive-rr": IsolationLevel.REPEATABLE_READ,
    "serializable": IsolationLevel.SERIALIZABLE,
    "for-update": IsolationLevel.READ_COMMITTED,
    "conditional": IsolationLevel.READ_COMMITTED,
}


def approve_once(strategy: str, conn, amount: int, who: str) -> tuple[bool, int]:
    match strategy:
        case "naive-rc" | "naive-rr":
            return approve.approve_naive(conn, "design", amount, who), 1
        case "serializable":
            return approve.approve_serializable(conn, "design", amount, who)
        case "for-update":
            return approve.approve_for_update(conn, "design", amount, who), 1
        case "conditional":
            return approve.approve_conditional(conn, "design", amount, who), 1
    raise ValueError(strategy)


def run(strategy: str) -> str:
    reset(BUDGET)
    jobs: queue.Queue = queue.Queue()
    rng = random.Random(42)
    for _ in range(APPROVALS):
        jobs.put(rng.randint(50, 500))

    results: list[tuple[bool, int]] = []

    def worker(n: int) -> None:
        conn = connect(f"w{n:02}", STRATEGIES[strategy])
        while True:
            try:
                amount = jobs.get_nowait()
            except queue.Empty:
                break
            results.append(approve_once(strategy, conn, amount, f"w{n:02}"))
        conn.close()

    threads = [threading.Thread(target=worker, args=(n,)) for n in range(WORKERS)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    with connect("check") as conn:
        total = conn.execute(
            "SELECT coalesce(sum(amount), 0) FROM expenses WHERE team_id = 'design'"
        ).fetchone()[0]
    approved = sum(ok for ok, _ in results)
    retries = sum(attempts - 1 for _, attempts in results)
    most = max(attempts for _, attempts in results)
    return (
        f"{strategy:<13} approved={approved:<3} retries={retries:<4} most_attempts={most:<2}"
        f" total={usd(total)} over_budget={usd(max(0, total - BUDGET))}"
    )


def main() -> None:
    approve.VERBOSE = False
    choice = sys.argv[1] if len(sys.argv) > 1 else "all"
    strategies = list(STRATEGIES) if choice == "all" else [choice]
    print(f"--- {WORKERS} workers, {APPROVALS} approvals, {usd(BUDGET)} budget ---")
    for strategy in strategies:
        print(run(strategy), flush=True)


if __name__ == "__main__":
    main()