Provide contextvars.Context with a ContextManager

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.

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