app/db/engine.py

"""Async SQLAlchemy engine + session factory.

Driver selection: production = `asyncpg` (Neon Postgres). Tests = `aiosqlite`
(in-process, fast). The DATABASE_URL is normalized so callers can supply
either a Neon connection string (with `?sslmode=require&channel_binding=require`)
or `sqlite+aiosqlite:///:memory:`.
"""
from __future__ import annotations

import os
import re
from contextlib import asynccontextmanager
from typing import AsyncIterator

from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool


def _normalize_url(url: str) -> str:
    # Neon hands out `postgresql://...` — SQLAlchemy needs the async driver explicit.
    if url.startswith("postgresql://"):
        url = url.replace("postgresql://", "postgresql+asyncpg://", 1)
    if url.startswith("postgres://"):
        url = url.replace("postgres://", "postgresql+asyncpg://", 1)
    # asyncpg doesn't understand libpq query flags Neon includes. Strip them
    # but keep them noted; ssl is requested via connect_args instead.
    if "+asyncpg://" in url:
        url = re.sub(r"[?&](sslmode|channel_binding|application_name)=[^&]*", "", url)
        url = re.sub(r"\?&", "?", url).rstrip("?&")
    return url


def _connect_args(url: str) -> dict:
    if "+asyncpg://" in url and "neon.tech" in url:
        return {"ssl": True}
    return {}


def get_database_url() -> str:
    url = os.environ.get("DATABASE_URL")
    if not url:
        raise RuntimeError("DATABASE_URL is not set")
    return _normalize_url(url)


_engine = None
_session_factory = None


def get_engine():
    global _engine
    if _engine is None:
        url = get_database_url()
        # SQLite (test) → NullPool: connection-per-checkout, no caching across
        # event loops. Production Postgres uses the default pool with pre-ping.
        if url.startswith("sqlite"):
            _engine = create_async_engine(url, poolclass=NullPool, connect_args=_connect_args(url))
        else:
            _engine = create_async_engine(
                url,
                pool_pre_ping=True,
                pool_size=int(os.environ.get("DB_POOL_SIZE", "5")),
                max_overflow=int(os.environ.get("DB_MAX_OVERFLOW", "5")),
                connect_args=_connect_args(url),
            )
    return _engine


def get_session_factory() -> async_sessionmaker[AsyncSession]:
    global _session_factory
    if _session_factory is None:
        _session_factory = async_sessionmaker(
            get_engine(),
            expire_on_commit=False,
            class_=AsyncSession,
        )
    return _session_factory


@asynccontextmanager
async def session_scope() -> AsyncIterator[AsyncSession]:
    factory = get_session_factory()
    async with factory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


def reset_engine_for_tests() -> None:
    """Test helper — discard cached engine so a new DATABASE_URL takes effect."""
    global _engine, _session_factory
    _engine = None
    _session_factory = None