"""MongoDB 异步连接器 — 使用 motor 驱动。 支持数据库变更工单的全生命周期: - 连接管理（支持认证、TLS） - 只读查询: find, findOne, countDocuments, distinct, aggregate - 变更操作: insertOne/Many, updateOne/Many, replaceOne, deleteOne/Many - 连接测试、库/collection 列表 - 安全验证：只读操作白名单 查询格式（JSON）: { "collection": "users", "operation": "find", "filter": {"age": {"$gt": 18}}, "projection": {"name": 1, "age": 1}, "sort": {"age": -1}, "limit": 100, "skip": 0 } 变更格式（JSON）: { "collection": "users", "operation": "insertMany", "documents": [{"name": "Alice", "age": 25}] } """ import json import time from typing import Any, Dict, List, Optional, Tuple from bson import json_util from motor.motor_asyncio import AsyncIOMotorClient from app.core.crypto import decrypt_password # ── 常量定义 ───────────────────────────────────────────────────── MONGODB_READONLY_OPS: set[str] = { "find", "findOne", "countDocuments", "estimatedDocumentCount", "distinct", "aggregate", } MONGODB_WRITE_OPS: set[str] = { "insertOne", "insertMany", "updateOne", "updateMany", "replaceOne", "deleteOne", "deleteMany", "drop", "createIndex", "dropIndex", "renameCollection", "findOneAndUpdate", "findOneAndReplace", "findOneAndDelete", "bulkWrite", } MAX_QUERY_LIMIT = 10000 # ── 安全验证 ───────────────────────────────────────────────────── def validate_operation(operation: str) -> Tuple[bool, str]: """验证操作是否为只读。 Args: operation: MongoDB 操作名称。 Returns: (is_readonly, error_message): 是否只读，若否附带错误原因。 """ if operation in MONGODB_READONLY_OPS: return True, "" if operation in MONGODB_WRITE_OPS: return False, f"MongoDB 操作 '{operation}' 是写操作，查询模式下禁止" return False, f"未知的 MongoDB 操作: {operation}" # ── 客户端缓存 ─────────────────────────────────────────────────── _mongo_clients: Dict[int, AsyncIOMotorClient] = {} class ConnectionResult: """MongoDB 连接测试结果。""" def __init__( self, success: bool, version: str = "", latency_ms: int = 0, message: str = "", ): self.success = success self.version = version self.latency_ms = latency_ms self.message = message def to_dict(self) -> dict: return { "success": self.success, "version": self.version, "latency_ms": self.latency_ms, "message": self.message, } class MongoConnector: """MongoDB 异步连接器。 提供连接池管理、只读查询、变更操作、元数据获取及安全验证。 """ # ── 连接管理 ──────────────────────────────────────────────── @staticmethod async def get_client(instance: Any) -> AsyncIOMotorClient: """获取（或创建）motor 客户端实例。 使用内部缓存避免重复创建连接池。密码通过 decrypt_password 解密。 Args: instance: Instance 模型对象，需包含 host, port, username, password 等字段。 Returns: AsyncIOMotorClient: 异步 MongoDB 客户端。 """ instance_id: int = int(instance.id) if instance_id in _mongo_clients: return _mongo_clients[instance_id] password = decrypt_password(instance.password) conn_str = MongoConnector._build_connection_string( host=instance.host, port=instance.port, username=instance.username or None, password=password or None, database=getattr(instance, "database", None), ) client = AsyncIOMotorClient(conn_str, serverSelectionTimeoutMS=5000) _mongo_clients[instance_id] = client return client @staticmethod def _build_connection_string( host: str, port: int, username: Optional[str] = None, password: Optional[str] = None, database: Optional[str] = None, ) -> str: """构建 MongoDB 连接字符串。 Args: host: 主机地址。 port: 端口号。 username: 用户名（可选）。 password: 密码（可选）。 database: 认证数据库（可选）。 Returns: MongoDB URI 字符串。 """ if username and password: auth_db = database or "admin" return ( f"mongodb://{username}:{password}@{host}:{port}/{auth_db}" f"?authSource={auth_db}" ) return f"mongodb://{host}:{port}" @staticmethod async def test_connection(instance: Any) -> ConnectionResult: """测试 MongoDB 连接并返回版本与延迟。 Args: instance: Instance 模型对象。 Returns: ConnectionResult: 包含 success, version, latency_ms, message。 """ start = time.time() client: AsyncIOMotorClient = await MongoConnector.get_client(instance) try: build_info = await client.admin.command("buildInfo") elapsed = int((time.time() - start) * 1000) version = build_info.get("version", "unknown") return ConnectionResult( success=True, version=version, latency_ms=elapsed, message="连接成功" ) except Exception as exc: elapsed = int((time.time() - start) * 1000) return ConnectionResult( success=False, version="", latency_ms=elapsed, message=f"连接失败: {exc}", ) @staticmethod async def close_client(instance_id: int) -> None: """关闭指定实例的 MongoDB 客户端并从缓存中移除。 Args: instance_id: 数据库实例 ID。 """ client = _mongo_clients.pop(instance_id, None) if client: client.close() # ── 只读查询 ──────────────────────────────────────────────── @staticmethod async def execute_find( instance: Any, database: str, collection: str, filter_doc: Optional[Dict[str, Any]] = None, projection: Optional[Dict[str, Any]] = None, sort: Optional[List[tuple]] = None, limit: int = 1000, skip: int = 0, ) -> List[Dict[str, Any]]: """执行 find 查询。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 filter_doc: 查询过滤器（MongoDB 查询文档）。 projection: 字段投影。 sort: 排序规则，可以是字典 {"field": 1/-1} 或列表 [("field", 1/-1)]。 limit: 最大返回文档数（上限 MAX_QUERY_LIMIT）。 skip: 跳过的文档数。 Returns: 文档列表，每个文档为字典。_id 被序列化为字符串。 Raises: ValueError: 当 limit 超过 MAX_QUERY_LIMIT 时。 """ if limit > MAX_QUERY_LIMIT: raise ValueError(f"limit 不能超过 {MAX_QUERY_LIMIT}") client = await MongoConnector.get_client(instance) coll = client[database][collection] cursor = coll.find(filter_doc or {}, projection=projection) if sort: if isinstance(sort, dict): sort_list = [(k, v) for k, v in sort.items()] cursor = cursor.sort(sort_list) else: cursor = cursor.sort(sort) if skip: cursor = cursor.skip(skip) if limit: cursor = cursor.limit(limit) docs = await cursor.to_list(length=limit or MAX_QUERY_LIMIT) return [MongoConnector._serialize_doc(doc) for doc in docs] @staticmethod async def execute_find_one( instance: Any, database: str, collection: str, filter_doc: Optional[Dict[str, Any]] = None, projection: Optional[Dict[str, Any]] = None, ) -> Optional[Dict[str, Any]]: """执行 findOne 查询，返回单个文档或 None。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 filter_doc: 查询过滤器。 projection: 字段投影。 Returns: 单个文档字典，或 None。_id 被序列化为字符串。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] doc = await coll.find_one(filter_doc or {}, projection=projection) if doc is None: return None return MongoConnector._serialize_doc(doc) @staticmethod async def execute_count( instance: Any, database: str, collection: str, filter_doc: Optional[Dict[str, Any]] = None, ) -> int: """执行 countDocuments 计数。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 filter_doc: 查询过滤器。 Returns: 符合条件的文档数量。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] return await coll.count_documents(filter_doc or {}) @staticmethod async def execute_distinct( instance: Any, database: str, collection: str, field: str, filter_doc: Optional[Dict[str, Any]] = None, ) -> List[Any]: """获取指定字段的去重值列表。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 field: 字段名。 filter_doc: 前置过滤器。 Returns: 去重后的值列表。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] return await coll.distinct(field, filter_doc or {}) @staticmethod async def execute_aggregate( instance: Any, database: str, collection: str, pipeline: List[Dict[str, Any]], ) -> List[Dict[str, Any]]: """执行聚合管道查询。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 pipeline: 聚合管道阶段列表。 Returns: 聚合结果文档列表。_id 被序列化为字符串。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] cursor = coll.aggregate(pipeline, allowDiskUse=True) docs = await cursor.to_list(length=MAX_QUERY_LIMIT) return [MongoConnector._serialize_doc(doc) for doc in docs] # ── 变更操作（工单执行用）─────────────────────────────────── @staticmethod async def execute_insert_one( instance: Any, database: str, collection: str, document: Dict[str, Any], ) -> str: """插入单个文档。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 document: 待插入的文档。 Returns: 插入文档的 ID 字符串。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] result = await coll.insert_one(document) return str(result.inserted_id) @staticmethod async def execute_insert_many( instance: Any, database: str, collection: str, documents: List[Dict[str, Any]], ) -> str: """批量插入文档。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 documents: 待插入的文档列表。 Returns: 插入数量的描述字符串。 Raises: ValueError: 文档列表为空时。 """ if not documents: raise ValueError("documents 不能为空列表") client = await MongoConnector.get_client(instance) coll = client[database][collection] result = await coll.insert_many(documents) return f"成功插入 {len(result.inserted_ids)} 条文档" @staticmethod async def execute_update_one( instance: Any, database: str, collection: str, filter_doc: Dict[str, Any], update_doc: Dict[str, Any], ) -> str: """更新单个文档。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 filter_doc: 查询过滤器。 update_doc: 更新操作文档（如 {"$set": {...}}）。 Returns: 影响行数描述字符串。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] result = await coll.update_one(filter_doc, update_doc) return f"匹配 {result.matched_count} 条，修改 {result.modified_count} 条" @staticmethod async def execute_update_many( instance: Any, database: str, collection: str, filter_doc: Dict[str, Any], update_doc: Dict[str, Any], ) -> str: """批量更新文档。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 filter_doc: 查询过滤器。 update_doc: 更新操作文档。 Returns: 影响行数描述字符串。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] result = await coll.update_many(filter_doc, update_doc) return f"匹配 {result.matched_count} 条，修改 {result.modified_count} 条" @staticmethod async def execute_replace_one( instance: Any, database: str, collection: str, filter_doc: Dict[str, Any], replacement: Dict[str, Any], ) -> str: """替换单个文档。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 filter_doc: 查询过滤器。 replacement: 替换用的新文档。 Returns: 影响行数描述字符串。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] result = await coll.replace_one(filter_doc, replacement) return f"匹配 {result.matched_count} 条，替换 {result.modified_count} 条" @staticmethod async def execute_delete_one( instance: Any, database: str, collection: str, filter_doc: Dict[str, Any], ) -> str: """删除单个文档。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 filter_doc: 查询过滤器。 Returns: 删除数量描述字符串。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] result = await coll.delete_one(filter_doc) return f"删除 {result.deleted_count} 条文档" @staticmethod async def execute_delete_many( instance: Any, database: str, collection: str, filter_doc: Dict[str, Any], ) -> str: """批量删除文档。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 filter_doc: 查询过滤器。 Returns: 删除数量描述字符串。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] result = await coll.delete_many(filter_doc) return f"删除 {result.deleted_count} 条文档" # ── 元数据 ────────────────────────────────────────────────── @staticmethod async def list_databases(instance: Any) -> List[str]: """列出 MongoDB 实例中所有数据库名称。 Args: instance: Instance 模型对象。 Returns: 数据库名称列表（排除系统库 admin, local, config）。 """ client = await MongoConnector.get_client(instance) db_info_list = await client.list_database_names() system_dbs = {"admin", "local", "config"} return [name for name in db_info_list if name not in system_dbs] @staticmethod async def list_collections(instance: Any, database: str) -> List[str]: """列出指定数据库中所有集合名称。 Args: instance: Instance 模型对象。 database: 数据库名称。 Returns: 集合名称列表。 """ client = await MongoConnector.get_client(instance) db = client[database] return await db.list_collection_names() @staticmethod async def get_collection_stats( instance: Any, database: str, collection: str ) -> Dict[str, Any]: """获取集合统计信息。 Args: instance: Instance 模型对象。 database: 数据库名称。 collection: 集合名称。 Returns: 包含 count, size, avgObjSize, storageSize, indexes, totalIndexSize 等字段的字典。 """ client = await MongoConnector.get_client(instance) coll = client[database][collection] stats = await coll.database.command("collStats", collection) return { "namespace": stats.get("ns", f"{database}.{collection}"), "document_count": stats.get("count", 0), "size_bytes": stats.get("size", 0), "avg_document_size_bytes": stats.get("avgObjSize", 0), "storage_size_bytes": stats.get("storageSize", 0), "index_count": stats.get("nindexes", 0), "index_size_bytes": stats.get("totalIndexSize", 0), "capped": stats.get("capped", False), } # ── 通用查询执行 ──────────────────────────────────────────── @staticmethod async def execute_query_json( instance: Any, database: str, query_json: Dict[str, Any], ) -> Any: """通过 JSON 格式执行 MongoDB 查询。 自动根据 query_json["operation"] 分发到对应方法。 仅支持只读操作，写操作会被拒绝。 Args: instance: Instance 模型对象。 database: 数据库名称。 query_json: 查询 JSON 字典，必须包含 collection 和 operation 字段。 Returns: 查询结果，类型取决于具体操作。 Raises: ValueError: 操作不合法或参数缺失时。 """ collection: str = query_json.get("collection") operation: str = query_json.get("operation") if not collection: raise ValueError("查询 JSON 缺少 'collection' 字段") if not operation: raise ValueError("查询 JSON 缺少 'operation' 字段") is_readonly, error_msg = validate_operation(operation) if not is_readonly: raise ValueError(error_msg) filter_doc = query_json.get("filter") or query_json.get("query") or {} projection = query_json.get("projection") sort = query_json.get("sort") limit = query_json.get("limit", 1000) skip = query_json.get("skip", 0) pipeline = query_json.get("pipeline", []) field = query_json.get("field", "") if operation == "find": return await MongoConnector.execute_find( instance, database, collection, filter_doc, projection, sort, limit, skip ) if operation == "findOne": return await MongoConnector.execute_find_one( instance, database, collection, filter_doc, projection ) if operation == "countDocuments": return await MongoConnector.execute_count( instance, database, collection, filter_doc ) if operation == "estimatedDocumentCount": client = await MongoConnector.get_client(instance) coll = client[database][collection] return await coll.estimated_document_count() if operation == "distinct": if not field: raise ValueError("distinct 操作需要提供 'field' 字段") return await MongoConnector.execute_distinct( instance, database, collection, field, filter_doc ) if operation == "aggregate": if not pipeline: raise ValueError("aggregate 操作需要提供 'pipeline' 字段") return await MongoConnector.execute_aggregate( instance, database, collection, pipeline ) raise ValueError(f"不支持的只读操作: {operation}") @staticmethod async def execute_change_json( instance: Any, database: str, change_json: Dict[str, Any], ) -> str: """通过 JSON 格式执行 MongoDB 变更操作。 自动根据 change_json["operation"] 分发到对应方法。 用于工单变更执行。 Args: instance: Instance 模型对象。 database: 数据库名称。 change_json: 变更 JSON 字典，必须包含 collection 和 operation 字段。 Returns: 操作结果描述字符串。 Raises: ValueError: 操作不合法或参数缺失时。 """ collection: str = change_json.get("collection") operation: str = change_json.get("operation") if not collection: raise ValueError("变更 JSON 缺少 'collection' 字段") if not operation: raise ValueError("变更 JSON 缺少 'operation' 字段") if operation not in MONGODB_WRITE_OPS: raise ValueError(f"不支持的变更操作: {operation}") filter_doc = change_json.get("filter") or change_json.get("query") or {} if operation == "insertOne": document = change_json.get("document") if document is None: raise ValueError("insertOne 需要提供 'document' 字段") return await MongoConnector.execute_insert_one( instance, database, collection, document ) if operation == "insertMany": documents = change_json.get("documents") if not documents: raise ValueError("insertMany 需要提供 'documents' 字段") return await MongoConnector.execute_insert_many( instance, database, collection, documents ) if operation == "updateOne": update_doc = change_json.get("update") if update_doc is None: raise ValueError("updateOne 需要提供 'update' 字段") return await MongoConnector.execute_update_one( instance, database, collection, filter_doc, update_doc ) if operation == "updateMany": update_doc = change_json.get("update") if update_doc is None: raise ValueError("updateMany 需要提供 'update' 字段") return await MongoConnector.execute_update_many( instance, database, collection, filter_doc, update_doc ) if operation == "replaceOne": replacement = change_json.get("replacement") if replacement is None: raise ValueError("replaceOne 需要提供 'replacement' 字段") return await MongoConnector.execute_replace_one( instance, database, collection, filter_doc, replacement ) if operation == "deleteOne": return await MongoConnector.execute_delete_one( instance, database, collection, filter_doc ) if operation == "deleteMany": return await MongoConnector.execute_delete_many( instance, database, collection, filter_doc ) raise ValueError(f"暂未实现的变更操作: {operation}") # ── 序列化工具 ────────────────────────────────────────────── @staticmethod def _serialize_doc(doc: Dict[str, Any]) -> Dict[str, Any]: """将 MongoDB 文档序列化为 JSON 可序列化字典。 使用 bson.json_util 处理 ObjectId、DateTime 等特殊类型。 Args: doc: 原始 MongoDB 文档。 Returns: JSON 可序列化的字典。 """ return json.loads(json_util.dumps(doc))