demand_server/app/db/hot_content_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

_hot_content_engine: Engine | None = None
_hot_content_session_local: sessionmaker[Session] | None = None


def _ensure_hot_content_engine() -> sessionmaker[Session]:
    global _hot_content_engine, _hot_content_session_local
    if not settings.hot_content_mysql_configured:
        raise RuntimeError(
            "新热内容库未配置，请在 .env 中设置 HOT_CONTENT_MYSQL_HOST、"
            "HOT_CONTENT_MYSQL_USER、HOT_CONTENT_MYSQL_PASSWORD、HOT_CONTENT_MYSQL_DATABASE"
        )
    if _hot_content_session_local is None:
        _hot_content_engine = create_engine(
            settings.hot_content_mysql_dsn,
            pool_pre_ping=True,
            pool_recycle=3600,
        )
        _hot_content_session_local = sessionmaker(
            bind=_hot_content_engine,
            autoflush=False,
            autocommit=False,
        )
    return _hot_content_session_local


def HotContentSessionLocal() -> Session:
    """打开新热内容库会话（独立 MySQL，用法同 SessionLocal()）。"""
    return _ensure_hot_content_engine()()


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


def hot_content_mysql_ping() -> bool:
    _ensure_hot_content_engine()
    assert _hot_content_engine is not None
    with _hot_content_engine.connect() as conn:
        conn.execute(text("SELECT 1"))
    return True