为 contextvars.Context 提供 ContextManager

Question

我正在尝试在我的数据库框架中管理事务（我使用 MongoDB 和 umongo 而不是 pymongo）。

\n

要使用事务，必须session沿整个调用链传递一个 kwarg。我想提供一个上下文管理器来隔离事务。只有调用链末尾的函数需要知道该session对象。

\n

我发现了上下文变量，并且我已经接近一些东西，但还没有完全实现。

\n

我想要什么：

\n

with Transaction():\n    #\xc2\xa0Do stuff\n    d = MyDocument.find_one()\n    d.attr = 12\n    d.commit()\n

\n

这是我现在想到的：

\n

s = ContextVar(\'session\', default=None)\n\nclass Transaction(AbstractContextManager):\n\n    def __init__(self):\n        self.ctx = copy_context()\n        # Create a new DB session\n        session = db.create_session()\n        # Set session in context\n        self.ctx.run(s.set, session)\n\n    def __exit__(self, *args, **kwargs):\n        pass\n\n    # Adding a run method for convenience\n    def run(self, func, *args, **kwargs):\n        self.ctx.run(func, *args, **kwargs)\n\ndef func():\n    d = MyDocument.find_one()\n    d.attr = 12\n    d.commit()\n\nwith Transaction() as t:\n    t.run(func)\n

\n

但我没有很好的上下文管理器语法。上下文管理器的要点是“其中的所有内容都应该在该上下文中运行”。

\n

我上面写的并不比仅仅使用函数更好：

\n

def run_transaction(func, *args, **kwargs):\n    ctx = copy_context()\n    session = 12\n    ctx.run(s.set, session)\n    ctx.run(func)\n\nrun_transaction(func)\n

\n

我是不是走错了路？

\n

我是否滥用了上下文变量？

\n

还有其他方法可以实现我想要做的事情吗？

\n

---

\n

基本上，我希望能够像上下文管理器一样打开上下文

\n

session = ContextVar(\'session\', default=None)\n\nwith copy_context() as ctx:\n    session = db.create_session()\n    # Do stuff\n    d = MyDocument.find_one()\n    d.attr = 12\n    d.commit()\n

\n

我将其嵌入到Transaction上下文管理器中来管理会话内容，并且仅在d用户代码中保留操作。

\n

Answer 1

您可以使用上下文管理器创建会话和事务，并将会话存储在 ContextVar 中以供其他函数使用。

from contextlib import contextmanager
from contextvars import ContextVar
import argparse
import pymongo


SESSION = ContextVar("session", default=None)


@contextmanager
def transaction(client):
    with client.start_session() as session:
        with session.start_transaction():
            t = SESSION.set(session)
            try:
                yield
            finally:
                SESSION.reset(t)


def insert1(client):
    client.test.txtest1.insert_one({"data": "insert1"}, session=SESSION.get())


def insert2(client):
    client.test.txtest2.insert_one({"data": "insert2"}, session=SESSION.get())


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", default="mongodb://localhost:27017")
    args = parser.parse_args()

    client = pymongo.MongoClient(args.url)

    # Create and lear collections, collections must be created outside the transaction
    insert1(client)
    client.test.txtest1.delete_many({})
    insert2(client)
    client.test.txtest2.delete_many({})

    with transaction(client):
        insert1(client)
        insert2(client)

    for doc in client.test.txtest1.find({}):
        print(doc)
    for doc in client.test.txtest2.find({}):
        print(doc)


if __name__ == "__main__":
    main()