如何让PyMongo事务装饰器的会话在仓储层代码中生效？

问题 #

需要实现类似Spring @Transactional的事务装饰器，但PyMongo事务要求将装饰器内创建的会话传入仓储层方法。现有装饰器无法让Service层获取到会话，需解决会话传递问题。

可行方案 #

方案1：将会话作为关键字参数传入被装饰函数 #

直接在装饰器中将session注入到被调用函数的关键字参数中，Service层方法通过参数接收会话后传递给仓储。

修改装饰器代码：

class MongoDB:
    def __init__(self):
        self.__client = MongoClient(
            DB_URL,
            tlsCAFile=certifi.where()
        )

    def close(self):
        self.__client.close()

    def transactional(self, origin_func):
        def wrapper_func(*args, **kwargs):
            with self.__client.start_session() as session:
                with session.start_transaction():
                    # 把会话加入关键字参数，传递给原函数
                    kwargs['session'] = session
                    return origin_func(*args, **kwargs)

        return wrapper_func

Service层使用示例：

class UserService:
    def __init__(self, mongo_user_repository: MongoUsersRepository):
        self.__mongo_user_repository = mongo_user_repository

    @mongodb.transactional
    def save(self, user_create_request: UserCreateRequestDto, session=None):
        mongo_user = MongoUsers.from_dto(user_create_request=user_create_request)
        self.__mongo_user_repository.save(mongo_user=mongo_user, session=session)

方案2：线程局部存储传递会话 #

利用threading.local()存储当前线程的会话，Service和仓储层可直接从MongoDB实例获取会话，无需修改方法参数。

修改MongoDB类：

import threading

class MongoDB:
    def __init__(self):
        self.__client = MongoClient(
            DB_URL,
            tlsCAFile=certifi.where()
        )
        # 线程局部变量，用于存放当前会话
        self._local = threading.local()

    def close(self):
        self.__client.close()

    @property
    def current_session(self):
        return getattr(self._local, 'session', None)

    def transactional(self, origin_func):
        def wrapper_func(*args, **kwargs):
            with self.__client.start_session() as session:
                with session.start_transaction():
                    self._local.session = session
                    try:
                        return origin_func(*args, **kwargs)
                    finally:
                        # 执行完毕清除会话，避免线程复用导致问题
                        del self._local.session

        return wrapper_func

Service层使用示例：

class UserService:
    def __init__(self, mongo_user_repository: MongoUsersRepository, mongodb: MongoDB):
        self.__mongo_user_repository = mongo_user_repository
        self.__mongodb = mongodb

    @mongodb.transactional
    def save(self, user_create_request: UserCreateRequestDto):
        mongo_user = MongoUsers.from_dto(user_create_request=user_create_request)
        self.__mongo_user_repository.save(
            mongo_user=mongo_user,
            session=self.__mongodb.current_session
        )

方案3：仓储层自动复用会话 #

进一步封装仓储层，让仓储方法自动从MongoDB实例获取当前会话，无需手动传递。

仓储层示例：

class MongoUsersRepository:
    def __init__(self, mongodb: MongoDB):
        self.__db = mongodb._MongoDB__client.get_database('your_db_name')
        self.__collection = self.__db.get_collection('users')
        self.__mongodb = mongodb

    def save(self, mongo_user, session=None):
        # 优先使用传入的session，否则取当前线程的会话
        session = session or self.__mongodb.current_session
        self.__collection.insert_one(mongo_user.dict(), session=session)

Service层简化使用：

@mongodb.transactional
def save(self, user_create_request: UserCreateRequestDto):
    mongo_user = MongoUsers.from_dto(user_create_request=user_create_request)
    self.__mongo_user_repository.save(mongo_user=mongo_user)

内容的提问来源于stack exchange，提问作者YM KIM