long_articles_agentic_supply/supply/infra/database/manager.py

"""数据库门面管理器 —— 依赖反转，接收已创建的 Backend 实例

与旧版 DatabaseManager 的核心区别：
  - 不再通过 config dict 内部创建后端，而是接收 Dict[str, BaseBackend]
  - 后端创建逻辑剥离为 create_backend() 工厂函数，由调用方（DI 容器 / app 工厂）负责
  - Manager 只关心生命周期编排与统一访问路由
"""

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

from supply.infra.database.ports import (
    BackendType,
    BaseBackend,
    SqlBackend,
)
from supply.infra.database.mysql import MySQLBackend
from supply.infra.database.postgresql import PgBackend
from supply.infra.database.milvus import MilvusBackend
from supply.infra.database.redis import RedisBackend

logger = logging.getLogger(__name__)


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


def _detect_backend_type(config: Any) -> BackendType:
    """从配置推断后端类型：优先读取显式 backend 字段，否则按类名匹配"""
    if hasattr(config, "backend") and config.backend:
        try:
            return BackendType(config.backend)
        except ValueError:
            pass
    type_name = type(config).__name__.lower()
    for bt in BackendType:
        if bt.value in type_name:
            return bt
    return BackendType.MYSQL


def create_backend(name: str, config: Any) -> BaseBackend:
    """工厂函数：根据配置创建对应后端实例

    由调用方（app 工厂 / DI 容器）使用，创建实例后传给 DatabaseManager。
    """
    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:
    """多后端数据库门面

    接收已创建的后端实例，负责：
      - 统一生命周期管理 (init_pools / close_pools)
      - 向后兼容 SQL 接口
      - 类型安全的后端直取访问器

    使用方式：
        backends = {
            "aigc": MySQLBackend("aigc", AigcDatabaseConfig()),
        }
        db = DatabaseManager(backends)
        await db.init_pools()
        rows = await db.async_fetch("SELECT * FROM t")
    """

    def __init__(
        self,
        backends: Dict[str, BaseBackend],
        *,
        default_db: str = "default",
    ):
        if not backends:
            raise ValueError("DatabaseManager 至少需要一个后端")
        self._backends = backends
        self._default_db = default_db

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

    async def init_pools(self) -> None:
        for name, backend in self._backends.items():
            await backend.init()
            logger.info("存储后端 [%s] 初始化完成 → %s", name, type(backend).__name__)

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

    # ==================== 向后兼容 SQL 接口 ====================

    async def async_fetch(
        self,
        query: str,
        *,
        params: Optional[tuple] = None,
        db_name: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        backend = self._get_sql(db_name)
        return await backend.fetch(query, params=params)

    async def async_fetch_one(
        self,
        query: str,
        *,
        params: Optional[tuple] = None,
        db_name: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        backend = self._get_sql(db_name)
        return await backend.fetch_one(query, params=params)

    async def async_save(
        self,
        query: str,
        *,
        params: Any = None,
        batch: bool = False,
        db_name: Optional[str] = None,
    ) -> int:
        backend = self._get_sql(db_name)
        return await backend.execute(query, params=params, batch=batch)

    # ==================== 类型安全的后端直取 ====================

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

    def pg(self, name: str) -> PgBackend:
        return self._get_typed(name, PgBackend)

    def milvus(self, name: str) -> MilvusBackend:
        return self._get_typed(name, MilvusBackend)

    def redis(self, name: str) -> RedisBackend:
        return self._get_typed(name, RedisBackend)

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

    def get_pool(self, db_name: str):
        """向后兼容：返回 MySQL 连接池对象"""
        backend = self.mysql(db_name)
        return backend._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