Skip to content

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

Description

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

Activity

  1. gregbrowndev commented on Jul 29, 2022

    @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!

  2. bl1nkker commented on Jan 19, 2023

    @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!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions