refactor: DB transactions context (#2513)

## Relevant issue(s)

Resolves #2512
Resolves #2516

## Description

This PR moves the db transactions to the context.

Notable Changes:
- moved `txn_db.go` to `store.go`
- ~replaced `explicitTxnDB` and `implicitTxnDb` with `store`~
- ~added `db/session.go` to simplify setting context~
- added `db/context.go` to manage context values

## Tasks

- [x] I made sure the code is well commented, particularly
hard-to-understand areas.
- [x] I made sure the repository-held documentation is changed
accordingly.
- [x] I made sure the pull request title adheres to the conventional
commit style (the subset used in the project can be found in
[tools/configs/chglog/config.yml](tools/configs/chglog/config.yml)).
- [x] I made sure to discuss its limitations such as threats to
validity, vulnerability to mistake and misuse, robustness to
invalidation of assumptions, resource requirements, ...

## How has this been tested?

`make test`

Specify the platform(s) on which this was tested:
- MacOS

cli/client.go

@@ -31,7 +31,7 @@ Execute queries, add schema types, obtain node info, etc.`,
 			if err := setContextTransaction(cmd, txID); err != nil {
 				return err
 			}
-			return setContextStore(cmd)
+			return setContextDB(cmd)
 		},
 	}
 	cmd.PersistentFlags().Uint64Var(&txID, "tx", 0, "Transaction ID")

cli/collection.go

@@ -17,7 +17,6 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/sourcenetwork/defradb/client"
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 func MakeCollectionCommand() *cobra.Command {
@@ -41,7 +40,7 @@ func MakeCollectionCommand() *cobra.Command {
 			if err := setContextTransaction(cmd, txID); err != nil {
 				return err
 			}
-			if err := setContextStore(cmd); err != nil {
+			if err := setContextDB(cmd); err != nil {
 				return err
 			}
 			store := mustGetContextStore(cmd)
@@ -71,10 +70,6 @@ func MakeCollectionCommand() *cobra.Command {
 			}
 			col := cols[0]
 
-			if tx, ok := cmd.Context().Value(txContextKey).(datastore.Txn); ok {
-				col = col.WithTxn(tx)
-			}
-
 			ctx := context.WithValue(cmd.Context(), colContextKey, col)
 			cmd.SetContext(ctx)
 			return nil

cli/index_create.go

@@ -14,7 +14,6 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/sourcenetwork/defradb/client"
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 func MakeIndexCreateCommand() *cobra.Command {
@@ -52,9 +51,6 @@ Example: create a named index for 'Users' collection on 'name' field:
 			if err != nil {
 				return err
 			}
-			if tx, ok := cmd.Context().Value(txContextKey).(datastore.Txn); ok {
-				col = col.WithTxn(tx)
-			}
 			desc, err = col.CreateIndex(cmd.Context(), desc)
 			if err != nil {
 				return err

cli/index_drop.go

@@ -12,8 +12,6 @@ package cli
 
 import (
 	"github.com/spf13/cobra"
-
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 func MakeIndexDropCommand() *cobra.Command {
@@ -34,9 +32,6 @@ Example: drop the index 'UsersByName' for 'Users' collection:
 			if err != nil {
 				return err
 			}
-			if tx, ok := cmd.Context().Value(txContextKey).(datastore.Txn); ok {
-				col = col.WithTxn(tx)
-			}
 			return col.DropIndex(cmd.Context(), nameArg)
 		},
 	}

cli/index_list.go

@@ -12,8 +12,6 @@ package cli
 
 import (
 	"github.com/spf13/cobra"
-
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 func MakeIndexListCommand() *cobra.Command {
@@ -38,9 +36,6 @@ Example: show all index for 'Users' collection:
 				if err != nil {
 					return err
 				}
-				if tx, ok := cmd.Context().Value(txContextKey).(datastore.Txn); ok {
-					col = col.WithTxn(tx)
-				}
 				indexes, err := col.GetIndexes(cmd.Context())
 				if err != nil {
 					return err

cli/schema_migration_down.go

@@ -17,8 +17,6 @@ import (
 
 	"github.com/sourcenetwork/immutable/enumerable"
 	"github.com/spf13/cobra"
-
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 func MakeSchemaMigrationDownCommand() *cobra.Command {
@@ -67,12 +65,7 @@ Example: migrate from stdin
 			if err := json.Unmarshal(srcData, &src); err != nil {
 				return err
 			}
-			lens := store.LensRegistry()
-			if tx, ok := cmd.Context().Value(txContextKey).(datastore.Txn); ok {
-				lens = lens.WithTxn(tx)
-			}
-
-			out, err := lens.MigrateDown(cmd.Context(), enumerable.New(src), collectionID)
+			out, err := store.LensRegistry().MigrateDown(cmd.Context(), enumerable.New(src), collectionID)
 			if err != nil {
 				return err
 			}

cli/schema_migration_reload.go

@@ -12,8 +12,6 @@ package cli
 
 import (
 	"github.com/spf13/cobra"
-
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 func MakeSchemaMigrationReloadCommand() *cobra.Command {
@@ -23,12 +21,7 @@ func MakeSchemaMigrationReloadCommand() *cobra.Command {
 		Long:  `Reload the schema migrations within DefraDB`,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			store := mustGetContextStore(cmd)
-
-			lens := store.LensRegistry()
-			if tx, ok := cmd.Context().Value(txContextKey).(datastore.Txn); ok {
-				lens = lens.WithTxn(tx)
-			}
-			return lens.ReloadLenses(cmd.Context())
+			return store.LensRegistry().ReloadLenses(cmd.Context())
 		},
 	}
 	return cmd

cli/schema_migration_up.go

@@ -17,8 +17,6 @@ import (
 
 	"github.com/sourcenetwork/immutable/enumerable"
 	"github.com/spf13/cobra"
-
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 func MakeSchemaMigrationUpCommand() *cobra.Command {
@@ -67,12 +65,7 @@ Example: migrate from stdin
 			if err := json.Unmarshal(srcData, &src); err != nil {
 				return err
 			}
-			lens := store.LensRegistry()
-			if tx, ok := cmd.Context().Value(txContextKey).(datastore.Txn); ok {
-				lens = lens.WithTxn(tx)
-			}
-
-			out, err := lens.MigrateUp(cmd.Context(), enumerable.New(src), collectionID)
+			out, err := store.LensRegistry().MigrateUp(cmd.Context(), enumerable.New(src), collectionID)
 			if err != nil {
 				return err
 			}

cli/utils.go

@@ -21,7 +21,7 @@ import (
 	"github.com/spf13/viper"
 
 	"github.com/sourcenetwork/defradb/client"
-	"github.com/sourcenetwork/defradb/datastore"
+	"github.com/sourcenetwork/defradb/db"
 	"github.com/sourcenetwork/defradb/http"
 )
 
@@ -32,17 +32,8 @@ var (
 	cfgContextKey = contextKey("cfg")
 	// rootDirContextKey is the context key for the root directory.
 	rootDirContextKey = contextKey("rootDir")
-	// txContextKey is the context key for the datastore.Txn
-	//
-	// This will only be set if a transaction id is specified.
-	txContextKey = contextKey("tx")
 	// dbContextKey is the context key for the client.DB
 	dbContextKey = contextKey("db")
-	// storeContextKey is the context key for the client.Store
-	//
-	// If a transaction exists, all operations will be executed
-	// in the current transaction context.
-	storeContextKey = contextKey("store")
 	// colContextKey is the context key for the client.Collection
 	//
 	// If a transaction exists, all operations will be executed
@@ -61,7 +52,7 @@ func mustGetContextDB(cmd *cobra.Command) client.DB {
 //
 // If a store is not set in the current context this function panics.
 func mustGetContextStore(cmd *cobra.Command) client.Store {
-	return cmd.Context().Value(storeContextKey).(client.Store)
+	return cmd.Context().Value(dbContextKey).(client.Store)
 }
 
 // mustGetContextP2P returns the p2p implementation for the current command context.
@@ -92,6 +83,18 @@ func tryGetContextCollection(cmd *cobra.Command) (client.Collection, bool) {
 	return col, ok
 }
 
+// setContextDB sets the db for the current command context.
+func setContextDB(cmd *cobra.Command) error {
+	cfg := mustGetContextConfig(cmd)
+	db, err := http.NewClient(cfg.GetString("api.address"))
+	if err != nil {
+		return err
+	}
+	ctx := context.WithValue(cmd.Context(), dbContextKey, db)
+	cmd.SetContext(ctx)
+	return nil
+}
+
 // setContextConfig sets teh config for the current command context.
 func setContextConfig(cmd *cobra.Command) error {
 	rootdir := mustGetContextRootDir(cmd)
@@ -115,24 +118,7 @@ func setContextTransaction(cmd *cobra.Command, txId uint64) error {
 	if err != nil {
 		return err
 	}
-	ctx := context.WithValue(cmd.Context(), txContextKey, tx)
-	cmd.SetContext(ctx)
-	return nil
-}
-
-// setContextStore sets the store for the current command context.
-func setContextStore(cmd *cobra.Command) error {
-	cfg := mustGetContextConfig(cmd)
-	db, err := http.NewClient(cfg.GetString("api.address"))
-	if err != nil {
-		return err
-	}
-	ctx := context.WithValue(cmd.Context(), dbContextKey, db)
-	if tx, ok := ctx.Value(txContextKey).(datastore.Txn); ok {
-		ctx = context.WithValue(ctx, storeContextKey, db.WithTxn(tx))
-	} else {
-		ctx = context.WithValue(ctx, storeContextKey, db)
-	}
+	ctx := db.SetContextTxn(cmd.Context(), tx)
 	cmd.SetContext(ctx)
 	return nil
 }

client/collection.go

@@ -14,8 +14,6 @@ import (
 	"context"
 
 	"github.com/sourcenetwork/immutable"
-
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 // Collection represents a defradb collection.
@@ -192,10 +190,6 @@ type Collection interface {
 		showDeleted bool,
 	) (*Document, error)
 
-	// WithTxn returns a new instance of the collection, with a transaction
-	// handle instead of a raw DB handle.
-	WithTxn(datastore.Txn) Collection
-
 	// GetAllDocIDs returns all the document IDs that exist in the collection.
 	GetAllDocIDs(ctx context.Context, identity immutable.Option[string]) (<-chan DocIDResult, error)
 

client/db.go

@@ -42,9 +42,6 @@ type DB interface {
 	// can safely operate on it concurrently.
 	NewConcurrentTxn(context.Context, bool) (datastore.Txn, error)
 
-	// WithTxn returns a new [client.Store] that respects the given transaction.
-	WithTxn(datastore.Txn) Store
-
 	// Root returns the underlying root store, within which all data managed by DefraDB is held.
 	Root() datastore.RootStore
 

client/lens.go

@@ -15,8 +15,6 @@ import (
 
 	"github.com/lens-vm/lens/host-go/config/model"
 	"github.com/sourcenetwork/immutable/enumerable"
-
-	"github.com/sourcenetwork/defradb/datastore"
 )
 
 // LensConfig represents the configuration of a Lens migration in Defra.
@@ -43,12 +41,6 @@ type LensConfig struct {
 // LensRegistry exposes several useful thread-safe migration related functions which may
 // be used to manage migrations.
 type LensRegistry interface {
-	// WithTxn returns a new LensRegistry scoped to the given transaction.
-	//
-	// WARNING: Currently this does not provide snapshot isolation, if other transactions are committed
-	// after this has been created, the results of those commits will be visible within this scope.
-	WithTxn(datastore.Txn) LensRegistry
-
 	// SetMigration caches the migration for the given collection ID. It does not persist the migration in long
 	// term storage, for that one should call [Store.SetMigration(ctx, cfg)].
 	//

db/backup.go

@@ -92,7 +92,7 @@ func (db *db) basicImport(ctx context.Context, txn datastore.Txn, filepath strin
 			}
 
 			// TODO-ACP: https://github.com/sourcenetwork/defradb/issues/2430 - Add identity ability to backup
-			err = col.WithTxn(txn).Create(ctx, acpIdentity.NoIdentity, doc)
+			err = col.Create(ctx, acpIdentity.NoIdentity, doc)
 			if err != nil {
 				return NewErrDocCreate(err)
 			}
@@ -104,7 +104,7 @@ func (db *db) basicImport(ctx context.Context, txn datastore.Txn, filepath strin
 					return NewErrDocUpdate(err)
 				}
 				// TODO-ACP: https://github.com/sourcenetwork/defradb/issues/2430 - Add identity ability to backup
-				err = col.WithTxn(txn).Update(ctx, acpIdentity.NoIdentity, doc)
+				err = col.Update(ctx, acpIdentity.NoIdentity, doc)
 				if err != nil {
 					return NewErrDocUpdate(err)
 				}
@@ -191,9 +191,8 @@ func (db *db) basicExport(ctx context.Context, txn datastore.Txn, config *client
 		if err != nil {
 			return err
 		}
-		colTxn := col.WithTxn(txn)
 		// TODO-ACP: https://github.com/sourcenetwork/defradb/issues/2430 - Add identity ability to export
-		docIDsCh, err := colTxn.GetAllDocIDs(ctx, acpIdentity.NoIdentity)
+		docIDsCh, err := col.GetAllDocIDs(ctx, acpIdentity.NoIdentity)
 		if err != nil {
 			return err
 		}
@@ -210,7 +209,7 @@ func (db *db) basicExport(ctx context.Context, txn datastore.Txn, config *client
 				}
 			}
 			// TODO-ACP: https://github.com/sourcenetwork/defradb/issues/2430 - Add identity ability to export
-			doc, err := colTxn.Get(ctx, acpIdentity.NoIdentity, docResultWithID.ID, false)
+			doc, err := col.Get(ctx, acpIdentity.NoIdentity, docResultWithID.ID, false)
 			if err != nil {
 				return err
 			}