Example of Go iterators for keyset pagination.

readme.md

I haven't seen a non-trivial example of Go's new iterators, so here's one of keyset pagination (from a mocked database).

It results in a clean API for the consumer to range over the results. The parameter—in this case 3—is the desired batch size.

for userBatch, err := range batchesOfUsers(3) {
   ...
}

Output:

go run main.go
Users: [Myra Guadalupe Wes], Cursor: Mg==
Users: [Angel Chong Shelia], Cursor: NQ==
Users: [Hunter], Cursor: <nil>

main.go

package main

import (
	"encoding/base64"
	"fmt"
	"iter"
	"strconv"
	"time"
)

type user struct {
	ID   int
	Name string
}

type userBatch struct {
	Users  []user
	Cursor *string
}

func (b userBatch) String() string {
	var users []string
	for _, user := range b.Users {
		users = append(users, user.Name)
	}
	cursorValue := "<nil>"
	if b.Cursor != nil {
		cursorValue = *b.Cursor
	}
	return fmt.Sprintf("Users: %v, Cursor: %v", users, cursorValue)
}

var allUsers []user

func main() {
	setupFakeUsersDB()

	// Iterate over all users in batches of 3.
	for userBatch, err := range batchesOfUsers(3) {
		if err != nil {
			panic(err)
		}
		fmt.Printf("%+v\n", userBatch)
	}
}

// batchesOfUsers returns an iterator that yields user batches of the given size.
func batchesOfUsers(batchSize int) iter.Seq2[userBatch, error] {
	return func(yield func(userBatch, error) bool) {
		var cursor string
		for {
			batch, err := getUserBatchFromDB(batchSize, cursor)
			if err != nil {
				yield(userBatch{}, err)
				return
			}
			if len(batch.Users) == 0 || !yield(batch, nil) || batch.Cursor == nil {
				return
			}
			cursor = *batch.Cursor
		}
	}
}

// getUserBatchFromDB simulates a database query that returns a batch of users
// of the given size, starting from the given cursor.
func getUserBatchFromDB(limit int, cursor string) (userBatch, error) {
	time.Sleep(time.Second) // Simulate DB latency
	var users []user
	var id int64
	if cursor == "" {
		id = -1
	} else {
		decodedCursor, err := base64.StdEncoding.DecodeString(cursor)
		if err != nil {
			return userBatch{}, err
		}
		id, err = strconv.ParseInt(string(decodedCursor), 10, 64)
		if err != nil {
			return userBatch{}, err
		}
	}
	for _, user := range allUsers {
		if user.ID <= int(id) || len(users) >= limit {
			continue
		}
		users = append(users, user)
	}
	var lastID *int
	var nextCursor *string
	if len(users) == limit {
		val := users[len(users)-1].ID
		lastID = &val
		encodedCursor := base64.StdEncoding.EncodeToString([]byte(strconv.Itoa(*lastID)))
		nextCursor = &encodedCursor
	}
	return userBatch{Users: users, Cursor: nextCursor}, nil
}

func setupFakeUsersDB() {
	allUsers = []user{
		{ID: 0, Name: "Myra"},
		{ID: 1, Name: "Guadalupe"},
		{ID: 2, Name: "Wes"},
		{ID: 3, Name: "Angel"},
		{ID: 4, Name: "Chong"},
		{ID: 5, Name: "Shelia"},
		{ID: 6, Name: "Hunter"},
	}
}