Understanding the Unit of Work Pattern in NestJS

Introduction

A Unit of Work is a behavioral design pattern that helps manage complex transactions and maintain data consistency. It ensures that a set of changes to your data is treated as a single unit, and either all changes are successfully applied or none are. This is especially critical in enterprise applications, where multiple operations on different data sources need to be coordinated and committed atomically.

Design patterns provide reusable solutions to common software design problems. They promote code organization, readability, and scalability, enabling developers to write maintainable and efficient systems.

This article will guide you through:

• Understanding the Unit of Work pattern.
• Implementing UoW in a NestJS application.
• Using a practical example to demonstrate the concept.
• Practical considerations and challenges for real-world usage.

What is the Unit of Work Pattern?

At its core, the Unit of Work pattern manages changes to your data as a single transaction. It helps track entities that need to be inserted, updated, or deleted and coordinates changes to ensure consistency and manage transactions efficiently.

In simpler terms, we can think of a Unit of Work as

• A "to-do list" for your database operations.
• A means to track all the actions (e.g., inserts, updates, deletes) you want to perform and ensure they are all applied together or none at all.
• A process to handle operation to the DB if something goes wrong during execution, the Unit of Work rolls back all changes to prevent partial updates and maintain data integrity.

Imagine you’re booking a flight ticket for a travel. You first select the flight going to your destination from your current location. Then you pick a seat after which you make payment. These steps are interdependent. If the payment fails, you wouldn’t want the flight to be booked without incomplete details or payment confirmation.
The Unit of Work ensures that all these steps are treated as one atomic operation:

• If all steps succeed, the booking is confirmed.
• If any step fails, the entire process is canceled.

Why Use the Unit of Work Pattern?

The UoW pattern reduces complexity by:

• Consolidating multiple database operations into a single transaction.
• Preventing partial updates that can lead to data corruption.
• Decoupling transaction management from business logic.

When to Use the Unit of Work Pattern

The UoW pattern is beneficial in scenarios where:

• Multiple entities are updated as part of a single operation.
• Transaction consistency is critical.
• You want to centralize transaction management for better maintainability.

Practical Example: Flight Booking System

Let’s use a Flight Booking System as an example. The system involves:

• Booking flight details.
• Saving passenger information.
• Processing payments.

All these operations should either succeed together or fail together.

Note that the examples provided in this article are simplified to explain the design process and core concepts of the Unit of Work pattern. In real-life scenarios, the service and architecture would require further expansion and adjustments to handle edge cases, scalability, and maintainability in production environments.

Without Unit of Work

In this approach, each database operation is handled independently, without transaction management.

@Injectable()
export class FlightBookingService {
  constructor(
    @InjectRepository(FlightBookingRepository)
    private readonly flightBookingRepository: FlightBookingRepository,
    @InjectRepository(PassengerRepository)
    private readonly passengerRepository: PassengerRepository,
    @InjectRepository(PaymentRepository)
    private readonly paymentRepository: PaymentRepository,
     ) {}

  async bookFlightWithoutUoW(
    bookingData: Partial<FlightBooking>,
    passengerData: Partial<Passenger>,
    paymentData: Partial<Payment>,
  ): Promise<FlightBooking> {
    try {
      // Step 1: Reserve the flight
      const booking = await this.flightBookingRepository.save(bookingData);

      // Step 2: Save passenger details
      const passenger = { ...passengerData, booking };
      await this.passengerRepository.save(passenger);

      // Step 3: Process payment
      const payment = { ...paymentData, booking };
      await this.paymentRepository.save(payment);

      return booking;
    } catch (error) {
      console.error('Error during booking:', error.message);
      throw error;
    }
  }
}

Problems Without Unit of Work

1. No Atomicity: If payment saving fails, the flight booking and passenger details are still saved, leaving the database in an inconsistent state.

2. Error Handling Complexity: You would need to implement manual rollback logic for each step, which is error-prone and repetitive.

3. Scattered Transaction Logic: Each repository call handles its operations separately, leading to duplication and harder maintainability.
   
   

With Unit of Work

Using the Unit of Work pattern, all operations are wrapped in a transaction, ensuring that they succeed or fail together.

First, we implement our unit of work service

import { Injectable } from '@nestjs/common';
import { DataSource, QueryRunner } from 'typeorm';

@Injectable()
export class UnitOfWork {
  private queryRunner: QueryRunner;

  constructor(private readonly dataSource: DataSource) {
    this.queryRunner = this.dataSource.createQueryRunner();
  }

  async startTransaction(): Promise<void> {
    await this.queryRunner.connect();
    await this.queryRunner.startTransaction();
  }

  async commitTransaction(): Promise<void> {
    await this.queryRunner.commitTransaction();
  }

  async rollbackTransaction(): Promise<void> {
    await this.queryRunner.rollbackTransaction();
  }

  async release(): Promise<void> {
    await this.queryRunner.release();
  }

  getManager() {
    return this.queryRunner.manager;
  }
}

Then, we apply it to the flight booking service

@Injectable()
export class FlightBookingService {
  constructor(private readonly unitOfWork: UnitOfWork) {}

  async bookFlightWithUoW(
    bookingData: Partial<FlightBooking>,
    passengerData: Partial<Passenger>,
    paymentData: Partial<Payment>,
  ): Promise<FlightBooking> {
    await this.unitOfWork.startTransaction();
    try {
      // Step 1: Reserve the flight
      const booking = await this.unitOfWork.getManager().save(FlightBooking, bookingData);

      // Step 2: Save passenger details
      const passenger = { ...passengerData, booking };
      await this.unitOfWork.getManager().save(Passenger, passenger);

      // Step 3: Process payment
      const payment = { ...paymentData, booking };
      await this.unitOfWork.getManager().save(Payment, payment);

      // Commit the transaction
      await this.unitOfWork.commitTransaction();
      return booking;
    } catch (error) {
      // Rollback the transaction if anything fails
      await this.unitOfWork.rollbackTransaction();
      throw new Error('Flight booking failed: ' + error.message);
    } finally {
      // Release the transaction resources
      await this.unitOfWork.release();
    }
  }
}

Testing the codes. Below are examples for both "without" and "with" Unit of Work approaches.

Testing Without Unit of Work

it('should book a flight without unit of work', async () => {
  const bookingData = { flightNumber: 'FL123', departureDate: new Date() };
  const savedBooking = { id: 1, ...bookingData };

  bookingRepository.save = jest.fn().mockResolvedValue(savedBooking);

  const result = await service.bookFlightWithoutUoW(bookingData, {}, {});
  expect(result).toEqual(savedBooking);
  expect(bookingRepository.save).toHaveBeenCalledWith(bookingData);
});

Testing With Unit of Work

it('should book a flight with unit of work', async () => {
  const bookingData = { flightNumber: 'FL123', departureDate: new Date() };
  const savedBooking = { id: 1, ...bookingData };

  const mockSave = jest.fn();
  unitOfWork.getManager = jest.fn().mockReturnValue({ save: mockSave });

  mockSave.mockResolvedValueOnce(savedBooking);

  const result = await service.bookFlightWithUoW(bookingData, {}, {});
  expect(result).toEqual(savedBooking);
  expect(unitOfWork.startTransaction).toHaveBeenCalled();
  expect(unitOfWork.commitTransaction).toHaveBeenCalled();
});

Benefits of Using Unit of Work

1. Atomicity: All operations are treated as a single unit. If any step fails, all previous changes are rolled back automatically.

2. Consistency: The database remains in a consistent state, even in case of errors.

3. Simplified Error Handling: No need to manually undo changes; the Unit of Work ensures the rollback of all operations.

4. Reusability: The Unit of Work service can be reused across different parts of the application, reducing boilerplate code.

Practical Considerations for Real-World Applications

When implementing the UoW pattern in real-world applications, keep these points in mind:

1. Error Handling: Design robust error handling mechanisms to ensure transactions are rolled back properly.
2. Scalability: Consider how transactions impact performance, especially in high-concurrency environments.
3. Framework-Specific Features: Leverage features like TypeORM’s Transaction decorator for simple use cases.

Trade-Offs and Considerations

While the Unit of Work pattern offers significant benefits, it’s not without challenges:

1. Performance Overhead: Managing multiple entities in memory can impact performance.
2. Scalability: Distributed transactions spanning multiple services require additional complexity.
3. Error Handling: Designing robust error recovery mechanisms is essential for production systems.

---

Key Takeaways

Key Takeaways

1. The Unit of Work pattern simplifies transaction management and ensures consistency.
2. It centralizes logic for committing or rolling back changes, reducing the risk of partial updates.
3. Implementing the UoW pattern in NestJS can streamline database operations and improve maintainability.
4. Real-world implementations require additional attention to scalability, error handling, and performance.

By adopting the Unit of Work pattern, you can build applications that are robust, maintainable, and resilient to failure.

View the code on Github