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()


if IS_POSTGRES:
    DB_PATH = None
    _COMMON = dict(pool_size=1, max_overflow=1, pool_pre_ping=True, pool_recycle=1800, future=True)
    # 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, **_COMMON)
    async_engine_ro = create_async_engine(
        DATABASE_URL, **_COMMON,
        connect_args={"server_settings": {"default_transaction_read_only": "on"}})
    sync_engine = create_engine(_sync_url(DATABASE_URL), pool_size=1, max_overflow=1,
                                pool_pre_ping=True, pool_recycle=1800, future=True)
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