Skip to content

Commit

Permalink
docstore/awsdynamodb: Fix Query for cases where the underlying query …
Browse files Browse the repository at this point in the history
…to AWS returns an empty set, but there's still more data
  • Loading branch information
vangent committed Apr 27, 2024
1 parent 3f641fd commit cb56d76
Showing 1 changed file with 18 additions and 12 deletions.
30 changes: 18 additions & 12 deletions docstore/awsdynamodb/query.go
Original file line number Diff line number Diff line change
Expand Up @@ -471,30 +471,36 @@ type documentIterator struct {

func (it *documentIterator) Next(ctx context.Context, doc driver.Document) error {
// Skip the first 'n' documents where 'n' is the offset.
if it.offset > 0 && it.count < it.offset {
it.curr++
it.count++
return it.Next(ctx, doc)
for it.count < it.offset {
if err := it.next(ctx, doc, false); err != nil {
return err
}
}
return it.next(ctx, doc, true)
}

func (it *documentIterator) next(ctx context.Context, doc driver.Document, decode bool) error {
// Only start counting towards the limit after the offset has been reached.
if it.limit > 0 && it.count >= it.offset+it.limit || it.curr >= len(it.items) && it.last == nil {
if it.limit > 0 && it.count >= it.offset+it.limit {
return io.EOF
}
if it.curr >= len(it.items) {
// it.items can be empty after a call to it.qr.run, but unless it.last is nil there may be more items.
for it.curr >= len(it.items) {
// Make a new query request at the end of this page.
if it.last == nil {
return io.EOF
}
var err error
it.items, it.last, it.asFunc, err = it.qr.run(ctx, it.last)
if err != nil {
return err
}
it.curr = 0
}
// If there are no more items, return EOF.
if len(it.items) == 0 {
return io.EOF
}
if err := decodeDoc(&dyn.AttributeValue{M: it.items[it.curr]}, doc); err != nil {
return err
if decode {
if err := decodeDoc(&dyn.AttributeValue{M: it.items[it.curr]}, doc); err != nil {
return err
}
}
it.curr++
it.count++
Expand Down

0 comments on commit cb56d76

Please sign in to comment.