garage-k2v-go/pager.go

package k2v

import (
	"context"
	"errors"
)

var StopScroll = errors.New("scroll canceled")

// ReadIndexResponseHandler is invoked for each batch of index read results.
//
// If an error is returned, scrolling is halted and the error is propagated.
// The sentinel value StopScroll can be returned to end iteration early without propagating an error.
type ReadIndexResponseHandler func(resp *ReadIndexResponse) error

type BatchSearchResultHandler func(result *BatchSearchResult) error

type IndexScroller interface {
	ReadIndex(ctx context.Context, b Bucket, q ReadIndexQuery) (*ReadIndexResponse, error)
}

var _ IndexScroller = &Client{}

type BatchSearchScroller interface {
	ReadBatch(ctx context.Context, b Bucket, q []BatchSearch) ([]*BatchSearchResult, error)
}

var _ BatchSearchScroller = &Client{}

// ScrollIndex calls the ReadIndex API serially, invoking the provided function for each response (batch) until there are no more results.
func ScrollIndex(ctx context.Context, client IndexScroller, b Bucket, q ReadIndexQuery, fn ReadIndexResponseHandler) error {
	for {
		resp, err := client.ReadIndex(ctx, b, q)
		if err != nil {
			return err
		}
		if err := fn(resp); err != nil {
			if errors.Is(err, StopScroll) {
				return nil
			}
			return err
		}
		if !resp.More || resp.NextStart == nil {
			break
		}
		q.Start = *resp.NextStart
	}
	return nil
}

func ScrollBatchSearch(ctx context.Context, client BatchSearchScroller, b Bucket, q []BatchSearch, fn BatchSearchResultHandler) error {
	for {
		results, err := client.ReadBatch(ctx, b, q)
		if err != nil {
			return err
		}
		var nextQ []BatchSearch
		for i := range results {
			if results[i].More && results[i].NextStart != nil {
				batch := q[i]
				batch.Start = *results[i].NextStart
				nextQ = append(nextQ, batch)
			}
			if err := fn(results[i]); err != nil {
				if errors.Is(err, StopScroll) {
					return nil
				}
				return err
			}
		}
		if len(nextQ) == 0 {
			break
		}
		q = nextQ
	}
	return nil
}