feat: add candidate-key pagination for selective transactions

cmd/serve.go

@@ -60,6 +60,8 @@ type ServeCommandConfig struct {
 	MaxPageSize            uint64 `mapstructure:"max-page-size"`
 	WorkerEnabled          bool   `mapstructure:"worker"`
 	WorkerAddress          string `mapstructure:"worker-grpc-address"`
+
+	ExperimentalTransactionCandidatePagination bool `mapstructure:"experimental-transaction-candidate-pagination"`
 }
 
 const (
@@ -75,6 +77,8 @@ const (
 	WorkerEnabledFlag     = "worker"
 	SemconvMetricsNames   = "semconv-metrics-names"
 	SchemaEnforcementMode = "schema-enforcement-mode"
+
+	ExperimentalTransactionCandidatePaginationFlag = "experimental-transaction-candidate-pagination"
 )
 
 func NewServeCommand() *cobra.Command {
@@ -106,6 +110,7 @@ func NewServeCommand() *cobra.Command {
 				storagefx.BunConnectModule(*connectionOptions, service.IsDebug(cmd)),
 				storage.NewFXModule(storage.ModuleConfig{
 					AutoUpgrade: cfg.AutoUpgrade,
+					ExperimentalTransactionCandidatePagination: cfg.ExperimentalTransactionCandidatePagination,
 				}),
 				drivers.NewFXModule(),
 				fx.Invoke(alldrivers.Register),
@@ -192,6 +197,7 @@ func NewServeCommand() *cobra.Command {
 	cmd.Flags().Uint64(DefaultPageSizeFlag, 15, "Default page size")
 	cmd.Flags().Bool(WorkerEnabledFlag, false, "Enable worker")
 	cmd.Flags().Bool(ExperimentalFeaturesFlag, false, "Enable features configurability")
+	cmd.Flags().Bool(ExperimentalTransactionCandidatePaginationFlag, false, "Enable experimental candidate-key pagination for selective transaction queries")
 	cmd.Flags().Bool(NumscriptInterpreterFlag, false, "Enable experimental numscript rewrite")
 	cmd.Flags().StringSlice(NumscriptInterpreterFlagsToPass, nil, "Feature flags to pass to the experimental numscript interpreter")
 	cmd.Flags().String(WorkerGRPCAddressFlag, "localhost:8081", "GRPC address")

internal/storage/common/resource.go

@@ -97,6 +97,11 @@ type RepositoryHandler[Opts any] interface {
 	Expand(query ResourceQuery[Opts], property string) (*bun.SelectQuery, *JoinCondition, error)
 }
 
+type CandidateKeyPaginationHandler[Opts any] interface {
+	UseCandidateKeyPagination(ResourceQuery[Opts], ColumnPaginatedQuery[Opts]) bool
+	CandidateKeyColumns(ResourceQuery[Opts], ColumnPaginatedQuery[Opts]) []string
+}
+
 type ResourceRepository[ResourceType, OptionsType any] struct {
 	resourceHandler RepositoryHandler[OptionsType]
 }
@@ -310,8 +315,10 @@ func (r *PaginatedResourceRepository[ResourceType, OptionsType]) Paginate(
 	}
 
 	var (
-		paginator     Paginator[ResourceType]
-		resourceQuery ResourceQuery[OptionsType]
+		paginator                Paginator[ResourceType]
+		resourceQuery            ResourceQuery[OptionsType]
+		columnPaginationQuery    *ColumnPaginatedQuery[OptionsType]
+		columnPaginationStrategy *columnPaginator[ResourceType, OptionsType]
 	)
 	switch v := any(paginationQuery).(type) {
 	case OffsetPaginatedQuery[OptionsType]:
@@ -322,12 +329,46 @@ func (r *PaginatedResourceRepository[ResourceType, OptionsType]) Paginate(
 		if field == nil {
 			return nil, fmt.Errorf("invalid property '%s' for pagination", v.Column)
 		}
-		paginator = newColumnPaginator[ResourceType, OptionsType](v, fieldName, field.Type)
+		cp := newColumnPaginator[ResourceType, OptionsType](v, fieldName, field.Type)
+		paginator = cp
+		columnPaginationQuery = &v
+		columnPaginationStrategy = &cp
 		resourceQuery = v.Options
 	default:
 		panic("should not happen")
 	}
 
+	if columnPaginationQuery != nil {
+		if candidateHandler, ok := r.resourceHandler.(CandidateKeyPaginationHandler[OptionsType]); ok &&
+			candidateHandler.UseCandidateKeyPagination(resourceQuery, *columnPaginationQuery) {
+			finalQuery, err := r.buildCandidateKeyPaginatedDataset(
+				resourceQuery,
+				*columnPaginationStrategy,
+				candidateHandler.CandidateKeyColumns(resourceQuery, *columnPaginationQuery),
+			)
+			if err != nil {
+				return nil, err
+			}
+
+			finalQuery, err = r.expand(finalQuery, resourceQuery)
+			if err != nil {
+				return nil, fmt.Errorf("expanding results: %w", err)
+			}
+
+			orderExpr := paginator.OrderExpression()
+			col, dir, _ := strings.Cut(orderExpr, " ")
+			finalQuery = finalQuery.Order(fmt.Sprintf("dataset.%s %s", col, dir))
+
+			ret := make([]ResourceType, 0)
+			err = finalQuery.Model(&ret).Scan(ctx)
+			if err != nil {
+				return nil, fmt.Errorf("scanning results: %w", err)
+			}
+
+			return paginator.BuildCursor(ret)
+		}
+	}
+
 	finalQuery, err := r.buildFilteredDataset(resourceQuery)
 	if err != nil {
 		return nil, fmt.Errorf("building filtered dataset: %w", err)
@@ -358,6 +399,56 @@ func (r *PaginatedResourceRepository[ResourceType, OptionsType]) Paginate(
 	return paginator.BuildCursor(ret)
 }
 
+func (r *PaginatedResourceRepository[ResourceType, OptionsType]) buildCandidateKeyPaginatedDataset(
+	resourceQuery ResourceQuery[OptionsType],
+	paginator columnPaginator[ResourceType, OptionsType],
+	keyColumns []string,
+) (*bun.SelectQuery, error) {
+	if len(keyColumns) == 0 {
+		return nil, fmt.Errorf("candidate key pagination requires at least one key column")
+	}
+
+	filteredDataset, err := r.buildFilteredDataset(resourceQuery)
+	if err != nil {
+		return nil, fmt.Errorf("building filtered candidate dataset: %w", err)
+	}
+
+	candidateKeys := filteredDataset.NewSelect().
+		ModelTableExpr("(?) candidate_source", filteredDataset)
+	for _, column := range keyColumns {
+		candidateKeys = candidateKeys.ColumnExpr("candidate_source." + column)
+	}
+
+	page := filteredDataset.NewSelect().
+		ModelTableExpr("candidate_keys")
+	for _, column := range keyColumns {
+		page = page.Column(column)
+	}
+	page, err = paginator.Paginate(page)
+	if err != nil {
+		return nil, fmt.Errorf("paginating candidate keys: %w", err)
+	}
+
+	unfilteredResourceQuery := resourceQuery
+	unfilteredResourceQuery.Builder = nil
+	unfilteredDataset, err := r.buildFilteredDataset(unfilteredResourceQuery)
+	if err != nil {
+		return nil, fmt.Errorf("building unfiltered page dataset: %w", err)
+	}
+
+	joinConditions := make([]string, 0, len(keyColumns))
+	for _, column := range keyColumns {
+		joinConditions = append(joinConditions, fmt.Sprintf("page.%s = dataset.%s", column, column))
+	}
+
+	return unfilteredDataset.NewSelect().
+		WithQuery(bun.NewWithQuery("candidate_keys", candidateKeys).Materialized()).
+		With("page", page).
+		ModelTableExpr("(?) dataset", unfilteredDataset).
+		ColumnExpr("dataset.*").
+		Join("join page on " + strings.Join(joinConditions, " and ")), nil
+}
+
 func NewPaginatedResourceRepository[ResourceType, OptionsType any](
 	handler RepositoryHandler[OptionsType],
 	defaultPaginationColumn string,

internal/storage/driver/module.go

@@ -16,7 +16,11 @@ import (
 	systemstore "github.com/formancehq/ledger/internal/storage/system"
 )
 
-func NewFXModule() fx.Option {
+type ModuleConfig struct {
+	ExperimentalTransactionCandidatePagination bool
+}
+
+func NewFXModule(config ModuleConfig) fx.Option {
 	return fx.Options(
 		fx.Provide(fx.Annotate(func(tracerProvider trace.TracerProvider) bucket.Factory {
 			return bucket.NewDefaultFactory(bucket.WithTracer(tracerProvider.Tracer("store")))
@@ -44,6 +48,9 @@ func NewFXModule() fx.Option {
 			if params.MeterProvider != nil {
 				options = append(options, ledgerstore.WithMeter(params.MeterProvider.Meter("store")))
 			}
+			if config.ExperimentalTransactionCandidatePagination {
+				options = append(options, ledgerstore.WithExperimentalTransactionCandidatePagination(true))
+			}
 			return ledgerstore.NewFactory(params.DB, options...)
 		}),
 		fx.Provide(func(

internal/storage/ledger/resource_transactions.go

@@ -6,6 +6,8 @@ import (
 
 	"github.com/uptrace/bun"
 
+	"github.com/formancehq/go-libs/v5/pkg/storage/bun/paginate"
+
 	"github.com/formancehq/ledger/internal/queries"
 	"github.com/formancehq/ledger/internal/storage/common"
 	"github.com/formancehq/ledger/pkg/features"
@@ -186,4 +188,31 @@ func (h transactionsResourceHandler) Expand(_ common.ResourceQuery[any], propert
 	}, nil
 }
 
+func (h transactionsResourceHandler) UseCandidateKeyPagination(
+	q common.ResourceQuery[any],
+	paginationQuery common.ColumnPaginatedQuery[any],
+) bool {
+	if !h.store.experimentalTransactionCandidatePagination {
+		return false
+	}
+	if paginationQuery.Column != "id" {
+		return false
+	}
+	if paginationQuery.Order == nil || *paginationQuery.Order != paginate.OrderDesc {
+		return false
+	}
+	if paginationQuery.Reverse {
+		return false
+	}
+	return hasPositiveSelectiveTransactionFilter(q.Builder)
+}
+
+func (h transactionsResourceHandler) CandidateKeyColumns(
+	common.ResourceQuery[any],
+	common.ColumnPaginatedQuery[any],
+) []string {
+	return []string{"ledger", "id"}
+}
+
 var _ common.RepositoryHandler[any] = transactionsResourceHandler{}
+var _ common.CandidateKeyPaginationHandler[any] = transactionsResourceHandler{}

internal/storage/ledger/store.go

@@ -34,6 +34,8 @@ type Store struct {
 	// (e.g. when a new ledger is created) immediately affects all stores.
 	aloneInBucket *atomic.Bool
 
+	experimentalTransactionCandidatePagination bool
+
 	tracer                             trace.Tracer
 	meter                              metric.Meter
 	checkBucketSchemaHistogram         metric.Int64Histogram
@@ -346,6 +348,12 @@ func WithTracer(tracer trace.Tracer) Option {
 	}
 }
 
+func WithExperimentalTransactionCandidatePagination(enabled bool) Option {
+	return func(s *Store) {
+		s.experimentalTransactionCandidatePagination = enabled
+	}
+}
+
 var defaultOptions = []Option{
 	WithMeter(noopmetrics.Meter{}),
 	WithTracer(nooptracer.Tracer{}),