Paginating fetcher against stubbed service

batch.go

package main

import (
	"io"
	"math/rand"
	"time"

	"github.com/sirupsen/logrus"
)

const batchSize = 10

type extSvcStubFunc func(string, string, int, int) ([]interface{}, error)

func main() {
	var (
		start = time.Now()
		end   = start.Add(time.Hour)
	)
	scenarios := []struct {
		desc       string
		batchSizes []int
	}{
		{
			desc:       "Two full batches and then an incomplete batch",
			batchSizes: []int{batchSize, batchSize, batchSize - 1},
		},
		{
			desc:       "Two full batches exactly",
			batchSizes: []int{batchSize, batchSize},
		},
		{
			desc:       "No data at all",
			batchSizes: nil,
		},
	}
	for _, scenario := range scenarios {
		logrus.WithField("scenario", scenario.desc).Info("starting scenario")
		f := mkExternalService(scenario.batchSizes)
		data, err := Fetch(start, end, f)
		if err != nil {
			logrus.WithError(err).Warn("fetch failed")
			continue
		}
		logrus.WithFields(logrus.Fields{
			"data_length": len(data),
			"data":        data,
		}).Info("fetch succeeded")
	}
}

func Fetch(start, end time.Time, extSvc extSvcStubFunc) ([]interface{}, error) {
	var (
		data   []interface{}
		offset = 0
	)
	logrus.Info("starting fetch")
	for {
		batch, err := fetchBatch(start, end, offset, batchSize, extSvc)
		if err != nil {
			if err == io.EOF {
				logrus.Info("empty batch received: exiting fetch loop")
				return data, nil
			}
			return nil, err
		}
		data = append(data, batch...)
		if len(batch) < batchSize {
			logrus.WithFields(logrus.Fields{
				"max_batch_size":      batchSize,
				"received_batch_size": len(batch),
			}).Info("received non-full batch: exiting fetch loop")
			return data, nil
		}
		offset += len(data)
	}
}

func fetchBatch(start, end time.Time, offset, limit int, extSvc extSvcStubFunc) ([]interface{}, error) {
	var (
		startStr = start.Format(time.RFC3339)
		endStr   = end.Format(time.RFC3339)
		startRow = offset
		endRow   = offset + limit - 1
	)
	data, err := extSvc(startStr, endStr, startRow, endRow)
	if err != nil {
		return nil, err
	}
	if len(data) == 0 {
		return nil, io.EOF
	}
	logrus.WithFields(logrus.Fields{
		"op":         "fetchBatch",
		"startTime":  startStr,
		"endTime":    endStr,
		"startRow":   startRow,
		"endRow":     endRow,
		"batch_size": len(data),
	}).Info("received batch")
	return data, nil
}

// Generate an external service stub that returns batches with sizes described by batchSizes.
func mkExternalService(batchSizes []int) extSvcStubFunc {
	batchIndex := 0
	return func(startTime, endTime string, startRow, endRow int) ([]interface{}, error) {
		if batchIndex >= len(batchSizes) {
			// There are no more batch sizes defined, so return empty batch.
			return nil, nil
		}
		batch := make([]interface{}, batchSizes[batchIndex])
		batchIndex++
		for i := range batch {
			batch[i] = rand.Intn(1 << 8)
		}
		return batch, nil
	}
}