fix: remove cache and use dynamic ledger count query

- Remove singleLedgerOptimization cache structure
- Add countLedgersInBucket function to Store
- Modify isSingleLedger to query count on each call
- Update applyLedgerFilter and getLedgerFilterSQL to accept context
- Inject countFunc via factory options in module.go
- Remove all UpdateSingleLedgerState calls from driver
- Add Ctx field to ResourceQuery for context propagation
- Update all call sites to pass context parameter (26 locations)

Author: flemzord

internal/storage/common/resource.go

@@ -209,6 +209,7 @@ func (r *ResourceRepository[ResourceType, OptionsType]) expand(dataset *bun.Sele
 }
 
 func (r *ResourceRepository[ResourceType, OptionsType]) GetOne(ctx context.Context, query ResourceQuery[OptionsType]) (*ResourceType, error) {
+	query.Ctx = ctx
 
 	finalQuery, err := r.buildFilteredDataset(query)
 	if err != nil {
@@ -235,6 +236,7 @@ func (r *ResourceRepository[ResourceType, OptionsType]) GetOne(ctx context.Conte
 }
 
 func (r *ResourceRepository[ResourceType, OptionsType]) Count(ctx context.Context, query ResourceQuery[OptionsType]) (int, error) {
+	query.Ctx = ctx
 
 	finalQuery, err := r.buildFilteredDataset(query)
 	if err != nil {
@@ -325,6 +327,7 @@ func (r *PaginatedResourceRepository[ResourceType, OptionsType]) Paginate(
 	default:
 		panic("should not happen")
 	}
+	resourceQuery.Ctx = ctx
 
 	finalQuery, err := r.buildFilteredDataset(resourceQuery)
 	if err != nil {
@@ -406,11 +409,12 @@ func NewPaginatedResourceRepositoryMapper[ToResourceType any, OriginalResourceTy
 }
 
 type ResourceQuery[Opts any] struct {
-	PIT     *time.Time    `json:"pit"`
-	OOT     *time.Time    `json:"oot"`
-	Builder query.Builder `json:"qb"`
-	Expand  []string      `json:"expand,omitempty"`
-	Opts    Opts          `json:"opts"`
+	Ctx     context.Context `json:"-"`
+	PIT     *time.Time      `json:"pit"`
+	OOT     *time.Time      `json:"oot"`
+	Builder query.Builder   `json:"qb"`
+	Expand  []string        `json:"expand,omitempty"`
+	Opts    Opts            `json:"opts"`
 }
 
 func (rq ResourceQuery[Opts]) UsePIT() bool {

internal/storage/driver/driver.go

@@ -81,10 +81,6 @@ func (d *Driver) CreateLedger(ctx context.Context, l *ledger.Ledger) (*ledgersto
 		return nil, postgres.ResolveError(err)
 	}
 
-	if err := ret.UpdateSingleLedgerState(ctx, d.systemStoreFactory.Create(d.db).CountLedgersInBucket); err != nil {
-		logging.FromContext(ctx).Debugf("Failed to update single-ledger state: %v", err)
-	}
-
 	return ret, nil
 }
 
@@ -97,10 +93,6 @@ func (d *Driver) OpenLedger(ctx context.Context, name string) (*ledgerstore.Stor
 
 	store := d.ledgerStoreFactory.Create(d.bucketFactory.Create(ret.Bucket), *ret)
 
-	if err := store.UpdateSingleLedgerState(ctx, d.systemStoreFactory.Create(d.db).CountLedgersInBucket); err != nil {
-		logging.FromContext(ctx).Debugf("Failed to update single-ledger state: %v", err)
-	}
-
 	return store, ret, err
 }
 

internal/storage/driver/module.go

@@ -29,12 +29,18 @@ func NewFXModule() fx.Option {
 				&ledger.Ledger{},
 			)
 		}),
+		fx.Provide(func(tracerProvider trace.TracerProvider) systemstore.StoreFactory {
+			return systemstore.NewStoreFactory(systemstore.WithTracer(
+				tracerProvider.Tracer("SystemStore"),
+			))
+		}),
 		fx.Provide(func(params struct {
 			fx.In
 
-			DB             *bun.DB
-			TracerProvider trace.TracerProvider `optional:"true"`
-			MeterProvider  metric.MeterProvider `optional:"true"`
+			DB                 *bun.DB
+			SystemStoreFactory systemstore.StoreFactory
+			TracerProvider     trace.TracerProvider `optional:"true"`
+			MeterProvider      metric.MeterProvider `optional:"true"`
 		}) ledgerstore.Factory {
 			options := make([]ledgerstore.Option, 0)
 			if params.TracerProvider != nil {
@@ -43,21 +49,25 @@ func NewFXModule() fx.Option {
 			if params.MeterProvider != nil {
 				options = append(options, ledgerstore.WithMeter(params.MeterProvider.Meter("store")))
 			}
+			options = append(options, ledgerstore.WithCountLedgersInBucketFunc(
+				func(ctx context.Context, bucketName string) (int, error) {
+					return params.SystemStoreFactory.Create(params.DB).CountLedgersInBucket(ctx, bucketName)
+				},
+			))
 			return ledgerstore.NewFactory(params.DB, options...)
 		}),
 		fx.Provide(func(
 			db *bun.DB,
 			bucketFactory bucket.Factory,
 			ledgerStoreFactory ledgerstore.Factory,
+			systemStoreFactory systemstore.StoreFactory,
 			tracerProvider trace.TracerProvider,
 		) (*Driver, error) {
 			return New(
 				db,
 				ledgerStoreFactory,
 				bucketFactory,
-				systemstore.NewStoreFactory(systemstore.WithTracer(
-					tracerProvider.Tracer("SystemStore"),
-				)),
+				systemStoreFactory,
 				WithTracer(tracerProvider.Tracer("StorageDriver")),
 			), nil
 		}),

internal/storage/ledger/accounts.go

@@ -91,7 +91,7 @@ func (store *Store) DeleteAccountMetadata(ctx context.Context, account, key stri
 				Set("metadata = metadata - ?", key).
 				Where("address = ?", account)
 
-			if filterSQL, filterArgs := store.getLedgerFilterSQL(); filterSQL != "" {
+			if filterSQL, filterArgs := store.getLedgerFilterSQL(ctx); filterSQL != "" {
 				query = query.Where(filterSQL, filterArgs...)
 			}
 

internal/storage/ledger/logs.go

@@ -124,7 +124,7 @@ func (store *Store) ReadLogWithIdempotencyKey(ctx context.Context, key string) (
 				Column("*").
 				Where("idempotency_key = ?", key).
 				Limit(1)
-			query = store.applyLedgerFilter(query, "logs")
+			query = store.applyLedgerFilter(ctx, query, "logs")
 			if err := query.Scan(ctx); err != nil {
 				return nil, postgres.ResolveError(err)
 			}

internal/storage/ledger/resource_accounts.go

@@ -29,7 +29,7 @@ func (h accountsResourceHandler) BuildDataset(opts common.RepositoryHandlerBuild
 	ret := h.store.db.NewSelect().
 		ModelTableExpr(h.store.GetPrefixedRelationName("accounts")).
 		Column("address", "address_array", "first_usage", "insertion_date", "updated_at")
-	ret = h.store.applyLedgerFilter(ret, "accounts")
+	ret = h.store.applyLedgerFilter(opts.Ctx, ret, "accounts")
 
 	if opts.PIT != nil && !opts.PIT.IsZero() {
 		ret = ret.Where("accounts.first_usage <= ?", opts.PIT)
@@ -42,7 +42,7 @@ func (h accountsResourceHandler) BuildDataset(opts common.RepositoryHandlerBuild
 			Column("accounts_address").
 			ColumnExpr("first_value(metadata) over (partition by accounts_address order by revision desc) as metadata").
 			Where("date <= ?", opts.PIT)
-		selectDistinctAccountMetadataHistories = h.store.applyLedgerFilter(selectDistinctAccountMetadataHistories, "accounts_metadata")
+		selectDistinctAccountMetadataHistories = h.store.applyLedgerFilter(opts.Ctx, selectDistinctAccountMetadataHistories, "accounts_metadata")
 
 		ret = ret.
 			Join(
@@ -77,12 +77,12 @@ func (h accountsResourceHandler) ResolveFilter(opts common.ResourceQuery[any], o
 				DistinctOn("asset").
 				ColumnExpr("first_value((post_commit_effective_volumes).inputs - (post_commit_effective_volumes).outputs) over (partition by (accounts_address, asset) order by effective_date desc, seq desc) as balance").
 				Where("effective_date <= ?", opts.PIT)
-			selectBalance = h.store.applyLedgerFilter(selectBalance, "moves")
+			selectBalance = h.store.applyLedgerFilter(opts.Ctx, selectBalance, "moves")
 		} else {
 			selectBalance = selectBalance.
 				ModelTableExpr(h.store.GetPrefixedRelationName("accounts_volumes")).
 				ColumnExpr("input - output as balance")
-			selectBalance = h.store.applyLedgerFilter(selectBalance, "accounts_volumes")
+			selectBalance = h.store.applyLedgerFilter(opts.Ctx, selectBalance, "accounts_volumes")
 		}
 
 		if balanceRegex.MatchString(property) {
@@ -130,7 +130,7 @@ func (h accountsResourceHandler) Expand(opts common.ResourceQuery[any], property
 			ModelTableExpr(h.store.GetPrefixedRelationName("moves")).
 			DistinctOn("accounts_address, asset").
 			Column("accounts_address", "asset")
-		selectRowsQuery = h.store.applyLedgerFilter(selectRowsQuery, "moves")
+		selectRowsQuery = h.store.applyLedgerFilter(opts.Ctx, selectRowsQuery, "moves")
 		if property == "volumes" {
 			selectRowsQuery = selectRowsQuery.
 				ColumnExpr("first_value(post_commit_volumes) over (partition by (accounts_address, asset) order by seq desc) as volumes").
@@ -145,7 +145,7 @@ func (h accountsResourceHandler) Expand(opts common.ResourceQuery[any], property
 			ModelTableExpr(h.store.GetPrefixedRelationName("accounts_volumes")).
 			Column("asset", "accounts_address").
 			ColumnExpr("(input, output)::"+h.store.GetPrefixedRelationName("volumes")+" as volumes")
-		selectRowsQuery = h.store.applyLedgerFilter(selectRowsQuery, "accounts_volumes")
+		selectRowsQuery = h.store.applyLedgerFilter(opts.Ctx, selectRowsQuery, "accounts_volumes")
 	}
 
 	return h.store.db.NewSelect().

internal/storage/ledger/resource_aggregated_balances.go

@@ -27,7 +27,7 @@ func (h aggregatedBalancesResourceRepositoryHandler) BuildDataset(query common.R
 			ModelTableExpr(h.store.GetPrefixedRelationName("moves")).
 			DistinctOn("accounts_address, asset").
 			Column("accounts_address", "asset")
-		ret = h.store.applyLedgerFilter(ret, "moves")
+		ret = h.store.applyLedgerFilter(query.Ctx, ret, "moves")
 		if query.Opts.UseInsertionDate {
 			if !h.store.ledger.HasFeature(features.FeatureMovesHistory, "ON") {
 				return nil, NewErrMissingFeature(features.FeatureMovesHistory)
@@ -53,7 +53,7 @@ func (h aggregatedBalancesResourceRepositoryHandler) BuildDataset(query common.R
 				TableExpr(h.store.GetPrefixedRelationName("accounts")).
 				Column("address_array").
 				Where("accounts.address = accounts_address")
-			subQuery = h.store.applyLedgerFilter(subQuery, "accounts")
+			subQuery = h.store.applyLedgerFilter(query.Ctx, subQuery, "accounts")
 
 			ret = ret.
 				ColumnExpr("accounts.address_array as accounts_address_array").
@@ -67,7 +67,7 @@ func (h aggregatedBalancesResourceRepositoryHandler) BuildDataset(query common.R
 				ColumnExpr("first_value(metadata) over (partition by accounts_address order by revision desc) as metadata").
 				Where("accounts_metadata.accounts_address = moves.accounts_address").
 				Where("date <= ?", query.PIT)
-			subQuery = h.store.applyLedgerFilter(subQuery, "accounts_metadata")
+			subQuery = h.store.applyLedgerFilter(query.Ctx, subQuery, "accounts_metadata")
 
 			ret = ret.
 				Join(`left join lateral (?) accounts_metadata on true`, subQuery).
@@ -80,7 +80,7 @@ func (h aggregatedBalancesResourceRepositoryHandler) BuildDataset(query common.R
 			ModelTableExpr(h.store.GetPrefixedRelationName("accounts_volumes")).
 			Column("asset", "accounts_address").
 			ColumnExpr("(input, output)::"+h.store.GetPrefixedRelationName("volumes")+" as volumes")
-		ret = h.store.applyLedgerFilter(ret, "accounts_volumes")
+		ret = h.store.applyLedgerFilter(query.Ctx, ret, "accounts_volumes")
 
 		if query.UseFilter("metadata") || query.UseFilter("address", func(value any) bool {
 			return isPartialAddress(value.(string))
@@ -89,7 +89,7 @@ func (h aggregatedBalancesResourceRepositoryHandler) BuildDataset(query common.R
 				TableExpr(h.store.GetPrefixedRelationName("accounts")).
 				Column("address").
 				Where("accounts.address = accounts_address")
-			subQuery = h.store.applyLedgerFilter(subQuery, "accounts")
+			subQuery = h.store.applyLedgerFilter(query.Ctx, subQuery, "accounts")
 
 			if query.UseFilter("address") {
 				subQuery = subQuery.ColumnExpr("address_array as accounts_address_array")

internal/storage/ledger/resource_logs.go

@@ -20,11 +20,11 @@ func (h logsResourceHandler) Schema() common.EntitySchema {
 	}
 }
 
-func (h logsResourceHandler) BuildDataset(_ common.RepositoryHandlerBuildContext[any]) (*bun.SelectQuery, error) {
+func (h logsResourceHandler) BuildDataset(opts common.RepositoryHandlerBuildContext[any]) (*bun.SelectQuery, error) {
 	ret := h.store.db.NewSelect().
 		ModelTableExpr(h.store.GetPrefixedRelationName("logs")).
 		ColumnExpr("*")
-	ret = h.store.applyLedgerFilter(ret, "logs")
+	ret = h.store.applyLedgerFilter(opts.Ctx, ret, "logs")
 	return ret, nil
 }
 

internal/storage/ledger/resource_transactions.go

@@ -45,7 +45,7 @@ func (h transactionsResourceHandler) BuildDataset(opts common.RepositoryHandlerB
 			"sources_arrays",
 			"destinations_arrays",
 		)
-	ret = h.store.applyLedgerFilter(ret, "transactions")
+	ret = h.store.applyLedgerFilter(opts.Ctx, ret, "transactions")
 
 	if slices.Contains(opts.Expand, "volumes") {
 		ret = ret.Column("post_commit_volumes")
@@ -62,7 +62,7 @@ func (h transactionsResourceHandler) BuildDataset(opts common.RepositoryHandlerB
 			Column("transactions_id", "metadata").
 			Order("transactions_id", "revision desc").
 			Where("date <= ?", opts.PIT)
-		selectDistinctTransactionMetadataHistories = h.store.applyLedgerFilter(selectDistinctTransactionMetadataHistories, "transactions_metadata")
+		selectDistinctTransactionMetadataHistories = h.store.applyLedgerFilter(opts.Ctx, selectDistinctTransactionMetadataHistories, "transactions_metadata")
 
 		ret = ret.
 			Join(
@@ -119,7 +119,7 @@ func (h transactionsResourceHandler) Project(_ common.ResourceQuery[any], select
 	return selectQuery.ColumnExpr("*"), nil
 }
 
-func (h transactionsResourceHandler) Expand(_ common.ResourceQuery[any], property string) (*bun.SelectQuery, *common.JoinCondition, error) {
+func (h transactionsResourceHandler) Expand(query common.ResourceQuery[any], property string) (*bun.SelectQuery, *common.JoinCondition, error) {
 	if property != "effectiveVolumes" {
 		return nil, nil, nil
 	}
@@ -130,7 +130,7 @@ func (h transactionsResourceHandler) Expand(_ common.ResourceQuery[any], propert
 		Column("transactions_id", "accounts_address", "asset").
 		ColumnExpr(`first_value(moves.post_commit_effective_volumes) over (partition by (transactions_id, accounts_address, asset) order by seq desc) as post_commit_effective_volumes`).
 		Where("transactions_id in (select id from dataset)")
-	innerMostQuery = h.store.applyLedgerFilter(innerMostQuery, "moves")
+	innerMostQuery = h.store.applyLedgerFilter(query.Ctx, innerMostQuery, "moves")
 
 	ret := h.store.db.NewSelect().
 		TableExpr(

internal/storage/ledger/resource_volumes.go

@@ -40,14 +40,14 @@ func (h volumesResourceHandler) BuildDataset(query common.RepositoryHandlerBuild
 			ColumnExpr("accounts_address as account").
 			ModelTableExpr(h.store.GetPrefixedRelationName("accounts_volumes")).
 			Order("accounts_address", "asset")
-		selectVolumes = h.store.applyLedgerFilter(selectVolumes, "accounts_volumes")
+		selectVolumes = h.store.applyLedgerFilter(query.Ctx, selectVolumes, "accounts_volumes")
 
 		if query.UseFilter("metadata") || query.UseFilter("first_usage") || needAddressSegments {
 			accountsQuery := h.store.db.NewSelect().
 				TableExpr(h.store.GetPrefixedRelationName("accounts")).
 				Column("address").
 				Where("accounts.address = accounts_address")
-			accountsQuery = h.store.applyLedgerFilter(accountsQuery, "accounts")
+			accountsQuery = h.store.applyLedgerFilter(query.Ctx, accountsQuery, "accounts")
 
 			if needAddressSegments {
 				accountsQuery = accountsQuery.ColumnExpr("address_array as account_array")
@@ -79,7 +79,7 @@ func (h volumesResourceHandler) BuildDataset(query common.RepositoryHandlerBuild
 			ModelTableExpr(h.store.GetPrefixedRelationName("moves")).
 			GroupExpr("accounts_address, asset").
 			Order("accounts_address", "asset")
-		selectVolumes = h.store.applyLedgerFilter(selectVolumes, "moves")
+		selectVolumes = h.store.applyLedgerFilter(query.Ctx, selectVolumes, "moves")
 
 		dateFilterColumn := "effective_date"
 		if query.Opts.UseInsertionDate {
@@ -98,7 +98,7 @@ func (h volumesResourceHandler) BuildDataset(query common.RepositoryHandlerBuild
 			accountsQuery := h.store.db.NewSelect().
 				TableExpr(h.store.GetPrefixedRelationName("accounts")).
 				Where("accounts.address = accounts_address")
-			accountsQuery = h.store.applyLedgerFilter(accountsQuery, "accounts")
+			accountsQuery = h.store.applyLedgerFilter(query.Ctx, accountsQuery, "accounts")
 
 			if needAddressSegments {
 				accountsQuery = accountsQuery.ColumnExpr("address_array")
@@ -117,7 +117,7 @@ func (h volumesResourceHandler) BuildDataset(query common.RepositoryHandlerBuild
 				ModelTableExpr(h.store.GetPrefixedRelationName("accounts_metadata")).
 				ColumnExpr("first_value(metadata) over (partition by accounts_address order by revision desc) as metadata").
 				Where("accounts_metadata.accounts_address = moves.accounts_address")
-			subQuery = h.store.applyLedgerFilter(subQuery, "accounts_metadata")
+			subQuery = h.store.applyLedgerFilter(query.Ctx, subQuery, "accounts_metadata")
 
 			selectVolumes = selectVolumes.
 				Join(`left join lateral (?) accounts_metadata on true`, subQuery).