Provide contextvars.Context with a ContextManager

Question

I'm trying to manage transactions in my DB framework (I use MongoDB with umongo over pymongo).

To use transaction, one must pass a session kwarg along the whole call chain. I would like to provide a context manager that would isolate the transaction. Only the function at the end of the call chain would need to be aware of the session object.

I found out about context variables and I'm close to something but not totally there.

What I would like to have:

with Transaction():
    # Do stuff
    d = MyDocument.find_one()
    d.attr = 12
    d.commit()

Here's what I came up with for now:

s = ContextVar('session', default=None)

class Transaction(AbstractContextManager):

    def __init__(self):
        self.ctx = copy_context()
        # Create a new DB session
        session = db.create_session()
        # Set session in context
        self.ctx.run(s.set, session)

    def __exit__(self, *args, **kwargs):
        pass

    # Adding a run method for convenience
    def run(self, func, *args, **kwargs):
        self.ctx.run(func, *args, **kwargs)

def func():
    d = MyDocument.find_one()
    d.attr = 12
    d.commit()

with Transaction() as t:
    t.run(func)

But I don't have the nice context manager syntax. The point of the context manager would be so say "everyting that's in there should be run in that context".

What I wrote above is not really better than just using a function:

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

run_transaction(func)

Am I on the wrong track?

Am I misusing context variables?

Any other way to achieve what I'm trying to do?

Basically, I'd like to be able to open a context like a context manager

session = ContextVar('session', default=None)

with copy_context() as ctx:
    session = db.create_session()
    # Do stuff
    d = MyDocument.find_one()
    d.attr = 12
    d.commit()

I'd embed this in a Transaction context manager to manage the session stuff and only keep operations on d in user code.

Jerome did you find the time to test this ? What final approach did you apply ? — Haidar Zeineddine, Feb 23 '21 at 13:08
I never went further but I think the accepted answer is the right (and nice) way. — Jérôme, Feb 23 '21 at 15:14
I opened a pull request with CPython to add context manager support: https://github.com/python/cpython/pull/99634 — Richard Hansen, Nov 21 '22 at 06:42

score 4 · Accepted Answer · answered Jan 25 '20 at 16:20

You can use a contextmanager to create the session and transaction and store the session in the ContextVar for use by other functions.


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()

Don't you need to use `copy_context` and `run` at some point? — Jérôme, Jan 26 '20 at 20:05
No, not unless you are doing something advanced and managing all this manually. `asyncio` and `futures` and so on all use `copy_context` and `run` as needed internally so that you don't need to. — rectalogic, Jan 27 '20 at 21:15
I guess I was overthinking it, then. I'll come back to upvote/accept once I get the time to try that. Thanks. — Jérôme, Jan 27 '20 at 21:56
I didn't get the time to try it but reading it now with a better understanding it makes total sense. Accepting this answer. Thanks. — Jérôme, Feb 23 '21 at 15:13

Provide contextvars.Context with a ContextManager

1 Answers1