agent-platform/services/model-gateway-service/app/db/session.py

from collections.abc import Generator

from core_db import DatabaseSettings, create_engine_from_settings, create_session_factory
from fastapi import Request
from sqlalchemy.orm import Session, sessionmaker

from app.bootstrap.settings import ModelGatewayServiceSettings


def build_session_factory(
    settings: ModelGatewayServiceSettings | None = None,
) -> sessionmaker[Session]:
    resolved_settings = settings or ModelGatewayServiceSettings()
    db_settings = DatabaseSettings(
        database_url=resolved_settings.database_url,
        echo_sql=resolved_settings.echo_sql,
    )
    engine = create_engine_from_settings(db_settings)
    return create_session_factory(engine)


def get_db(request: Request) -> Generator[Session, None, None]:
    session_factory: sessionmaker[Session] = request.app.state.session_factory
    session = session_factory()
    try:
        yield session
    finally:
        session.close()