初始化提交

mongodb_db.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+MongoDB 数据库操作模块
+支持基本的文档操作
+"""
+
+import sys
+import json
+
+try:
+    from pymongo import MongoClient
+except ImportError as e:
+    print(f"错误: 无法导入 pymongo 库，请确保已安装: pip install pymongo")
+    print(f"详细错误: {e}")
+    sys.exit(1)
+
+
+class MongoDatabase:
+    """MongoDB 数据库连接和操作类"""
+    
+    def __init__(self, host='localhost', port=27017, database='test', username=None, password=None):
+        """
+        初始化 MongoDB 连接参数
+        :param host: 数据库主机地址
+        :param port: 数据库端口
+        :param database: 数据库名
+        :param username: 用户名（可选）
+        :param password: 密码（可选）
+        """
+        self.host = host
+        self.port = port
+        self.database_name = database
+        self.username = username
+        self.password = password
+        self.client = None
+        self.db = None
+    
+    def connect(self):
+        """建立数据库连接"""
+        try:
+            # 构建连接参数
+            if self.username and self.password:
+                # 使用用户名密码认证
+                self.client = MongoClient(
+                    host=self.host,
+                    port=self.port,
+                    username=self.username,
+                    password=self.password,
+                    authSource=self.database_name,  # 认证数据库
+                    serverSelectionTimeoutMS=5000
+                )
+            else:
+                # 无需认证
+                self.client = MongoClient(
+                    host=self.host,
+                    port=self.port,
+                    serverSelectionTimeoutMS=5000
+                )
+            
+            # 测试连接
+            self.client.admin.command('ping')
+            self.db = self.client[self.database_name]
+            return True, f"MongoDB 连接成功 (数据库: {self.database_name})"
+        except Exception as e:
+            return False, f"MongoDB 连接失败: {str(e)}"
+    
+    def execute(self, command):
+        """
+        执行 MongoDB 命令
+        :param command: MongoDB 命令字符串（JSON 格式或 shell 命令）
+        示例: {"collection": "users", "operation": "find", "query": {}}
+        或: show dbs, show collections, use <database>
+        """
+        if self.db is None:
+            return False, "未连接到数据库"
+        
+        # 处理 MongoDB shell 命令
+        cmd_lower = command.strip().lower()
+        
+        # show dbs - 列出所有数据库
+        if cmd_lower in ['show dbs', 'show databases']:
+            try:
+                db_list = self.client.list_database_names()
+                result = []
+                for db_name in db_list:
+                    db_stats = self.client[db_name].command('dbstats')
+                    size_mb = db_stats.get('dataSize', 0) / (1024 * 1024)
+                    result.append(f"{db_name}: {size_mb:.2f} MB")
+                return True, result
+            except Exception as e:
+                return False, f"执行失败: {str(e)}"
+        
+        # show collections - 列出当前数据库的所有集合
+        elif cmd_lower in ['show collections', 'show tables']:
+            try:
+                collections = self.db.list_collection_names()
+                return True, collections
+            except Exception as e:
+                return False, f"执行失败: {str(e)}"
+        
+        # use <database> - 切换数据库
+        elif cmd_lower.startswith('use '):
+            try:
+                db_name = command.strip().split(maxsplit=1)[1]
+                self.db = self.client[db_name]
+                self.database_name = db_name
+                return True, f"已切换到数据库: {db_name}"
+            except Exception as e:
+                return False, f"切换数据库失败: {str(e)}"
+        
+        try:
+            # 解析命令
+            if isinstance(command, str):
+                cmd = json.loads(command)
+            else:
+                cmd = command
+            
+            collection_name = cmd.get('collection')
+            operation = cmd.get('operation')
+            
+            if not collection_name or not operation:
+                return False, "命令必须包含 collection 和 operation 字段"
+            
+            collection = self.db[collection_name]
+            
+            # 根据操作类型执行不同的命令
+            if operation == 'find':
+                query = cmd.get('query', {})
+                limit = cmd.get('limit', 0)
+                result = list(collection.find(query).limit(limit))
+                # 转换 ObjectId 为字符串
+                for doc in result:
+                    if '_id' in doc:
+                        doc['_id'] = str(doc['_id'])
+                return True, result
+            
+            elif operation == 'insert':
+                document = cmd.get('document')
+                if not document:
+                    return False, "插入操作需要 document 字段"
+                result = collection.insert_one(document)
+                return True, f"插入成功, ID: {result.inserted_id}"
+            
+            elif operation == 'update':
+                query = cmd.get('query', {})
+                update = cmd.get('update')
+                if not update:
+                    return False, "更新操作需要 update 字段"
+                result = collection.update_many(query, update)
+                return True, f"更新了 {result.modified_count} 条记录"
+            
+            elif operation == 'delete':
+                query = cmd.get('query', {})
+                result = collection.delete_many(query)
+                return True, f"删除了 {result.deleted_count} 条记录"
+            
+            elif operation == 'count':
+                query = cmd.get('query', {})
+                count = collection.count_documents(query)
+                return True, f"共有 {count} 条记录"
+            
+            else:
+                return False, f"不支持的操作: {operation}"
+                
+        except json.JSONDecodeError:
+            return False, "命令格式错误，请使用 JSON 格式"
+        except Exception as e:
+            return False, f"执行失败: {str(e)}"
+    
+    def close(self):
+        """关闭数据库连接"""
+        if self.client:
+            self.client.close()
+            self.client = None
+            self.db = None
+            return "MongoDB 连接已关闭"
+        return "连接已经关闭"