fix(query): filtered collections pagination (backport #16905)

CHANGELOG.md

@@ -57,6 +57,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ### Bug Fixes
 
 * (x/bank) [#16841](https://github.com/cosmos/cosmos-sdk/pull/16841) Correctly process legacy `DenomAddressIndex` values.
+* (types/query) [#16905](https://github.com/cosmos/cosmos-sdk/pull/16905) Collections Pagination now applies proper count when filtering results.
 
 ### API Breaking Changes
 

simapp/go.mod

@@ -12,8 +12,8 @@ require (
 	cosmossdk.io/store v1.0.0-alpha.1
 	cosmossdk.io/tools/confix v0.0.0-20230713160716-d4e95eec9f29
 	cosmossdk.io/tools/rosetta v0.2.1-0.20230713160716-d4e95eec9f29
-	cosmossdk.io/x/circuit v0.0.0-20230713160716-d4e95eec9f29
-	cosmossdk.io/x/evidence v0.0.0-20230713160716-d4e95eec9f29
+	cosmossdk.io/x/circuit v0.0.0-20230713220914-c0369a888135
+	cosmossdk.io/x/evidence v0.0.0-20230713220914-c0369a888135
 	cosmossdk.io/x/feegrant v0.0.0-20230713160716-d4e95eec9f29
 	cosmossdk.io/x/nft v0.0.0-20230713160716-d4e95eec9f29
 	cosmossdk.io/x/tx v0.9.1

simapp/go.sum

@@ -209,10 +209,10 @@ cosmossdk.io/tools/confix v0.0.0-20230713160716-d4e95eec9f29 h1:uIWFA16Db7llUhWZ
 cosmossdk.io/tools/confix v0.0.0-20230713160716-d4e95eec9f29/go.mod h1:b/bU9v699JldsEnnsDFGAAem+L1dHA+EWOI4l3Ik1cg=
 cosmossdk.io/tools/rosetta v0.2.1-0.20230713160716-d4e95eec9f29 h1:7c/8Q6icB72IUkOO47ua2o7/6+GxdmTUmaVbgoP5gOA=
 cosmossdk.io/tools/rosetta v0.2.1-0.20230713160716-d4e95eec9f29/go.mod h1:jdr/6CIOCmyPd27m3/GyN9I1lDuSmZx+zpsieG3Hdbc=
-cosmossdk.io/x/circuit v0.0.0-20230713160716-d4e95eec9f29 h1:DQqRU/fT8TM8HykDLRCBEsbOgjRe6wPI5SrbWmGvsNE=
-cosmossdk.io/x/circuit v0.0.0-20230713160716-d4e95eec9f29/go.mod h1:oo7//3TLAxU3sHyCxUeiUmUw0a6VFoXiJs7TR0Te0+4=
-cosmossdk.io/x/evidence v0.0.0-20230713160716-d4e95eec9f29 h1:xzwV67rZifJRCCUSzKlN5s3xPGoj/mlRQW5NLOJ9uCQ=
-cosmossdk.io/x/evidence v0.0.0-20230713160716-d4e95eec9f29/go.mod h1:SUYYupIrpF2id6W+1VMcDpAfmeglZ3cTIfc0/PLtAHo=
+cosmossdk.io/x/circuit v0.0.0-20230713220914-c0369a888135 h1:Qfywx4oCgJx3TzMiU5MM0R5zQO5dR2m29r6Y3EOAaf0=
+cosmossdk.io/x/circuit v0.0.0-20230713220914-c0369a888135/go.mod h1:jxt8hWrbvtpNAxC1bIh5gojBIqvX2w83wJVquVwsOKM=
+cosmossdk.io/x/evidence v0.0.0-20230713220914-c0369a888135 h1:SWjE56GlSCe5RSMUuC+peWhnkZa4Lbev0fJip3WflbA=
+cosmossdk.io/x/evidence v0.0.0-20230713220914-c0369a888135/go.mod h1:WiGXY9/0GKp0FFKPmNr0OtMxE5MOoRkDeE6HWZE41Pw=
 cosmossdk.io/x/feegrant v0.0.0-20230713160716-d4e95eec9f29 h1:oqmjaQAAYjNo9Sj77IsCuTuuKNN4Q+Py3zwclSDsroU=
 cosmossdk.io/x/feegrant v0.0.0-20230713160716-d4e95eec9f29/go.mod h1:izdG3ENxO7Z6kSRbZorbRJHarV2efZM63+Xyk27KksU=
 cosmossdk.io/x/nft v0.0.0-20230713160716-d4e95eec9f29 h1:Orne21OeV4lMN7zdbKQs4/vQAiTpR1IOeX1tFcezXuw=

tests/go.mod

@@ -12,7 +12,7 @@ require (
 	cosmossdk.io/math v1.0.1
 	cosmossdk.io/simapp v0.0.0-20230620040119-e078f1a49e8b
 	cosmossdk.io/store v1.0.0-alpha.1
-	cosmossdk.io/x/evidence v0.0.0-20230713160716-d4e95eec9f29
+	cosmossdk.io/x/evidence v0.0.0-20230713220914-c0369a888135
 	cosmossdk.io/x/feegrant v0.0.0-20230713160716-d4e95eec9f29
 	cosmossdk.io/x/nft v0.0.0-20230713160716-d4e95eec9f29 // indirect
 	cosmossdk.io/x/tx v0.9.1
@@ -39,7 +39,7 @@ require (
 	cloud.google.com/go/iam v1.1.0 // indirect
 	cloud.google.com/go/storage v1.30.1 // indirect
 	cosmossdk.io/client/v2 v2.0.0-20230713160716-d4e95eec9f29 // indirect
-	cosmossdk.io/x/circuit v0.0.0-20230713160716-d4e95eec9f29 // indirect
+	cosmossdk.io/x/circuit v0.0.0-20230713220914-c0369a888135 // indirect
 	filippo.io/edwards25519 v1.0.0 // indirect
 	github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 // indirect
 	github.com/99designs/keyring v1.2.1 // indirect

tests/go.sum

@@ -205,10 +205,10 @@ cosmossdk.io/math v1.0.1 h1:Qx3ifyOPaMLNH/89WeZFH268yCvU4xEcnPLu3sJqPPg=
 cosmossdk.io/math v1.0.1/go.mod h1:Ygz4wBHrgc7g0N+8+MrnTfS9LLn9aaTGa9hKopuym5k=
 cosmossdk.io/store v1.0.0-alpha.1 h1:/151XxAgm0tiKuYrtJzMG61lf6enpPuP+D/hIN8cRjQ=
 cosmossdk.io/store v1.0.0-alpha.1/go.mod h1:ejgU9GhRGMNBduVnDwC3RyhOmu4uKlNQlTiJgPmbDkI=
-cosmossdk.io/x/circuit v0.0.0-20230713160716-d4e95eec9f29 h1:DQqRU/fT8TM8HykDLRCBEsbOgjRe6wPI5SrbWmGvsNE=
-cosmossdk.io/x/circuit v0.0.0-20230713160716-d4e95eec9f29/go.mod h1:oo7//3TLAxU3sHyCxUeiUmUw0a6VFoXiJs7TR0Te0+4=
-cosmossdk.io/x/evidence v0.0.0-20230713160716-d4e95eec9f29 h1:xzwV67rZifJRCCUSzKlN5s3xPGoj/mlRQW5NLOJ9uCQ=
-cosmossdk.io/x/evidence v0.0.0-20230713160716-d4e95eec9f29/go.mod h1:SUYYupIrpF2id6W+1VMcDpAfmeglZ3cTIfc0/PLtAHo=
+cosmossdk.io/x/circuit v0.0.0-20230713220914-c0369a888135 h1:Qfywx4oCgJx3TzMiU5MM0R5zQO5dR2m29r6Y3EOAaf0=
+cosmossdk.io/x/circuit v0.0.0-20230713220914-c0369a888135/go.mod h1:jxt8hWrbvtpNAxC1bIh5gojBIqvX2w83wJVquVwsOKM=
+cosmossdk.io/x/evidence v0.0.0-20230713220914-c0369a888135 h1:SWjE56GlSCe5RSMUuC+peWhnkZa4Lbev0fJip3WflbA=
+cosmossdk.io/x/evidence v0.0.0-20230713220914-c0369a888135/go.mod h1:WiGXY9/0GKp0FFKPmNr0OtMxE5MOoRkDeE6HWZE41Pw=
 cosmossdk.io/x/feegrant v0.0.0-20230713160716-d4e95eec9f29 h1:oqmjaQAAYjNo9Sj77IsCuTuuKNN4Q+Py3zwclSDsroU=
 cosmossdk.io/x/feegrant v0.0.0-20230713160716-d4e95eec9f29/go.mod h1:izdG3ENxO7Z6kSRbZorbRJHarV2efZM63+Xyk27KksU=
 cosmossdk.io/x/nft v0.0.0-20230713160716-d4e95eec9f29 h1:Orne21OeV4lMN7zdbKQs4/vQAiTpR1IOeX1tFcezXuw=

types/query/collections_pagination.go

@@ -35,24 +35,40 @@ type Collection[K, V any] interface {
 	KeyCodec() collcodec.KeyCodec[K]
 }
 
-// CollectionPaginate follows the same behavior as Paginate but works on a Collection.
-func CollectionPaginate[K, V any, C Collection[K, V]](
+// CollectionPaginate follows the same logic as Paginate but for collection types.
+// transformFunc is used to transform the result to a different type.
+func CollectionPaginate[K, V any, C Collection[K, V], T any](
 	ctx context.Context,
 	coll C,
 	pageReq *PageRequest,
-) ([]collections.KeyValue[K, V], *PageResponse, error) {
-	return CollectionFilteredPaginate[K, V](ctx, coll, pageReq, nil)
+	transformFunc func(key K, value V) (T, error),
+	opts ...func(opt *CollectionsPaginateOptions[K]),
+) ([]T, *PageResponse, error) {
+	return CollectionFilteredPaginate(
+		ctx,
+		coll,
+		pageReq,
+		nil,
+		transformFunc,
+		opts...,
+	)
 }
 
-// CollectionFilteredPaginate works in the same way as FilteredPaginate but for collection types.
+// CollectionFilteredPaginate works in the same way as CollectionPaginate but allows to filter
+// results using a predicateFunc.
 // A nil predicateFunc means no filtering is applied and results are collected as is.
-func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
+// TransformFunc is applied only to results which are in range of the pagination and allow
+// to convert the result to a different type.
+// NOTE: do not collect results using the values/keys passed to predicateFunc as they are not
+// guaranteed to be in the pagination range requested.
+func CollectionFilteredPaginate[K, V any, C Collection[K, V], T any](
 	ctx context.Context,
 	coll C,
 	pageReq *PageRequest,
 	predicateFunc func(key K, value V) (include bool, err error),
+	transformFunc func(key K, value V) (T, error),
 	opts ...func(opt *CollectionsPaginateOptions[K]),
-) ([]collections.KeyValue[K, V], *PageResponse, error) {
+) (results []T, pageRes *PageResponse, err error) {
 	pageReq = initPageRequestDefaults(pageReq)
 
 	offset := pageReq.Offset
@@ -65,12 +81,6 @@ func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 		return nil, nil, fmt.Errorf("invalid request, either offset or key is expected, got both")
 	}
 
-	var (
-		results []collections.KeyValue[K, V]
-		pageRes *PageResponse
-		err     error
-	)
-
 	opt := new(CollectionsPaginateOptions[K])
 	for _, o := range opts {
 		o(opt)
@@ -85,9 +95,9 @@ func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 	}
 
 	if len(key) != 0 {
-		results, pageRes, err = collFilteredPaginateByKey(ctx, coll, prefix, key, reverse, limit, predicateFunc)
+		results, pageRes, err = collFilteredPaginateByKey(ctx, coll, prefix, key, reverse, limit, predicateFunc, transformFunc)
 	} else {
-		results, pageRes, err = collFilteredPaginateNoKey(ctx, coll, prefix, reverse, offset, limit, countTotal, predicateFunc)
+		results, pageRes, err = collFilteredPaginateNoKey(ctx, coll, prefix, reverse, offset, limit, countTotal, predicateFunc, transformFunc)
 	}
 	// invalid iter error is ignored to retain Paginate behavior
 	if errors.Is(err, collections.ErrInvalidIterator) {
@@ -102,7 +112,7 @@ func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
 
 // collFilteredPaginateNoKey applies the provided pagination on the collection when the starting key is not set.
 // If predicateFunc is nil no filtering is applied.
-func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
+func collFilteredPaginateNoKey[K, V any, C Collection[K, V], T any](
 	ctx context.Context,
 	coll C,
 	prefix []byte,
@@ -111,7 +121,8 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 	limit uint64,
 	countTotal bool,
 	predicateFunc func(K, V) (bool, error),
-) ([]collections.KeyValue[K, V], *PageResponse, error) {
+	transformFunc func(K, V) (T, error),
+) ([]T, *PageResponse, error) {
 	iterator, err := getCollIter[K, V](ctx, coll, prefix, nil, reverse)
 	if err != nil {
 		return nil, nil, err
@@ -125,7 +136,7 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 	var (
 		count   uint64
 		nextKey []byte
-		results []collections.KeyValue[K, V]
+		results []T
 	)
 
 	for ; iterator.Valid(); iterator.Next() {
@@ -138,18 +149,28 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 			}
 			// if no predicate function is specified then we just include the result
 			if predicateFunc == nil {
-				results = append(results, kv)
+				transformed, err := transformFunc(kv.Key, kv.Value)
+				if err != nil {
+					return nil, nil, err
+				}
+				results = append(results, transformed)
+				count++
+
 				// if predicate function is defined we check if the result matches the filtering criteria
 			} else {
 				include, err := predicateFunc(kv.Key, kv.Value)
 				if err != nil {
 					return nil, nil, err
 				}
 				if include {
-					results = append(results, kv)
+					transformed, err := transformFunc(kv.Key, kv.Value)
+					if err != nil {
+						return nil, nil, err
+					}
+					results = append(results, transformed)
+					count++
 				}
 			}
-			count++
 		// second case, we found all the objects specified within the limit
 		case count == limit:
 			key, err := iterator.Key()
@@ -172,12 +193,31 @@ func collFilteredPaginateNoKey[K, V any, C Collection[K, V]](
 		// but we need to count how many possible results exist in total.
 		// so we keep increasing the count until the iterator is fully consumed.
 		case count > limit:
-			count++
+			if predicateFunc == nil {
+				count++
+
+				// if predicate function is defined we check if the result matches the filtering criteria
+			} else {
+				kv, err := iterator.KeyValue()
+				if err != nil {
+					return nil, nil, err
+				}
+
+				include, err := predicateFunc(kv.Key, kv.Value)
+				if err != nil {
+					return nil, nil, err
+				}
+				if include {
+					count++
+				}
+			}
 		}
 	}
+
 	resp := &PageResponse{
 		NextKey: nextKey,
 	}
+
 	if countTotal {
 		resp.Total = count + offset
 	}
@@ -200,15 +240,16 @@ func advanceIter[I interface {
 
 // collFilteredPaginateByKey paginates a collection when a starting key
 // is provided in the PageRequest. Predicate is applied only if not nil.
-func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
+func collFilteredPaginateByKey[K, V any, C Collection[K, V], T any](
 	ctx context.Context,
 	coll C,
 	prefix []byte,
 	key []byte,
 	reverse bool,
 	limit uint64,
-	predicateFunc func(K, V) (bool, error),
-) ([]collections.KeyValue[K, V], *PageResponse, error) {
+	predicateFunc func(key K, value V) (bool, error),
+	transformFunc func(key K, value V) (transformed T, err error),
+) (results []T, pageRes *PageResponse, err error) {
 	iterator, err := getCollIter[K, V](ctx, coll, prefix, key, reverse)
 	if err != nil {
 		return nil, nil, err
@@ -218,7 +259,6 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
 	var (
 		count   uint64
 		nextKey []byte
-		results []collections.KeyValue[K, V]
 	)
 
 	for ; iterator.Valid(); iterator.Next() {
@@ -243,7 +283,11 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
 		}
 		// if no predicate is specified then we just append the result
 		if predicateFunc == nil {
-			results = append(results, kv)
+			transformed, err := transformFunc(kv.Key, kv.Value)
+			if err != nil {
+				return nil, nil, err
+			}
+			results = append(results, transformed)
 			// if predicate is applied we execute the predicate function
 			// and append only if predicateFunc yields true.
 		} else {
@@ -252,7 +296,11 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V]](
 				return nil, nil, err
 			}
 			if include {
-				results = append(results, kv)
+				transformed, err := transformFunc(kv.Key, kv.Value)
+				if err != nil {
+					return nil, nil, err
+				}
+				results = append(results, transformed)
 			}
 		}
 		count++

types/query/collections_pagination_test.go

@@ -100,14 +100,15 @@ func TestCollectionPagination(t *testing.T) {
 				Limit: 3,
 			},
 			expResp: &PageResponse{
-				NextKey: encodeKey(3),
+				NextKey: encodeKey(5),
 			},
 			filter: func(key, value uint64) (bool, error) {
 				return key%2 == 0, nil
 			},
 			expResults: []collections.KeyValue[uint64, uint64]{
 				{Key: 0, Value: 0},
 				{Key: 2, Value: 2},
+				{Key: 4, Value: 4},
 			},
 		},
 		"filtered with key": {
@@ -131,7 +132,15 @@ func TestCollectionPagination(t *testing.T) {
 	for name, tc := range tcs {
 		tc := tc
 		t.Run(name, func(t *testing.T) {
-			gotResults, gotResponse, err := CollectionFilteredPaginate(ctx, m, tc.req, tc.filter)
+			gotResults, gotResponse, err := CollectionFilteredPaginate(
+				ctx,
+				m,
+				tc.req,
+				tc.filter,
+				func(key, value uint64) (collections.KeyValue[uint64, uint64], error) {
+					return collections.KeyValue[uint64, uint64]{Key: key, Value: value}, nil
+				},
+			)
 			if tc.wantErr != nil {
 				require.ErrorIs(t, err, tc.wantErr)
 				return

x/auth/keeper/grpc_query.go

@@ -47,15 +47,14 @@ func (s queryServer) Accounts(ctx context.Context, req *types.QueryAccountsReque
 		return nil, status.Error(codes.InvalidArgument, "empty request")
 	}
 
-	var accounts []*codectypes.Any
-	_, pageRes, err := query.CollectionFilteredPaginate(ctx, s.k.Accounts, req.Pagination, func(_ sdk.AccAddress, value sdk.AccountI) (include bool, err error) {
-		accountAny, err := codectypes.NewAnyWithValue(value)
-		if err != nil {
-			return false, err
-		}
-		accounts = append(accounts, accountAny)
-		return false, nil // we don't include it since we're already appending the account
-	})
+	accounts, pageRes, err := query.CollectionPaginate(
+		ctx,
+		s.k.Accounts,
+		req.Pagination,
+		func(_ sdk.AccAddress, value sdk.AccountI) (*codectypes.Any, error) {
+			return codectypes.NewAnyWithValue(value)
+		},
+	)
 
 	return &types.QueryAccountsResponse{Accounts: accounts, Pagination: pageRes}, err
 }