311 lines
10 KiB
Python
311 lines
10 KiB
Python
"""PostgreSQL connection, async SQLAlchemy ORM and session management.
|
|
|
|
Configuration comes from the environment, with `.env` read from `backend/`
|
|
(`DB_USERNAME`, `DB_PASSWORD`, `DB_HOST`, `DB_PORT`, `DB_NAME`, `PROD_ENV`, and
|
|
the `DB_*` tuning fields below). There is no repo-root `.env`. 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 logging
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
from typing import Annotated, Any, AsyncIterator, Sequence
|
|
|
|
from dotenv import load_dotenv
|
|
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 sqlmodel import SQLModel
|
|
|
|
load_dotenv()
|
|
|
|
logger = logging.getLogger("db")
|
|
BASE_DIR = Path(__file__).resolve().parent
|
|
|
|
_TRUE = {"1", "true", "yes", "on"}
|
|
_LOOPBACK_HOSTS = frozenset({"localhost", "127.0.0.1"})
|
|
|
|
|
|
def _running_in_docker() -> bool:
|
|
"""True inside a container (/.dockerenv) or when Compose sets IN_DOCKER=1."""
|
|
return Path("/.dockerenv").exists() or os.environ.get("IN_DOCKER", "").strip().lower() in _TRUE
|
|
|
|
|
|
class Settings(BaseSettings):
|
|
"""Every field is overridden by an environment variable of the same name."""
|
|
|
|
model_config = SettingsConfigDict(
|
|
env_file=BASE_DIR / ".env",
|
|
extra="ignore",
|
|
)
|
|
|
|
database_url: str = "" # full DSN; wins over the DB_* parts below
|
|
db_username: str = ""
|
|
db_password: str = ""
|
|
db_host: str = "localhost"
|
|
db_port: int = 5432
|
|
db_name: str = ""
|
|
db_sslmode: str = "" # blank = derive from PROD_ENV (require on RDS, off locally)
|
|
prod_env: bool = False # true → RDS (SSL); false → local psql over asyncpg
|
|
|
|
db_schemas: Annotated[list[str], NoDecode] = "app"
|
|
db_default_schema: str = "app"
|
|
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
|
|
db_autogenerate: bool = True
|
|
db_model_modules: Annotated[list[str], NoDecode] = []
|
|
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
|
|
|
|
@field_validator("prod_env", mode="before")
|
|
@classmethod
|
|
def _bool(cls, value: Any) -> Any:
|
|
if isinstance(value, str):
|
|
return value.strip().lower() in _TRUE
|
|
return value
|
|
|
|
def url(self, *, async_driver: bool = True) -> URL:
|
|
"""DSN with the driver forced; `sslmode` is mapped to asyncpg's `ssl` mode name.
|
|
|
|
Local Docker: `DB_HOST=localhost` means the container itself, so rewrite to
|
|
`host.docker.internal` for the connection URL only (Settings.db_host unchanged).
|
|
Prod never rewrites — RDS hostname is used as-is.
|
|
"""
|
|
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,
|
|
)
|
|
)
|
|
if (
|
|
not self.prod_env
|
|
and _running_in_docker()
|
|
and (url.host or "").lower() in _LOOPBACK_HOSTS
|
|
):
|
|
url = url.set(host="host.docker.internal")
|
|
|
|
query = dict(url.query)
|
|
|
|
# PROD_ENV=true → RDS needs SSL. Local psql talks plain asyncpg (no SSL).
|
|
sslmode = self.db_sslmode.strip() if self.db_sslmode else ("require" if self.prod_env else "")
|
|
if sslmode:
|
|
query.setdefault("sslmode", sslmode)
|
|
else:
|
|
query.pop("sslmode", None)
|
|
|
|
if async_driver and (mode := query.pop("sslmode", None)) is not None:
|
|
query["ssl"] = mode
|
|
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,
|
|
)
|
|
|
|
|
|
# SQLModel keeps its own registry, so `SQLModel` tables would be invisible to the
|
|
# Alembic autogenerate in alembic_setup.py, which diffs `Base.metadata` alone.
|
|
# Pointing SQLModel at the same MetaData gives both styles one registry, and lets
|
|
# SQLModel tables inherit the naming convention and the default schema above.
|
|
SQLModel.metadata = Base.metadata
|
|
|
|
|
|
_engine: AsyncEngine | None = None
|
|
_sessionmaker: async_sessionmaker[AsyncSession] | None = None
|
|
|
|
|
|
def _connect_args(settings: Settings) -> dict:
|
|
"""UTC session + SSL for RDS. `require` encrypts without verifying the CA."""
|
|
import ssl as ssl_mod
|
|
|
|
args: dict = {
|
|
"server_settings": {"timezone": "UTC", "application_name": settings.app_name}
|
|
}
|
|
mode = (settings.db_sslmode or "").strip().lower()
|
|
if mode and mode not in ("disable", "allow", "prefer"):
|
|
ctx = ssl_mod.create_default_context()
|
|
if mode == "require":
|
|
ctx.check_hostname = False
|
|
ctx.verify_mode = ssl_mod.CERT_NONE
|
|
args["ssl"] = ctx
|
|
return args
|
|
|
|
|
|
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=_connect_args(s),
|
|
)
|
|
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
|
|
s = get_settings()
|
|
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 [PROD_ENV=%s]",
|
|
database_url(hide_password=True),
|
|
s.prod_env,
|
|
)
|
|
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())
|