179 lines
6.5 KiB
Python
179 lines
6.5 KiB
Python
#!/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 "连接已经关闭"
|