fix: publish scan mutations and RETURNING atomically (#16)

README.md

@@ -520,6 +520,12 @@ UPDATE accounts SET balance = balance + 100 WHERE name = 'Bob';
 COMMIT;
 ```
 
+Multi-row `UPDATE` and `DELETE`, primary-key moves, and DML with `RETURNING`
+are staged before one atomic commit. An evaluation or validation failure leaves
+no partial statement effects; in an explicit transaction, earlier successful
+statements remain staged. Mutations exceeding a single storage transaction's
+limits fail rather than being split into independently committed chunks.
+
 **Indexes:**
 ```sql
 CREATE INDEX idx_users_email ON users(email);

pkg/executor/atomic_mutation_test.go

@@ -0,0 +1,52 @@
+package executor
+
+import (
+	"reflect"
+	"testing"
+)
+
+func checkAtomicMutations(t *testing.T, e *Executor) {
+	t.Helper()
+	execMust(t, e, "CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT, j TEXT)")
+	execMust(t, e, `INSERT INTO t VALUES (1,'old','{"x":"new"}'),(2,'old','invalid')`)
+	want := [][]interface{}{{int64(1), "old"}, {int64(2), "old"}}
+	for _, sql := range []string{
+		"UPDATE t SET v=json_extract(j,'$.x')",
+		"UPDATE t SET v='new' WHERE json_extract(j,'$.x')='new'",
+		"DELETE FROM t WHERE json_extract(j,'$.x')='new'",
+		"UPDATE t SET v='new' RETURNING json_extract(j,'$.x')",
+		"DELETE FROM t RETURNING json_extract(j,'$.x')",
+		"UPDATE t SET id=CASE WHEN id=1 THEN 3 ELSE 3 END",
+	} {
+		if _, err := execSQL(e, sql); err == nil {
+			t.Fatalf("%s: expected failure", sql)
+		}
+		if got := execMust(t, e, "SELECT id,v FROM t ORDER BY id").Rows; !reflect.DeepEqual(got, want) {
+			t.Fatalf("%s left partial effects: %v", sql, got)
+		}
+	}
+	execMust(t, e, "BEGIN")
+	execMust(t, e, "INSERT INTO t VALUES (4,'before','{}')")
+	if _, err := execSQL(e, "UPDATE t SET v=json_extract(j,'$.x') WHERE id<4"); err == nil {
+		t.Fatal("expected failed in-transaction update")
+	}
+	execMust(t, e, "COMMIT")
+	if got := execMust(t, e, "SELECT id,v FROM t WHERE id<4 ORDER BY id").Rows; !reflect.DeepEqual(got, want) {
+		t.Fatalf("failed statement committed partial overlay: %v", got)
+	}
+	if got := execMust(t, e, "SELECT id FROM t WHERE id=4").RowCount; got != 1 {
+		t.Fatal("earlier transaction work was lost")
+	}
+}
+
+func TestAtomicMutationEvaluationAndReturning(t *testing.T) {
+	_, schema, table := newTestDB(t)
+	checkAtomicMutations(t, newExec(schema, table))
+}
+
+func TestAtomicMutationRealKV(t *testing.T) {
+	e := newOptExec(t, "atomic_mutation_regressions")
+	execSQL(e, "DROP TABLE IF EXISTS t")
+	t.Cleanup(func() { execSQL(e, "DROP TABLE IF EXISTS t") })
+	checkAtomicMutations(t, e)
+}

pkg/executor/atomic_storage_test.go

@@ -0,0 +1,33 @@
+package executor
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/danfragoso/pizzasql-next/pkg/storage"
+)
+
+func TestStorageScanUpdateCommitConflictHasNoPartialEffects(t *testing.T) {
+	_, schema, table := newTestDB(t)
+	e := newExec(schema, table)
+	other := newExec(schema, table)
+	execMust(t, e, "CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)")
+	execMust(t, e, "INSERT INTO t VALUES (1,'old'),(2,'old')")
+	raced := false
+	count, err := table.UpdateFunc("t", func(storage.Row) (storage.Row, error) {
+		if !raced {
+			raced = true
+			execMust(t, other, "UPDATE t SET v='external' WHERE id=2")
+		}
+		return storage.Row{"v": "staged"}, nil
+	}, nil)
+	if !errors.Is(err, storage.ErrSerialization) || count != 0 {
+		t.Fatalf("update returned count=%d err=%v", count, err)
+	}
+	if got := execMust(t, e, "SELECT v FROM t WHERE id=1").Rows[0][0]; got != "old" {
+		t.Fatalf("conflicted batch published earlier row: %v", got)
+	}
+	if got := execMust(t, e, "SELECT v FROM t WHERE id=2").Rows[0][0]; got != "external" {
+		t.Fatalf("conflicted batch overwrote concurrent writer: %v", got)
+	}
+}

pkg/executor/executor.go

@@ -255,11 +255,20 @@ func (e *Executor) Execute(stmt parser.Statement) (*Result, error) {
 	case *parser.SelectStmt:
 		return e.executeSelect(s)
 	case *parser.InsertStmt:
-		return e.executeInsert(s)
+		if len(s.Returning) == 0 {
+			return e.executeInsert(s)
+		}
+		return e.executeAtomicMutation(func() (*Result, error) { return e.executeInsert(s) })
 	case *parser.UpdateStmt:
-		return e.executeUpdate(s)
+		if len(s.Returning) == 0 && e.isPointMutation(s.Table.Name, s.Where, s.Set) {
+			return e.executeUpdate(s)
+		}
+		return e.executeAtomicMutation(func() (*Result, error) { return e.executeUpdate(s) })
 	case *parser.DeleteStmt:
-		return e.executeDelete(s)
+		if len(s.Returning) == 0 && e.isPointMutation(s.Table.Name, s.Where, nil) {
+			return e.executeDelete(s)
+		}
+		return e.executeAtomicMutation(func() (*Result, error) { return e.executeDelete(s) })
 	case *parser.CreateTableStmt:
 		return e.executeCreateTable(s)
 	case *parser.DropTableStmt:
@@ -2968,8 +2977,13 @@ func (e *Executor) runDMLAtomic(tableName string, apply func() error) error {
 // transaction even when the table has no UNIQUE index. INSERT OR REPLACE needs
 // it because a primary-key replacement must delete and insert atomically.
 func (e *Executor) runDMLAtomicForce(tableName string, force bool, apply func() error) error {
-	if e.inTransaction {
-		return apply()
+	if e.session.InTx() {
+		pos := e.session.Snapshot()
+		if err := apply(); err != nil {
+			e.session.RollbackTo(pos)
+			return err
+		}
+		return nil
 	}
 	uniq, err := e.table.HasUniqueIndex(tableName)
 	if err != nil {
@@ -2988,6 +3002,54 @@ func (e *Executor) runDMLAtomicForce(tableName string, force bool, apply func()
 	return e.session.Commit()
 }
 
+// executeAtomicMutation includes evaluation and RETURNING in the statement's
+// atomic boundary. Nested calls use mutation-log savepoints, not nested commits.
+func (e *Executor) executeAtomicMutation(apply func() (*Result, error)) (*Result, error) {
+	owned := !e.session.InTx()
+	if owned {
+		if err := e.session.Begin(); err != nil {
+			return nil, err
+		}
+	}
+	pos := e.session.Snapshot()
+	changes, total, last := e.changes, e.totalChanges, e.lastInsertRowID
+	result, err := apply()
+	if err == nil && owned {
+		var lsn uint64
+		lsn, err = e.session.CommitWithLSN()
+		if err == nil {
+			result.LSN = lsn
+		}
+	}
+	if err != nil {
+		if owned {
+			_ = e.session.Rollback()
+		} else {
+			e.session.RollbackTo(pos)
+		}
+		e.changes, e.totalChanges, e.lastInsertRowID = changes, total, last
+		return nil, err
+	}
+	return result, nil
+}
+
+func (e *Executor) isPointMutation(table string, where parser.Expr, assignments []parser.Assignment) bool {
+	column, value, equality := e.extractIndexableCondition(where)
+	if !equality || !safePrimaryKeyLiteral(value) {
+		return false
+	}
+	schema, err := e.schema.GetSchema(table)
+	if err != nil || !strings.EqualFold(column, schema.PrimaryKey) {
+		return false
+	}
+	for _, assignment := range assignments {
+		if strings.EqualFold(assignment.Column, schema.PrimaryKey) {
+			return false
+		}
+	}
+	return true
+}
+
 func containsFold(values []string, target string) bool {
 	for _, value := range values {
 		if strings.EqualFold(value, target) {
@@ -3096,12 +3158,14 @@ func (e *Executor) executeUpdate(stmt *parser.UpdateStmt) (*Result, error) {
 		return e.executeUpdateFrom(stmt, schema)
 	}
 
-	// Build filter
+	// Keep predicate errors inside the statement rollback boundary.
+	var filterErr error
 	var filter func(storage.Row) bool
 	if stmt.Where != nil {
 		filter = func(row storage.Row) bool {
 			val, err := e.evalExpr(stmt.Where, row)
 			if err != nil {
+				filterErr = err
 				return false
 			}
 			return toBool(val)
@@ -3153,7 +3217,7 @@ func (e *Executor) executeUpdate(stmt *parser.UpdateStmt) (*Result, error) {
 				break
 			}
 		}
-		if equality && strings.EqualFold(column, schema.PrimaryKey) && !updatesPrimaryKey {
+		if equality && safePrimaryKeyLiteral(value) && strings.EqualFold(column, schema.PrimaryKey) && !updatesPrimaryKey {
 			pk := fmt.Sprintf("%v", value)
 			_, updated, err := e.session.UpdateByPK(tableName, pk, updateFn)
 			if err != nil {
@@ -3199,6 +3263,9 @@ func (e *Executor) executeUpdate(stmt *parser.UpdateStmt) (*Result, error) {
 		if err != nil {
 			return err
 		}
+		if filterErr != nil {
+			return filterErr
+		}
 		count = n
 		return nil
 	})
@@ -3481,12 +3548,13 @@ func (e *Executor) executeDelete(stmt *parser.DeleteStmt) (*Result, error) {
 		return nil, err
 	}
 
-	// Build filter
+	var filterErr error
 	var filter func(storage.Row) bool
 	if stmt.Where != nil {
 		filter = func(row storage.Row) bool {
 			val, err := e.evalExpr(stmt.Where, row)
 			if err != nil {
+				filterErr = err
 				return false
 			}
 			return toBool(val)
@@ -3499,7 +3567,7 @@ func (e *Executor) executeDelete(stmt *parser.DeleteStmt) (*Result, error) {
 
 	if stmt.Where != nil {
 		column, value, equality := e.extractIndexableCondition(stmt.Where)
-		if equality && strings.EqualFold(column, schema.PrimaryKey) {
+		if equality && safePrimaryKeyLiteral(value) && strings.EqualFold(column, schema.PrimaryKey) {
 			deletedRow, deleted, err := e.session.DeleteByPK(tableName, fmt.Sprintf("%v", value))
 			if err != nil {
 				return nil, err
@@ -3531,6 +3599,9 @@ func (e *Executor) executeDelete(stmt *parser.DeleteStmt) (*Result, error) {
 	if err != nil {
 		return nil, err
 	}
+	if filterErr != nil {
+		return nil, filterErr
+	}
 
 	return e.finishDML("DELETE", count, 0, false, affected, stmt.Returning, schema)
 }

pkg/pgserver/features_test.go

@@ -74,6 +74,18 @@ func dataRowValues(t *testing.T, msg *Message) []string {
 	return values
 }
 
+// A durable implicit commit may report its LSN between result completion and
+// ReadyForQuery. Keep result-shape assertions independent of status updates.
+func resultMessages(messages []*Message) []*Message {
+	var result []*Message
+	for _, message := range messages {
+		if message.Type != MsgParameterStatus {
+			result = append(result, message)
+		}
+	}
+	return result
+}
+
 // commandTag extracts the NUL-terminated tag from a CommandComplete message.
 func commandTag(t *testing.T, msg *Message) string {
 	t.Helper()
@@ -87,7 +99,7 @@ func TestSimpleQueryInsertReturning(t *testing.T) {
 	c, client := newDataConnection(t)
 	runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("CREATE TABLE t (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT)"), 0)})
 
-	msgs := runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("INSERT INTO t (name) VALUES ('alice') RETURNING id, name"), 0)})
+	msgs := resultMessages(runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("INSERT INTO t (name) VALUES ('alice') RETURNING id, name"), 0)}))
 	if len(msgs) != 4 {
 		t.Fatalf("got %d messages, want RowDescription+DataRow+CommandComplete+ReadyForQuery: %v", len(msgs), msgs)
 	}
@@ -139,7 +151,7 @@ func TestCommandTagSelectAndUpdate(t *testing.T) {
 		t.Fatalf("select tag = %q", tag)
 	}
 
-	msgs = runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("UPDATE t SET v='b' RETURNING id"), 0)})
+	msgs = resultMessages(runQuery(t, c, client, &Message{Type: MsgQuery, Data: append([]byte("UPDATE t SET v='b' RETURNING id"), 0)}))
 	if msgs[0].Type != MsgRowDescription {
 		t.Fatalf("update returning first message = %c", msgs[0].Type)
 	}

pkg/storage/table.go

@@ -1002,94 +1002,21 @@ func (m *TableManager) UpdateFunc(table string, updateFn func(Row) (Row, error),
 	return m.updateLocked(table, updateFn, filter)
 }
 
-// updateLocked is the shared scan-based update implementation. It applies
-// per-row updates, revalidates unique indexes, and moves a row when its primary
-// key changes (deleting the old key and refusing to overwrite an occupied new
-// key) so an UPDATE of the primary key does not orphan or clobber rows.
+// updateLocked stages the complete update before one compare-batch publication.
+// Evaluation, encoding, validation or I/O errors must never leave earlier rows
+// committed, including when a primary key moves to a different durable key.
 func (m *TableManager) updateLocked(table string, updateFn func(Row) (Row, error), filter func(Row) bool) (int, error) {
-	tl := m.tableLock(strings.ToLower(table))
-	tl.Lock()
-	defer tl.Unlock()
-	rows, err := m.selectRows(table, filter)
-	if err != nil {
+	session := NewSession(m.schema, m)
+	if err := session.Begin(); err != nil {
 		return 0, err
 	}
-	schema, err := m.schema.GetSchema(table)
+	defer session.Rollback()
+	count, err := session.UpdateFunc(table, updateFn, filter)
 	if err != nil {
 		return 0, err
 	}
-
-	count := 0
-	for _, row := range rows {
-		oldRow := cloneRow(row)
-		oldKey := m.dataKey(table, fmt.Sprintf("%v", oldRow[schema.PrimaryKey]))
-		if err := m.updateIndexesForRow(table, row, false); err != nil {
-			return count, err
-		}
-
-		updates, err := updateFn(row)
-		if err != nil {
-			_ = m.updateIndexesForRow(table, oldRow, true)
-			return count, err
-		}
-		for k, v := range updates {
-			for _, col := range schema.Columns {
-				if strings.EqualFold(k, col.Name) {
-					row[col.Name] = v
-					break
-				}
-			}
-		}
-
-		newKey := m.dataKey(table, fmt.Sprintf("%v", row[schema.PrimaryKey]))
-		data, err := encodeRow(row)
-		if err != nil {
-			_ = m.updateIndexesForRow(table, oldRow, true)
-			continue
-		}
-
-		excluded := map[string]bool{oldKey: true}
-		if err := m.validateUniqueRows(table, []Row{row}, excluded); err != nil {
-			_ = m.updateIndexesForRow(table, oldRow, true)
-			return count, err
-		}
-
-		if newKey != oldKey {
-			// Refuse to move onto an existing primary key.
-			if _, _, gerr := m.getByPKWithLSN(table, fmt.Sprintf("%v", row[schema.PrimaryKey])); gerr == nil {
-				_ = m.updateIndexesForRow(table, oldRow, true)
-				return count, fmt.Errorf("duplicate primary key: %v", row[schema.PrimaryKey])
-			} else if gerr != ErrKeyNotFound {
-				_ = m.updateIndexesForRow(table, oldRow, true)
-				return count, gerr
-			}
-			if err := m.pool.WithClient(func(c *KVClient) error {
-				_, derr := c.Del([]byte(oldKey))
-				return derr
-			}); err != nil {
-				_ = m.updateIndexesForRow(table, oldRow, true)
-				continue
-			}
-		}
-
-		if err := m.pool.WithClient(func(c *KVClient) error {
-			_, perr := c.Put([]byte(newKey), data)
-			return perr
-		}); err != nil {
-			_ = m.updateIndexesForRow(table, oldRow, true)
-			continue
-		}
-		if err := m.updateIndexesForRow(table, row, true); err != nil {
-			return count, err
-		}
-		if err := m.bumpIndexPredicates(table, oldRow, row); err != nil {
-			return count, err
-		}
-		count++
-	}
-
-	if count > 0 {
-		m.bumpGeneration(table)
+	if err := session.Commit(); err != nil {
+		return 0, err
 	}
 	return count, nil
 }
@@ -1170,49 +1097,18 @@ func (m *TableManager) UpdateByPK(table, pk string, updateFn func(Row) (Row, err
 
 // Delete deletes rows matching the filter.
 func (m *TableManager) Delete(table string, filter func(Row) bool) (int, error) {
-	tl := m.tableLock(strings.ToLower(table))
-	tl.Lock()
-	defer tl.Unlock()
-	rows, err := m.selectRows(table, filter)
-	if err != nil {
+	session := NewSession(m.schema, m)
+	if err := session.Begin(); err != nil {
 		return 0, err
 	}
-	schema, err := m.schema.GetSchema(table)
+	defer session.Rollback()
+	count, err := session.Delete(table, filter)
 	if err != nil {
 		return 0, err
 	}
-	wasInit := m.countInitialized(table)
-
-	count := 0
-	for _, row := range rows {
-		// Remove index entries before deleting row
-		if err := m.updateIndexesForRow(table, row, false); err != nil {
-			return count, err
-		}
-
-		pkValue := row[schema.PrimaryKey]
-		pk := fmt.Sprintf("%v", pkValue)
-		key := m.dataKey(table, pk)
-
-		err = m.pool.WithClient(func(c *KVClient) error {
-			_, err := c.Del([]byte(key))
-			return err
-		})
-		if err == nil {
-			if err := m.bumpIndexPredicates(table, row); err != nil {
-				return count, err
-			}
-			count++
-		} else {
-			// Restore the index entries removed above.
-			_ = m.updateIndexesForRow(table, row, true)
-		}
-	}
-
-	if count > 0 {
-		m.bumpGeneration(table)
+	if err := session.Commit(); err != nil {
+		return 0, err
 	}
-	m.incrCount(table, schema.CreatedAt, -count, wasInit)
 	return count, nil
 }