diff --git a/types/query/collections_pagination.go b/types/query/collections_pagination.go index 2877df26396c..902d68b1dde3 100644 --- a/types/query/collections_pagination.go +++ b/types/query/collections_pagination.go @@ -265,53 +265,59 @@ func collFilteredPaginateByKey[K, V any, C Collection[K, V], T any]( defer iterator.Close() var ( - count uint64 - nextKey []byte + count uint64 + nextKey []byte + transformed T ) for ; iterator.Valid(); iterator.Next() { - // if we reached the specified limit - // then we get the next key, and we exit the iteration. - if count == limit { - concreteKey, err := iterator.Key() - if err != nil { - return nil, nil, err - } - - nextKey, err = encodeCollKey[K, V](coll, concreteKey) - if err != nil { - return nil, nil, err - } - break - } - kv, err := iterator.KeyValue() if err != nil { return nil, nil, err } + + include := false // if no predicate is specified then we just append the result if predicateFunc == nil { - transformed, err := transformFunc(kv.Key, kv.Value) + transformed, err = transformFunc(kv.Key, kv.Value) if err != nil { return nil, nil, err } - results = append(results, transformed) + include = true // if predicate is applied we execute the predicate function // and append only if predicateFunc yields true. } else { - include, err := predicateFunc(kv.Key, kv.Value) + include, err = predicateFunc(kv.Key, kv.Value) if err != nil { return nil, nil, err } if include { - transformed, err := transformFunc(kv.Key, kv.Value) + transformed, err = transformFunc(kv.Key, kv.Value) if err != nil { return nil, nil, err } - results = append(results, transformed) } } - count++ + + if include { + // if we reached the specified limit + // then we get the next key, and we exit the iteration. + if count == limit { + concreteKey, err := iterator.Key() + if err != nil { + return nil, nil, err + } + + nextKey, err = encodeCollKey[K, V](coll, concreteKey) + if err != nil { + return nil, nil, err + } + break + } + + results = append(results, transformed) + count++ + } } return results, &PageResponse{ diff --git a/types/query/collections_pagination_test.go b/types/query/collections_pagination_test.go index 25dc267a4481..6e4beb8c753c 100644 --- a/types/query/collections_pagination_test.go +++ b/types/query/collections_pagination_test.go @@ -127,7 +127,7 @@ func TestCollectionPagination(t *testing.T) { Limit: 3, }, expResp: &PageResponse{ - NextKey: encodeKey(5), + NextKey: encodeKey(8), }, filter: func(key, value uint64) (bool, error) { return key%2 == 0, nil @@ -135,6 +135,21 @@ func TestCollectionPagination(t *testing.T) { expResults: []collections.KeyValue[uint64, uint64]{ {Key: 2, Value: 2}, {Key: 4, Value: 4}, + {Key: 6, Value: 6}, + }, + }, + "filtered with key and empty next key in response": { + req: &PageRequest{ + Key: encodeKey(295), + }, + expResp: &PageResponse{ + NextKey: nil, + }, + filter: func(key, value uint64) (bool, error) { + return key%5 == 0, nil + }, + expResults: []collections.KeyValue[uint64, uint64]{ + {Key: 295, Value: 295}, }, }, }