Model/infrastructure/postgres/engine.py

from collections.abc import Iterator
from contextlib import contextmanager

from sqlalchemy.engine import Engine
from sqlmodel import Session, create_engine

from infrastructure.postgres.config import PostgresConfig


def make_engine(config: PostgresConfig) -> Engine:
    return create_engine(
        config.url(),
        pool_size=config.pool_size,
        max_overflow=config.max_overflow,
        pool_pre_ping=config.pool_pre_ping,
        pool_recycle=config.pool_recycle,
    )


def make_session(engine: Engine) -> Session:
    return Session(engine)


@contextmanager  # pyright: ignore[reportDeprecated]
def transactional_session(engine: Engine) -> Iterator[Session]:
    """Yield a session whose lifecycle owns the transaction.

    On clean exit the session commits; on any exception it rolls back and
    re-raises. Either way the session is closed. Callers in the application
    layer can do their work inside the ``with`` block without ever invoking
    ``.commit()`` / ``.rollback()`` themselves -- transaction semantics stay
    in the infrastructure layer.
    """
    session = Session(engine)
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager  # pyright: ignore[reportDeprecated]
def commit_scope(session: Session) -> Iterator[Session]:
    """Commit a caller-owned session on clean exit; roll back on error.

    Like ``transactional_session`` but for a session the caller already holds
    and will close itself. Use it to keep slow, non-DB work *outside* the
    transaction: build the session, run the slow work, then enter
    ``commit_scope`` only for the persistence -- so a connection is checked out
    (SQLModel sessions are lazy) for the shortest possible window.
    """
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise