feat: Add unit-of-work identity map with nested transaction savepoints

README.md

@@ -27,12 +27,18 @@ Most ORM tutorials stop at `findMany`. This repo goes the other way. It builds t
 - Hash join in the client: `IN (...)` then group by foreign key
 - Identity map so two children share one parent object
 - JOIN cartesian product vs two-query nesting (parent rows are not duplicated)
+- Unit of work: batches inserts/updates/deletes and flushes them as one transaction
+- Session-scoped identity map so repeated loads of the same row return one object
+- Snapshot-based dirty checking (diff the live object against its load-time snapshot, no proxies)
+- Nested transactions via SQL `SAVEPOINT`/`RELEASE`/`ROLLBACK TO`, since real nested `BEGIN` isn't a thing
+- Partial rollback: an inner failure unwinds to its savepoint without discarding the outer transaction
 
 ## What's implemented
 
 - Type-safe SELECT/WHERE/JOIN query builder with prepared-statement binding
 - Declarative schema DSL + migration runner with up/down and a diff generator
 - hasMany/belongsTo relations with eager loading to kill the N+1 problem
+- Unit-of-work / identity map for change tracking, with nested transactions via savepoints
 
 ## Usage
 
@@ -91,6 +97,23 @@ const graph = eagerHasMany(session, userOrders, parents, {
 
 A JOIN of users to orders repeats every user column once per order. The two-query path keeps parent rows unique and attaches `orders: []` when a user has none.
 
+A `UnitOfWork` tracks entities you load or create, batches the resulting SQL, and flushes them in one transaction. Loading the same row twice through `attach()` returns the same object instead of a second copy, and a plain field write is enough to mark it dirty: flush diffs the object against the snapshot it took at load time and only sends the columns that changed.
+
+```ts
+import { TransactionManager, UnitOfWork } from 'orm-workbench'
+
+const tx = new TransactionManager(session)
+const uow = new UnitOfWork(session, 'sqlite', tx)
+
+const pat = uow.attach(users, 'id', session.all('SELECT * FROM users WHERE id = ?', [1])[0])
+pat.email = 'pat@new.com'
+
+const draft = uow.registerNew(orders, 'id', { id: 42, userId: 1, totalCents: 900 })
+uow.flush() // one BEGIN/COMMIT: an UPDATE for pat, an INSERT for the order
+```
+
+`TransactionManager.run()` nests correctly: SQLite and Postgres don't support a second real `BEGIN`, so past the outermost call it issues `SAVEPOINT`/`RELEASE SAVEPOINT` instead. If the inner block throws, `ROLLBACK TO SAVEPOINT` undoes only that block; the outer transaction is still live and can catch the error and keep going, or commit what came before it.
+
 ```bash
 pnpm install
 pnpm run typecheck

src/index.ts

@@ -3,3 +3,4 @@ export * from './query'
 export * from './ddl'
 export * from './migrate'
 export * from './relations'
+export * from './unit-of-work'

src/unit-of-work.ts

@@ -0,0 +1,205 @@
+import type { SqlSession } from './migrate'
+import type { Dialect } from './query'
+import { quoteIdent, type AnyTable } from './schema'
+
+export class TransactionManager {
+  private depth = 0
+
+  constructor(private readonly session: SqlSession) {}
+
+  get isOpen(): boolean {
+    return this.depth > 0
+  }
+
+  begin(): void {
+    if (this.depth === 0) this.session.exec('BEGIN')
+    else this.session.exec(`SAVEPOINT ${this.savepointName(this.depth)}`)
+    this.depth += 1
+  }
+
+  commit(): void {
+    if (this.depth === 0) throw new Error('commit() called with no open transaction')
+    this.depth -= 1
+    if (this.depth === 0) this.session.exec('COMMIT')
+    else this.session.exec(`RELEASE SAVEPOINT ${this.savepointName(this.depth)}`)
+  }
+
+  rollback(): void {
+    if (this.depth === 0) throw new Error('rollback() called with no open transaction')
+    this.depth -= 1
+    if (this.depth === 0) {
+      this.session.exec('ROLLBACK')
+      return
+    }
+    const name = this.savepointName(this.depth)
+    this.session.exec(`ROLLBACK TO SAVEPOINT ${name}`)
+    this.session.exec(`RELEASE SAVEPOINT ${name}`)
+  }
+
+  run<T>(fn: () => T): T {
+    this.begin()
+    try {
+      const result = fn()
+      this.commit()
+      return result
+    } catch (err) {
+      this.rollback()
+      throw err
+    }
+  }
+
+  private savepointName(depth: number): string {
+    return `sp_${depth}`
+  }
+}
+
+type EntityState = 'clean' | 'new' | 'removed'
+
+type Tracked<Row extends Record<string, unknown>> = {
+  readonly table: AnyTable
+  readonly pkColumn: string
+  readonly row: Row
+  snapshot: Row
+  state: EntityState
+}
+
+function entityKey(table: AnyTable, pkValue: unknown): string {
+  return `${table.$name}:${String(pkValue)}`
+}
+
+function valuesEqual(a: unknown, b: unknown): boolean {
+  if (a instanceof Date && b instanceof Date) return a.getTime() === b.getTime()
+  return Object.is(a, b)
+}
+
+function encodeValue(value: unknown, dialect: Dialect): unknown {
+  if (dialect !== 'sqlite') return value
+  if (typeof value === 'boolean') return value ? 1 : 0
+  if (value instanceof Date) return value.toISOString()
+  return value
+}
+
+function requirePkColumn(table: AnyTable, pkColumn: string): void {
+  if (!(pkColumn in table.$columns)) {
+    throw new Error(`${table.$name} has no column ${JSON.stringify(pkColumn)}`)
+  }
+}
+
+function requireRowHasPk<Row extends Record<string, unknown>>(table: AnyTable, pkColumn: string, row: Row): void {
+  if (!(pkColumn in row) || row[pkColumn] === null || row[pkColumn] === undefined) {
+    throw new Error(`Row for ${table.$name} is missing primary key ${pkColumn}`)
+  }
+}
+
+export class UnitOfWork {
+  private readonly identityMap = new Map<string, Tracked<Record<string, unknown>>>()
+
+  constructor(
+    private readonly session: SqlSession,
+    private readonly dialect: Dialect,
+    private readonly tx: TransactionManager = new TransactionManager(session),
+  ) {}
+
+  attach<Row extends Record<string, unknown>>(table: AnyTable, pkColumn: string, row: Row): Row {
+    requirePkColumn(table, pkColumn)
+    requireRowHasPk(table, pkColumn, row)
+    const key = entityKey(table, row[pkColumn])
+    const existing = this.identityMap.get(key)
+    if (existing) {
+      if (existing.state === 'removed') throw new Error(`${key} is scheduled for removal`)
+      return existing.row as Row
+    }
+    this.identityMap.set(key, { table, pkColumn, row, snapshot: { ...row }, state: 'clean' })
+    return row
+  }
+
+  registerNew<Row extends Record<string, unknown>>(table: AnyTable, pkColumn: string, row: Row): Row {
+    requirePkColumn(table, pkColumn)
+    requireRowHasPk(table, pkColumn, row)
+    const key = entityKey(table, row[pkColumn])
+    if (this.identityMap.has(key)) throw new Error(`${key} is already tracked`)
+    this.identityMap.set(key, { table, pkColumn, row, snapshot: { ...row }, state: 'new' })
+    return row
+  }
+
+  remove(table: AnyTable, pkColumn: string, pkValue: unknown): void {
+    const key = entityKey(table, pkValue)
+    const tracked = this.identityMap.get(key)
+    if (!tracked) throw new Error(`Cannot remove untracked entity ${key}`)
+    // A 'new' entity was never flushed, so there's no row to DELETE; just drop it.
+    if (tracked.state === 'new') this.identityMap.delete(key)
+    else tracked.state = 'removed'
+  }
+
+  isTracked(table: AnyTable, pkValue: unknown): boolean {
+    return this.identityMap.has(entityKey(table, pkValue))
+  }
+
+  flush(): number {
+    const inserts: Tracked<Record<string, unknown>>[] = []
+    const updates: { tracked: Tracked<Record<string, unknown>>; columns: string[] }[] = []
+    const deletes: Tracked<Record<string, unknown>>[] = []
+    for (const tracked of this.identityMap.values()) {
+      if (tracked.state === 'new') inserts.push(tracked)
+      else if (tracked.state === 'removed') deletes.push(tracked)
+      else {
+        const columns = this.dirtyColumns(tracked)
+        if (columns.length > 0) updates.push({ tracked, columns })
+      }
+    }
+    if (inserts.length === 0 && updates.length === 0 && deletes.length === 0) return 0
+    return this.tx.run(() => {
+      for (const tracked of inserts) this.insert(tracked)
+      for (const { tracked, columns } of updates) this.update(tracked, columns)
+      // Reverse registration order is a cheap FK-safety heuristic, not a real dependency graph.
+      for (const tracked of [...deletes].reverse()) this.delete(tracked)
+      for (const tracked of inserts) {
+        tracked.state = 'clean'
+        tracked.snapshot = { ...tracked.row }
+      }
+      for (const { tracked } of updates) tracked.snapshot = { ...tracked.row }
+      for (const tracked of deletes) this.identityMap.delete(entityKey(tracked.table, tracked.row[tracked.pkColumn]))
+      return inserts.length + updates.length + deletes.length
+    })
+  }
+
+  private dirtyColumns(tracked: Tracked<Record<string, unknown>>): string[] {
+    const changed: string[] = []
+    for (const key of Object.keys(tracked.row)) {
+      if (key === tracked.pkColumn) continue
+      if (!valuesEqual(tracked.row[key], tracked.snapshot[key])) changed.push(key)
+    }
+    return changed
+  }
+
+  private ph(n: number): string {
+    return this.dialect === 'sqlite' ? '?' : `$${n}`
+  }
+
+  private insert(tracked: Tracked<Record<string, unknown>>): void {
+    const columns = Object.keys(tracked.row)
+    const placeholders = columns.map((_, i) => this.ph(i + 1))
+    const values = columns.map((c) => encodeValue(tracked.row[c], this.dialect))
+    this.session.exec(
+      `INSERT INTO ${quoteIdent(tracked.table.$name)} (${columns.map(quoteIdent).join(', ')}) VALUES (${placeholders.join(', ')})`,
+      values,
+    )
+  }
+
+  private update(tracked: Tracked<Record<string, unknown>>, columns: string[]): void {
+    const setSql = columns.map((c, i) => `${quoteIdent(c)} = ${this.ph(i + 1)}`).join(', ')
+    const values = columns.map((c) => encodeValue(tracked.row[c], this.dialect))
+    values.push(encodeValue(tracked.row[tracked.pkColumn], this.dialect))
+    this.session.exec(
+      `UPDATE ${quoteIdent(tracked.table.$name)} SET ${setSql} WHERE ${quoteIdent(tracked.pkColumn)} = ${this.ph(columns.length + 1)}`,
+      values,
+    )
+  }
+
+  private delete(tracked: Tracked<Record<string, unknown>>): void {
+    this.session.exec(
+      `DELETE FROM ${quoteIdent(tracked.table.$name)} WHERE ${quoteIdent(tracked.pkColumn)} = ${this.ph(1)}`,
+      [encodeValue(tracked.row[tracked.pkColumn], this.dialect)],
+    )
+  }
+}