Loading content...
Having understood why we need the Unit of Work pattern and where transaction boundaries should be drawn, we're now ready to build a complete implementation. While production applications often use ORM-provided Unit of Work implementations (Entity Framework, Hibernate, Prisma), understanding how to build one from scratch gives you the insight to use these tools effectively and troubleshoot issues when they arise.
This page walks through building a Unit of Work implementation step by step, covering all the essential components: the entity state machine, identity mapping, change tracking, dependency resolution, and commit orchestration.
By the end of this page, you will be able to implement a fully functional Unit of Work that tracks entity states, manages an identity map, orders operations by dependencies, and commits all changes atomically within a database transaction.
At the core of any Unit of Work is the concept of entity state—tracking whether an entity is new, modified, deleted, or unchanged. Understanding this state machine is fundamental.
Entity States:
| State | Description | On Commit |
|---|---|---|
| Transient | Newly created, not yet tracked by Unit of Work | Nothing (not tracked) |
| Added (New) | New entity registered for insertion | INSERT into database |
| Clean (Unchanged) | Loaded from DB, no modifications since | Nothing (no changes) |
| Dirty (Modified) | Existing entity with modified properties | UPDATE in database |
| Deleted | Entity marked for removal | DELETE from database |
| Detached | Removed from Unit of Work tracking | Nothing (ignored) |
12345678910111213141516171819202122232425262728293031323334
// Entity state enumerationenum EntityState { Transient, // Not tracked Added, // New, pending INSERT Clean, // Loaded, unchanged Dirty, // Modified, pending UPDATE Deleted // Pending DELETE} // Entity entry - wrapper that tracks stateinterface EntityEntry<T = unknown> { entity: T; state: EntityState; originalValues: Map<string, unknown> | null; // For dirty checking entityType: EntityType;} // Entity type metadatainterface EntityType { name: string; tableName: string; keyProperty: string; properties: PropertyMetadata[]; relationships: RelationshipMetadata[];} // State tracker interfaceinterface StateTracker { getEntry<T>(entity: T): EntityEntry<T> | undefined; getState<T>(entity: T): EntityState; setState<T>(entity: T, state: EntityState): void; getEntriesByState(state: EntityState): EntityEntry[]; hasChanges(): boolean;}Not all state transitions are valid. For example, you can't go from Deleted directly to Added—the entity must be detached first, then re-added. The Unit of Work enforces valid transitions and throws exceptions for invalid ones.
The Identity Map is a crucial component that ensures each entity is represented by exactly one object instance within a Unit of Work scope. Without it, you could have two different object instances representing the same database row, leading to data corruption when both are modified.
Why Identity Map Matters:
12345678910111213141516171819202122232425262728293031323334
// ❌ WITHOUT Identity Map - Disaster waiting to happenasync function processOrder(orderId: string) { // First query: get order for display const orderForDisplay = await orderRepository.findById(orderId); console.log(`Order total: ${orderForDisplay.total}`); // Second query: get same order for modification const orderForUpdate = await orderRepository.findById(orderId); orderForUpdate.addItem(newItem); orderForUpdate.recalculateTotal(); // Now total = 150 await orderRepository.save(orderForUpdate); // Problem: orderForDisplay.total is still 100 (stale) // Even worse: if we save orderForDisplay later, we overwrite! // orderForDisplay and orderForUpdate are DIFFERENT objects console.log(orderForDisplay === orderForUpdate); // false!} // ✅ WITH Identity Map - Same object alwaysasync function processOrderSafe(orderId: string, unitOfWork: IUnitOfWork) { const orderRepo = unitOfWork.orderRepository; // First query const orderForDisplay = await orderRepo.findById(orderId); // Second query - returns SAME INSTANCE from identity map const orderForUpdate = await orderRepo.findById(orderId); console.log(orderForDisplay === orderForUpdate); // true! orderForUpdate.addItem(newItem); // Modifies the single instance // orderForDisplay sees the change too - same object!}Identity Map Implementation:
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
class IdentityMap { // Map: EntityTypeName -> Map<EntityId, EntityEntry> private entities: Map<string, Map<string, EntityEntry>> = new Map(); /** * Create a unique key for an entity */ private getKey(entityType: string, id: unknown): string { // Handle composite keys if (Array.isArray(id)) { return id.map(String).join('::'); } return String(id); } /** * Check if entity exists in the map */ contains(entityType: string, id: unknown): boolean { const typeMap = this.entities.get(entityType); if (!typeMap) return false; return typeMap.has(this.getKey(entityType, id)); } /** * Get entity from the map */ get<T>(entityType: string, id: unknown): EntityEntry<T> | undefined { const typeMap = this.entities.get(entityType); if (!typeMap) return undefined; return typeMap.get(this.getKey(entityType, id)) as EntityEntry<T> | undefined; } /** * Add entity to the map */ add<T>(entityType: string, id: unknown, entry: EntityEntry<T>): void { let typeMap = this.entities.get(entityType); if (!typeMap) { typeMap = new Map(); this.entities.set(entityType, typeMap); } const key = this.getKey(entityType, id); if (typeMap.has(key)) { throw new Error( `Entity of type ${entityType} with id ${key} already exists in identity map` ); } typeMap.set(key, entry as EntityEntry); } /** * Remove entity from the map */ remove(entityType: string, id: unknown): boolean { const typeMap = this.entities.get(entityType); if (!typeMap) return false; return typeMap.delete(this.getKey(entityType, id)); } /** * Get all entries of a given state */ getByState(state: EntityState): EntityEntry[] { const results: EntityEntry[] = []; for (const typeMap of this.entities.values()) { for (const entry of typeMap.values()) { if (entry.state === state) { results.push(entry); } } } return results; } /** * Get all tracked entities */ getAll(): EntityEntry[] { const results: EntityEntry[] = []; for (const typeMap of this.entities.values()) { for (const entry of typeMap.values()) { results.push(entry); } } return results; } /** * Clear all tracked entities */ clear(): void { this.entities.clear(); }}Repositories must check the Identity Map before querying the database. If an entity with the requested ID already exists in the map, return that instance instead of creating a new one from the database results. This is the key to preventing duplicate instances.
The Change Tracker is the heart of the Unit of Work—it manages entity entries, detects modifications, and maintains the identity map. Let's build a complete implementation:
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210
class ChangeTracker { private identityMap: IdentityMap = new IdentityMap(); private metadata: EntityMetadataRegistry; constructor(metadata: EntityMetadataRegistry) { this.metadata = metadata; } /** * Track a new entity (will be INSERTed on commit) */ trackNew<T extends object>(entity: T): void { const entityType = this.metadata.getTypeFor(entity); const id = this.extractId(entity, entityType); // For new entities, id might be null/undefined (auto-generated) // Use a temporary id until commit const trackingId = id ?? this.generateTempId(); const entry: EntityEntry<T> = { entity, state: EntityState.Added, originalValues: null, entityType }; this.identityMap.add(entityType.name, trackingId, entry); } /** * Track an entity loaded from database (starts as Clean) */ trackLoaded<T extends object>(entity: T): void { const entityType = this.metadata.getTypeFor(entity); const id = this.extractId(entity, entityType); // Check if already tracked if (this.identityMap.contains(entityType.name, id)) { return; // Already tracking this entity } const entry: EntityEntry<T> = { entity, state: EntityState.Clean, originalValues: this.takeSnapshot(entity, entityType), entityType }; this.identityMap.add(entityType.name, id, entry); } /** * Mark entity for deletion */ trackDeleted<T extends object>(entity: T): void { const entityType = this.metadata.getTypeFor(entity); const id = this.extractId(entity, entityType); const entry = this.identityMap.get<T>(entityType.name, id); if (!entry) { throw new Error('Cannot delete entity that is not tracked'); } if (entry.state === EntityState.Added) { // New entity being deleted - just remove from tracking this.identityMap.remove(entityType.name, id); } else { entry.state = EntityState.Deleted; } } /** * Take a snapshot of entity property values */ private takeSnapshot(entity: object, entityType: EntityType): Map<string, unknown> { const snapshot = new Map<string, unknown>(); for (const prop of entityType.properties) { const value = (entity as any)[prop.name]; // Deep clone for objects/arrays snapshot.set(prop.name, this.cloneValue(value)); } return snapshot; } /** * Detect changes by comparing current values to snapshot */ detectChanges(): void { for (const entry of this.identityMap.getByState(EntityState.Clean)) { if (this.hasChanges(entry)) { entry.state = EntityState.Dirty; } } } /** * Check if entity has changed since loaded */ private hasChanges(entry: EntityEntry): boolean { if (!entry.originalValues) return false; for (const prop of entry.entityType.properties) { const currentValue = (entry.entity as any)[prop.name]; const originalValue = entry.originalValues.get(prop.name); if (!this.valuesEqual(currentValue, originalValue)) { return true; } } return false; } /** * Get all pending changes grouped by operation type */ getPendingChanges(): PendingChanges { this.detectChanges(); return { inserts: this.identityMap.getByState(EntityState.Added), updates: this.identityMap.getByState(EntityState.Dirty), deletes: this.identityMap.getByState(EntityState.Deleted) }; } /** * Mark all entities as clean after successful commit */ acceptChanges(): void { for (const entry of this.identityMap.getAll()) { if (entry.state === EntityState.Deleted) { // Remove deleted entities const id = this.extractId(entry.entity as object, entry.entityType); this.identityMap.remove(entry.entityType.name, id); } else if (entry.state === EntityState.Added || entry.state === EntityState.Dirty) { // Reset to clean entry.state = EntityState.Clean; entry.originalValues = this.takeSnapshot( entry.entity as object, entry.entityType ); } } } /** * Discard all changes and restore original values */ rejectChanges(): void { for (const entry of this.identityMap.getAll()) { if (entry.state === EntityState.Added) { // Remove new entities const id = this.extractId(entry.entity as object, entry.entityType); this.identityMap.remove(entry.entityType.name, id); } else if (entry.state === EntityState.Dirty) { // Restore original values this.restoreSnapshot(entry); entry.state = EntityState.Clean; } else if (entry.state === EntityState.Deleted) { // Undelete entry.state = EntityState.Clean; } } } private restoreSnapshot(entry: EntityEntry): void { if (!entry.originalValues) return; for (const [prop, value] of entry.originalValues) { (entry.entity as any)[prop] = this.cloneValue(value); } } // Helper methods private extractId(entity: object, entityType: EntityType): unknown { return (entity as any)[entityType.keyProperty]; } private generateTempId(): string { return `temp_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`; } private cloneValue(value: unknown): unknown { if (value === null || value === undefined) return value; if (typeof value !== 'object') return value; if (value instanceof Date) return new Date(value.getTime()); return JSON.parse(JSON.stringify(value)); } private valuesEqual(a: unknown, b: unknown): boolean { if (a === b) return true; if (a instanceof Date && b instanceof Date) { return a.getTime() === b.getTime(); } if (typeof a === 'object' && typeof b === 'object') { return JSON.stringify(a) === JSON.stringify(b); } return false; }} interface PendingChanges { inserts: EntityEntry[]; updates: EntityEntry[]; deletes: EntityEntry[];}Taking snapshots via JSON serialization is simple but expensive for large object graphs. Production ORMs use more sophisticated techniques: property-level tracking, lazy snapshot creation, or compile-time code generation. For most applications, the simpler approach suffices; optimize only when profiling reveals a bottleneck.
Database foreign key constraints require operations to be executed in a specific order. New parent records must be inserted before child records that reference them. Deleted child records must be removed before deleting their parents. This is the dependency ordering problem.
The Challenge:
1234567891011121314151617181920212223242526
// Domain modelclass Order { id: string; customerId: string; // FK to Customer items: OrderItem[];} class OrderItem { id: string; orderId: string; // FK to Order productId: string; // FK to Product} // If we try to insert Order before Customer exists: FK violation!// If we try to insert OrderItem before Order exists: FK violation! // Correct INSERT order:// 1. Customer (if new)// 2. Order (depends on Customer)// 3. Product (if new, no dependencies)// 4. OrderItem (depends on Order AND Product) // Correct DELETE order (opposite!):// 1. OrderItem (children first)// 2. Order// 3. Customer (only if no other orders)Topological Sort for Dependency Resolution:
The solution is to build a dependency graph of entity types and perform a topological sort to determine the correct execution order:
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128
class DependencyResolver { private metadata: EntityMetadataRegistry; constructor(metadata: EntityMetadataRegistry) { this.metadata = metadata; } /** * Order entity entries for INSERT operations * Parents must be inserted before children */ orderForInsert(entries: EntityEntry[]): EntityEntry[] { const graph = this.buildDependencyGraph(entries); return this.topologicalSort(graph, entries); } /** * Order entity entries for DELETE operations * Children must be deleted before parents (reverse of insert) */ orderForDelete(entries: EntityEntry[]): EntityEntry[] { const insertOrder = this.orderForInsert(entries); return insertOrder.reverse(); } /** * Build a graph of type dependencies */ private buildDependencyGraph(entries: EntityEntry[]): Map<string, Set<string>> { const graph = new Map<string, Set<string>>(); // Initialize nodes for (const entry of entries) { const typeName = entry.entityType.name; if (!graph.has(typeName)) { graph.set(typeName, new Set()); } } // Add edges based on foreign key relationships for (const entry of entries) { const typeName = entry.entityType.name; for (const rel of entry.entityType.relationships) { if (rel.type === 'many-to-one' || rel.type === 'one-to-one') { // This type depends on the related type const dependsOn = rel.targetType; if (graph.has(dependsOn)) { graph.get(typeName)!.add(dependsOn); } } } } return graph; } /** * Kahn's algorithm for topological sort */ private topologicalSort( graph: Map<string, Set<string>>, entries: EntityEntry[] ): EntityEntry[] { // Calculate in-degrees const inDegree = new Map<string, number>(); for (const [node, deps] of graph) { if (!inDegree.has(node)) { inDegree.set(node, 0); } for (const dep of deps) { inDegree.set(dep, (inDegree.get(dep) || 0)); } } for (const [node, deps] of graph) { for (const dep of deps) { // node depends on dep, so node's in-degree increases // Actually, we want nodes with NO dependencies first } } // Nodes with no dependencies come first const queue: string[] = []; for (const [node] of graph) { const deps = graph.get(node)!; if (deps.size === 0) { queue.push(node); } } const ordered: string[] = []; const processed = new Set<string>(); while (queue.length > 0) { const current = queue.shift()!; if (processed.has(current)) continue; processed.add(current); ordered.push(current); // Find nodes that depend on current for (const [node, deps] of graph) { if (deps.has(current) && !processed.has(node)) { // Check if all dependencies are now processed const allDepsProcessed = [...deps].every(d => processed.has(d)); if (allDepsProcessed) { queue.push(node); } } } } // Detect cycles if (processed.size !== graph.size) { throw new Error('Circular dependency detected in entity graph'); } // Map type order to entry order const typeOrder = new Map(ordered.map((type, index) => [type, index])); return entries.sort((a, b) => { const orderA = typeOrder.get(a.entityType.name) ?? 0; const orderB = typeOrder.get(b.entityType.name) ?? 0; return orderA - orderB; }); }}Self-referencing entities (like a tree structure with parent_id referencing the same table) require special handling. Common solutions: defer FK checks until after all inserts, insert with NULL parent_id then UPDATE, or process level-by-level from root to leaves.
Now let's bring all the components together into a complete Unit of Work implementation:
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196
interface IUnitOfWork { // Repository access readonly orders: IOrderRepository; readonly products: IProductRepository; readonly customers: ICustomerRepository; // Transaction control begin(): Promise<void>; commit(): Promise<void>; rollback(): Promise<void>; // Change tracking registerNew<T extends object>(entity: T): void; registerDirty<T extends object>(entity: T): void; registerDeleted<T extends object>(entity: T): void; // Status hasChanges(): boolean;} class UnitOfWork implements IUnitOfWork { private connection: DatabaseConnection; private transaction: DatabaseTransaction | null = null; private changeTracker: ChangeTracker; private dependencyResolver: DependencyResolver; private sqlGenerator: SqlGenerator; private isDisposed: boolean = false; // Lazy-initialized repositories private _orders: OrderRepository | null = null; private _products: ProductRepository | null = null; private _customers: CustomerRepository | null = null; constructor( connectionFactory: () => Promise<DatabaseConnection>, metadata: EntityMetadataRegistry ) { this.connection = connectionFactory(); this.changeTracker = new ChangeTracker(metadata); this.dependencyResolver = new DependencyResolver(metadata); this.sqlGenerator = new SqlGenerator(metadata); } // Repository accessors - all share the same change tracker get orders(): IOrderRepository { if (!this._orders) { this._orders = new OrderRepository(this.connection, this.changeTracker); } return this._orders; } get products(): IProductRepository { if (!this._products) { this._products = new ProductRepository(this.connection, this.changeTracker); } return this._products; } get customers(): ICustomerRepository { if (!this._customers) { this._customers = new CustomerRepository(this.connection, this.changeTracker); } return this._customers; } // Change registration (delegates to change tracker) registerNew<T extends object>(entity: T): void { this.ensureNotDisposed(); this.changeTracker.trackNew(entity); } registerDirty<T extends object>(entity: T): void { this.ensureNotDisposed(); // Note: With snapshot-based tracking, this might not be needed // The change tracker can auto-detect dirty entities } registerDeleted<T extends object>(entity: T): void { this.ensureNotDisposed(); this.changeTracker.trackDeleted(entity); } hasChanges(): boolean { const changes = this.changeTracker.getPendingChanges(); return changes.inserts.length > 0 || changes.updates.length > 0 || changes.deletes.length > 0; } // Transaction management async begin(): Promise<void> { this.ensureNotDisposed(); if (this.transaction) { throw new Error('Transaction already in progress'); } this.transaction = await this.connection.beginTransaction(); } async commit(): Promise<void> { this.ensureNotDisposed(); try { // Get all pending changes const changes = this.changeTracker.getPendingChanges(); // Order operations for dependency constraints const orderedInserts = this.dependencyResolver.orderForInsert(changes.inserts); const orderedDeletes = this.dependencyResolver.orderForDelete(changes.deletes); // Execute in correct order: deletes first (optional), then inserts, then updates // Some strategies do inserts first, then updates, then deletes // Order depends on your constraint handling preference // Execute INSERTs for (const entry of orderedInserts) { const sql = this.sqlGenerator.generateInsert(entry); const result = await this.executeWithTransaction(sql); // If ID was auto-generated, update the entity if (result.insertedId && !this.hasId(entry.entity)) { this.setId(entry.entity, entry.entityType, result.insertedId); } } // Execute UPDATEs for (const entry of changes.updates) { const sql = this.sqlGenerator.generateUpdate(entry); await this.executeWithTransaction(sql); } // Execute DELETEs for (const entry of orderedDeletes) { const sql = this.sqlGenerator.generateDelete(entry); await this.executeWithTransaction(sql); } // Commit transaction if we started one if (this.transaction) { await this.transaction.commit(); this.transaction = null; } // Mark all entities as clean this.changeTracker.acceptChanges(); } catch (error) { // Rollback on any failure await this.rollback(); throw error; } } async rollback(): Promise<void> { if (this.transaction) { await this.transaction.rollback(); this.transaction = null; } // Restore in-memory state this.changeTracker.rejectChanges(); } private async executeWithTransaction(sql: SqlStatement): Promise<QueryResult> { if (this.transaction) { return await this.transaction.execute(sql.text, sql.parameters); } return await this.connection.execute(sql.text, sql.parameters); } private ensureNotDisposed(): void { if (this.isDisposed) { throw new Error('Unit of Work has been disposed'); } } private hasId(entity: object): boolean { // Check if entity has an ID set return (entity as any).id != null; } private setId(entity: object, entityType: EntityType, id: unknown): void { (entity as any)[entityType.keyProperty] = id; } // Cleanup async dispose(): Promise<void> { if (this.isDisposed) return; if (this.transaction) { await this.rollback(); } await this.connection.close(); this.isDisposed = true; }}This implementation doesn't show concurrency handling. Production implementations typically add version columns (optimistic locking) or row locks (pessimistic locking). On UPDATE, compare the current version with the expected version and throw ConcurrencyException if they differ.
Repositories work seamlessly with the Unit of Work by sharing the same change tracker and connection. Here's how a repository integrates:
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
class OrderRepository implements IOrderRepository { private connection: DatabaseConnection; private changeTracker: ChangeTracker; constructor(connection: DatabaseConnection, changeTracker: ChangeTracker) { this.connection = connection; this.changeTracker = changeTracker; } async findById(id: string): Promise<Order | null> { // CRITICAL: Check identity map first! const tracked = this.changeTracker.getTracked<Order>('Order', id); if (tracked) { return tracked; } // Not in identity map, query database const result = await this.connection.query( 'SELECT * FROM orders WHERE id = $1', [id] ); if (result.rows.length === 0) { return null; } // Map row to entity const order = this.mapRowToOrder(result.rows[0]); // Register as clean (loaded from DB) this.changeTracker.trackLoaded(order); return order; } async findByCustomer(customerId: string): Promise<Order[]> { const result = await this.connection.query( 'SELECT * FROM orders WHERE customer_id = $1', [customerId] ); return result.rows.map(row => { // Check identity map for each row const existing = this.changeTracker.getTracked<Order>('Order', row.id); if (existing) { return existing; } const order = this.mapRowToOrder(row); this.changeTracker.trackLoaded(order); return order; }); } add(order: Order): void { // Register as new with change tracker this.changeTracker.trackNew(order); } remove(order: Order): void { this.changeTracker.trackDeleted(order); } private mapRowToOrder(row: any): Order { return new Order({ id: row.id, customerId: row.customer_id, total: parseFloat(row.total), status: row.status, createdAt: row.created_at }); }}We've built a complete Unit of Work implementation from first principles. Let's consolidate the key components and concepts:
What's Next:
With a complete Unit of Work implementation in hand, the next page explores how to integrate the Unit of Work with the Repository pattern for a cohesive persistence architecture. We'll examine patterns like Repository Factory, the relationship between aggregates and repositories, and testing strategies.
You now have a deep understanding of how to implement the Unit of Work pattern, including entity state management, identity mapping, change tracking, dependency resolution, and commit orchestration. This knowledge enables you to effectively use ORM-provided implementations and troubleshoot persistence issues.