pizzasql/pkg/storage/tx.go

package storage

import (
	"fmt"
	"sort"
	"strings"
	"sync"
)

// ErrSerialization is returned when a transaction's optimistic validation fails
// because a concurrent transaction committed a conflicting change.
var ErrSerialization = fmt.Errorf("serialization failure: concurrent transaction modified the database")

// Session is a buffered SQL transaction/session wrapper around a TableManager.
// While in a transaction, writes are staged in memory, reads observe that
// staged overlay, and COMMIT issues a single atomic compare-and-swap batch
// write. Rollback simply discards the staged changes (no compensating durable
// writes). A Session is used for both autocommit statements (where operations
// delegate straight to the durable TableManager) and buffered transactions.
type Session struct {
	schema *SchemaManager
	table  *TableManager

	mu sync.Mutex

	inTx    bool
	aborted bool

	// overlay holds staged writes keyed by lowercased table then data key.
	overlay map[string]map[string]*overlayEntry

	// log records mutations in order so savepoints can roll back.
	log []txMutation

	// reads records the durable LSN of every key the transaction observed
	// (0 means the key was observed absent). It becomes the compare set at
	// commit and also captures the base LSN of every written key.
	reads map[string]uint64

	// scanGens records the per-table generation captured by the first scan of
	// each table, validated at commit to detect phantoms. Indexed equality reads
	// use predicateGens so writes to other index values do not cause conflicts.
	scanGens      map[string]uint64
	predicateGens map[string]predicateRead
}

type predicateRead struct {
	table string
	gen   uint64
}

type overlayEntry struct {
	row    Row
	absent bool
}

type txMutation struct {
	table string
	key   string
	prev  *overlayEntry
}

// NewSession creates a session wrapping the given schema and table managers.
func NewSession(schema *SchemaManager, table *TableManager) *Session {
	return &Session{
		schema:        schema,
		table:         table,
		overlay:       make(map[string]map[string]*overlayEntry),
		reads:         make(map[string]uint64),
		scanGens:      make(map[string]uint64),
		predicateGens: make(map[string]predicateRead),
	}
}

// Begin starts a buffered transaction.
func (s *Session) Begin() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.inTx {
		return fmt.Errorf("cannot start a transaction within a transaction")
	}
	s.inTx = true
	s.aborted = false
	s.overlay = make(map[string]map[string]*overlayEntry)
	s.log = nil
	s.reads = make(map[string]uint64)
	s.scanGens = make(map[string]uint64)
	s.predicateGens = make(map[string]predicateRead)
	return nil
}

// InTx reports whether a transaction is in progress.
func (s *Session) InTx() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.inTx
}

// Abort marks the transaction as aborted without discarding state.
func (s *Session) Abort() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.inTx {
		s.aborted = true
	}
}

// Snapshot returns the current mutation-log position for a savepoint.
func (s *Session) Snapshot() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.log)
}

// RollbackTo discards mutations after the given savepoint position.
func (s *Session) RollbackTo(pos int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := len(s.log) - 1; i >= pos && i >= 0; i-- {
		m := s.log[i]
		if m.prev == nil {
			delete(s.overlay[m.table], m.key)
			if len(s.overlay[m.table]) == 0 {
				delete(s.overlay, m.table)
			}
		} else {
			if s.overlay[m.table] == nil {
				s.overlay[m.table] = make(map[string]*overlayEntry)
			}
			s.overlay[m.table][m.key] = m.prev
		}
	}
	if pos < len(s.log) {
		s.log = s.log[:pos]
	}
	s.aborted = false
}

// Rollback discards the transaction without writing anything durable.
func (s *Session) Rollback() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		return fmt.Errorf("cannot rollback: no transaction in progress")
	}
	s.resetLocked()
	return nil
}

// Commit validates and durably applies the staged transaction in one atomic
// compare-and-swap batch write. It discards the commit LSN; callers that need a
// ordering cursor use CommitWithLSN.
func (s *Session) Commit() error {
	_, err := s.CommitWithLSN()
	return err
}

// CommitWithLSN behaves like Commit and additionally returns the LSN assigned
// to the durable batch. A read-only or empty transaction returns zero because
// no new durable state was published.
func (s *Session) CommitWithLSN() (uint64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		return 0, fmt.Errorf("cannot commit: no transaction in progress")
	}
	if s.aborted {
		s.resetLocked()
		return 0, fmt.Errorf("current transaction is aborted")
	}

	// Referential integrity spans tables: a commit must validate the final
	// child rows against the final parent rows, including rows another table
	// staged in this same transaction. Take the foreign-key write gate before
	// any table gate (the canonical lock order) and register the constrained
	// component as scan reads so acquireCommitLocks holds it exclusively while
	// the validation scans run. A concurrent write on either side would bump one
	// of those generations and fail the commit with 40001 instead of committing
	// a dangling reference.
	closure, err := s.foreignKeyClosure()
	if err != nil {
		s.resetLocked()
		return 0, err
	}
	if len(closure) > 0 {
		s.table.foreignKeyMu.Lock()
		defer s.table.foreignKeyMu.Unlock()
	}
	for table := range closure {
		if _, ok := s.scanGens[table]; !ok {
			s.scanGens[table] = s.table.generation(table)
		}
	}

	// Acquire the per-table gates, escalating any table whose UNIQUE index
	// appeared after the initial unlocked probe. See acquireCommitLocks.
	_, _, unlockAll, uniqueTables, err := s.acquireCommitLocks()
	if err != nil {
		s.resetLocked()
		return 0, err
	}
	defer unlockAll()

	// A staged row was normalized against the schema as of its INSERT or UPDATE,
	// which a concurrent ADD COLUMN can predate. Re-apply the schema that is
	// current under the held gates before any derived state is computed: the CAS
	// check cannot reject such a row (a new key is only checked for absence), so
	// a NOT NULL column published after the row image was built would otherwise
	// be violated durably, and a UNIQUE index would not see the default the row
	// is about to receive. Run it before uniqueness and cache planning so both
	// observe the final values.
	if err := s.applyOverlaySchemaDefaults(); err != nil {
		s.resetLocked()
		return 0, err
	}

	// Validate scan generations for phantom detection.
	for t, gen := range s.scanGens {
		if s.table.generation(t) != gen {
			s.resetLocked()
			return 0, ErrSerialization
		}
	}
	for key, predicate := range s.predicateGens {
		if s.table.predicateGeneration(key) != predicate.gen {
			s.resetLocked()
			return 0, ErrSerialization
		}
	}

	// Validate the final overlay against every foreign key in the constrained
	// component. The per-table gates acquired above exclude concurrent writers,
	// so this sees a stable state; checking the final state (rather than each
	// statement) is what lets a transaction delete a parent and its children in
	// either order as long as the committed state stays consistent.
	if err := s.validateForeignKeyFinalState(closure); err != nil {
		s.resetLocked()
		return 0, err
	}

	// Validate the final overlay of every unique-indexed written table against
	// durable rows. This runs under the exclusive gate acquired above, so it
	// cannot race a concurrent writer. Duplicate unique values fail the commit
	// with a "UNIQUE constraint failed" error rather than ErrSerialization.
	uniquePlans := make(map[string]uniqueCachePlan, len(uniqueTables))
	for t := range uniqueTables {
		plan, err := s.table.prepareUniqueRows(t, s.overlay[t])
		if err != nil {
			s.resetLocked()
			return 0, err
		}
		uniquePlans[t] = plan
	}

	exclusiveTables := make(map[string]bool)
	for table := range uniqueTables {
		exclusiveTables[table] = true
	}
	for table := range s.scanGens {
		exclusiveTables[table] = true
	}
	for _, predicate := range s.predicateGens {
		exclusiveTables[predicate.table] = true
	}
	cachePlans := make(map[string]*tableCachePlan)
	for table, entries := range s.overlay {
		if exclusiveTables[table] {
			plan, err := s.table.prepareTableCache(table, entries, s.reads)
			if err != nil {
				s.resetLocked()
				return 0, err
			}
			cachePlans[table] = plan
		}
	}

	// Build the compare set from every observed key.
	checks := make([]CompareCheck, 0, len(s.reads))
	for key, lsn := range s.reads {
		checks = append(checks, CompareCheck{Key: []byte(key), LSN: lsn})
	}
	sort.Slice(checks, func(i, j int) bool { return string(checks[i].Key) < string(checks[j].Key) })

	// Build the batch ops from the staged overlay.
	ops := make([]BatchOp, 0)
	for _, entries := range s.overlay {
		for key, e := range entries {
			if e.absent {
				ops = append(ops, BatchOp{Op: batchDelete, Key: []byte(key)})
			} else {
				data, err := encodeRow(e.row)
				if err != nil {
					s.resetLocked()
					return 0, err
				}
				ops = append(ops, BatchOp{Op: batchPut, Key: []byte(key), Value: data})
			}
		}
	}
	sort.Slice(ops, func(i, j int) bool { return string(ops[i].Key) < string(ops[j].Key) })

	// A read-only transaction has nothing to write; commit trivially.
	if len(ops) == 0 {
		s.resetLocked()
		return 0, nil
	}

	var committed bool
	var commitLSN uint64
	err = s.table.pool.WithClient(func(c *KVClient) error {
		lsn, ok, err := c.CompareBatchWrite(checks, ops, nil)
		commitLSN = lsn
		committed = ok
		return err
	})
	if err != nil || !committed {
		// A transport error can hide a durable commit; a compare conflict can
		// also expose an out-of-band writer. Rebuild derived state next time.
		for table := range s.overlay {
			s.table.invalidateWriteFailure(table)
		}
	}
	if err != nil {
		s.resetLocked()
		return 0, err
	}
	if !committed {
		s.resetLocked()
		return 0, ErrSerialization
	}

	// Publish derived state before generations while retaining the table gates.
	for t := range s.overlay {
		if plan := cachePlans[t]; plan != nil {
			s.table.applyTableCache(plan)
		} else {
			s.table.invalidateCache(t)
		}
		s.table.applyUniquePlan(uniquePlans[t], nil)
		s.table.bumpIndexPredicateWildcard(t)
		s.table.bumpGeneration(t)
	}

	s.resetLocked()
	return commitLSN, nil
}

// acquireCommitLocks acquires the per-table gates needed to commit the staged
// transaction, returning the held locks, whether each is exclusive, an unlock
// function, and the set of written tables that carry a UNIQUE index. A written
// table with a UNIQUE index is locked exclusively so its final overlay can be
// validated against durable rows. The uniqueness probe is re-checked under the
// held gates so a concurrent CREATE UNIQUE INDEX cannot slip in after the probe
// and leave a duplicate unvalidated.
func (s *Session) acquireCommitLocks() (locks []*sync.RWMutex, exclusive []bool, unlockAll func(), uniqueTables map[string]bool, err error) {
	unlock := func(ls []*sync.RWMutex, ex []bool) {
		for i := len(ls) - 1; i >= 0; i-- {
			if ex[i] {
				ls[i].Unlock()
			} else {
				ls[i].RUnlock()
			}
		}
	}
	for {
		// Tables read through scans or predicates need an exclusive validation
		// gate. Tables that are only written use the shared publication gate,
		// allowing disjoint optimistic commits to proceed concurrently while
		// still excluding scanner commits.
		affected := make(map[string]bool)
		for t := range s.overlay {
			affected[t] = false
		}
		for t := range s.scanGens {
			affected[t] = true
		}
		for _, predicate := range s.predicateGens {
			affected[predicate.table] = true
		}

		// Tables written in this transaction that carry a UNIQUE index need the
		// exclusive gate so their final overlay can be validated against durable
		// rows without racing another writer.
		uniqueTables = make(map[string]bool)
		for t := range s.overlay {
			uniq, err := s.table.hasUniqueIndex(t)
			if err != nil {
				return nil, nil, nil, nil, err
			}
			if uniq {
				uniqueTables[t] = true
				affected[t] = true
			}
		}

		tables := make([]string, 0, len(affected))
		for t := range affected {
			tables = append(tables, t)
		}
		sort.Strings(tables)

		locks = make([]*sync.RWMutex, len(tables))
		exclusive = make([]bool, len(tables))
		for i, t := range tables {
			locks[i] = s.table.tableLock(t)
			exclusive[i] = affected[t]
		}
		for i, lock := range locks {
			if exclusive[i] {
				lock.Lock()
			} else {
				lock.RLock()
			}
		}

		// Re-check under the held gates. A concurrent CREATE UNIQUE INDEX can
		// only have completed before we acquired the gate (it needs the exclusive
		// gate), so this probe is authoritative; escalate and retry if a written
		// table gained a unique index.
		retry := false
		for i, t := range tables {
			if exclusive[i] {
				continue
			}
			if _, written := s.overlay[t]; !written {
				continue
			}
			uniq, err := s.table.hasUniqueIndex(t)
			if err != nil {
				unlock(locks, exclusive)
				return nil, nil, nil, nil, err
			}
			if uniq {
				retry = true
				break
			}
		}
		if retry {
			unlock(locks, exclusive)
			continue
		}

		unlockAll = func() { unlock(locks, exclusive) }
		return locks, exclusive, unlockAll, uniqueTables, nil
	}
}

// applyOverlaySchemaDefaults re-applies the current schema to every staged row
// (the caller holds the per-table gates). Columns published after a row image
// was built receive their catalog default, and required columns are validated,
// so a commit cannot persist a row that violates the schema now in effect.
func (s *Session) applyOverlaySchemaDefaults() error {
	for table, entries := range s.overlay {
		schema, err := s.schema.GetSchema(table)
		if err != nil {
			return err
		}
		for _, entry := range entries {
			if entry.absent {
				continue
			}
			if err := s.schema.applySchemaDefaults(schema, entry.row); err != nil {
				return err
			}
			if err := s.schema.validateChecks(schema, entry.row); err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *Session) resetLocked() {
	s.inTx = false
	s.aborted = false
	s.overlay = make(map[string]map[string]*overlayEntry)
	s.log = nil
	s.reads = make(map[string]uint64)
	s.scanGens = make(map[string]uint64)
	s.predicateGens = make(map[string]predicateRead)
}

func (s *Session) stagePut(table, key string, row Row) {
	tl := strings.ToLower(table)
	if s.overlay[tl] == nil {
		s.overlay[tl] = make(map[string]*overlayEntry)
	}
	s.log = append(s.log, txMutation{table: tl, key: key, prev: s.overlay[tl][key]})
	s.overlay[tl][key] = &overlayEntry{row: cloneRow(row)}
	if _, ok := s.reads[key]; !ok {
		s.reads[key] = 0
	}
}

func (s *Session) stageDelete(table, key string, row Row) {
	tl := strings.ToLower(table)
	if s.overlay[tl] == nil {
		s.overlay[tl] = make(map[string]*overlayEntry)
	}
	s.log = append(s.log, txMutation{table: tl, key: key, prev: s.overlay[tl][key]})
	// Keep the deleted row so COMMIT can exempt its rowid from the unique-index
	// scan and release its unique value.
	s.overlay[tl][key] = &overlayEntry{row: cloneRow(row), absent: true}
	if _, ok := s.reads[key]; !ok {
		s.reads[key] = 0
	}
}

// GetByPK retrieves a row by primary key, observing the staged overlay in a
// transaction and recording the observed LSN for validation.
func (s *Session) GetByPK(table, pk string) (Row, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		return s.table.GetByPK(table, pk)
	}
	key := s.table.dataKey(table, pk)
	tl := strings.ToLower(table)
	if e, ok := s.overlay[tl][key]; ok {
		if e.absent {
			return nil, ErrKeyNotFound
		}
		return cloneRow(e.row), nil
	}
	row, lsn, err := s.table.getByPKWithLSN(table, pk)
	if err != nil {
		if err == ErrKeyNotFound {
			s.reads[key] = 0
			return nil, ErrKeyNotFound
		}
		return nil, err
	}
	s.reads[key] = lsn
	return row, nil
}

// Select scans a table, merging the staged overlay so a transaction sees its
// own writes, and records per-row LSNs plus the table generation.
func (s *Session) Select(table string, filter func(Row) bool) ([]Row, error) {
	s.mu.Lock()
	if !s.inTx {
		s.mu.Unlock()
		return s.table.Select(table, filter)
	}
	defer s.mu.Unlock()
	return s.selectLocked(table, filter)
}

// SelectLimited bounds an autocommit scan. Explicit transactions retain the
// complete overlay/read-set behavior until bounded transactional scans have a
// separate phantom/overlay contract; their caller still applies LIMIT later.
func (s *Session) SelectLimited(table string, filter func(Row) bool, maxRows int) ([]Row, error) {
	s.mu.Lock()
	if s.inTx {
		defer s.mu.Unlock()
		return s.selectLocked(table, filter)
	}
	s.mu.Unlock()
	if maxRows == 0 {
		return []Row{}, nil
	}
	return s.table.SelectWithLimit(table, filter, maxRows, 0)
}

func (s *Session) selectLocked(table string, filter func(Row) bool) ([]Row, error) {
	schema, err := s.schema.GetSchema(table)
	if err != nil {
		return nil, err
	}
	tl := strings.ToLower(table)
	tableLock := s.table.tableLock(tl)
	tableLock.RLock()
	defer tableLock.RUnlock()
	if _, ok := s.scanGens[tl]; !ok {
		s.scanGens[tl] = s.table.generation(tl)
	}
	overlay := s.overlay[tl]

	var rows []Row
	err = s.table.scanRowsWithLSN(table, func(row Row, lsn uint64) (bool, error) {
		pk, err := schema.PrimaryKeyForRow(row)
		if err != nil {
			return true, err
		}
		key := s.table.dataKey(table, pk)
		if _, ok := overlay[key]; ok {
			return false, nil
		}
		s.reads[key] = lsn
		if filter == nil || filter(row) {
			rows = append(rows, row)
		}
		return false, nil
	})
	if err != nil {
		return nil, err
	}
	for _, e := range overlay {
		if e.absent {
			continue
		}
		if filter == nil || filter(e.row) {
			rows = append(rows, cloneRow(e.row))
		}
	}
	return rows, nil
}

// SelectByIndex reads only matching durable rows while capturing their LSNs,
// then merges the transaction overlay. The table generation protects against
// matching rows being inserted or removed after the lookup.
func (s *Session) SelectByIndex(table, indexName string, colValue interface{}) ([]Row, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		return s.table.SelectByIndex(table, indexName, colValue)
	}
	tableKey := strings.ToLower(table)
	index, err := s.schema.GetIndex(indexName)
	if err != nil {
		return nil, err
	}
	col := index.Columns[0].Name
	want := formatIndexValue(colValue)
	versions, predicate, err := s.table.selectByIndexWithLSN(table, indexName, colValue)
	if err != nil {
		return nil, err
	}
	if _, ok := s.predicateGens[predicate.valueKey]; !ok {
		s.predicateGens[predicate.valueKey] = predicateRead{table: tableKey, gen: predicate.valueGen}
	}
	if _, ok := s.predicateGens[predicate.wildcardKey]; !ok {
		s.predicateGens[predicate.wildcardKey] = predicateRead{table: tableKey, gen: predicate.wildcardGen}
	}
	overlay := s.overlay[tableKey]
	rows := make([]Row, 0, len(versions)+len(overlay))
	for _, version := range versions {
		if _, staged := overlay[version.key]; staged {
			continue
		}
		s.reads[version.key] = version.lsn
		rows = append(rows, version.row)
	}
	for _, entry := range overlay {
		if !entry.absent && formatIndexValue(entry.row[col]) == want {
			rows = append(rows, cloneRow(entry.row))
		}
	}
	return rows, nil
}

// SelectByIndexKey retrieves the rows whose index key equals valueKey, merging
// the staged overlay so a transaction sees its own writes. valueKey is the
// formatted key from TableManager.IndexRowKey, which supports plain, composite,
// and expression indexes uniformly.
func (s *Session) SelectByIndexKey(table string, idx *Index, valueKey string) ([]Row, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		versions, _, err := s.table.selectByIndexKeyWithLSN(table, idx.Name, valueKey)
		if err != nil {
			return nil, err
		}
		rows := make([]Row, len(versions))
		for i := range versions {
			rows[i] = versions[i].row
		}
		return rows, nil
	}

	tableKey := strings.ToLower(table)
	versions, predicate, err := s.table.selectByIndexKeyWithLSN(table, idx.Name, valueKey)
	if err != nil {
		return nil, err
	}
	if _, ok := s.predicateGens[predicate.valueKey]; !ok {
		s.predicateGens[predicate.valueKey] = predicateRead{table: tableKey, gen: predicate.valueGen}
	}
	if _, ok := s.predicateGens[predicate.wildcardKey]; !ok {
		s.predicateGens[predicate.wildcardKey] = predicateRead{table: tableKey, gen: predicate.wildcardGen}
	}
	overlay := s.overlay[tableKey]
	rows := make([]Row, 0, len(versions)+len(overlay))
	for _, version := range versions {
		if _, staged := overlay[version.key]; staged {
			continue
		}
		s.reads[version.key] = version.lsn
		rows = append(rows, version.row)
	}
	for _, entry := range overlay {
		if entry.absent {
			continue
		}
		key, kerr := s.table.IndexRowKey(idx, entry.row)
		if kerr != nil {
			return nil, kerr
		}
		if key == valueKey {
			rows = append(rows, cloneRow(entry.row))
		}
	}
	return rows, nil
}

// CountFast returns the exact row count, observing the staged overlay in a
// transaction.
func (s *Session) CountFast(table string) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		return s.table.CountFast(table)
	}
	rows, err := s.selectLocked(table, nil)
	if err != nil {
		return 0, err
	}
	return len(rows), nil
}

// Insert stages an insert in a transaction, or performs a durable autocommit
// insert otherwise.
func (s *Session) Insert(table string, row Row) error {
	_, err := s.InsertWithRowID(table, row)
	return err
}

// InsertWithRowID stages an insert in a transaction, or performs a durable
// autocommit insert otherwise, returning the actual generated ROWID.
func (s *Session) InsertWithRowID(table string, row Row) (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		return s.table.InsertWithRowID(table, row)
	}
	return s.insertLocked(table, row)
}

// InsertBulk stages or durably bulk-inserts multiple rows.
func (s *Session) InsertBulk(table string, rows []Row) (int, error) {
	count, _, err := s.InsertBulkWithLastRowID(table, rows)
	return count, err
}

// InsertBulkWithLastRowID is InsertBulk, additionally returning the ROWID of the
// last staged/persisted row (0 when nothing was inserted).
func (s *Session) InsertBulkWithLastRowID(table string, rows []Row) (int, int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		return s.table.InsertBulkWithLastRowID(table, rows)
	}
	var lastRowID int64
	count := 0
	for _, row := range rows {
		rid, err := s.insertLocked(table, row)
		if err != nil {
			return count, lastRowID, err
		}
		lastRowID = rid
		count++
	}
	return count, lastRowID, nil
}

// insertLocked is the transaction insert helper (caller holds s.mu).
func (s *Session) insertLocked(table string, row Row) (int64, error) {
	nr, key, err := s.table.prepareInsert(table, row)
	if err != nil {
		return 0, err
	}
	rowid, _ := rowIDFromRow(nr)
	tl := strings.ToLower(table)
	// prepareInsert returned the durable key; duplicate reporting shows the same
	// encoded primary-key component the key carries.
	pk := strings.TrimPrefix(key, s.table.dataPrefix(table))

	if e, ok := s.overlay[tl][key]; ok {
		if !e.absent {
			return 0, &PrimaryKeyViolationError{Table: table, Key: pk}
		}
	} else {
		_, lsn, err := s.table.getByPKWithLSN(table, pk)
		if err == nil {
			s.reads[key] = lsn
			return 0, &PrimaryKeyViolationError{Table: table, Key: pk}
		}
		if err != ErrKeyNotFound {
			return 0, err
		}
		s.reads[key] = 0
	}

	s.stagePut(table, key, nr)
	return rowid, nil
}

// UpdateByPK stages or durably applies a single-row update.
func (s *Session) UpdateByPK(table, pk string, updateFn func(Row) (Row, error)) (Row, bool, error) {
	s.mu.Lock()
	if !s.inTx {
		s.mu.Unlock()
		return s.table.UpdateByPK(table, pk, updateFn)
	}
	defer s.mu.Unlock()

	schema, err := s.schema.GetSchema(table)
	if err != nil {
		return nil, false, err
	}
	key := s.table.dataKey(table, pk)

	row, err := s.getByPKLocked(table, pk)
	if err == ErrKeyNotFound {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}

	oldRow := cloneRow(row)
	updates, err := s.runUpdateFn(updateFn, row)
	if err != nil {
		return nil, false, err
	}
	for name, value := range updates {
		for _, column := range schema.Columns {
			if strings.EqualFold(name, column.Name) {
				row[column.Name] = value
				break
			}
		}
	}

	if err := s.schema.validateRowConstraints(schema, row); err != nil {
		return nil, false, err
	}
	s.stagePut(table, key, row)
	return oldRow, true, nil
}

// runUpdateFn releases the session lock while invoking the update callback so the
// callback can run nested reads (e.g. a scalar subquery in SET) through the same
// session without deadlocking on s.mu. A connection executes one statement at a
// time, so the staged overlay cannot change while the callback runs. The lock is
// re-acquired before returning.
func (s *Session) runUpdateFn(updateFn func(Row) (Row, error), row Row) (Row, error) {
	s.mu.Unlock()
	defer s.mu.Lock()
	return updateFn(row)
}

// DeleteByPK stages or durably applies a single-row delete.
func (s *Session) DeleteByPK(table, pk string) (Row, bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if !s.inTx {
		return s.table.DeleteByPK(table, pk)
	}

	key := s.table.dataKey(table, pk)
	row, err := s.getByPKLocked(table, pk)
	if err == ErrKeyNotFound {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}
	s.stageDelete(table, key, row)
	return row, true, nil
}

// getByPKLocked reads a row observing the overlay (caller holds s.mu).
func (s *Session) getByPKLocked(table, pk string) (Row, error) {
	key := s.table.dataKey(table, pk)
	tl := strings.ToLower(table)
	if e, ok := s.overlay[tl][key]; ok {
		if e.absent {
			return nil, ErrKeyNotFound
		}
		return cloneRow(e.row), nil
	}
	row, lsn, err := s.table.getByPKWithLSN(table, pk)
	if err != nil {
		if err == ErrKeyNotFound {
			s.reads[key] = 0
			return nil, ErrKeyNotFound
		}
		return nil, err
	}
	s.reads[key] = lsn
	return row, nil
}

// UpdateFunc stages or durably applies a scan-based update.
func (s *Session) UpdateFunc(table string, updateFn func(Row) (Row, error), filter func(Row) bool) (int, error) {
	s.mu.Lock()
	if !s.inTx {
		s.mu.Unlock()
		return s.table.UpdateFunc(table, updateFn, filter)
	}
	defer s.mu.Unlock()

	schema, err := s.schema.GetSchema(table)
	if err != nil {
		return 0, err
	}
	rows, err := s.selectLocked(table, filter)
	if err != nil {
		return 0, err
	}
	count := 0
	for _, row := range rows {
		oldRow := cloneRow(row)
		oldPK, err := schema.PrimaryKeyForRow(row)
		if err != nil {
			return count, err
		}
		oldKey := s.table.dataKey(table, oldPK)
		updates, err := s.runUpdateFn(updateFn, row)
		if err != nil {
			return count, err
		}
		for name, value := range updates {
			for _, column := range schema.Columns {
				if strings.EqualFold(name, column.Name) {
					row[column.Name] = value
					break
				}
			}
		}
		if err := s.schema.validateRowConstraints(schema, row); err != nil {
			return count, err
		}
		newPK, err := schema.PrimaryKeyForRow(row)
		if err != nil {
			return count, err
		}
		newKey := s.table.dataKey(table, newPK)
		if newKey != oldKey {
			// The primary key changed: move the staged row instead of leaving a
			// stale copy under the old key, and refuse to overwrite another row.
			if _, gerr := s.getByPKLocked(table, newPK); gerr == nil {
				return count, &PrimaryKeyViolationError{Table: table, Key: newPK}
			} else if gerr != ErrKeyNotFound {
				return count, gerr
			}
			s.stageDelete(table, oldKey, oldRow)
		}
		s.stagePut(table, newKey, row)
		count++
	}
	return count, nil
}

// Delete stages or durably applies a scan-based delete.
func (s *Session) Delete(table string, filter func(Row) bool) (int, error) {
	s.mu.Lock()
	if !s.inTx {
		s.mu.Unlock()
		return s.table.Delete(table, filter)
	}
	defer s.mu.Unlock()

	schema, err := s.schema.GetSchema(table)
	if err != nil {
		return 0, err
	}
	rows, err := s.selectLocked(table, filter)
	if err != nil {
		return 0, err
	}
	count := 0
	for _, row := range rows {
		pk, err := schema.PrimaryKeyForRow(row)
		if err != nil {
			return count, err
		}
		s.stageDelete(table, s.table.dataKey(table, pk), row)
		count++
	}
	return count, nil
}

// ClearIndex passes through to the durable table manager.
func (s *Session) ClearIndex(indexName, tableName string, columns []string) error {
	return s.table.ClearIndex(indexName, tableName, columns)
}

// InvalidateCache passes through to the durable table manager.
func (s *Session) InvalidateCache(table string) {
	s.table.InvalidateCache(table)
}

// BuildIndex passes through to the durable table manager.
func (s *Session) BuildIndex(indexName, tableName string, columns []string) error {
	return s.table.BuildIndex(indexName, tableName, columns)
}