feat: emulate Firestore read-after-write transactions

adapters/dalgo2memory/database.go

@@ -41,6 +41,9 @@ type database struct {
 	// mu without attempting an RWMutex lock upgrade.
 	enginesMu sync.Mutex
 	schema    *memorySchema
+	// noReadsAfterWritesInTransaction emulates Firestore's transaction ordering
+	// rule for databases created with WithNoReadsAfterWritesInTransaction.
+	noReadsAfterWritesInTransaction bool
 	// schemaRefBreaking is the schema-wide columnar fidelity default (faithful
 	// unless WithoutSchemaRefBreaking was used). NewDB initializes it to true.
 	schemaRefBreaking bool
@@ -67,7 +70,9 @@ func (db *database) RunReadonlyTransaction(ctx context.Context, f dal.ROTxWorker
 func (db *database) RunReadwriteTransaction(ctx context.Context, f dal.RWTxWorker, _ ...dal.TransactionOption) error {
 	db.mu.Lock()
 	defer db.mu.Unlock()
-	return f(ctx, session{db: db})
+	return f(ctx, session{db: db, txState: &transactionState{
+		noReadsAfterWrites: db.noReadsAfterWritesInTransaction,
+	}})
 }
 
 func (db *database) Exists(ctx context.Context, key *dal.Key) (bool, error) {
@@ -157,7 +162,30 @@ func (db *database) UpdateMulti(ctx context.Context, keys []*dal.Key, updates []
 var _ dal.DB = (*database)(nil)
 
 type session struct {
-	db *database
+	db      *database
+	txState *transactionState
+}
+
+type transactionState struct {
+	noReadsAfterWrites bool
+	hasWritten         bool
+}
+
+// ErrReadAfterWriteInTransaction matches the ordering error returned by
+// Firestore transactions when a read follows a queued write.
+var ErrReadAfterWriteInTransaction = errors.New("firestore: read after write in transaction")
+
+func (s session) allowRead() error {
+	if s.txState != nil && s.txState.noReadsAfterWrites && s.txState.hasWritten {
+		return ErrReadAfterWriteInTransaction
+	}
+	return nil
+}
+
+func (s session) markWrite() {
+	if s.txState != nil {
+		s.txState.hasWritten = true
+	}
 }
 
 func (s session) ID() string {
@@ -169,10 +197,17 @@ func (s session) Options() dal.TransactionOptions {
 }
 
 func (s session) Exists(_ context.Context, key *dal.Key) (bool, error) {
+	if err := s.allowRead(); err != nil {
+		return false, err
+	}
 	return s.db.engine(key.Collection()).exists(keyID(key)), nil
 }
 
 func (s session) Get(_ context.Context, record dal.Record) error {
+	if err := s.allowRead(); err != nil {
+		record.SetError(err)
+		return err
+	}
 	if err := s.db.guardCollection(record.Key().Collection()); err != nil {
 		record.SetError(err)
 		return err
@@ -186,6 +221,9 @@ func (s session) Get(_ context.Context, record dal.Record) error {
 }
 
 func (s session) GetMulti(ctx context.Context, records []dal.Record) error {
+	if err := s.allowRead(); err != nil {
+		return err
+	}
 	for _, record := range records {
 		if err := s.Get(ctx, record); err != nil && !dal.IsNotFound(err) {
 			return err
@@ -195,7 +233,11 @@ func (s session) GetMulti(ctx context.Context, records []dal.Record) error {
 }
 
 func (s session) Set(_ context.Context, record dal.Record) error {
-	return s.save(record, true)
+	if err := s.save(record, true); err != nil {
+		return err
+	}
+	s.markWrite()
+	return nil
 }
 
 func (s session) SetMulti(ctx context.Context, records []dal.Record) error {
@@ -223,7 +265,7 @@ func (s session) Insert(ctx context.Context, record dal.Record, opts ...dal.Inse
 		gen = dal.NewInsertOptions(dal.WithRandomStringKey(dal.DefaultRandomStringIDLength, 5)).IDGenerator()
 	}
 	if gen != nil {
-		return dal.InsertWithIdGenerator(ctx, record, gen, insertWithGeneratorMaxAttempts,
+		err := dal.InsertWithIdGenerator(ctx, record, gen, insertWithGeneratorMaxAttempts,
 			func(key *dal.Key) error {
 				if s.db.engine(key.Collection()).exists(keyID(key)) {
 					return nil // id is taken: signal "exists" so generation retries
@@ -234,8 +276,16 @@ func (s session) Insert(ctx context.Context, record dal.Record, opts ...dal.Inse
 				return s.save(r, false)
 			},
 		)
+		if err == nil {
+			s.markWrite()
+		}
+		return err
 	}
-	return s.save(record, false)
+	if err := s.save(record, false); err != nil {
+		return err
+	}
+	s.markWrite()
+	return nil
 }
 
 func (s session) InsertMulti(ctx context.Context, records []dal.Record, opts ...dal.InsertOption) error {
@@ -249,6 +299,7 @@ func (s session) InsertMulti(ctx context.Context, records []dal.Record, opts ...
 
 func (s session) Delete(_ context.Context, key *dal.Key) error {
 	s.db.engine(key.Collection()).delete(keyID(key))
+	s.markWrite()
 	return nil
 }
 
@@ -269,7 +320,11 @@ func (s session) UpdateRecord(_ context.Context, record dal.Record, updates []up
 	if err := s.db.guardCollection(collectionName); err != nil {
 		return err
 	}
-	return s.db.engine(collectionName).update(keyID(record.Key()), updates)
+	if err := s.db.engine(collectionName).update(keyID(record.Key()), updates); err != nil {
+		return err
+	}
+	s.markWrite()
+	return nil
 }
 
 func (s session) UpdateMulti(ctx context.Context, keys []*dal.Key, updates []update.Update, preconditions ...dal.Precondition) error {
@@ -282,6 +337,9 @@ func (s session) UpdateMulti(ctx context.Context, keys []*dal.Key, updates []upd
 }
 
 func (s session) ExecuteQueryToRecordsReader(_ context.Context, query dal.Query) (dal.RecordsReader, error) {
+	if err := s.allowRead(); err != nil {
+		return nil, err
+	}
 	q, ok := query.(dal.StructuredQuery)
 	if !ok {
 		return nil, dal.ErrNotSupported

adapters/dalgo2memory/database_test.go

@@ -198,6 +198,98 @@ func TestQueryEmptyCollection(t *testing.T) {
 	require.ErrorIs(t, err, dal.ErrNoMoreRecords)
 }
 
+func TestNoReadsAfterWritesInTransaction(t *testing.T) {
+	ctx := context.Background()
+	key := dal.NewKeyWithID("Things", "existing")
+
+	newDB := func() *database {
+		db := NewDB(WithNoReadsAfterWritesInTransaction()).(*database)
+		require.NoError(t, db.Set(ctx, dal.NewRecordWithData(key, &thing{Name: "before", Count: 1})))
+		return db
+	}
+
+	t.Run("read then write then reads fail", func(t *testing.T) {
+		db := newDB()
+		err := db.RunReadwriteTransaction(ctx, func(ctx context.Context, tx dal.ReadwriteTransaction) error {
+			got := &thing{}
+			if err := tx.Get(ctx, dal.NewRecordWithData(key, got)); err != nil {
+				return err
+			}
+			if err := tx.Set(ctx, dal.NewRecordWithData(dal.NewKeyWithID("Things", "new"), &thing{Name: "after"})); err != nil {
+				return err
+			}
+
+			_, err := tx.Exists(ctx, key)
+			require.ErrorIs(t, err, ErrReadAfterWriteInTransaction)
+			err = tx.Get(ctx, dal.NewRecordWithData(key, &thing{}))
+			require.ErrorIs(t, err, ErrReadAfterWriteInTransaction)
+			err = tx.GetMulti(ctx, []dal.Record{dal.NewRecordWithData(key, &thing{})})
+			require.ErrorIs(t, err, ErrReadAfterWriteInTransaction)
+			q := dal.From(dal.NewRootCollectionRef("Things", "")).NewQuery().SelectKeysOnly(reflect.String)
+			_, err = tx.ExecuteQueryToRecordsReader(ctx, q)
+			require.ErrorIs(t, err, ErrReadAfterWriteInTransaction)
+			_, err = tx.ExecuteQueryToRecordsetReader(ctx, q)
+			require.ErrorIs(t, err, ErrReadAfterWriteInTransaction)
+			return nil
+		})
+		require.NoError(t, err)
+	})
+
+	t.Run("each write operation blocks subsequent reads", func(t *testing.T) {
+		writes := map[string]func(dal.ReadwriteTransaction) error{
+			"Set": func(tx dal.ReadwriteTransaction) error {
+				return tx.Set(ctx, dal.NewRecordWithData(dal.NewKeyWithID("Things", "set"), &thing{}))
+			},
+			"SetMulti": func(tx dal.ReadwriteTransaction) error {
+				return tx.SetMulti(ctx, []dal.Record{dal.NewRecordWithData(dal.NewKeyWithID("Things", "set-multi"), &thing{})})
+			},
+			"Insert": func(tx dal.ReadwriteTransaction) error {
+				return tx.Insert(ctx, dal.NewRecordWithData(dal.NewKeyWithID("Things", "insert"), &thing{}))
+			},
+			"InsertMulti": func(tx dal.ReadwriteTransaction) error {
+				return tx.InsertMulti(ctx, []dal.Record{dal.NewRecordWithData(dal.NewKeyWithID("Things", "insert-multi"), &thing{})})
+			},
+			"Update": func(tx dal.ReadwriteTransaction) error {
+				return tx.Update(ctx, key, []update.Update{update.ByFieldName("Count", 2)})
+			},
+			"UpdateRecord": func(tx dal.ReadwriteTransaction) error {
+				return tx.UpdateRecord(ctx, dal.NewRecordWithData(key, &thing{}), []update.Update{update.ByFieldName("Count", 2)})
+			},
+			"UpdateMulti": func(tx dal.ReadwriteTransaction) error {
+				return tx.UpdateMulti(ctx, []*dal.Key{key}, []update.Update{update.ByFieldName("Count", 2)})
+			},
+			"Delete":      func(tx dal.ReadwriteTransaction) error { return tx.Delete(ctx, key) },
+			"DeleteMulti": func(tx dal.ReadwriteTransaction) error { return tx.DeleteMulti(ctx, []*dal.Key{key}) },
+		}
+
+		for name, write := range writes {
+			t.Run(name, func(t *testing.T) {
+				db := newDB()
+				err := db.RunReadwriteTransaction(ctx, func(ctx context.Context, tx dal.ReadwriteTransaction) error {
+					if err := write(tx); err != nil {
+						return err
+					}
+					_, err := tx.Exists(ctx, key)
+					require.ErrorIs(t, err, ErrReadAfterWriteInTransaction)
+					return nil
+				})
+				require.NoError(t, err)
+			})
+		}
+	})
+
+	t.Run("default remains permissive", func(t *testing.T) {
+		db := NewDB().(*database)
+		err := db.RunReadwriteTransaction(ctx, func(ctx context.Context, tx dal.ReadwriteTransaction) error {
+			if err := tx.Set(ctx, dal.NewRecordWithData(key, &thing{Name: "stored"})); err != nil {
+				return err
+			}
+			return tx.Get(ctx, dal.NewRecordWithData(key, &thing{}))
+		})
+		require.NoError(t, err)
+	})
+}
+
 // TestConcurrentReadonlyQueriesInitializeEnginesSafely verifies that queries
 // against previously unseen collections can initialize their storage engines
 // concurrently. It is intended to run under the race detector.

adapters/dalgo2memory/recordset.go

@@ -23,6 +23,9 @@ import (
 // (any) columns, matching the adapter's schemaless JSON storage. A non-
 // structured query is not supported.
 func (s session) ExecuteQueryToRecordsetReader(ctx context.Context, query dal.Query, options ...recordset.Option) (dal.RecordsetReader, error) {
+	if err := s.allowRead(); err != nil {
+		return nil, err
+	}
 	q, ok := query.(dal.StructuredQuery)
 	if !ok {
 		return nil, dal.ErrNotSupported

adapters/dalgo2memory/schema.go

@@ -10,6 +10,17 @@ import (
 // Option configures an in-memory database created by NewDB.
 type Option func(*database)
 
+// WithNoReadsAfterWritesInTransaction enables Firestore-compatible transaction
+// ordering for this in-memory database. In a read-write transaction, every
+// read after the first successful write returns ErrReadAfterWriteInTransaction.
+// It is intended for tests that need to catch Firestore-only ordering errors;
+// the default in-memory behavior remains permissive.
+func WithNoReadsAfterWritesInTransaction() Option {
+	return func(db *database) {
+		db.noReadsAfterWritesInTransaction = true
+	}
+}
+
 // collectionDef describes a single collection in an in-memory schema.
 // It is produced by WithCollection and consumed by WithSchema.
 type collectionDef struct {