"""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" 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())