Skip to content

Commit

Permalink
fix: collection filtered pagination
Browse files Browse the repository at this point in the history
  • Loading branch information
akhilkumarpilli committed Dec 19, 2024
1 parent a68f032 commit 4d651c5
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 24 deletions.
52 changes: 29 additions & 23 deletions types/query/collections_pagination.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down
17 changes: 16 additions & 1 deletion types/query/collections_pagination_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -127,14 +127,29 @@ 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
},
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},
},
},
}
Expand Down

0 comments on commit 4d651c5

Please sign in to comment.