depager/depager.go

/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

package depager

import (
	"encoding/json"
	"fmt"
	"net/url"
)

/*
The `Page` interface must wrap server responses. This
allows pagers to calculate page sizes and iterate over
response aggregates.

The underlying implementation must be a pointer to a
struct containing the desired response fields, all tagged
appropriately. Any fields corresponding to
platform-specific error responses should also be included.
*/
type Page[T any] interface {
	// Count must return the total number of items to be paged
	Count() uint64

	// Elems must return the items from the current page
	Elems() []T
}

// Exposes the part of the client that depager understands.
type Client interface {
	// Get must return the response body and/or an error
	Get(uri url.URL) ([]byte, error)
}

// Maps depager's parameters into the API's parameters.
type PagedURI interface {
	// PageURI generates a URI to request the indicated items
	PageURI(limit, offset uint64) url.URL
}

type Pager[T any] interface {
	// Iter is intended to be used in a for-range loop
	Iter() <-chan T

	// LastErr must return the first error encountered, if any
	LastErr() error
}

func NewPager[T any](
	u PagedURI,
	c Client,
	pageSize uint64,
	newPage func() Page[T],
) Pager[T] {
	return &pager[T]{
		client:  c,
		uri:     u,
		m:       0,
		n:       pageSize,
		p:       4,
		newPage: newPage,
	}
}

/*
Retrieve n items in the range [m*n, m*n + n - 1], inclusive.
Keep p pages buffered.
*/
type pager[T any] struct {
	client  Client
	uri     PagedURI
	m       uint64
	n       uint64
	err     error
	p       int
	newPage func() Page[T]
}

func (p *pager[T]) iteratePages() <-chan Page[T] {
	ch := make(chan Page[T], p.p)
	go func() {
		defer close(ch)
		for {
			uri := p.uri.PageURI(p.n, p.m*p.n)

			body, err := p.client.Get(uri)
			if err != nil {
				// assume the client knows better than we do
				p.err = err
				return
			}

			page := p.newPage()

			// TODO: eventually allow other codecs
			err = json.Unmarshal(body, page)
			if err != nil {
				p.err = fmt.Errorf("pager: iterate pages: unmarshal response: %v", err)
				return
			}
			ch <- page

			if (p.m*p.n + p.n) >= page.Count() {
				return
			}
			p.m++
		}
	}()
	return ch
}

func (p *pager[T]) Iter() <-chan T {
	ch := make(chan T, p.n)
	go func() {
		defer close(ch)
		for page := range p.iteratePages() {
			for _, i := range page.Elems() {
				ch <- i
			}
			if p.err != nil {
				p.err = fmt.Errorf("pager: iterate items: %s", p.err)
				return
			}
		}
	}()
	return ch
}

func (p *pager[T]) LastErr() error {
	return p.err
}