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
– Haidar Zeineddine, Commented Feb 23, 2021 at 13:08
I never went further but I think the accepted answer is the right (and nice) way. — Jérôme
– Jérôme, Commented Feb 23, 2021 at 15:14
I opened a pull request with CPython to add context manager support: github.com/python/cpython/pull/99634 — Richard Hansen
– Richard Hansen, Commented Nov 21, 2022 at 6:42

rectalogic · Accepted Answer · 2020-01-25 16:20:21Z

7

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

answered Jan 25, 2020 at 16:20

rectalogic

1,36311 silver badges12 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

Jérôme Over a year ago

Don't you need to use copy_context and run at some point?

rectalogic Over a year ago

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.

Jérôme Over a year ago

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 Over a year ago

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.

Collectives™ on Stack Overflow

Provide contextvars.Context with a ContextManager

1 Answer 1

4 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

4 Comments

Your Answer

Sign up or log in

Post as a guest

Related