demand_server/app/db/supply_mysql.py

from collections.abc import Generator

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker

from app.core.config import settings

_supply_engine: Engine | None = None
_supply_session_local: sessionmaker[Session] | None = None


def _ensure_supply_engine() -> sessionmaker[Session]:
    global _supply_engine, _supply_session_local
    if not settings.supply_mysql_configured:
        raise RuntimeError(
            "供需内容库未配置，请在 .env 中设置 SUPPLY_MYSQL_HOST、"
            "SUPPLY_MYSQL_USER、SUPPLY_MYSQL_PASSWORD、SUPPLY_MYSQL_DATABASE"
        )
    if _supply_session_local is None:
        _supply_engine = create_engine(
            settings.supply_mysql_dsn,
            pool_pre_ping=True,
            pool_recycle=3600,
        )
        _supply_session_local = sessionmaker(
            bind=_supply_engine,
            autoflush=False,
            autocommit=False,
        )
    return _supply_session_local


def SupplySessionLocal() -> Session:
    return _ensure_supply_engine()()


def get_supply_db_session() -> Generator[Session, None, None]:
    db = SupplySessionLocal()
    try:
        yield db
    finally:
        db.close()