Chapter 6 Excercise for the Reader: Separate UoW and Context Manager

[chrisgrf]

experiment with separating the UoW (whose responsibilities are
commit() , rollback() , and providing the .batches repository) from the
context manager, whose job is to initialize things, and then do the commit
or rollback on exit.

Could someone give a solution for this exercise?

Following is my code example. Would there be a point in splitting FakeUnitOfWork?

src/allocation/service_layer/unit_of_work.py

class AbstractUnitOfWork(abc.ABC):
    batches: repository.AbstractRepository
    
    @abc.abstractmethod
    def commit(self):
        raise NotImplementedError

    @abc.abstractmethod
    def rollback(self):
        raise NotImplementedError


DEFAULT_SESSION_FACTORY = sessionmaker(
    bind=create_engine(
        config.get_postgres_uri(),
    )
)


class SQLAlchemyUoWContextManager:
    def __init__(self, session_factory=DEFAULT_SESSION_FACTORY):
        self.session_factory = session_factory

    def __enter__(self):
        self.session: Session = self.session_factory()
        self.uow = SqlAlchemyUnitOfWork(session=self.session)
        return self.uow

    def __exit__(self, *args):
        self.uow.rollback()
        self.session.close()


class SqlAlchemyUnitOfWork(AbstractUnitOfWork):
    def __init__(self, session: Session):
        self.session = session
        self.batches = repository.SqlAlchemyRepository(self.session)

    def commit(self):
        self.session.commit()

    def rollback(self):
        self.session.rollback()

[gregbrowndev]

@chrisgrf, here's how I've implemented this pattern in a few projects at my work:

service_layer/ports.py:

import typing as t


class IUnitOfWorkContext(t.Protocol):
    batches_repository: IBatchesRepository
    products_repository: IProductsRepository


class IUnitOfWork(t.ContextManager[IUnitOfWorkContext], t.Protocol):
  
    def commit(self) -> None:
        """
        Commit unit of work

        :raises UnitOfWorkNotStarted if called outside the context manager scope
        """
        ...

    def rollback(self) -> None:
        """
        Rollback unit of work

        :raises UnitOfWorkNotStarted if called outside the context manager scope
        """
        ...

Note: the protocol conveniently defines the base type t.ContextManager[IUnitOfWorkContext] using ContextManager from the typing module. This defines the __enter__ and __exit__ methods for you, i.e. that __enter__ returns a IUnitOfWorkContext.

This is the concrete implementation:

class PostgresUnitOfWorkContext(IUnitOfWorkContext):
    session: Session

    def __init__(self, session: Session):
        self.session = session
        self.batches_repository = BatchesRepository(session)
        self.products_repository = ProductsRepository(session)


class PostgresUnitOfWork(IUnitOfWork):
    session_factory: sessionmaker
    context: PostgresUnitOfWorkContext | None

    def __init__(self, session_factory: sessionmaker):
        self.session_factory = session_factory
        self.context = None

    def __enter__(self) -> "PostgresUnitOfWorkContext":
        session = self.session_factory()
        self.context = PostgresUnitOfWorkContext(session)
        return self.context

    def __exit__(
        self,
        exc_type: Type[BaseException] | None,
        exc_val: BaseException | None,  # noqa: F841
        exc_tb: TracebackType | None,  # noqa: F841
    ) -> None:
        if not self.context:
            raise UnitOfWorkNotStarted()
        if exc_type is None:
            # Commit on exit by default
            self.commit()
        else:
            self.rollback()
        self.context = None

    def commit(self) -> None:
        if not self.context:
            raise UnitOfWorkNotStarted()

        self.context.session.commit()

    def rollback(self) -> None:
        if not self.context:
            raise UnitOfWorkNotStarted()

        self.context.session.rollback()

It can be used like this:

def some_command_handler(uow: IUnitOfWork, product_id: int):
    with uow as ctx:
        product = ctx.product_repository.get_by_id(product_id)
        ...

One of the reasons I implemented it like this was that it allows you to centralise the UoW responsibility within the message bus instead of in the command handler which allows you to execute multiple commands within the same UoW (bends the rules but this might be useful). So the definition of your command handlers becomes:

def some_command_handler(ctx: IUnitOfWorkContext, product_id: int):
    product = ctx.product_repository.get_by_id(product_id)
    ...

However, I've changed this a lot since I started using this pattern. Here are some thoughts/notes:

• Recently, I've actually dropped the IUnitOfWorkContext object altogether in favour of injecting the repositories directly into the command handlers rather than accessing them through the UoW. The problem is the UoW/context tends to become a single god object that provides access to all adapters, not just the repositories. E.g. something like:

  def some_command_handler(
      batches_repository: BatchesRepository, 
      products_repository: IProductsRepository, 
      product_id: int
  ):
      product = product_repository.get_by_id(product_id)
      ...

  However, this requires you re-think the relationship between the session, UoW and repositories since the repos still need access to the session. Instead, you must create the session in the bootstrap function and pass it into the UoW and repositories rather than the UoW being the sessionmaker.

  Note: I've always found the relationship of the session and UoW to be a bit mysterious. I've come to the conclusion that they aren't the same thing, and it's OK to create a session to pass around to different adapters as required. There's an example in Vaughn Vernon's red book in Chapter 12: "Repositories", Section: "Considerations for a TopLink Implementation" that shows a Java ORM called TopLink that has a API where you get the UoW from the session, e.g.:

  UnitOfWork unitOfWork = session.acquireUnitOfWork();`

• The naming of the two classes, IUnitOfWork and IUnitOfWorkContext, is a bit weird. You might want to name them something like, IUnitOfWorkManager and IUnitOfWork, respectively.

• I've also changed around where the commit and rollback methods live. The benefit of putting them on (what I called) the IUnitOfWorkContext object is that it eliminates the UnitOfWorkNotStarted error state since you have to use the context manager to get the ctx back. This is I think what the exercise was getting at. However, I chose to keep them in IUnitOfWork mainly because the command handlers don't need those methods (my team wanted the UoW to commit by default and throw exceptions for errors - all of this is handled in the PostgresUnitOfWork as you can see.)

Hope that is helpful!

[bl1nkker]

@gregbrowndev Thanks for a very detailed and interesting answer!

However, our task was to separate the UoW form the context manager. How good is your version, given that the __enter__ and __exit__ dunder methods are still inside UoW?

My solution looks like this (unit_of_work.py):

class AbstractUnitOfWorkManager():
    def __enter__(self):
        raise NotImplementedError

    def __exit__(self):
        raise NotImplementedError


class UnitOfWorkManager(AbstractUnitOfWorkManager):

    def __init__(self,  session_factory=DEFAULT_SESSION_FACTORY):
        self.session_factory = session_factory

    def __enter__(self):
        self.session = self.session_factory()
        self.uow = SqlAlchemyUnitOfWork(self.session)
        return self.uow

    def __exit__(self, *args):
        self.session.close()


class SqlAlchemyUnitOfWork:

    def __init__(self, session) -> None:
        self.session = session
        self.batches = repository.SqlAlchemyRepository(session=self.session)

    def commit(self):
        self.session.commit()

    def rollback(self):
        self.session.rollback()

and its usage (services.py):

...

def allocate(
    orderid: str,
    sku: str,
    qty: int,
    uow_manager: unit_of_work.AbstractUnitOfWorkManager
) -> str:
    line = OrderLine(orderid, sku, qty)
    with uow_manager as uow:
        batches = uow.batches.list()
        if not is_valid_sku(line.sku, batches):
            raise InvalidSku(f"Invalid sku {line.sku}")
        batchref = model.allocate(line, batches)
        uow.commit()
    return batchref

...

In my opinion, this solution most explicitly shows that UoW only has the necessary responsibilities.
If I am wrong, I will be glad to hear comments and contrarguments!