HR-ATS-Portal/backend/db_setup.py

"""PostgreSQL connection, async SQLAlchemy ORM and session management.

Configuration comes from the environment, with `.env` read from the repo root or
from `backend/` (`Db_USERNAME`, `Db_PASSWORD`, `Db_HOST`, `Db_PORT`, `Db_NAME`, and
the `DB_*` tuning fields below). Alembic lives in `alembic_setup.py`; `init_db()`
calls into it.

    app = FastAPI(lifespan=lifespan)                          # migrate on startup
    async def endpoint(db: AsyncSession = Depends(get_session)): ...
    async with session_scope() as db: ...                     # workers and scripts
"""

from __future__ import annotations

import asyncio
import os
import logging
from contextlib import asynccontextmanager
from functools import lru_cache
from pathlib import Path
from typing import Annotated, Any, AsyncIterator, Sequence

from pydantic import field_validator
from pydantic_settings import BaseSettings, NoDecode, SettingsConfigDict
from sqlalchemy import MetaData, text
from sqlalchemy.engine import URL, make_url
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import DeclarativeBase
from dotenv import load_dotenv

load_dotenv()

logger = logging.getLogger("db")
BASE_DIR = Path(__file__).resolve().parent


class Settings(BaseSettings):
    """Every field is overridden by an environment variable of the same name."""

    model_config = SettingsConfigDict(
        env_file=(BASE_DIR.parent / ".env", BASE_DIR / ".env"), extra="ignore"
    )

    database_url: str = ""  # full DSN; wins over the Db_* parts below
    db_username: str = os.getenv("DB_USERNAME")
    db_password: str = os.getenv("DB_PASSWORD")
    db_host: str = os.getenv("DB_HOST")
    db_port: int = int(os.getenv("DB_PORT"))
    db_name: str = os.getenv("DB_NAME")
    db_sslmode: str = ""  # e.g. "require" on Azure


    db_schemas: Annotated[list[str], NoDecode] = ["app", "ref", "audit", "ai", "staging"]
    db_default_schema: str = "app"  # schema for models that declare none
    db_echo: bool = False
    db_pool_size: int = 5
    db_max_overflow: int = 10
    db_pool_recycle: int = 1800
    db_connect_retries: int = 10
    db_auto_migrate: bool = True  # run `upgrade head` on startup
    db_autogenerate: bool = True  # write a revision when models drift from the schema
    db_model_modules: Annotated[list[str], NoDecode] = []  # empty means auto-discover
    app_name: str = "hr-ats-portal"

    @field_validator("db_schemas", "db_model_modules", mode="before")
    @classmethod
    def _csv(cls, value: Any) -> Any:
        if isinstance(value, str):
            return [item.strip() for item in value.split(",") if item.strip()]
        return value

    def url(self, *, async_driver: bool = True) -> URL:
        """DSN with the driver forced; `sslmode` is translated to asyncpg's `ssl`."""
        url = (
            make_url(self.database_url)
            if self.database_url
            else URL.create(
                "postgresql",
                self.db_username,
                self.db_password,
                self.db_host,
                self.db_port,
                self.db_name,
            )
        )
        query = dict(url.query)
        if self.db_sslmode:
            query.setdefault("sslmode", self.db_sslmode)
        if async_driver and query.pop("sslmode", None) not in (None, "disable", "allow", "prefer"):
            query["ssl"] = "true"
        driver = "asyncpg" if async_driver else "psycopg2"
        return url.set(drivername=f"postgresql+{driver}", query=query)


@lru_cache(maxsize=1)
def get_settings() -> Settings:
    return Settings()


def database_url(*, async_driver: bool = True, hide_password: bool = False) -> str:
    return get_settings().url(async_driver=async_driver).render_as_string(hide_password=hide_password)


NAMING_CONVENTION = {
    "ix": "ix_%(table_name)s_%(column_0_N_name)s",
    "uq": "uq_%(table_name)s_%(column_0_N_name)s",
    "ck": "ck_%(table_name)s_%(constraint_name)s",
    "fk": "fk_%(table_name)s_%(column_0_N_name)s_%(referred_table_name)s",
    "pk": "pk_%(table_name)s",
}


class Base(DeclarativeBase):
    """Declarative base every model must inherit from."""

    metadata = MetaData(
        naming_convention=NAMING_CONVENTION,
        schema=get_settings().db_default_schema or None,
    )


_engine: AsyncEngine | None = None
_sessionmaker: async_sessionmaker[AsyncSession] | None = None


def get_engine() -> AsyncEngine:
    """The process-wide AsyncEngine, created on first use."""
    global _engine
    if _engine is None:
        s = get_settings()
        _engine = create_async_engine(
            s.url(),
            echo=s.db_echo,
            pool_pre_ping=True,
            pool_size=s.db_pool_size,
            max_overflow=s.db_max_overflow,
            pool_recycle=s.db_pool_recycle,
            connect_args={
                "server_settings": {"timezone": "UTC", "application_name": s.app_name}
            },
        )
    return _engine


def get_sessionmaker() -> async_sessionmaker[AsyncSession]:
    global _sessionmaker
    if _sessionmaker is None:
        _sessionmaker = async_sessionmaker(
            bind=get_engine(), class_=AsyncSession, expire_on_commit=False, autoflush=False
        )
    return _sessionmaker


async def get_session() -> AsyncIterator[AsyncSession]:
    """FastAPI dependency. Rolls back on error; committing is the caller's job."""
    async with get_sessionmaker()() as session:
        try:
            yield session
        except Exception:
            await session.rollback()
            raise


@asynccontextmanager
async def session_scope() -> AsyncIterator[AsyncSession]:
    """Transactional session for scripts and workers: commits, or rolls back on error."""
    async with get_sessionmaker()() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


async def close_db() -> None:
    """Dispose of the pool and reset the cached factories."""
    global _engine, _sessionmaker
    if _engine is not None:
        await _engine.dispose()
        logger.info("connection pool closed")
    _engine, _sessionmaker = None, None


async def check_connection(retries: int | None = None, delay: float = 1.0) -> None:
    """Wait for Postgres to answer `SELECT 1`, retrying with a capped backoff."""
    attempts = get_settings().db_connect_retries if retries is None else retries
    for attempt in range(1, max(attempts, 1) + 1):
        try:
            async with get_engine().connect() as conn:
                await conn.execute(text("SELECT 1"))
            logger.info("connected to %s", database_url(hide_password=True))
            return
        except Exception as exc:
            if attempt >= attempts:
                raise RuntimeError(f"cannot reach {database_url(hide_password=True)}") from exc
            logger.warning("database not ready (%s/%s): %s", attempt, attempts, exc)
            await asyncio.sleep(delay)
            delay = min(delay * 2, 10.0)


async def create_schemas(schemas: Sequence[str] | None = None) -> None:
    """`CREATE SCHEMA IF NOT EXISTS` for every configured schema."""
    names = list(get_settings().db_schemas if schemas is None else schemas)
    if not names:
        return
    async with get_engine().begin() as conn:
        for name in names:
            await conn.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{name}"'))
    logger.info("schemas ensured: %s", ", ".join(names))


async def init_db(*, migrate: bool | None = None, autogen: bool | None = None) -> None:
    """Connect, create the schemas, then bring migrations up to head."""
    should_migrate = get_settings().db_auto_migrate if migrate is None else migrate
    await check_connection()
    await create_schemas()
    if should_migrate:
        from alembic_setup import migrate as run_migrations  # local import: avoids a cycle

        await run_migrations(autogen=autogen)


@asynccontextmanager
async def lifespan(app: Any = None) -> AsyncIterator[None]:
    """FastAPI lifespan: `app = FastAPI(lifespan=lifespan)`."""
    await init_db()
    try:
        yield
    finally:
        await close_db()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)-8s %(name)s: %(message)s")

    async def _main() -> None:
        try:
            await init_db()
        finally:
            await close_db()

    asyncio.run(_main())