python操作mongodb底层封装并验证pymongo是否应该关闭连接

一、pymongo简介

github地址：https://github.com/mongodb/mongo-python-driver

mymongo安装命令：pip install pymongo==4.7.2

mymongo接口文档：PyMongo 4.7.2 Documentation

PyMongo发行版包含Python与MongoDB数据库交互的工具。bson包是用于Python的bson格式的实现。pymongo包是MongoDB的原生Python驱动程序。gridfs包是pymongo之上的一个gridfs实现。pymongo是python与mongodb交互的首选方式。

二、封装接口

import uuid
from datetime import datetime, timezone
from pydantic import BaseModel,ConfigDict
from pycommon.utils.mongo.mongo_collection_pool import mongo_collection_pool
from pymongo.cursor import Cursor
from pymongo.collection import Collectionclass BaseDbService:def __init__(self, mongo_settings: MongoSettings, encryptor_key: str) -> None:self.mongo_settings = mongo_settingsself.encryptor_key = encryptor_keydef create(self, data: T, user_id: str, tenant_id: str = None, collection_name: str = None) -> str:collection: Collection = mongo_collection_pool.get_collection(self.mongo_settings, collection_name, tenant_id)if not get_attr_value(data, 'id'):set_attr_value(data, 'id', f"{uuid.uuid1()}")set_attr_value(data, 'created_by', user_id)set_attr_value(data, 'created_time', datetime.now(timezone.utc))request = data.model_dump(by_alias = True)id = collection.insert_one(request).inserted_idlogger.log_info(f'{type(data).__name__} id {id} is created')return iddef create_list(self, datas: list[T], user_id: str, tenant_id: str = None, collection_name: str = None) -> list[str]:collection = mongo_collection_pool.get_collection(self.mongo_settings, collection_name, tenant_id)request_list = []for data in datas:if not get_attr_value(data, 'id'):set_attr_value(data, 'id', f"{uuid.uuid1()}")set_attr_value(data, 'created_by', user_id)set_attr_value(data, 'created_time', datetime.now(timezone.utc))request = data.model_dump(by_alias = True)request_list.append(request)ids = collection.insert_many(request_list).inserted_idslogger.log_info(f'{type(data).__name__} ids {ids} are created')return idsdef update(self, data: T, user_id: str, tenant_id: str = None, collection_name: str = None):if not get_attr_value(data, 'id'):raise ValueError(f'{type(data).__name__} id is null or empty when update, user_id={user_id}, tenant_id={tenant_id}, collection_name={collection_name}')set_attr_value(data, 'updated_by', user_id)set_attr_value(data, 'updated_time', datetime.now(timezone.utc))collection = mongo_collection_pool.get_collection(self.mongo_settings, collection_name, tenant_id)# response is dict typeoriginal_data = collection.find_one(data.id)if original_data:if "createdBy" in original_data:set_attr_value(data, "created_by", original_data["createdBy"])if "createdTime" in original_data:set_attr_value(data, "created_time", original_data["createdTime"])request = data.model_dump(by_alias = True)collection.find_one_and_replace({"_id": data.id}, request)logger.log_info(f'{type(data).__name__} id {id} is updated')else:logger.log_info(f'{type(data).__name__} id {id} does not exist when update')def get(self, id: str, responseType: TResponse, tenant_id: str = None, collection_name: str = None) -> TResponse:if not id:logger.log_info(f'{type(responseType).__name__} id is null or empty when get')return Nonecollection = mongo_collection_pool.get_collection(self.mongo_settings, collection_name, tenant_id)result = collection.find_one(id)return responseType.model_validate(result)def get_list(self, ids: list[str], entityType: T, tenant_id: str = None, collection_name: str = None) -> list[T]:if not ids:logger.log_info(f'{type(entityType).__name__} ids is null or empty when get_list')return Nonecollection = mongo_collection_pool.get_collection(self.mongo_settings, collection_name, tenant_id)response = list(collection.find({"_id": { '$in': ids }}))if not response:return []result = []for item in response:result_item = entityType.model_validate(item)result.append(result_item)return result

整体思路是创建BaseDbService实例时指定目标数据库，进行具体的CRUD操作时再指定表名以及查询结果要转换成的类型（pymongo查询出结果后默认是dict类型），这样底层的数据库服务就能操作任意库的所有表了，而且最终获取的数据是具体类型的对象。

实际上用户大致会按照下述方式使用BaseDbService：

if __name__ == '__main__':mongo_settings = MongoSettings(connection_string="mongodb://localhost:27017/", database_name="test-activity")db_service = BaseDbService(mongo_settings, "bWRHdzkzVDFJbWNB")collection_name = "pycommon"tenant_id = "test-tenant-id"result = db_service.get("74bae4fe-13ec-11ef-be60-7404f152b5a1", ActivityQueryItem, tenant_id, collection_name)

三、数据库连接池机制

第二步中对数据库表操作之前都要先执行mongo_collection_pool.get_collection(self.mongo_settings, collection_name, tenant_id)来获取collection的连接，mongo_collection_pool是一个单例对象，内部有一个缓存机制来存储所有已经操作过的client、database和collection，连接从不关闭，这样可以保证连接不会重复的创建销毁从而提升性能。

看到这里有些细心的朋友可能会提问连接从不关闭，不会造成当前用户过多导致新用户无法连接的情况吗？本节先展示数据库连接池，大家的疑问咱们在下边有验证

from pymongo import MongoClient
from pymongo.database import Database
from pymongo.collection import Collection
from pycommon.models.mongo.mongo_settings import MongoSettingsclass MongoCollectionPool:def __init__(self):self.mongo_client_cache = {}self.mongo_database_cache = {}self.mongo_collection_cache = {}def get_collection(self, mongo_settings: MongoSettings, collection_name: str, tenant_id: str) -> Collection:client_cache_key = mongo_settings.connection_stringdatabase_name = mongo_settings.database_namedatabase_cache_key = client_cache_key + "_" + mongo_settings.database_namecollection_name = collection_name + "_" + tenant_id if tenant_id else collection_namecollection_cache_key = client_cache_key + "_" + mongo_settings.database_name + "_" + collection_namemongo_client: MongoClient = Noneif client_cache_key in self.mongo_client_cache:mongo_client = self.mongo_client_cache[client_cache_key]else:mongo_client = MongoClient(mongo_settings.connection_string)self.mongo_client_cache[client_cache_key] = mongo_clientdatabase: Database = Noneif database_cache_key in self.mongo_database_cache:database = self.mongo_database_cache[database_cache_key]else:database = mongo_client[database_name]self.mongo_database_cache[database_cache_key] = databaseif collection_cache_key in self.mongo_collection_cache:return self.mongo_collection_cache[collection_cache_key]else:collection = database[collection_name]self.mongo_collection_cache[collection_cache_key] = collectionreturn collection# 单例对象
mongo_collection_pool = MongoCollectionPool()

from typing import Optional
from pycommon.models.base_core_model import BaseCoreModelclass MongoSettings(BaseCoreModel):connection_string: Optional[str] = Nonedatabase_name: Optional[str] = None