plm-lottery/app/db/base.py

from sqlalchemy import event
from sqlalchemy.ext.asyncio import AsyncEngine, 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
# five concurrent background tasks (scheduler, confirmation poller, RBF bumper,
# two reconcilers) plus HTTP handlers briefly overlapping a write.
_SQLITE_BUSY_TIMEOUT_MS = 5000


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`.

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

    @event.listens_for(target_engine.sync_engine, "connect")
    def _set_sqlite_pragmas(dbapi_connection, connection_record) -> None:
        cursor = dbapi_connection.cursor()
        try:
            cursor.execute("PRAGMA journal_mode=WAL")
            cursor.execute("PRAGMA synchronous=NORMAL")
            cursor.execute(f"PRAGMA busy_timeout={_SQLITE_BUSY_TIMEOUT_MS}")
        finally:
            cursor.close()


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


class Base(DeclarativeBase):
    pass