db.py

"""Database configuration and session management using SQLAlchemy 2.x."""
from __future__ import annotations

import os
from typing import Generator

from sqlalchemy import create_engine
from sqlalchemy.orm import DeclarativeBase, sessionmaker

from api.config import get_settings


class Base(DeclarativeBase):
    """Base class for SQLAlchemy models."""


def _normalize_postgres_url(url: str) -> str:
    """Normalize Postgres URL to use the psycopg driver.

    Accepts legacy prefixes and ensures the SQLAlchemy URL includes the
    ``+psycopg`` dialect when using PostgreSQL.
    """
    if url.startswith("postgres://"):
        return url.replace("postgres://", "postgresql+psycopg://", 1)
    if url.startswith("postgresql://") and "+psycopg" not in url:
        return url.replace("postgresql://", "postgresql+psycopg://", 1)
    return url


def get_engine_url() -> str:
    """Resolve the database URL from env or settings with sane defaults."""
    # Use the cached settings instance instead of instantiating ``Settings``
    # directly so configuration is read a single time across the app.
    settings = get_settings()
    url = os.getenv("DATABASE_URL", settings.database_url)
    if url.startswith("sqlite"):
        return url
    return _normalize_postgres_url(url)


ENGINE_URL = get_engine_url()
engine = create_engine(ENGINE_URL, pool_pre_ping=True, future=True)

SessionLocal = sessionmaker(
    bind=engine,
    autoflush=False,
    autocommit=False,
    future=True,
)

# Ensure models are registered on Base.metadata for create_all/drop_all
import api.models  # noqa: E402,F401


def get_db() -> Generator:
    """FastAPI dependency to provide a session per request."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()