""" Async SQLAlchemy engine and session factory for AnemiaLens. Supports SQLite (dev) and PostgreSQL (production) via DATABASE_URL. Performance optimizations: - Connection pooling with configurable pool size - pool_pre_ping for stale connection detection - pool_recycle for automatic connection refresh - Query result caching for frequently-accessed data - Indexed columns for common query patterns """ from __future__ import annotations import logging import os from contextlib import asynccontextmanager from pathlib import Path from dotenv import load_dotenv from sqlalchemy import event from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase from sqlalchemy.pool import AsyncAdaptedQueuePool BACKEND_ROOT = Path(__file__).resolve().parents[1] load_dotenv(BACKEND_ROOT / ".env") log = logging.getLogger("anemialens.db") DEFAULT_SQLITE_DATABASE_URL = "sqlite+aiosqlite:///./anemialens.db" # --------------------------------------------------------------------------- # Connection pooling configuration # --------------------------------------------------------------------------- POOL_SIZE = int(os.getenv("DB_POOL_SIZE", "10")) MAX_OVERFLOW = int(os.getenv("DB_MAX_OVERFLOW", "20")) POOL_RECYCLE = int(os.getenv("DB_POOL_RECYCLE", "1800")) POOL_TIMEOUT = int(os.getenv("DB_POOL_TIMEOUT", "30")) POOL_PRE_PING = os.getenv("DB_POOL_PRE_PING", "true").lower() == "true" SLOW_QUERY_THRESHOLD_MS = int(os.getenv("SLOW_QUERY_THRESHOLD_MS", "0")) def _normalize_database_url(database_url: str) -> str: url = database_url.strip() if url.startswith("postgres://"): return url.replace("postgres://", "postgresql+asyncpg://", 1) if url.startswith("postgresql://") and "+asyncpg" not in url: return url.replace("postgresql://", "postgresql+asyncpg://", 1) return url def _resolve_database_url() -> str: return _normalize_database_url( os.getenv("DATABASE_URL", DEFAULT_SQLITE_DATABASE_URL) ) def _runtime_environment() -> str: return ( os.getenv("ANEMIALENS_ENVIRONMENT") or os.getenv("ENVIRONMENT") or "development" ).strip().lower() def _truthy_env(value: str | None) -> bool | None: if value is None: return None return value.strip().lower() in {"1", "true", "yes", "on"} def _development_database_fallback_enabled() -> bool: explicit = _truthy_env(os.getenv("ANEMIALENS_ENABLE_DEV_DB_FALLBACK")) if explicit is not None: return explicit return _runtime_environment() != "production" def _development_fallback_database_url() -> str: return _normalize_database_url( os.getenv("ANEMIALENS_DEV_DATABASE_URL", DEFAULT_SQLITE_DATABASE_URL) ) def _build_connect_args(database_url: str) -> dict[str, object]: connect_args: dict[str, object] = {} if "sqlite" in database_url: connect_args["check_same_thread"] = False elif "postgres" in database_url: connect_args["ssl"] = "require" return connect_args def _build_engine_kwargs(database_url: str) -> dict[str, object]: """Build engine configuration based on database type.""" connect_args = _build_connect_args(database_url) base: dict[str, object] = { "echo": False, "pool_pre_ping": POOL_PRE_PING, "connect_args": connect_args, } if "sqlite" in database_url: base["connect_args"] = {**connect_args, "timeout": 30} base["pool_size"] = 1 base["max_overflow"] = 0 elif "postgres" in database_url: base["pool_size"] = POOL_SIZE base["max_overflow"] = MAX_OVERFLOW base["pool_recycle"] = POOL_RECYCLE base["pool_timeout"] = POOL_TIMEOUT base["poolclass"] = AsyncAdaptedQueuePool return base def _attach_slow_query_logging(target_engine) -> None: if SLOW_QUERY_THRESHOLD_MS <= 0: return @event.listens_for(target_engine.sync_engine, "before_cursor_execute") def _before_cursor_execute(conn, cursor, statement, parameters, context, executemany): conn.info.setdefault("query_start_time", []).append(__import__("time").time()) @event.listens_for(target_engine.sync_engine, "after_cursor_execute") def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany): start_times = conn.info.get("query_start_time", []) if start_times: start = start_times.pop() elapsed_ms = (__import__("time").time() - start) * 1000 if elapsed_ms > SLOW_QUERY_THRESHOLD_MS: log.warning( "SLOW QUERY (%.1fms): %s", elapsed_ms, statement[:200], ) def _build_engine_and_session(database_url: str): current_engine = create_async_engine( database_url, **_build_engine_kwargs(database_url), ) _attach_slow_query_logging(current_engine) session_factory = async_sessionmaker( current_engine, class_=AsyncSession, expire_on_commit=False, ) return current_engine, session_factory DATABASE_URL = _resolve_database_url() _configured_database_url = DATABASE_URL engine, async_session_factory = _build_engine_and_session(DATABASE_URL) async def _rebind_engine(database_url: str) -> None: global DATABASE_URL, _configured_database_url, engine, async_session_factory current_engine = engine DATABASE_URL = database_url _configured_database_url = database_url engine, async_session_factory = _build_engine_and_session(database_url) os.environ["DATABASE_URL"] = database_url if current_engine is not engine: await current_engine.dispose() database_kind = "sqlite" if "sqlite" in database_url else "postgresql" log.info("Database engine rebound for %s runtime.", database_kind) def _ensure_engine_current() -> None: current_database_url = _resolve_database_url() if current_database_url == _configured_database_url: return globals()["DATABASE_URL"] = current_database_url globals()["_configured_database_url"] = current_database_url globals()["engine"], globals()["async_session_factory"] = _build_engine_and_session( current_database_url ) database_kind = "sqlite" if "sqlite" in current_database_url else "postgresql" log.info("Database engine rebound for %s runtime.", database_kind) class Base(DeclarativeBase): """Declarative base for all ORM models.""" pass @asynccontextmanager async def get_db_session(): """ Async context manager for database sessions. Preferred over the dependency version for service-layer code. """ _ensure_engine_current() async with async_session_factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() async def get_db(): """FastAPI dependency that yields an async database session.""" _ensure_engine_current() async with async_session_factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise finally: await session.close() async def create_tables() -> None: """Create all ORM tables - called once during app startup.""" _ensure_engine_current() primary_database_url = DATABASE_URL try: await _create_all_tables_for_url(primary_database_url) log.info("Database tables created/verified successfully.") except Exception as exc: log.error("Database table creation FAILED: %s", exc, exc_info=True) fallback_url = await _activate_development_database_fallback(exc) if fallback_url is not None: await _create_all_tables_for_url(fallback_url) log.info("Database tables created/verified successfully using SQLite fallback.") return log.warning("Continuing startup despite table creation error.") async def close_engine() -> None: """Dispose of the engine - called during app shutdown.""" _ensure_engine_current() await engine.dispose() log.info("Database engine disposed.") async def _create_all_tables_for_url(database_url: str) -> None: ddl_url = database_url if "pooler.supabase.com:6543" in ddl_url: ddl_url = ddl_url.replace(":6543/", ":5432/") log.info("Using session pooler (port 5432) for DDL operations.") ddl_engine = create_async_engine( ddl_url, echo=False, pool_pre_ping=True, connect_args={"ssl": "require"} if "postgres" in ddl_url else {}, ) try: async with ddl_engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) finally: await ddl_engine.dispose() async def _activate_development_database_fallback(exc: Exception) -> str | None: if "sqlite" in DATABASE_URL: return None if not _development_database_fallback_enabled(): return None fallback_url = _development_fallback_database_url() if fallback_url == DATABASE_URL: return None log.warning( "Primary database unavailable in %s mode; falling back to local SQLite runtime. " "Original error: %s", _runtime_environment(), exc, ) await _rebind_engine(fallback_url) return fallback_url # --------------------------------------------------------------------------- # Query optimization utilities # --------------------------------------------------------------------------- def paginated_query(query, page: int = 1, page_size: int = 20): """Apply pagination to a SQLAlchemy query.""" offset = (page - 1) * page_size return query.offset(offset).limit(page_size) async def cached_count(session: AsyncSession, model, cache_key: str | None = None, ttl: int = 60): """ Get cached row count for a model. For large tables, count(*) is expensive; this caches the result. """ from app.services.cache import response_cache from sqlalchemy import func, select key = cache_key or f"count:{model.__tablename__}" cached = await response_cache.get(key) if cached is not None: return cached result = await session.execute(select(func.count()).select_from(model)) count = result.scalar() or 0 await response_cache.set(key, count, ttl_seconds=ttl) return count