Design Patterns in Python: Unit of Work Pattern
The Unit of Work Design Pattern Explained and Implemented in Python
The Unit of Work (UoW) design pattern helps manage transactional data in an object-oriented system. It is commonly used in applications that require database transactions, such as web applications, and provides a way to handle database operations in a way that is both efficient and easy to maintain.
The UoW and Repository patterns are often used together to provide a complete solution for managing data in an application. While they are related, they serve different purposes and have distinct responsibilities.
The Repository pattern serves as an abstraction layer for persistent storage, while the UoW pattern serves as an abstraction layer for atomic operations.
The UoW pattern is responsible for managing transactions and coordinating data operations within a single transactional boundary.
The Repository pattern, on the other hand, is responsible for encapsulating the logic for accessing and querying data from a data store. It acts as an abstraction layer between the application and the data store, allowing the application to access and manipulate data without needing to know the underlying details of the data store.
Python Implementation
UnitOfWork is responsible for managing the database transaction and providing a context manager interface.
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
class UnitOfWork:
def __init__(self):
self.engine = create_engine('postgresql://user:password@host:port/database')
self.Session = sessionmaker(bind=self.engine)
self.session = None
def __enter__(self):
self.session = self.Session()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is None:
self.session.commit()
else:
self.session.rollback()
self.session.close()
def add(self, entity):
self.session.add(entity)
def delete(self, entity):
self.session.delete(entity)When a block of code is executed within a with statement, the __enter__ method of the context manager is called, which typically returns an object that will be used within the block. Once the block is complete, the __exit__ method of the context manager is called to clean up any resources that were used.
In our case, the __enter__() method is called at the beginning of the transaction, and the __exit__() method is called at the end of the transaction. If no exceptions are raised within the transaction, the changes are committed; otherwise, the changes are rolled back.
The add() and delete() methods are provided as convenience methods for adding and deleting entities from the database within the transactional boundary.
with UnitOfWork() as uow:
customer = Customer(name='John Doe', email='johndoe@example.com')
uow.add(customer)
order = Order(customer=customer, total=100.00)
uow.add(order)UoW ensures that all changes made within the context manager are persisted in the database as a single atomic operation.
from sqlalchemy.orm import sessionmaker
from myapp.models import Order, OrderLineItem
class OrderService:
def __init__(self, db_engine):
Session = sessionmaker(bind=db_engine)
self.session = Session()
def place_order(self, customer_id, order_items):
# Create a Unit of Work instance
with self.session.begin():
# Create a new Order instance
order = Order(customer_id=customer_id)
self.session.add(order)
# Add order line items to the Order instance
for item in order_items:
line_item = OrderLineItem(
order_id=order.id,
product_id=item['product_id'],
quantity=item['quantity'],
price=item['price']
)
self.session.add(line_item)
# Commit the changes to the database
self.session.commit()We call self.session.commit() to commit the changes to the database. If any errors occur during this process, the changes will be rolled back to their previous state, ensuring that the transaction is completed atomically and transactional consistency is maintained.
By using the UoW pattern in this way, we can perform multiple database operations within a single transaction, improving performance by reducing the number of round-trips to the database. This can be especially important when dealing with high-concurrency scenarios or when dealing with large amounts of data.
Suppose we have an application that needs to support multiple data sources, such as a main database and a backup database. We can use the Unit of Work pattern to abstract away the details of database persistence and allow our application to easily switch between data sources.
class CustomerService:
def __init__(self, primary_db_engine, backup_db_engine):
self.primary_unit_of_work = UnitOfWork(primary_db_engine)
self.backup_unit_of_work = UnitOfWork(backup_db_engine)
def get_customer_by_id(self, customer_id):
with self.primary_unit_of_work as primary_uow:
customer = primary_uow.session.query(Customer).get(customer_id)
if not customer:
with self.backup_unit_of_work as backup_uow:
customer = backup_uow.session.query(Customer).get(customer_id)
return customer
def create_customer(self, name, email):
with self.primary_unit_of_work as uow:
customer = Customer(name=name, email=email)
uow.session.add(customer)
uow.save_changes()Within the CustomerService class, we define methods for getting a customer by ID and creating a new customer. For the get_customer_by_id method, we first try to retrieve the customer from the primary database. If the customer does not exist in the primary database, we retrieve it from the backup database. This allows us to easily switch between data sources if the primary database is unavailable.
By using the UoW pattern in this way, we can easily support multiple data sources and switch between them as needed. We can abstract away the details of database persistence and focus on the business logic of our application.
We can use UoW with a repository.
class CustomerRepository:
def __init__(self, unit_of_work):
self.unit_of_work = unit_of_work
def get_by_id(self, customer_id):
with self.unit_of_work as uow:
customer = uow.session.query(Customer).get(customer_id)
return customer
def add(self, customer):
with self.unit_of_work as uow:
uow.session.add(customer)
uow.save_changes()class CustomerService:
def __init__(self, db_engine):
self.db_engine = db_engine
def get_customer(self, customer_id):
with UnitOfWork(self.db_engine) as uow:
customer_repo = CustomerRepository(uow)
customer = customer_repo.get_by_id(customer_id)
return customer
def add_customer(self, name, email):
customer = Customer(name=name, email=email)
with UnitOfWork(self.db_engine) as uow:
customer_repo = CustomerRepository(uow)
customer_repo.add(customer)
return customer.idWe define a CustomerService class that provides a simple interface for retrieving and adding customer data. Within the get_customer method, we use the UnitOfWork and CustomerRepository classes to retrieve a Customer object by ID.
In conclusion, UoW is a powerful tool for simplifying database interactions in applications. By abstracting away the details of database persistence and providing a simple interface for performing atomic operations, UoW can help ensure transactional consistency, simplify code, and support multiple data sources.
