""" MongoDB Tool - Quản lý MongoDB qua pymongo. Author: Hieu Louis (2026) Operations: find, insert, update, delete, aggregate, create_index. """ from __future__ import annotations import json from datetime import datetime, date from typing import Any, Dict, List, Optional from .base import Tool, ToolResult, ToolContext, ToolCategory, ToolSafety # Operations ghi/xoá // write ops WRITE_OPS = {"insert", "update", "delete", "create_index"} def _json_default(o: Any) -> Any: """Serializer cho BSON ObjectId / datetime // BSON fallback.""" try: from bson import ObjectId # type: ignore if isinstance(o, ObjectId): return str(o) except Exception: pass if isinstance(o, (datetime, date)): return o.isoformat() if isinstance(o, bytes): try: return o.decode("utf-8") except Exception: return o.hex() return str(o) class MongoTool(Tool): """Quản lý MongoDB: find, insert, update, delete, aggregate, create_index.""" category = ToolCategory.DATABASE safety = ToolSafety.DANGEROUS requires_confirmation = True @property def name(self) -> str: return "mongo" @property def description(self) -> str: return ( "Quản lý MongoDB qua pymongo: find, insert (one/many), update (one/many), " "delete (one/many), aggregate, create_index." ) @property def parameters(self) -> Dict[str, Any]: return { "type": "object", "properties": { "connection_string": {"type": "string", "description": "MongoDB URL (mongodb://host:27017 hoặc mongodb+srv://...)"}, "database": {"type": "string", "description": "Tên database"}, "collection": {"type": "string", "description": "Tên collection"}, "operation": { "type": "string", "enum": ["find", "insert", "update", "delete", "aggregate", "create_index"], "description": "Mongo operation", }, "filter": {"type": "object", "description": "Query filter (find/update/delete)"}, "doc": {"type": "object", "description": "Document (insert one)"}, "docs": {"type": "array", "items": {"type": "object"}, "description": "Documents (insert many)"}, "update": {"type": "object", "description": "Update spec (vd: {$set: {..}})"}, "pipeline": {"type": "array", "items": {}, "description": "Aggregation pipeline (aggregate)"}, "index": { "oneOf": [{"type": "string"}, {"type": "object"}], "description": "Index spec (create_index, vd: {\"email\": 1})", }, "many": {"type": "boolean", "description": "Update/Delete nhiều (default false)"}, "limit": {"type": "integer", "description": "Giới hạn kết quả (default 100)"}, }, "required": ["connection_string", "database", "collection", "operation"], } def validate_args(self, args: Dict[str, Any]) -> Optional[str]: for k in ("connection_string", "database", "collection", "operation"): if not args.get(k): return f"Missing required arg: {k}" op = args["operation"] if op == "insert" and not (args.get("doc") or args.get("docs")): return "Operation 'insert' requires 'doc' or 'docs' arg" if op == "update" and not args.get("update"): return "Operation 'update' requires 'update' arg" if op == "aggregate" and not args.get("pipeline"): return "Operation 'aggregate' requires 'pipeline' arg" if op == "create_index" and not args.get("index"): return "Operation 'create_index' requires 'index' arg" return None def execute(self, args: Dict[str, Any], context: ToolContext) -> ToolResult: cs: str = args["connection_string"] db_name: str = args["database"] coll_name: str = args["collection"] op: str = args["operation"] # Lazy import pymongo // lazy import try: from pymongo import MongoClient, ASCENDING # type: ignore except ImportError as e: return ToolResult( success=False, error=f"pymongo not installed: {e}. Cài: pip install pymongo", return_code=127, ) # Dry-run cho write ops // dry-run if context.dry_run and op in WRITE_OPS: return ToolResult( success=True, output=f"[dry-run] Would run {op} on {db_name}.{coll_name}", metadata={"dry_run": True, "operation": op, "database": db_name, "collection": coll_name}, ) try: client = MongoClient(cs, serverSelectionTimeoutMS=context.timeout * 1000) db = client[db_name] coll = db[coll_name] if op == "find": flt = args.get("filter") or {} limit = int(args.get("limit") or 100) docs = list(coll.find(flt).limit(limit)) payload = json.dumps(docs, default=_json_default, ensure_ascii=False, indent=2) return ToolResult(success=True, output=payload, metadata={"count": len(docs)}) if op == "insert": if args.get("docs"): result = coll.insert_many(args["docs"]) return ToolResult( success=True, output=json.dumps({"inserted": len(result.inserted_ids)}, default=_json_default), metadata={"inserted_ids": [str(i) for i in result.inserted_ids]}, ) result = coll.insert_one(args["doc"]) return ToolResult( success=True, output=json.dumps({"inserted_id": str(result.inserted_id)}, default=_json_default), metadata={"inserted_id": str(result.inserted_id)}, ) if op == "update": flt = args.get("filter") or {} upd = args["update"] if args.get("many"): result = coll.update_many(flt, upd) else: result = coll.update_one(flt, upd) return ToolResult( success=True, output=json.dumps({"matched": result.matched_count, "modified": result.modified_count}), metadata={"matched": result.matched_count, "modified": result.modified_count}, ) if op == "delete": flt = args.get("filter") or {} if args.get("many"): result = coll.delete_many(flt) else: result = coll.delete_one(flt) return ToolResult( success=True, output=json.dumps({"deleted": result.deleted_count}), metadata={"deleted": result.deleted_count}, ) if op == "aggregate": pipeline = args["pipeline"] docs = list(coll.aggregate(pipeline)) payload = json.dumps(docs, default=_json_default, ensure_ascii=False, indent=2) return ToolResult(success=True, output=payload, metadata={"count": len(docs)}) if op == "create_index": spec = args["index"] if isinstance(spec, str): spec = [(spec, ASCENDING)] name = coll.create_index(spec) return ToolResult( success=True, output=f"Index created: {name}", metadata={"index_name": name}, ) return ToolResult(success=False, error=f"Unknown operation: {op}", return_code=1) except Exception as e: return ToolResult(success=False, error=str(e), return_code=1) finally: try: client.close() except Exception: pass