From e75239a15cf31af502c92e61eef8ce4deb67f529 Mon Sep 17 00:00:00 2001 From: Thomas Hart Date: Wed, 16 Sep 2026 18:00:30 +0000 Subject: [PATCH] feat: Add unit-of-work identity map with nested transaction savepoints Co-Authored-By: Claude Sonnet 5 --- README.md | 23 +++++ src/index.ts | 1 + src/unit-of-work.ts | 205 ++++++++++++++++++++++++++++++++++++++ test/unit-of-work.test.ts | 178 +++++++++++++++++++++++++++++++++ 4 files changed, 407 insertions(+) create mode 100644 src/unit-of-work.ts create mode 100644 test/unit-of-work.test.ts diff --git a/README.md b/README.md index 79712ec..97a9c6e 100644 --- a/README.md +++ b/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 diff --git a/src/index.ts b/src/index.ts index 31a6e86..0d66c9b 100644 --- a/src/index.ts +++ b/src/index.ts @@ -3,3 +3,4 @@ export * from './query' export * from './ddl' export * from './migrate' export * from './relations' +export * from './unit-of-work' diff --git a/src/unit-of-work.ts b/src/unit-of-work.ts new file mode 100644 index 0000000..1de491d --- /dev/null +++ b/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(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> = { + 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>(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>>() + + constructor( + private readonly session: SqlSession, + private readonly dialect: Dialect, + private readonly tx: TransactionManager = new TransactionManager(session), + ) {} + + attach>(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>(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>[] = [] + const updates: { tracked: Tracked>; columns: string[] }[] = [] + const deletes: Tracked>[] = [] + 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>): 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>): 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>, 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>): void { + this.session.exec( + `DELETE FROM ${quoteIdent(tracked.table.$name)} WHERE ${quoteIdent(tracked.pkColumn)} = ${this.ph(1)}`, + [encodeValue(tracked.row[tracked.pkColumn], this.dialect)], + ) + } +} diff --git a/test/unit-of-work.test.ts b/test/unit-of-work.test.ts new file mode 100644 index 0000000..1f5ea7d --- /dev/null +++ b/test/unit-of-work.test.ts @@ -0,0 +1,178 @@ +import { expect, test } from 'vitest' +import { boolean, defineTable, integer, text, TransactionManager, UnitOfWork, type SqlSession } from '../src/index' + +const accounts = defineTable('accounts', { id: integer(), name: text(), balanceCents: integer(), active: boolean() }) + +type SqliteDb = { + exec: (sql: string) => void + prepare: (sql: string) => { all: (...p: unknown[]) => unknown; run: (...p: unknown[]) => unknown } + close: () => void +} + +function sqlite(): { session: SqlSession; queries: () => number; rows: () => Record[]; close: () => void } { + const loaded = ( + globalThis as { process?: { getBuiltinModule?: (id: string) => { DatabaseSync: new (path: string) => SqliteDb } } } + ).process?.getBuiltinModule?.('node:sqlite') + if (!loaded) throw new Error('node:sqlite is not available in this runtime') + const db = new loaded.DatabaseSync(':memory:') + db.exec(`CREATE TABLE "accounts" ("id" INTEGER PRIMARY KEY, "name" TEXT NOT NULL, "balanceCents" INTEGER NOT NULL, "active" INTEGER NOT NULL)`) + let queries = 0 + const session: SqlSession = { + exec: (sql, params) => { + queries += 1 + if (!params || params.length === 0) db.exec(sql) + else db.prepare(sql).run(...params) + }, + all: (sql, params) => { + queries += 1 + return db.prepare(sql).all(...(params ?? [])) as Record[] + }, + } + return { + session, + queries: () => queries, + rows: () => db.prepare('SELECT * FROM "accounts" ORDER BY "id"').all() as Record[], + close: () => db.close(), + } +} + +test('identity map returns the same object for repeated attach() of one row', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + const row = { id: 1, name: 'pat', balanceCents: 500, active: true } + const first = uow.attach(accounts, 'id', row) + const second = uow.attach(accounts, 'id', { id: 1, name: 'stale-copy', balanceCents: 0, active: false }) + expect(second).toBe(first) + expect(second.name).toBe('pat') + db.close() +}) + +test('attach and registerNew reject rows missing their primary key', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + expect(() => uow.attach(accounts, 'id', { name: 'x' } as never)).toThrow(/missing primary key/) + expect(() => uow.registerNew(accounts, 'id', { id: null, name: 'x' } as never)).toThrow(/missing primary key/) + expect(() => uow.attach(accounts, 'ssn', { id: 1 } as never)).toThrow(/no column/) + db.close() +}) + +test('registerNew rejects a pk already tracked; remove rejects an untracked pk', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + uow.registerNew(accounts, 'id', { id: 1, name: 'a', balanceCents: 0, active: true }) + expect(() => uow.registerNew(accounts, 'id', { id: 1, name: 'b', balanceCents: 0, active: true })).toThrow(/already tracked/) + expect(() => uow.remove(accounts, 'id', 99)).toThrow(/untracked/) + db.close() +}) + +test('flush() with nothing pending never touches the session', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + uow.attach(accounts, 'id', { id: 1, name: 'a', balanceCents: 0, active: true }) + const before = db.queries() + expect(uow.flush()).toBe(0) + expect(db.queries()).toBe(before) + db.close() +}) + +test('flush() inserts new entities inside a transaction', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + uow.registerNew(accounts, 'id', { id: 1, name: 'pat', balanceCents: 500, active: true }) + uow.registerNew(accounts, 'id', { id: 2, name: 'lee', balanceCents: 0, active: false }) + expect(uow.flush()).toBe(2) + expect(db.rows()).toEqual([ + { id: 1, name: 'pat', balanceCents: 500, active: 1 }, + { id: 2, name: 'lee', balanceCents: 0, active: 0 }, + ]) +}) + +test('flush() only updates columns that actually changed', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + const row = uow.registerNew(accounts, 'id', { id: 1, name: 'pat', balanceCents: 500, active: true }) + uow.flush() + row.balanceCents = 400 + const updateCount = uow.flush() + expect(updateCount).toBe(1) + expect(db.rows()).toEqual([{ id: 1, name: 'pat', balanceCents: 400, active: 1 }]) + expect(uow.flush()).toBe(0) +}) + +test('remove() deletes on flush and forgets the entity so a fresh attach starts clean', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + uow.registerNew(accounts, 'id', { id: 1, name: 'pat', balanceCents: 500, active: true }) + uow.flush() + uow.remove(accounts, 'id', 1) + expect(uow.flush()).toBe(1) + expect(db.rows()).toEqual([]) + expect(uow.isTracked(accounts, 1)).toBe(false) + const revived = uow.attach(accounts, 'id', { id: 1, name: 'new-pat', balanceCents: 0, active: false }) + expect(revived.name).toBe('new-pat') +}) + +test('remove() on a never-flushed new entity just drops it, no DELETE is issued', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + uow.registerNew(accounts, 'id', { id: 1, name: 'pat', balanceCents: 500, active: true }) + uow.remove(accounts, 'id', 1) + expect(uow.isTracked(accounts, 1)).toBe(false) + const before = db.queries() + expect(uow.flush()).toBe(0) + expect(db.queries()).toBe(before) + db.close() +}) + +test('a failing flush rolls back every write and leaves entities pending for retry', () => { + const db = sqlite() + const uow = new UnitOfWork(db.session, 'sqlite') + uow.registerNew(accounts, 'id', { id: 1, name: 'ok', balanceCents: 0, active: true }) + db.session.exec(`INSERT INTO "accounts" ("id", "name", "balanceCents", "active") VALUES (2, 'seed', 0, 0)`) + uow.registerNew(accounts, 'id', { id: 2, name: 'collides-with-existing-row', balanceCents: 0, active: true }) + expect(() => uow.flush()).toThrow() + expect(db.rows().map((r) => r.id)).toEqual([2]) + expect(db.rows()[0]?.name).toBe('seed') + uow.remove(accounts, 'id', 2) + expect(uow.flush()).toBe(1) + expect(db.rows().map((r) => r.name).sort()).toEqual(['ok', 'seed']) + db.close() +}) + +test('TransactionManager rejects commit/rollback with no open transaction', () => { + const db = sqlite() + const tx = new TransactionManager(db.session) + expect(() => tx.commit()).toThrow(/no open transaction/) + expect(() => tx.rollback()).toThrow(/no open transaction/) + db.close() +}) + +test('nested run() uses a savepoint so an inner failure does not undo the outer transaction', () => { + const db = sqlite() + const tx = new TransactionManager(db.session) + tx.run(() => { + db.session.exec(`INSERT INTO "accounts" ("id", "name", "balanceCents", "active") VALUES (1, 'a', 0, 1)`) + expect(() => + tx.run(() => { + db.session.exec(`INSERT INTO "accounts" ("id", "name", "balanceCents", "active") VALUES (2, 'b', 0, 1)`) + throw new Error('inner boom') + }), + ).toThrow('inner boom') + db.session.exec(`INSERT INTO "accounts" ("id", "name", "balanceCents", "active") VALUES (3, 'c', 0, 1)`) + }) + expect(db.rows().map((r) => r.id)).toEqual([1, 3]) + expect(tx.isOpen).toBe(false) +}) + +test('UnitOfWork.flush() nested inside a manual transaction is undone by an outer rollback', () => { + const db = sqlite() + const tx = new TransactionManager(db.session) + const uow = new UnitOfWork(db.session, 'sqlite', tx) + tx.begin() + uow.registerNew(accounts, 'id', { id: 1, name: 'pat', balanceCents: 500, active: true }) + expect(uow.flush()).toBe(1) + expect(db.rows()).toHaveLength(1) + tx.rollback() + expect(db.rows()).toEqual([]) + expect(tx.isOpen).toBe(false) +})