"""数据库管理器 —— 连接封装 + 生命周期 + 后端路由 不包含具体数据操作逻辑(操作在各子包 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 = "vector"): """获取 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