group_xinghuo_jinrong/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() 清缓存。

角色（D20 · 架构设计-基金转换交易 §11.1）：ro / rw / admin 三档，
由 `get_engine` 的 `role` 参数选择账号；账号未配置时回退 mysql_user。
"""

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()

_VALID_ROLES = ("ro", "rw", "admin")


def _resolve_credentials(database: str, role: str) -> tuple[str, str]:
    """解析该 (库, 角色) 应使用的 (user, password)。

    - `admin`：恒用 `mysql_user`（DDL / 测试 teardown 等需 DELETE 的场景）
    - `ro` / `rw` + core 库：分别用 `mysql_core_ro_*` / `mysql_core_rw_*`
    - `ro` / `rw` + agent 库：统一用 `mysql_agent_*`（agent 库无独立只读账号）
    - 对应账号为空串 → 回退 `mysql_user`（渐进启用，不阻塞本地开发）
    """
    if role == "admin":
        return settings.mysql_user, settings.mysql_password
    if role not in _VALID_ROLES:
        raise ValueError(f"unknown db role: {role!r}（应为 ro / rw / admin）")

    if database == settings.mysql_core_database:
        if role == "ro":
            user, password = settings.mysql_core_ro_user, settings.mysql_core_ro_password
        else:
            user, password = settings.mysql_core_rw_user, settings.mysql_core_rw_password
    else:
        user, password = settings.mysql_agent_user, settings.mysql_agent_password

    if not user:
        return settings.mysql_user, settings.mysql_password
    return user, password


def get_engine(database: str, role: str = "rw") -> Engine:
    """取指定 (库名, 角色) 的单例 Engine（进程内共享连接池）。

    ⚠️ 调用点必须显式声明角色（F-11 清单共 13 处），不要依赖默认值表达意图：
    - `role="ro"`：只读账号 → `core_ro` / `core_tools` / deps 归属校验 /
      risk 扫描 / **gateway 的校验读**
    - `role="rw"`：读写账号 → `gateway_repository`（core 写）、
      `risk_repository` / `session_repository`（agent 读写）
    - `role="admin"`：`mysql_user` → 建库/DDL、`tests/conftest.py` teardown（需 DELETE）

    账号未配置时回退 `mysql_user`，与单账号时代行为一致。
    """
    if role not in _VALID_ROLES:
        raise ValueError(f"unknown db role: {role!r}（应为 ro / rw / admin）")

    cache_key = f"{database}:{role}"
    with _engines_lock:
        engine = _engines.get(cache_key)
        if engine is None:
            user, password = _resolve_credentials(database, role)
            auth = f"{user}:{password}" if password else 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[cache_key] = 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()