plm-lottery/app/db/base.py

from sqlalchemy import event
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase

from app.config import settings

# How long a writer waits for a lock held by another writer before SQLite raises
# "database is locked" (B-39). A few seconds is enough to ride out this app's own
# six concurrent background tasks (Electrum listener, scheduler, confirmation
# poller, RBF bumper, two reconcilers) plus HTTP handlers briefly overlapping a
# write.
_SQLITE_BUSY_TIMEOUT_MS = 5000


# Execution option read by the `begin` event below: which BEGIN a connection's next
# transaction opens with. Set through `begin_immediate`, never by hand.
_BEGIN_MODE_OPTION = "sqlite_begin_mode"


def _register_sqlite_pragmas(target_engine: AsyncEngine) -> None:
    """Without WAL, SQLite's default (rollback-journal) mode lets a writer block
    every reader for the duration of its transaction, and a second writer arriving
    while one is already active fails immediately rather than waiting at all —
    realistic under this app's concurrency, and nothing previously handled it.
    WAL lets readers and writers proceed without blocking each other, and
    busy_timeout gives a second writer a real window to wait for the first
    instead of an instant `OperationalError`.

    B-88: two more things, both about the code's transactions being what it
    assumes. `synchronous=FULL`, because with WAL `NORMAL` can lose the last
    commits on power loss — a registration, and with it a derivation index that
    would then be handed to the next user. And real transaction control: pysqlite's
    legacy mode emits `BEGIN` only at the first INSERT/UPDATE/DELETE, so every read
    before it ran outside any transaction and the lock wait surfaced as `database is
    locked` at an arbitrary later statement. `isolation_level = None` turns the
    driver's implicit BEGIN off and the `begin` event emits it ourselves — a plain
    `BEGIN` by default, `BEGIN IMMEDIATE` for a transaction opened through
    `begin_immediate`, which takes the write lock *before* the first read so that
    what a money-moving transaction reads is what it then writes on top of.

    No-op for any dialect other than sqlite (e.g. a future PostgreSQL
    DATABASE_URL), which neither needs nor understands any of this.
    """
    if target_engine.dialect.name != "sqlite":
        return

    @event.listens_for(target_engine.sync_engine, "connect")
    def _set_sqlite_pragmas(dbapi_connection, connection_record) -> None:
        # Before the pragmas: journal_mode cannot change inside a transaction, and in
        # legacy mode the driver may have one open.
        dbapi_connection.isolation_level = None
        cursor = dbapi_connection.cursor()
        try:
            cursor.execute("PRAGMA journal_mode=WAL")
            cursor.execute("PRAGMA synchronous=FULL")  # B-88
            cursor.execute(f"PRAGMA busy_timeout={_SQLITE_BUSY_TIMEOUT_MS}")
        finally:
            cursor.close()

    @event.listens_for(target_engine.sync_engine, "begin")
    def _emit_begin(conn) -> None:  # B-88
        mode = conn.get_execution_options().get(_BEGIN_MODE_OPTION, "DEFERRED")
        conn.exec_driver_sql("BEGIN IMMEDIATE" if mode == "IMMEDIATE" else "BEGIN")


async def begin_immediate(session: AsyncSession) -> None:
    """B-88: start the session's next transaction with `BEGIN IMMEDIATE`, i.e. take
    SQLite's write lock now, before the first read.

    Every transaction that moves money or fixes its state calls this first thing, so
    that what it reads (balance, free UTXOs, participants, the round's status) cannot
    change under it before its own writes, and so that the wait for the lock happens
    here — at the very start, before anything has been broadcast — instead of at the
    first write. Raises `sqlalchemy.exc.OperationalError` ("database is locked") when
    the lock is not granted within `busy_timeout`; the caller decides whether that
    means "refuse" (before a broadcast) or "leave it to the reconciler" (after).

    A transaction the session already has open is committed first — a lock cannot be
    upgraded in place, and the callers open the transaction with nothing but reads
    behind them (typically the authentication lookup). Pending, unflushed changes
    make that unsafe to do silently, so they are an error.

    Harmless on a session whose engine did not go through `_register_sqlite_pragmas`
    (the option is simply never read) and on a non-sqlite dialect.
    """
    if session.new or session.dirty or session.deleted:
        raise RuntimeError("begin_immediate called with unflushed changes on the session")
    if session.in_transaction():
        await session.commit()
    await session.connection(execution_options={_BEGIN_MODE_OPTION: "IMMEDIATE"})


engine = create_async_engine(settings.database_url)
_register_sqlite_pragmas(engine)
AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False)


class Base(DeclarativeBase):
    pass