Chapter 6 Excercise for the Reader: Separate UoW and Context Manager
chrisgrf opened this issue · 2 comments
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()
@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 thesession
. Instead, you must create thesession
in thebootstrap
function and pass it into theUoW
and repositories rather than theUoW
being thesessionmaker
.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
andIUnitOfWorkContext
, is a bit weird. You might want to name them something like,IUnitOfWorkManager
andIUnitOfWork
, respectively. -
I've also changed around where the
commit
androllback
methods live. The benefit of putting them on (what I called) theIUnitOfWorkContext
object is that it eliminates theUnitOfWorkNotStarted
error state since you have to use the context manager to get thectx
back. This is I think what the exercise was getting at. However, I chose to keep them inIUnitOfWork
mainly because the command handlers don't need those methods (my team wanted theUoW
to commit by default and throw exceptions for errors - all of this is handled in thePostgresUnitOfWork
as you can see.)
Hope that is helpful!
@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!