XingHuo/app/utils/db.py

"""SQLAlchemy 引擎工厂（B7 · 开发计划挂账⑥）。

按库名缓存单例 Engine：deps / risk / simulate 路由每请求实例化 Repository 时
复用同一连接池，不再每次 create_engine（B6 复审 P3：实例化点泄漏）；应用
shutdown 经 dispose_engines 统一释放（B7 评审 P1-1：显式 Engine.dispose()，
不依赖 GC 兜底）。测试直传 engine= 参数的用法不受影响；
monkeypatch settings 后须先 dispose_engines() 清缓存。
"""

from __future__ import annotations

import threading

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine

from app.config.settings import settings

_engines: dict[str, Engine] = {}
_engines_lock = threading.Lock()


def get_engine(database: str) -> Engine:
    """取指定库的单例 Engine（进程内共享连接池）。"""
    with _engines_lock:
        engine = _engines.get(database)
        if engine is None:
            pwd = settings.mysql_password
            auth = f"{settings.mysql_user}:{pwd}" if pwd else settings.mysql_user
            url = (
                f"mysql+pymysql://{auth}@{settings.mysql_host}:{settings.mysql_port}"
                f"/{database}?charset=utf8mb4"
            )
            engine = create_engine(url, pool_pre_ping=True)
            _engines[database] = engine
        return engine


def dispose_engines() -> None:
    """释放全部缓存 Engine（lifespan shutdown / 测试隔离）。

    逐个 Engine.dispose() 关闭连接池空闲连接后清空缓存；调用方仍持有的
    checked-out 连接不受影响，归还时由旧池关闭。dispose 后再 get_engine
    按当前 settings 重建（测试 monkeypatch 配置依赖此语义）。
    """
    with _engines_lock:
        for engine in _engines.values():
            engine.dispose()
        _engines.clear()