SQUASH if it works

client/lens.go

@@ -15,6 +15,8 @@ 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.
@@ -38,9 +40,17 @@ type LensConfig struct {
 	model.Lens
 }
 
+// TxnSource represents an object capable of constructing the transactions that
+// implicit-transaction registries need internally.
+type TxnSource interface {
+	NewTxn(context.Context, bool) (datastore.Txn, error)
+}
+
 // LensRegistry exposes several useful thread-safe migration related functions which may
 // be used to manage migrations.
 type LensRegistry interface {
+	Init(TxnSource)
+
 	// 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)].
 	//

http/client_lens.go

@@ -35,6 +35,8 @@ type setMigrationRequest struct {
 	Config       model.Lens
 }
 
+func (w *LensRegistry) Init(txnSource client.TxnSource) {}
+
 func (c *LensRegistry) SetMigration(ctx context.Context, collectionID uint32, config model.Lens) error {
 	methodURL := c.http.baseURL.JoinPath("lens", "registry")
 

internal/db/db.go

@@ -109,6 +109,10 @@ func newDB(
 		opt(db)
 	}
 
+	if lens != nil {
+		lens.Init(db)
+	}
+
 	err = db.initialize(ctx)
 	if err != nil {
 		return nil, err

internal/lens/registry.go

@@ -70,20 +70,13 @@ func newTxnCtx(txn datastore.Txn) *txnContext {
 	}
 }
 
-// TxnSource represents an object capable of constructing the transactions that
-// implicit-transaction registries need internally.
-type TxnSource interface {
-	NewTxn(context.Context, bool) (datastore.Txn, error)
-}
-
 // DefaultPoolSize is the default size of the lens pool for each schema version.
 const DefaultPoolSize int = 5
 
 // NewRegistry instantiates a new registery.
 //
 // It will be of size 5 (per schema version) if a size is not provided.
 func NewRegistry(
-	db TxnSource,
 	poolSize int,
 	runtime module.Runtime,
 ) client.LensRegistry {
@@ -97,7 +90,6 @@ func NewRegistry(
 	}
 
 	return &implicitTxnLensRegistry{
-		db:       db,
 		registry: registry,
 	}
 }

internal/lens/txn_registry.go

@@ -22,7 +22,7 @@ import (
 
 type implicitTxnLensRegistry struct {
 	registry *lensRegistry
-	db       TxnSource
+	db       client.TxnSource
 }
 
 type explicitTxnLensRegistry struct {
@@ -33,13 +33,12 @@ type explicitTxnLensRegistry struct {
 var _ client.LensRegistry = (*implicitTxnLensRegistry)(nil)
 var _ client.LensRegistry = (*explicitTxnLensRegistry)(nil)
 
-func (r *implicitTxnLensRegistry) WithTxn(txn datastore.Txn) client.LensRegistry {
-	return &explicitTxnLensRegistry{
-		registry: r.registry,
-		txn:      txn,
-	}
+func (r *implicitTxnLensRegistry) Init(txnSource client.TxnSource) {
+	r.db = txnSource
 }
 
+func (r *explicitTxnLensRegistry) Init(txnSource client.TxnSource) {}
+
 func (r *explicitTxnLensRegistry) WithTxn(txn datastore.Txn) client.LensRegistry {
 	return &explicitTxnLensRegistry{
 		registry: r.registry,

node/lens.go

@@ -16,7 +16,6 @@ import (
 	"github.com/lens-vm/lens/host-go/engine/module"
 
 	"github.com/sourcenetwork/defradb/client"
-	"github.com/sourcenetwork/defradb/datastore"
 	"github.com/sourcenetwork/defradb/internal/lens"
 )
 
@@ -63,7 +62,6 @@ func WithLensPoolSize(size int) Option {
 
 func NewLens(
 	ctx context.Context,
-	txnSource *txnSource,
 	opts ...LenOpt,
 ) (client.LensRegistry, error) {
 	options := DefaultLensOptions()
@@ -79,18 +77,7 @@ func NewLens(
 	}
 
 	return lens.NewRegistry(
-		txnSource,
 		options.lensPoolSize,
 		runtime,
 	), nil
 }
-
-type txnSource struct {
-	db *client.DB
-}
-
-var _ lens.TxnSource = (*txnSource)(nil)
-
-func (s txnSource) NewTxn(ctx context.Context, readonly bool) (datastore.Txn, error) {
-	return (*s.db).NewTxn(ctx, readonly)
-}

node/node.go

@@ -128,8 +128,7 @@ func NewNode(ctx context.Context, opts ...Option) (*Node, error) {
 		return nil, err
 	}
 
-	txnSource := &txnSource{}
-	lens, err := NewLens(ctx, txnSource, lensOpts...)
+	lens, err := NewLens(ctx, lensOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -138,7 +137,6 @@ func NewNode(ctx context.Context, opts ...Option) (*Node, error) {
 	if err != nil {
 		return nil, err
 	}
-	txnSource.db = &db
 
 	var node *net.Node
 	if !options.disableP2P {

tests/clients/cli/wrapper_lens.go

@@ -28,6 +28,8 @@ type LensRegistry struct {
 	cmd *cliWrapper
 }
 
+func (w *LensRegistry) Init(txnSource client.TxnSource) {}
+
 func (w *LensRegistry) SetMigration(ctx context.Context, collectionID uint32, config model.Lens) error {
 	args := []string{"client", "schema", "migration", "set-registry"}