long_articles_agentic_supply/src/infra/database/manager.py

"""数据库管理器 —— 连接封装 + 生命周期 + 后端路由

不包含具体数据操作逻辑（操作在各子包 backend.py 实现类内部）。
子门面 MysqlManager / PgManager / RedisManager 在各子包 manager.py 中。
"""

import logging
from typing import Any, Dict, List, Optional

from .ports import (
    BackendType,
    BaseBackend,
    SqlBackend,
)

logger = logging.getLogger(__name__)


# ==================== 后端工厂函数 ====================


def _detect_backend_type(config: Any) -> BackendType:
    """从配置读取后端类型。要求显式声明 backend 字段。"""
    backend = getattr(config, "backend", None)
    if not backend:
        raise ValueError(
            f"配置 {type(config).__name__} 缺少 backend 字段，"
            "请显式设置为 mysql/postgresql/milvus/redis"
        )
    try:
        return BackendType(str(backend).lower())
    except ValueError as exc:
        allowed = ", ".join(bt.value for bt in BackendType)
        raise ValueError(
            f"配置 {type(config).__name__} 的 backend={backend!r} 非法，可选: {allowed}"
        ) from exc


def create_backend(name: str, config: Any) -> BaseBackend:
    """工厂函数：根据配置创建对应后端实例"""
    from .mysql.backend import MySQLBackend
    from .postgresql.backend import PgBackend
    from .milvus.backend import MilvusBackend
    from .redis.backend import RedisBackend

    bt = _detect_backend_type(config)
    if bt == BackendType.MYSQL:
        return MySQLBackend(name, config)
    elif bt == BackendType.POSTGRESQL:
        return PgBackend(name, config)
    elif bt == BackendType.MILVUS:
        return MilvusBackend(name, config)
    elif bt == BackendType.REDIS:
        return RedisBackend(name, config)
    raise ValueError(f"未知的存储后端类型: {bt}")


# ==================== 管理器 ====================


class DatabaseManager:
    """数据库连接管理器 —— 封装多后端实例的生命周期、路由、子门面工厂

    职责: 连接池初始化/关闭、后端直取、子门面创建
    不负责: 具体 SQL/KV 操作（操作在各子门面及后端实现类内部）

    使用方式:
        backends = {"aigc": MySQLBackend("aigc", config)}
        db = DatabaseManager(backends)
        await db.init_pools()

        # 通过子门面操作
        mysql = db.mysql_manager("long_articles")
        rows = await mysql.fetch("SELECT * FROM t")
    """

    def __init__(
        self,
        backends: Dict[str, BaseBackend],
        *,
        default_db: Optional[str] = None,
    ):
        if not backends:
            raise ValueError("DatabaseManager 至少需要一个后端")
        self._backends = backends
        if default_db is None:
            self._default_db = next(iter(backends))
        else:
            if default_db not in backends:
                raise ValueError(
                    f"default_db={default_db!r} 未注册，可用: {list(backends.keys())}"
                )
            self._default_db = default_db

    # ==================== 子门面工厂 ====================

    def mysql_manager(self, db_name: str = "long_articles"):
        """获取 MySQL 专用门面"""
        from .mysql.manager import MysqlManager
        return MysqlManager(self._get_sql(db_name))

    def pg_manager(self, db_name: str = "default"):
        """获取 PostgreSQL 专用门面"""
        from .postgresql.manager import PgManager
        return PgManager(self._get_sql(db_name))

    def redis_manager(self, db_name: str = "default"):
        """获取 Redis 专用门面"""
        from .redis.backend import RedisBackend
        from .redis.manager import RedisManager
        return RedisManager(self._get_typed(db_name, RedisBackend))

    # ==================== 生命周期 ====================

    async def init_pools(self) -> None:
        initialized: list[tuple[str, BaseBackend]] = []
        try:
            for name, backend in self._backends.items():
                await backend.init()
                initialized.append((name, backend))
                logger.info("存储后端 [%s] 初始化完成 → %s", name, type(backend).__name__)
        except Exception as exc:
            logger.exception("存储后端初始化失败，开始回滚已初始化连接: %s", exc)
            for name, backend in reversed(initialized):
                try:
                    await backend.close()
                    logger.info("存储后端 [%s] 回滚关闭完成", name)
                except Exception:
                    logger.exception("存储后端 [%s] 回滚关闭失败", name)
            raise

    async def close_pools(self) -> None:
        for name, backend in self._backends.items():
            await backend.close()
            logger.info("存储后端 [%s] 已关闭", name)

    # ==================== 后端直取 ====================

    def mysql(self, name: Optional[str] = None):
        from .mysql.backend import MySQLBackend
        return self._get_typed(name or self._default_db, MySQLBackend)

    def pg(self, name: str):
        from .postgresql.backend import PgBackend
        return self._get_typed(name, PgBackend)

    def milvus(self, name: str):
        from .milvus.backend import MilvusBackend
        return self._get_typed(name, MilvusBackend)

    def redis(self, name: str):
        from .redis.backend import RedisBackend
        return self._get_typed(name, RedisBackend)

    # ==================== 辅助 ====================

    def get_pool(self, db_name: str):
        """向后兼容：返回 MySQL 连接池对象（不建议新增调用）"""
        return self.mysql(db_name)._pool

    def list_databases(self) -> List[str]:
        return list(self._backends.keys())

    def list_backends(self) -> Dict[str, str]:
        return {n: type(b).__name__ for n, b in self._backends.items()}

    # ==================== 内部 ====================

    def _get_backend(self, db_name: Optional[str] = None) -> BaseBackend:
        name = db_name or self._default_db
        backend = self._backends.get(name)
        if backend is None:
            raise KeyError(f"存储后端 '{name}' 未注册，可用: {self.list_databases()}")
        return backend

    def _get_sql(self, db_name: Optional[str] = None) -> SqlBackend:
        backend = self._get_backend(db_name)
        if not isinstance(backend, SqlBackend):
            raise TypeError(
                f"后端 '{db_name}' 类型为 {type(backend).__name__}，不支持 SQL 操作"
            )
        return backend

    def _get_typed(self, name: str, expected_type: type) -> BaseBackend:
        backend = self._get_backend(name)
        if not isinstance(backend, expected_type):
            raise TypeError(
                f"后端 '{name}' 实际类型为 {type(backend).__name__}，"
                f"期望 {expected_type.__name__}"
            )
        return backend