fix: recover balances query performance

internal/storage/bucket/bucket.go

@@ -9,26 +9,27 @@ import (
 	"go.opentelemetry.io/otel/trace/noop"
 )
 
+//go:generate mockgen -write_source_comment=false -write_package_comment=false -source bucket.go -destination bucket_generated_test.go -package bucket . Bucket
 type Bucket interface {
-	Migrate(ctx context.Context, opts ...migrations.Option) error
-	AddLedger(ctx context.Context, ledger ledger.Ledger) error
-	HasMinimalVersion(ctx context.Context) (bool, error)
-	IsUpToDate(ctx context.Context) (bool, error)
-	GetMigrationsInfo(ctx context.Context) ([]migrations.Info, error)
-	IsInitialized(context.Context) (bool, error)
+	Migrate(ctx context.Context, db bun.IDB, opts ...migrations.Option) error
+	AddLedger(ctx context.Context, db bun.IDB, ledger ledger.Ledger) error
+	HasMinimalVersion(ctx context.Context, db bun.IDB) (bool, error)
+	IsUpToDate(ctx context.Context, db bun.IDB) (bool, error)
+	GetMigrationsInfo(ctx context.Context, db bun.IDB) ([]migrations.Info, error)
+	IsInitialized(context.Context, bun.IDB) (bool, error)
 }
 
 type Factory interface {
-	Create(name string, db bun.IDB) Bucket
+	Create(name string) Bucket
 	GetMigrator(b string, db bun.IDB) *migrations.Migrator
 }
 
 type DefaultFactory struct {
 	tracer trace.Tracer
 }
 
-func (f *DefaultFactory) Create(name string, db bun.IDB) Bucket {
-	return NewDefault(db, f.tracer, name)
+func (f *DefaultFactory) Create(name string) Bucket {
+	return NewDefault(f.tracer, name)
 }
 
 func (f *DefaultFactory) GetMigrator(b string, db bun.IDB) *migrations.Migrator {

internal/storage/bucket/default_bucket.go

@@ -19,12 +19,11 @@ const MinimalSchemaVersion = 12
 
 type DefaultBucket struct {
 	name string
-	db   bun.IDB
 	tracer trace.Tracer
 }
 
-func (b *DefaultBucket) IsInitialized(ctx context.Context) (bool, error) {
-	_, err := GetMigrator(b.db, b.name).GetLastVersion(ctx)
+func (b *DefaultBucket) IsInitialized(ctx context.Context, db bun.IDB) (bool, error) {
+	_, err := GetMigrator(db, b.name).GetLastVersion(ctx)
 	if err == nil {
 		return true, nil
 	}
@@ -34,32 +33,32 @@ func (b *DefaultBucket) IsInitialized(ctx context.Context) (bool, error) {
 	return false, err
 }
 
-func (b *DefaultBucket) IsUpToDate(ctx context.Context) (bool, error) {
-	return GetMigrator(b.db, b.name).IsUpToDate(ctx)
+func (b *DefaultBucket) IsUpToDate(ctx context.Context, db bun.IDB) (bool, error) {
+	return GetMigrator(db, b.name).IsUpToDate(ctx)
 }
 
-func (b *DefaultBucket) Migrate(ctx context.Context, options ...migrations.Option) error {
-	return runMigrate(ctx, b.tracer, b.db, b.name, append(options, migrations.WithTracer(b.tracer))...)
+func (b *DefaultBucket) Migrate(ctx context.Context, db bun.IDB, options ...migrations.Option) error {
+	return runMigrate(ctx, b.tracer, db, b.name, append(options, migrations.WithTracer(b.tracer))...)
 }
 
-func (b *DefaultBucket) HasMinimalVersion(ctx context.Context) (bool, error) {
-	lastVersion, err := b.GetLastVersion(ctx)
+func (b *DefaultBucket) HasMinimalVersion(ctx context.Context, db bun.IDB) (bool, error) {
+	lastVersion, err := b.GetLastVersion(ctx, db)
 	if err != nil {
 		return false, err
 	}
 
 	return lastVersion >= MinimalSchemaVersion, nil
 }
 
-func (b *DefaultBucket) GetLastVersion(ctx context.Context) (int, error) {
-	return GetMigrator(b.db, b.name).GetLastVersion(ctx)
+func (b *DefaultBucket) GetLastVersion(ctx context.Context, db bun.IDB) (int, error) {
+	return GetMigrator(db, b.name).GetLastVersion(ctx)
 }
 
-func (b *DefaultBucket) GetMigrationsInfo(ctx context.Context) ([]migrations.Info, error) {
-	return GetMigrator(b.db, b.name).GetMigrations(ctx)
+func (b *DefaultBucket) GetMigrationsInfo(ctx context.Context, db bun.IDB) ([]migrations.Info, error) {
+	return GetMigrator(db, b.name).GetMigrations(ctx)
 }
 
-func (b *DefaultBucket) AddLedger(ctx context.Context, l ledger.Ledger) error {
+func (b *DefaultBucket) AddLedger(ctx context.Context, db bun.IDB, l ledger.Ledger) error {
 
 	for _, setup := range ledgerSetups {
 		if l.Features.Match(setup.requireFeatures) {
@@ -69,7 +68,7 @@ func (b *DefaultBucket) AddLedger(ctx context.Context, l ledger.Ledger) error {
 				return fmt.Errorf("executing template: %w", err)
 			}
 
-			_, err := b.db.ExecContext(ctx, buf.String())
+			_, err := db.ExecContext(ctx, buf.String())
 			if err != nil {
 				return fmt.Errorf("executing sql: %w", err)
 			}
@@ -79,9 +78,8 @@ func (b *DefaultBucket) AddLedger(ctx context.Context, l ledger.Ledger) error {
 	return nil
 }
 
-func NewDefault(db bun.IDB, tracer trace.Tracer, name string) *DefaultBucket {
+func NewDefault(tracer trace.Tracer, name string) *DefaultBucket {
 	return &DefaultBucket{
-		db:   db,
 		name: name,
 		tracer: tracer,
 	}

internal/storage/bucket/default_bucket_test.go

@@ -30,6 +30,6 @@ func TestBuckets(t *testing.T) {
 
 	require.NoError(t, system.Migrate(ctx, db))
 
-	b := bucket.NewDefault(db, noop.Tracer{}, name)
-	require.NoError(t, b.Migrate(ctx))
+	b := bucket.NewDefault(noop.Tracer{}, name)
+	require.NoError(t, b.Migrate(ctx, db))
 }

internal/storage/driver/adapters.go

@@ -20,7 +20,7 @@ func (d *DefaultStorageDriverAdapter) OpenLedger(ctx context.Context, name strin
 		return nil, nil, err
 	}
 
-	isUpToDate, err := store.GetBucket().IsUpToDate(ctx)
+	isUpToDate, err := store.GetBucket().IsUpToDate(ctx, d.db)
 	if err != nil {
 		return nil, nil, fmt.Errorf("checking if bucket is up to date: %w", err)
 	}