feat(collections): implement pagination (#14468)
Co-authored-by: testinginprod <testinginprod@somewhere.idk>
This commit is contained in:
testinginprod
co-authored by
testinginprod
testinginprod
parent
512953cd68
commit
f771f20da4
@@ -0,0 +1,223 @@
|
||||
package query
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"cosmossdk.io/collections"
|
||||
)
|
||||
|
||||
// Collection defines the minimum required API of a collection
|
||||
// to work with pagination.
|
||||
type Collection[K, V any] interface {
|
||||
// IterateRaw allows to iterate over a raw set of byte keys.
|
||||
IterateRaw(ctx context.Context, start, end []byte, order collections.Order) (collections.Iterator[K, V], error)
|
||||
// KeyCodec exposes the KeyCodec of a collection, required to encode a collection key from and to bytes
|
||||
// for pagination request and response.
|
||||
KeyCodec() collections.KeyCodec[K]
|
||||
}
|
||||
|
||||
// CollectionPaginate follows the same behaviour as Paginate but works on a Collection.
|
||||
func CollectionPaginate[K, V any, C Collection[K, V]](
|
||||
ctx context.Context,
|
||||
coll C,
|
||||
pageReq *PageRequest,
|
||||
) ([]collections.KeyValue[K, V], *PageResponse, error) {
|
||||
return CollectionFilteredPaginate[K, V](ctx, coll, pageReq, nil)
|
||||
}
|
||||
|
||||
// CollectionFilteredPaginate works in the same way as FilteredPaginate but for collection types.
|
||||
// A nil predicateFunc means no filtering is applied and results are collected as is.
|
||||
func CollectionFilteredPaginate[K, V any, C Collection[K, V]](
|
||||
ctx context.Context,
|
||||
coll C,
|
||||
pageReq *PageRequest,
|
||||
predicateFunc func(key K, value V) (include bool),
|
||||
) ([]collections.KeyValue[K, V], *PageResponse, error) {
|
||||
if pageReq == nil {
|
||||
pageReq = &PageRequest{}
|
||||
}
|
||||
|
||||
offset := pageReq.Offset
|
||||
key := pageReq.Key
|
||||
limit := pageReq.Limit
|
||||
countTotal := pageReq.CountTotal
|
||||
reverse := pageReq.Reverse
|
||||
|
||||
if offset > 0 && key != nil {
|
||||
return nil, nil, fmt.Errorf("invalid request, either offset or key is expected, got both")
|
||||
}
|
||||
|
||||
if limit == 0 {
|
||||
limit = DefaultLimit
|
||||
countTotal = true
|
||||
}
|
||||
|
||||
if len(key) != 0 {
|
||||
return collFilteredPaginateByKey(ctx, coll, key, reverse, limit, predicateFunc)
|
||||
}
|
||||
|
||||
return collFilteredPaginateNoKey(ctx, coll, reverse, offset, limit, countTotal, predicateFunc)
|
||||
}
|
||||
|
||||
// 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]](
|
||||
ctx context.Context,
|
||||
coll C,
|
||||
reverse bool,
|
||||
offset uint64,
|
||||
limit uint64,
|
||||
countTotal bool,
|
||||
predicateFunc func(K, V) bool,
|
||||
) ([]collections.KeyValue[K, V], *PageResponse, error) {
|
||||
iterator, err := getCollIter[K, V](ctx, coll, nil, reverse)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
defer iterator.Close()
|
||||
// we advance the iter equal to the provided offset
|
||||
if !advanceIter(iterator, offset) {
|
||||
return nil, nil, collections.ErrInvalidIterator
|
||||
}
|
||||
|
||||
var (
|
||||
count uint64
|
||||
nextKey []byte
|
||||
results []collections.KeyValue[K, V]
|
||||
)
|
||||
|
||||
for ; iterator.Valid(); iterator.Next() {
|
||||
switch {
|
||||
// first case, we still haven't found all the results up to the limit
|
||||
case count < limit:
|
||||
kv, err := iterator.KeyValue()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
// if no predicate function is specified then we just include the result
|
||||
if predicateFunc == nil {
|
||||
results = append(results, kv)
|
||||
count++
|
||||
// if predicate function is defined we check if the result matches the filtering criteria
|
||||
} else if predicateFunc(kv.Key, kv.Value) {
|
||||
results = append(results, kv)
|
||||
count++
|
||||
}
|
||||
// second case, we found all the objects specified within the limit
|
||||
case count == limit:
|
||||
key, err := iterator.Key()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
nextKey, err = encodeCollKey[K, V](coll, key)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
// if count total was not specified, we return the next key only
|
||||
if !countTotal {
|
||||
return results, &PageResponse{
|
||||
NextKey: nextKey,
|
||||
}, nil
|
||||
}
|
||||
// otherwise we fallthrough the third case
|
||||
fallthrough
|
||||
// this is the case in which we found all the required results
|
||||
// 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++
|
||||
}
|
||||
}
|
||||
return results, &PageResponse{
|
||||
NextKey: nextKey,
|
||||
Total: count + offset,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func advanceIter[I interface {
|
||||
Next()
|
||||
Valid() bool
|
||||
}](iter I, offset uint64,
|
||||
) bool {
|
||||
for i := uint64(0); i < offset; i++ {
|
||||
if !iter.Valid() {
|
||||
return false
|
||||
}
|
||||
iter.Next()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// 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]](
|
||||
ctx context.Context,
|
||||
coll C,
|
||||
key []byte,
|
||||
reverse bool,
|
||||
limit uint64,
|
||||
predicateFunc func(K, V) bool,
|
||||
) ([]collections.KeyValue[K, V], *PageResponse, error) {
|
||||
iterator, err := getCollIter[K, V](ctx, coll, key, reverse)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
defer iterator.Close()
|
||||
|
||||
var (
|
||||
count uint64
|
||||
nextKey []byte
|
||||
results []collections.KeyValue[K, V]
|
||||
)
|
||||
|
||||
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
|
||||
}
|
||||
// if no predicate is specified then we just append the result
|
||||
if predicateFunc == nil {
|
||||
results = append(results, kv)
|
||||
count++
|
||||
// if predicate is applied we execute the predicate function
|
||||
// and append only if predicateFunc yields true.
|
||||
} else if predicateFunc(kv.Key, kv.Value) {
|
||||
results = append(results, kv)
|
||||
count++
|
||||
}
|
||||
}
|
||||
|
||||
return results, &PageResponse{
|
||||
NextKey: nextKey,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// todo maybe move to collections?
|
||||
func encodeCollKey[K, V any](coll Collection[K, V], key K) ([]byte, error) {
|
||||
buffer := make([]byte, coll.KeyCodec().Size(key))
|
||||
_, err := coll.KeyCodec().Encode(buffer, key)
|
||||
return buffer, err
|
||||
}
|
||||
|
||||
func getCollIter[K, V any](ctx context.Context, coll Collection[K, V], start []byte, reverse bool) (collections.Iterator[K, V], error) {
|
||||
if reverse {
|
||||
return coll.IterateRaw(ctx, nil, start, collections.OrderDescending)
|
||||
}
|
||||
return coll.IterateRaw(ctx, start, nil, collections.OrderAscending)
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
package query
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"cosmossdk.io/collections"
|
||||
"cosmossdk.io/core/store"
|
||||
db "github.com/cosmos/cosmos-db"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCollectionPagination(t *testing.T) {
|
||||
sk, ctx := deps()
|
||||
sb := collections.NewSchemaBuilder(sk)
|
||||
m := collections.NewMap(sb, collections.NewPrefix(0), "_", collections.Uint64Key, collections.Uint64Value)
|
||||
|
||||
for i := uint64(0); i < 300; i++ {
|
||||
require.NoError(t, m.Set(ctx, i, i))
|
||||
}
|
||||
|
||||
createResults := func(from, to uint64) []collections.KeyValue[uint64, uint64] {
|
||||
var res []collections.KeyValue[uint64, uint64]
|
||||
if from <= to {
|
||||
for i := from; i <= to; i++ {
|
||||
res = append(res, collections.KeyValue[uint64, uint64]{
|
||||
Key: i,
|
||||
Value: i,
|
||||
})
|
||||
}
|
||||
} else {
|
||||
for i := from; i >= to; i-- {
|
||||
res = append(res, collections.KeyValue[uint64, uint64]{
|
||||
Key: i,
|
||||
Value: i,
|
||||
})
|
||||
}
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
encodeKey := func(key uint64) []byte {
|
||||
b, err := encodeCollKey[uint64, uint64](m, key)
|
||||
require.NoError(t, err)
|
||||
return b
|
||||
}
|
||||
|
||||
type test struct {
|
||||
req *PageRequest
|
||||
expResp *PageResponse
|
||||
filter func(key uint64, value uint64) bool
|
||||
expResults []collections.KeyValue[uint64, uint64]
|
||||
wantErr error
|
||||
}
|
||||
|
||||
tcs := map[string]test{
|
||||
"nil pagination": {
|
||||
req: nil,
|
||||
expResp: &PageResponse{
|
||||
NextKey: encodeKey(100),
|
||||
Total: 300,
|
||||
},
|
||||
expResults: createResults(0, 99),
|
||||
},
|
||||
"with key and limit": {
|
||||
req: &PageRequest{
|
||||
Key: encodeKey(100),
|
||||
Limit: 149,
|
||||
},
|
||||
expResp: &PageResponse{
|
||||
NextKey: encodeKey(249),
|
||||
},
|
||||
expResults: createResults(100, 248),
|
||||
},
|
||||
"with reverse": {
|
||||
req: &PageRequest{
|
||||
Reverse: true,
|
||||
},
|
||||
expResp: &PageResponse{
|
||||
NextKey: encodeKey(199),
|
||||
Total: 300,
|
||||
},
|
||||
expResults: createResults(299, 200),
|
||||
},
|
||||
"with offset and count total": {
|
||||
req: &PageRequest{
|
||||
Offset: 50,
|
||||
Limit: 100,
|
||||
CountTotal: true,
|
||||
},
|
||||
expResp: &PageResponse{
|
||||
NextKey: encodeKey(150),
|
||||
Total: 300,
|
||||
},
|
||||
expResults: createResults(50, 149),
|
||||
},
|
||||
"filtered no key": {
|
||||
req: &PageRequest{
|
||||
Limit: 3,
|
||||
},
|
||||
expResp: &PageResponse{
|
||||
NextKey: encodeKey(5),
|
||||
},
|
||||
filter: func(key uint64, value uint64) bool {
|
||||
return key%2 == 0
|
||||
},
|
||||
expResults: []collections.KeyValue[uint64, uint64]{
|
||||
{Key: 0, Value: 0},
|
||||
{Key: 2, Value: 2},
|
||||
{Key: 4, Value: 4},
|
||||
},
|
||||
},
|
||||
"filtered with key": {
|
||||
req: &PageRequest{
|
||||
Key: encodeKey(2),
|
||||
Limit: 3,
|
||||
},
|
||||
expResp: &PageResponse{
|
||||
NextKey: encodeKey(7),
|
||||
},
|
||||
filter: func(key uint64, value uint64) bool {
|
||||
return key%2 == 0
|
||||
},
|
||||
expResults: []collections.KeyValue[uint64, uint64]{
|
||||
{Key: 2, Value: 2},
|
||||
{Key: 4, Value: 4},
|
||||
{Key: 6, Value: 6},
|
||||
},
|
||||
},
|
||||
"error offset > total items": {
|
||||
req: &PageRequest{Offset: 500},
|
||||
wantErr: collections.ErrInvalidIterator,
|
||||
},
|
||||
}
|
||||
|
||||
for name, tc := range tcs {
|
||||
tc := tc
|
||||
t.Run(name, func(t *testing.T) {
|
||||
gotResults, gotResponse, err := CollectionFilteredPaginate(ctx, m, tc.req, tc.filter)
|
||||
if tc.wantErr != nil {
|
||||
require.ErrorIs(t, err, tc.wantErr)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.expResults, gotResults)
|
||||
require.Equal(t, tc.expResp, gotResponse)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type testStore struct {
|
||||
db db.DB
|
||||
}
|
||||
|
||||
func (t testStore) OpenKVStore(ctx context.Context) store.KVStore {
|
||||
return t
|
||||
}
|
||||
|
||||
func (t testStore) Get(key []byte) []byte {
|
||||
res, err := t.db.Get(key)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
func (t testStore) Has(key []byte) bool {
|
||||
res, err := t.db.Has(key)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
func (t testStore) Set(key, value []byte) {
|
||||
err := t.db.Set(key, value)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
func (t testStore) Delete(key []byte) {
|
||||
err := t.db.Delete(key)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
func (t testStore) Iterator(start, end []byte) store.Iterator {
|
||||
res, err := t.db.Iterator(start, end)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
func (t testStore) ReverseIterator(start, end []byte) store.Iterator {
|
||||
res, err := t.db.ReverseIterator(start, end)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
var _ store.KVStore = testStore{}
|
||||
|
||||
func deps() (store.KVStoreService, context.Context) {
|
||||
kv := db.NewMemDB()
|
||||
return &testStore{kv}, context.Background()
|
||||
}
|
||||
Reference in New Issue
Block a user