MongoDB

类

from pymongo import MongoClient
from bson import ObjectId
from typing import Union, List

class MongoDBHelper:
    def __init__(self, host: str, port: int):
        """
        初始化 MongoDBHelper 类

        Args:
            host (str): MongoDB 主机名
            port (int): MongoDB 端口号
        """
        self.host = host
        self.port = port
        self.client = MongoClient(host, port)

    def get_database(self, db_name: str):
        """
        获取指定名称的数据库对象

        Args:
            db_name (str): 数据库名称

        Returns:
            Database: 数据库对象
        """
        return self.client[db_name]

    def get_collection(self, db_name: str, collection_name: str):
        """
        获取指定数据库中的集合对象

        Args:
            db_name (str): 数据库名称
            collection_name (str): 集合名称

        Returns:
            Collection: 集合对象
        """
        db = self.get_database(db_name)
        return db[collection_name]

    def insert_document(self, db_name: str, collection_name: str, document: dict) -> Union[str, None]:
        """
        在指定集合中插入文档

        Args:
            db_name (str): 数据库名称
            collection_name (str): 集合名称
            document (dict): 要插入的文档

        Returns:
            Union[str, None]: 插入文档的 ObjectID，如果插入失败则返回 None
        """
        collection = self.get_collection(db_name, collection_name)
        result = collection.insert_one(document)
        return str(result.inserted_id) if result.acknowledged else None

    def find_documents(self, db_name: str, collection_name: str, query: dict = None) -> List[dict]:
        """
        在指定集合中查询文档

        Args:
            db_name (str): 数据库名称
            collection_name (str): 集合名称
            query (dict, optional): 查询条件，默认为 None

        Returns:
            List[dict]: 查询到的文档列表
        """
        collection = self.get_collection(db_name, collection_name)
        cursor = collection.find(query) if query else collection.find()
        return list(cursor)

    def find_document_by_id(self, db_name: str, collection_name: str, document_id: str) -> Union[dict, None]:
        """
        根据文档 ID 在指定集合中查询文档

        Args:
            db_name (str): 数据库名称
            collection_name (str): 集合名称
            document_id (str): 文档 ID

        Returns:
            Union[dict, None]: 查询到的文档，如果未找到则返回 None
        """
        collection = self.get_collection(db_name, collection_name)
        result = collection.find_one({"_id": ObjectId(document_id)})
        return result if result else None

    def update_document(self, db_name: str, collection_name: str, document_id: str, update: dict) -> bool:
        """
        根据文档 ID 更新指定集合中的文档

        Args:
            db_name (str): 数据库名称
            collection_name (str): 集合名称
            document_id (str): 文档 ID
            update (dict): 更新内容

        Returns:
            bool: 更新是否成功
        """
        collection = self.get_collection(db_name, collection_name)
        result = collection.update_one({"_id": ObjectId(document_id)}, {"$set": update})
        return result.acknowledged

    def delete_document(self, db_name: str, collection_name: str, document_id: str) -> bool:
        """
        根据文档 ID 删除指定集合中的文档

        Args:
            db_name (str): 数据库名称
            collection_name (str): 集合名称
            document_id (str): 文档 ID

        Returns:
            bool: 删除是否成功
        """
        collection = self.get_collection(db_name, collection_name)
        result = collection.delete_one({"_id": ObjectId(document_id)})
        return result.acknowledged

使用示例

# 使用示例
# 初始化 MongoDBHelper
mongo_helper = MongoDBHelper("localhost", 27017)

# 插入文档
document = {
    "name": "John Doe",
    "age": 30,
    "email": "johndoe@example.com"
}
document_id = mongo_helper.insert_document("mydb", "mycollection", document)
print(f"Inserted document ID: {document_id}")

# 查询所有文档
documents = mongo_helper.find_documents("mydb", "mycollection")
print("All documents:")
for doc in documents:
    print(doc)

# 根据文档 ID 查询文档
document_id = "615af6e6799e9f9dfc4c0e7c"  # Assuming a valid document ID
document = mongo_helper.find_document_by_id("mydb", "mycollection", document_id)
if document:
    print(f"Found document with ID {document_id}:")
    print(document)
else:
    print(f"No document found with ID {document_id}")

# 更新文档
document_id = "615af6e6799e9f9dfc4c0e7c"  # Assuming a valid document ID
update = {
    "age": 35
}
if mongo_helper.update_document("mydb", "mycollection", document_id, update):
    print("Document updated successfully")
else:
    print("Failed to update document")

# 删除文档
document_id = "615af6e6799e9f9dfc4c0e7c"  # Assuming a valid document ID
if mongo_helper.delete_document("mydb", "mycollection", document_id):
    print("Document deleted successfully")
else:
    print("Failed to delete document")