"""Async SQLAlchemy engine, session factory and transaction helpers. Spec ยง7: *"all multi-step writes inside DB transactions"*. The only sanctioned way to touch the database is `transaction()` (write) or `session_scope()` (read), so a half-applied pipeline stage cannot be committed. """ from __future__ import annotations import ssl from collections.abc import AsyncIterator from contextlib import asynccontextmanager from typing import Any from sqlalchemy.ext.asyncio import ( AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine, ) from app.core.logging import get_logger from app.core.settings import Settings, get_settings logger = get_logger(__name__) _engine: AsyncEngine | None = None _session_factory: async_sessionmaker[AsyncSession] | None = None def _connect_args(settings: Settings) -> dict[str, Any]: args: dict[str, Any] = { "server_settings": { "application_name": "ledgerlens-api", "statement_timeout": str(settings.db_statement_timeout_ms), }, # A suspended serverless compute (Neon, Aurora Serverless) takes several # seconds to wake; measured cold connects to Neon run ~7 s. 10 s left no # margin. `wait_for_database` retries on top of this. "timeout": 20.0, } if settings.database_requires_tls: # Neon terminates TLS; asyncpg needs an SSLContext rather than `sslmode=`. args["ssl"] = ssl.create_default_context() return args def create_engine(settings: Settings | None = None) -> AsyncEngine: """Build a fresh engine. Used by the app factory and by tests.""" settings = settings or get_settings() return create_async_engine( settings.sqlalchemy_url, pool_size=settings.db_pool_size, max_overflow=settings.db_max_overflow, pool_pre_ping=True, pool_recycle=280, # below Neon's idle cutoff connect_args=_connect_args(settings), echo=False, future=True, ) def init_engine(settings: Settings | None = None) -> AsyncEngine: """Initialise the process-wide engine and session factory (idempotent).""" global _engine, _session_factory if _engine is None: _engine = create_engine(settings) _session_factory = async_sessionmaker(_engine, expire_on_commit=False, class_=AsyncSession) logger.info("db_engine_initialised") return _engine async def dispose_engine() -> None: """Close every pooled connection. Called on application shutdown.""" global _engine, _session_factory if _engine is not None: await _engine.dispose() logger.info("db_engine_disposed") _engine = None _session_factory = None def get_session_factory() -> async_sessionmaker[AsyncSession]: if _session_factory is None: init_engine() assert _session_factory is not None return _session_factory @asynccontextmanager async def session_scope() -> AsyncIterator[AsyncSession]: """Read-oriented session. Rolls back on exit; never commits implicitly.""" async with get_session_factory()() as session: try: yield session finally: await session.rollback() @asynccontextmanager async def transaction() -> AsyncIterator[AsyncSession]: """Write session wrapped in a single transaction: commit on success, rollback on error.""" async with get_session_factory()() as session: try: yield session except BaseException: await session.rollback() raise else: await session.commit() async def get_db() -> AsyncIterator[AsyncSession]: """FastAPI dependency yielding a read session.""" async with session_scope() as session: yield session