thermograph/accounts/db.py

"""Authoritative account database — SQLAlchemy over PostgreSQL (or SQLite).

This is deliberately a *separate* database from the disposable derived cache
(data/store.py). Accounts, subscriptions and notifications have no source to
recompute from, so they live with stricter rules: foreign keys enforced, errors
surface rather than get swallowed, and the data must be backed up (it is not
regenerable).

Backend is chosen by ``THERMOGRAPH_DATABASE_URL``:

* **PostgreSQL** (``postgresql+asyncpg://…``) in container/prod. Two async engines
  over the one primary give the read / read-write split the deployment asks for —
  a read-only engine (``default_transaction_read_only``) for pure-GET endpoints
  and a read-write engine for everything else (nearly every account op is a
  SELECT-then-write in one transaction and must be RW). A separate **sync**
  engine (psycopg) drives the background notifier thread, which has no event loop.
* **SQLite** (default / tests) keeps the previous single-file behaviour with WAL
  so the async web path and the sync notifier coexist; the RO "engine" is just an
  alias of the RW one (SQLite has no read-only transaction mode to enforce).
"""
import os

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

import paths

# The async SQLAlchemy URL. When unset we fall back to the legacy SQLite file
# (THERMOGRAPH_ACCOUNTS_DB overrides its path; tests point it at a temp file).
DATABASE_URL = os.environ.get("THERMOGRAPH_DATABASE_URL", "").strip()
IS_POSTGRES = DATABASE_URL.startswith("postgresql")


def _sync_url(async_url: str) -> str:
    """The sync driver URL for the notifier: asyncpg is async-only, so the sync
    engine uses psycopg (``postgresql+psycopg://``); SQLite drops the aiosqlite
    driver suffix."""
    return (async_url.replace("+asyncpg", "+psycopg")
            if async_url.startswith("postgresql")
            else async_url.replace("+aiosqlite", ""))


class Base(DeclarativeBase):
    """Declarative base shared by every account-domain table (models.py)."""


def _apply_sqlite_pragmas(dbapi_conn, _rec):
    # SQLite-only: WAL for concurrent async/sync access, NORMAL sync, and
    # foreign_keys ON so the ON DELETE CASCADE relationships are enforced
    # (SQLite defaults them off; Postgres always enforces FKs).
    cur = dbapi_conn.cursor()
    cur.execute("PRAGMA journal_mode=WAL")
    cur.execute("PRAGMA synchronous=NORMAL")
    cur.execute("PRAGMA foreign_keys=ON")
    cur.close()


# Postgres-only pool sizing + sync-path timeouts, defined at module scope (not
# only inside `if IS_POSTGRES` below) so the values themselves are unit-testable
# without a real Postgres connection.
#
# pool_size=1/max_overflow=1 (the original setting on both) meant a 3rd
# concurrent request per worker had to wait out pool_timeout (~30s default) for a
# slot — easy to hit with 4 uvicorn workers x real traffic. 5+5 gives each async
# engine headroom for a burst of concurrent account requests without coming close
# to Postgres's own connection ceiling (two async engines here plus the sync
# engine below, times N workers, all share that budget).
_ASYNC_POOL_KWARGS = dict(pool_size=5, max_overflow=5, pool_pre_ping=True,
                          pool_recycle=1800, future=True)
# Smaller than the async engines: only the notifier leader (one process
# cluster-wide, see core/singleton.py) uses the sync engine, driving its periodic
# sweep from a plain thread with no event loop. connect_timeout + a
# statement_timeout (via psycopg's `options`) bound how long a wedged Postgres
# can hang the notifier — it runs for the lifetime of that process's leader
# flock, so an indefinite hang here would silently stop every subscription
# notification until the process is restarted.
_SYNC_POOL_KWARGS = dict(pool_size=2, max_overflow=3, pool_pre_ping=True,
                         pool_recycle=1800, future=True,
                         connect_args={"connect_timeout": 5,
                                       "options": "-c statement_timeout=30000"})

if IS_POSTGRES:
    DB_PATH = None
    # Read-write engine (default) + a read-only engine that pins every session to a
    # read-only transaction on the server, so a pure-read endpoint can't accidentally
    # write and a long read stays off the write connection.
    async_engine = create_async_engine(DATABASE_URL, **_ASYNC_POOL_KWARGS)
    async_engine_ro = create_async_engine(
        DATABASE_URL, **_ASYNC_POOL_KWARGS,
        connect_args={"server_settings": {"default_transaction_read_only": "on"}})
    sync_engine = create_engine(_sync_url(DATABASE_URL), **_SYNC_POOL_KWARGS)
else:
    DB_PATH = os.environ.get("THERMOGRAPH_ACCOUNTS_DB") or os.path.join(paths.DATA_DIR, "accounts.sqlite")
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    async_engine = create_async_engine(f"sqlite+aiosqlite:///{DB_PATH}", future=True)
    event.listen(async_engine.sync_engine, "connect", _apply_sqlite_pragmas)
    async_engine_ro = async_engine   # SQLite has no read-only transaction mode
    sync_engine = create_engine(f"sqlite:///{DB_PATH}", future=True)
    event.listen(sync_engine, "connect", _apply_sqlite_pragmas)


async_session_maker = async_sessionmaker(async_engine, expire_on_commit=False)
async_session_ro_maker = async_sessionmaker(async_engine_ro, expire_on_commit=False)
sync_session_maker = sessionmaker(sync_engine, expire_on_commit=False)


async def get_async_session() -> AsyncSession:
    """FastAPI dependency: a read-write AsyncSession per request (the default)."""
    async with async_session_maker() as session:
        yield session


async def get_read_session() -> AsyncSession:
    """FastAPI dependency for genuinely read-only endpoints — served by the
    read-only engine so reads prefer the read connection. On Postgres the session
    runs in a read-only transaction; on SQLite it's the same engine as RW."""
    async with async_session_ro_maker() as session:
        yield session


async def create_db_and_tables() -> None:
    """Create any missing tables. Called once at startup (app lifespan). On
    Postgres, retry briefly so an app boot that races the DB container recovers
    instead of crashing."""
    from accounts import models  # noqa: F401  (registers tables on Base.metadata)

    import asyncio

    last = None
    for attempt in range(10):
        try:
            async with async_engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
            return
        except Exception as e:  # noqa: BLE001 - transient "DB not ready yet" on container start
            last = e
            if not IS_POSTGRES:
                raise
            await asyncio.sleep(min(2 ** attempt, 5))
    raise last