# Unit of Work Pattern

## Problem

When multiple domain objects are modified during a business transaction, each change could trigger an individual database call. This leads to partial updates on failure, poor performance from excessive round-trips, and difficulty maintaining consistency across related changes.

## Solution

The **Unit of Work** keeps track of all objects that have been created, modified, or deleted during a business transaction. When the transaction is complete, it commits all changes in a single atomic operation. If anything fails, all changes are rolled back.

## When to Use

- Multiple entities must be persisted atomically in a single transaction.
- You want to batch database writes for performance.
- You need rollback capability across multiple repository operations.
- Used alongside the Repository pattern to coordinate persistence.

## When to Avoid

- Single-entity CRUD operations where atomicity is trivial.
- Event-driven architectures where eventual consistency is acceptable.
- Simple scripts or pipelines with no transactional requirements.

## Pseudocode

```
CLASS UnitOfWork:
    new_objects = []
    dirty_objects = []
    removed_objects = []

    METHOD register_new(obj):
        APPEND obj TO new_objects

    METHOD register_dirty(obj):
        APPEND obj TO dirty_objects

    METHOD register_removed(obj):
        APPEND obj TO removed_objects

    METHOD commit():
        INSERT all new_objects into database
        UPDATE all dirty_objects in database
        DELETE all removed_objects from database
        CLEAR all tracking lists

    METHOD rollback():
        CLEAR all tracking lists
```

## Python

```python
from dataclasses import dataclass


@dataclass
class Product:
    id: int
    name: str
    price: float


class UnitOfWork:
    def __init__(self) -> None:
        self._new: list[Product] = []
        self._dirty: list[Product] = []
        self._removed: list[Product] = []

    def register_new(self, obj: Product) -> None:
        self._new.append(obj)
        print(f"  Registered new: {obj}")

    def register_dirty(self, obj: Product) -> None:
        self._dirty.append(obj)
        print(f"  Registered dirty: {obj}")

    def register_removed(self, obj: Product) -> None:
        self._removed.append(obj)
        print(f"  Registered removed: {obj}")

    def commit(self) -> None:
        print("  Committing transaction...")
        for obj in self._new:
            print(f"    INSERT {obj}")
        for obj in self._dirty:
            print(f"    UPDATE {obj}")
        for obj in self._removed:
            print(f"    DELETE {obj}")
        self._new.clear()
        self._dirty.clear()
        self._removed.clear()
        print("  Transaction committed.")

    def rollback(self) -> None:
        count = len(self._new) + len(self._dirty) + len(self._removed)
        self._new.clear()
        self._dirty.clear()
        self._removed.clear()
        print(f"  Rolled back {count} pending operations.")


class ProductRepository:
    def __init__(self, uow: UnitOfWork) -> None:
        self._uow = uow
        self._store: dict[int, Product] = {}

    def add(self, product: Product) -> None:
        self._store[product.id] = product
        self._uow.register_new(product)

    def update(self, product: Product) -> None:
        self._store[product.id] = product
        self._uow.register_dirty(product)

    def remove(self, product_id: int) -> None:
        product = self._store.pop(product_id, None)
        if product:
            self._uow.register_removed(product)


def main() -> None:
    uow = UnitOfWork()
    repo = ProductRepository(uow)

    print("--- Register changes ---")
    repo.add(Product(1, "Widget", 9.99))
    repo.add(Product(2, "Gadget", 24.99))

    p = Product(1, "Widget Pro", 14.99)
    repo.update(p)

    repo.remove(2)

    print("\n--- Commit ---")
    uow.commit()

    print("\n--- Register then rollback ---")
    repo.add(Product(3, "Doohickey", 4.99))
    uow.rollback()


if __name__ == "__main__":
    main()
```

**Output:**
```
--- Register changes ---
  Registered new: Product(id=1, name='Widget', price=9.99)
  Registered new: Product(id=2, name='Gadget', price=24.99)
  Registered dirty: Product(id=1, name='Widget Pro', price=14.99)
  Registered removed: Product(id=2, name='Gadget', price=24.99)

--- Commit ---
  Committing transaction...
    INSERT Product(id=1, name='Widget', price=9.99)
    INSERT Product(id=2, name='Gadget', price=24.99)
    UPDATE Product(id=1, name='Widget Pro', price=14.99)
    DELETE Product(id=2, name='Gadget', price=24.99)
  Transaction committed.

--- Register then rollback ---
  Registered new: Product(id=3, name='Doohickey', price=4.99)
  Rolled back 1 pending operations.
```

## Go

```go
package main

import "fmt"

type Product struct {
	ID    int
	Name  string
	Price float64
}

type UnitOfWork struct {
	newObjs     []Product
	dirtyObjs   []Product
	removedObjs []Product
}

func NewUnitOfWork() *UnitOfWork {
	return &UnitOfWork{}
}

func (u *UnitOfWork) RegisterNew(p Product) {
	u.newObjs = append(u.newObjs, p)
	fmt.Printf("  Registered new: %+v\n", p)
}

func (u *UnitOfWork) RegisterDirty(p Product) {
	u.dirtyObjs = append(u.dirtyObjs, p)
	fmt.Printf("  Registered dirty: %+v\n", p)
}

func (u *UnitOfWork) RegisterRemoved(p Product) {
	u.removedObjs = append(u.removedObjs, p)
	fmt.Printf("  Registered removed: %+v\n", p)
}

func (u *UnitOfWork) Commit() {
	fmt.Println("  Committing transaction...")
	for _, p := range u.newObjs {
		fmt.Printf("    INSERT %+v\n", p)
	}
	for _, p := range u.dirtyObjs {
		fmt.Printf("    UPDATE %+v\n", p)
	}
	for _, p := range u.removedObjs {
		fmt.Printf("    DELETE %+v\n", p)
	}
	u.newObjs = nil
	u.dirtyObjs = nil
	u.removedObjs = nil
	fmt.Println("  Transaction committed.")
}

func (u *UnitOfWork) Rollback() {
	count := len(u.newObjs) + len(u.dirtyObjs) + len(u.removedObjs)
	u.newObjs = nil
	u.dirtyObjs = nil
	u.removedObjs = nil
	fmt.Printf("  Rolled back %d pending operations.\n", count)
}

type ProductRepository struct {
	uow   *UnitOfWork
	store map[int]Product
}

func NewProductRepository(uow *UnitOfWork) *ProductRepository {
	return &ProductRepository{uow: uow, store: make(map[int]Product)}
}

func (r *ProductRepository) Add(p Product) {
	r.store[p.ID] = p
	r.uow.RegisterNew(p)
}

func (r *ProductRepository) Update(p Product) {
	r.store[p.ID] = p
	r.uow.RegisterDirty(p)
}

func (r *ProductRepository) Remove(id int) {
	if p, ok := r.store[id]; ok {
		delete(r.store, id)
		r.uow.RegisterRemoved(p)
	}
}

func main() {
	uow := NewUnitOfWork()
	repo := NewProductRepository(uow)

	fmt.Println("--- Register changes ---")
	repo.Add(Product{1, "Widget", 9.99})
	repo.Add(Product{2, "Gadget", 24.99})
	repo.Update(Product{1, "Widget Pro", 14.99})
	repo.Remove(2)

	fmt.Println("\n--- Commit ---")
	uow.Commit()

	fmt.Println("\n--- Register then rollback ---")
	repo.Add(Product{3, "Doohickey", 4.99})
	uow.Rollback()
}
```

**Output:**
```
--- Register changes ---
  Registered new: {ID:1 Name:Widget Price:9.99}
  Registered new: {ID:2 Name:Gadget Price:24.99}
  Registered dirty: {ID:1 Name:Widget Pro Price:14.99}
  Registered removed: {ID:2 Name:Gadget Price:24.99}

--- Commit ---
  Committing transaction...
    INSERT {ID:1 Name:Widget Price:9.99}
    INSERT {ID:2 Name:Gadget Price:24.99}
    UPDATE {ID:1 Name:Widget Pro Price:14.99}
    DELETE {ID:2 Name:Gadget Price:24.99}
  Transaction committed.

--- Register then rollback ---
  Registered new: {ID:3 Name:Doohickey Price:4.99}
  Rolled back 1 pending operations.
```

## Related Patterns

- **Repository** — repositories register changes with the unit of work; the unit of work commits or rolls them back atomically.
- **Identity Map** — the identity map is often embedded in the unit of work to prevent loading the same entity twice within a transaction.
- **Transaction Script** — transaction script puts all logic in a procedure; unit of work is used when objects need to coordinate persistence.
