【Python】封装一个操作MongoDB的工具类

封装操作MongoDB

from pymongo.errors import ConnectionFailure
from typing import Optional, Dict, List, Any
from pymongo.database import Database
from pymongo import MongoClient
from loguru import logger


class MongoDBManager:
    """
    MongoDB管理
    """

    def __init__(self, uri: str, db_name: str):
        """
        初始化连接信息
        :param uri: 连接URI
        :param db_name: 数据库名
        """
        self.uri = uri
        self.db_name = db_name
        self.client: Optional[MongoClient] = None
        self.db: Optional[Database] = None

    def connect(self):
        """
        连接数据库
        :return:
        """
        try:
            if self.client is None:
                self.client = MongoClient(self.uri)
                self.db = self.client[self.db_name]
                # 进行ping测试以确保连接成功
                self.db.command('ping')
                logger.info('MongoDB Connection Successful.')
                return True
        except ConnectionFailure:
            logger.info('MongoDB Connection Failed.')
            self.client = None
            self.db = None
            return False
        except Exception as e:
            logger.error(f"Unexpected error: {e}")
            self.client = None
            self.db = None
        return False

    def close(self):
        """
        关闭数据库连接
        :return:
        """
        try:
            self.client.close()
            logger.info('MongoDB Connection Close Successful.')
        except ConnectionFailure:
            logger.info('MongoDB Connection Failed.')

    def insert_document(self, collection: str, document: dict):
        """
        插入单个文档
        :param collection: 集合名
        :param document: 文档
        :return: MongoDB Object ID
        """
        result = self.db[collection].insert_one(document)
        return result.inserted_id

    def insert_many_document(self, collection: str, documents: list):
        """
        插入多个文档
        :param collection: 集合名
        :param documents: 文档列表
        :return: MongoDB Object ID List
        """
        result = self.db[collection].insert_many(documents)
        return result.inserted_ids

    def find_document(self, collection: str, query: dict):
        """
        查找单个文档
        :param collection: 集合名
        :param query: 查找条件
        :return: 返回文档
        """
        result = self.db[collection].find_one(query)
        return result

    def find_many_document(self, collection: str, query: dict):
        """
        查找多个文档
        :param collection: 集合名
        :param query: 查找条件
        :return: 返回文档列表
        """
        cursor = self.db[collection].find(query)
        documents = list(cursor)  # 将cursor对象转换为列表
        return documents

    def find_paginate(self, collection: str, page: int = 1, page_size: int = 10,
                      sort: Optional[Dict[str, int]] = None) -> List[Dict[str, Any]]:
        """
        分页查询文档
        :param collection: 集合名
        :param page: 起始页
        :param page_size: 分页大小
        :param sort: 排序条件
        :return: 文档列表
        """
        skip = (page - 1) * page_size  # 计算跳过的文档数量
        try:
            # 执行查询
            cursor = self.db[collection].find().skip(skip).limit(page_size)
            if sort:  # 应用排序
                cursor = cursor.sort(sort)
            result = list(cursor)
            return result
        except Exception as e:
            print(f"分页查询异常: {e}")
            return []

    def update_document(self, collection: str, query: dict, new_document: dict, upsert: bool = False):
        """
        更新文档
        :param collection: 集合名
        :param query: 查找条件
        :param new_document: 新文档
        :param upsert: 查询条件未匹配到任何现有文档时，是否应该插入一个新的文档
        :return: modified_count
        """
        result = self.db[collection].update_one(query, {"$set": new_document}, upsert=upsert)
        return result.modified_count

    def update_many_document(self, collection: str, query: dict, new_documents: list, upsert: bool = False):
        """
        更新多个文档
        :param collection: 集合名
        :param query: 查找条件
        :param new_documents: 新文档
        :param upsert: 查询条件未匹配到任何现有文档时，是否应该插入一个新的文档
        :return: modified_count
        """
        result = self.db[collection].update_many(query, {"$set": new_documents}, upsert=upsert)
        return result.modified_count

    def delete_document(self, collection: str, query: dict):
        """
        删除文档
        :param collection: 集合名
        :param query: 查找条件
        :return: deleted_count
        """
        result = self.db[collection].delete_one(query)
        return result.deleted_count

    def delete_many_document(self, collection: str, query: dict):
        """
        删除多个文档
        :param collection: 集合
        :param query: 查找条件
        :return: deleted_count
        """
        result = self.db[collection].delete_many(query)
        return result.deleted_count


if __name__ == '__main__':
    from config import Config
    from urllib.parse import quote_plus
    
    db_host = '192.168.11.174'
    db_port = 27017
    db_name = 'uaa'
    user = quote_plus('root')
    password = quote_plus('mqjc@123')
	uri = f"mongodb://{user}:{password}@{db_host}:{db_port}?authSource=admin"
	
	# 创建连接对象 
    db = MongoDBManager(uri, db_name)

    db.connect()
    db.close()