Add pagination support for gcs backend

lib/backend/gcsbackend/client.go

@@ -175,23 +175,29 @@ func (c *Client) List(prefix string, opts ...backend.ListOption) (*backend.ListR
 		opt(options)
 	}
 
-	if options.Paginated {
-		return nil, errors.New("pagination not supported")
-	}
-
 	var names []string
 
 	absPrefix := path.Join(c.pather.BasePath(), prefix)
 	pageIterator := c.gcs.GetObjectIterator(absPrefix)
-	// Start at beginning.
-	objectsPage := iterator.NewPager(pageIterator, c.config.ListMaxKeys, "")
-	_, err := objectsPage.NextPage(&names)
+
+	maxKeys := c.config.ListMaxKeys
+	paginationToken := ""
+	if options.Paginated {
+		maxKeys = options.MaxKeys
+		paginationToken = options.ContinuationToken
+	}
+	objectsPage := iterator.NewPager(pageIterator, maxKeys, paginationToken)
+	continuationToken, err := objectsPage.NextPage(&names)
 	if err != nil {
 		return nil, err
 	}
+	if !options.Paginated {
+		continuationToken = ""
+	}
 
 	return &backend.ListResult{
-		Names: names,
+		Names:             names,
+		ContinuationToken: continuationToken,
 	}, nil
 }
 

lib/backend/gcsbackend/client_test.go

@@ -16,11 +16,13 @@ package gcsbackend
 import (
 	"bytes"
 	"fmt"
+	"math/rand"
 	"strconv"
 	"strings"
 	"testing"
 
 	"github.com/uber/kraken/core"
+	"github.com/uber/kraken/lib/backend"
 	"github.com/uber/kraken/mocks/lib/backend/gcsbackend"
 	"github.com/uber/kraken/utils/mockutil"
 	"github.com/uber/kraken/utils/randutil"
@@ -142,46 +144,72 @@ func TestClientUpload(t *testing.T) {
 	require.NoError(client.Upload(core.NamespaceFixture(), "test", dataReader))
 }
 
-func Alphabets() *AlphaIterator {
-	it := &AlphaIterator{}
+func Alphabets(t *testing.T, maxIterate int) *AlphaIterator {
+	it := &AlphaIterator{assert: require.New(t), maxIterate: maxIterate}
 	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
 		it.next,
 		func() int { return len(it.elems) },
 		func() interface{} { e := it.elems; it.elems = nil; return e })
 	return it
 }
 
+// Iterates from 0-maxIterate
 type AlphaIterator struct {
-	pageInfo *iterator.PageInfo
-	nextFunc func() error
-	elems    []string
+	assert     *require.Assertions
+	pageInfo   *iterator.PageInfo
+	nextFunc   func() error
+	elems      []string
+	maxIterate int
 }
 
 func (it *AlphaIterator) PageInfo() *iterator.PageInfo {
 	return it.pageInfo
 }
 
 func (it *AlphaIterator) next(pageSize int, pageToken string) (string, error) {
-	// A simple implementation that ignores the pageToken
-	for i := 0; i < pageSize; i++ {
+	i := 0
+	if pageToken != "" {
+		var err error
+		i, err = strconv.Atoi(pageToken)
+		it.assert.NoError(err)
+	}
+	endCount := i + pageSize
+	for ; i < endCount && i < it.maxIterate; i++ {
 		it.elems = append(it.elems, "test/"+strconv.Itoa(i))
 	}
-	return "", nil
+	if i == it.maxIterate {
+		return "", nil
+	}
+	return strconv.Itoa(i), nil
 }
 
 func TestClientList(t *testing.T) {
 	require := require.New(t)
+	maxIterate := 100
 
 	mocks, cleanup := newClientMocks(t)
 	defer cleanup()
 
 	client := mocks.new()
 
-	mocks.gcs.EXPECT().GetObjectIterator(
-		"/root/test",
-	).Return(Alphabets())
-
-	result, err := client.List("test")
-	require.NoError(err)
-	require.Equal([]string{"test/0", "test/1", "test/2", "test/3", "test/4"}, result.Names)
+	contToken := ""
+	for i := 0; i < maxIterate; {
+		mocks.gcs.EXPECT().GetObjectIterator(
+			"/root/test",
+		).Return(Alphabets(t, maxIterate))
+
+		count := (rand.Int() % 10) + 1
+		result, err := client.List("test", backend.ListWithPagination(),
+			backend.ListWithMaxKeys(count),
+			backend.ListWithContinuationToken(contToken))
+		require.NoError(err)
+		var expected []string
+		for j := i; j < (i+count) && j < maxIterate; j++ {
+			expected = append(expected, "test/"+strconv.Itoa(j))
+		}
+		require.Equal(expected, result.Names)
+		contToken = result.ContinuationToken
+		i += count
+	}
+	require.Equal(contToken, "")
 }

lib/backend/options.go

@@ -17,15 +17,15 @@ package backend
 // when listing names. It is used to enable pagination in list requests.
 type ListOptions struct {
 	Paginated bool
-	MaxKeys int64
+	MaxKeys int
 	ContinuationToken string
 }
 
 // DefaultListOptions defines the defaults for list operations.
 func DefaultListOptions() *ListOptions {
 	return &ListOptions{
 		Paginated: false,
-		MaxKeys: int64(DefaultListMaxKeys),
+		MaxKeys: DefaultListMaxKeys,
 		ContinuationToken: "",
 	}
 }
@@ -42,7 +42,7 @@ func ListWithPagination() ListOption {
 
 // ListWithMaxKeys configures the list command to return a max
 // number of keys if pagination is enabled.
-func ListWithMaxKeys(max int64) ListOption {
+func ListWithMaxKeys(max int) ListOption {
 	return func(opt *ListOptions) {
 		opt.MaxKeys = max
 	}

lib/backend/s3backend/client.go

@@ -244,7 +244,7 @@ func (c *Client) List(prefix string, opts ...backend.ListOption) (*backend.ListR
 	if options.Paginated {
 		listInput := &s3.ListObjectsV2Input{
 			Bucket:  aws.String(c.config.Bucket),
-			MaxKeys: aws.Int64(options.MaxKeys),
+			MaxKeys: aws.Int64(int64(options.MaxKeys)),
 			Prefix:  aws.String(path.Join(c.pather.BasePath(), prefix)[1:]),
 		}