Support pagination for tagserver

build-index/tagserver/server.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"path"
 	"strconv"
 	"strings"
@@ -44,6 +45,11 @@ import (
 	"github.com/uber-go/tally"
 )
 
+const (
+	limitQ  string = "limit"
+	offsetQ string = "offset"
+)
+
 // Server provides tag operations for the build-index.
 type Server struct {
 	config            Config
@@ -63,6 +69,16 @@ type Server struct {
 	depResolver tagtype.DependencyResolver
 }
 
+// List Response with pagination.
+type ListResponse struct {
+	Links struct {
+		Next string `json:"next"`
+		Self string `json:"self"`
+	}
+	Size   int      `json:"size"`
+	Result []string `json:"result"`
+}
+
 // New creates a new Server.
 func New(
 	config Config,
@@ -244,11 +260,23 @@ func (s *Server) listHandler(w http.ResponseWriter, r *http.Request) error {
 	if err != nil {
 		return handler.Errorf("backend manager: %s", err)
 	}
-	result, err := client.List(prefix)
+
+	opts, err := buildPaginationOptions(r.URL)
+	if err != nil {
+		return err
+	}
+
+	result, err := client.List(prefix, opts...)
+	if err != nil {
+		return handler.Errorf("error listing from backend: %s", err)
+	}
+
+	resp, err := buildPaginationResponse(r.URL, result.ContinuationToken,
+		result.Names)
 	if err != nil {
 		return err
 	}
-	if err := json.NewEncoder(w).Encode(&result.Names); err != nil {
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
 		return handler.Errorf("json encode: %s", err)
 	}
 	return nil
@@ -265,10 +293,17 @@ func (s *Server) listRepositoryHandler(w http.ResponseWriter, r *http.Request) e
 	if err != nil {
 		return handler.Errorf("backend manager: %s", err)
 	}
-	result, err := client.List(path.Join(repo, "_manifests/tags"))
+
+	opts, err := buildPaginationOptions(r.URL)
 	if err != nil {
 		return err
 	}
+
+	result, err := client.List(path.Join(repo, "_manifests/tags"), opts...)
+	if err != nil {
+		return handler.Errorf("error listing from backend: %s", err)
+	}
+
 	var tags []string
 	for _, name := range result.Names {
 		// Strip repo prefix.
@@ -279,7 +314,12 @@ func (s *Server) listRepositoryHandler(w http.ResponseWriter, r *http.Request) e
 		}
 		tags = append(tags, parts[1])
 	}
-	if err := json.NewEncoder(w).Encode(&tags); err != nil {
+
+	resp, err := buildPaginationResponse(r.URL, result.ContinuationToken, tags)
+	if err != nil {
+		return err
+	}
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
 		return handler.Errorf("json encode: %s", err)
 	}
 	return nil
@@ -405,3 +445,69 @@ func (s *Server) replicateTag(tag string, d core.Digest, deps core.DigestList) e
 	}
 	return nil
 }
+
+func buildPaginationOptions(u *url.URL) ([]backend.ListOption, error) {
+	var opts []backend.ListOption
+	q := u.Query()
+	for k, v := range q {
+		if len(v) != 1 {
+			return nil, handler.Errorf(
+				"invalid query %s:%s", k, v).Status(http.StatusBadRequest)
+		}
+		switch k {
+		case limitQ:
+			limitCount, err := strconv.Atoi(v[0])
+			if err != nil {
+				return nil, handler.Errorf(
+					"invalid limit %s: %s", v, err).Status(http.StatusBadRequest)
+			}
+			if limitCount == 0 {
+				return nil, handler.Errorf(
+					"invalid limit %d", limitCount).Status(http.StatusBadRequest)
+			}
+			opts = append(opts, backend.ListWithMaxKeys(limitCount))
+		case offsetQ:
+			opts = append(opts, backend.ListWithContinuationToken(v[0]))
+		default:
+			return nil, handler.Errorf(
+				"invalid query %s", k).Status(http.StatusBadRequest)
+		}
+	}
+	if len(opts) > 0 {
+		// Enable pagination if either or both of the query param exists.
+		opts = append(opts, backend.ListWithPagination())
+	}
+
+	return opts, nil
+}
+
+func buildPaginationResponse(u *url.URL, continuationToken string,
+	result []string) (interface{}, error) {
+
+	if continuationToken == "" {
+		return result, nil
+	}
+
+	// Deep copy url.
+	nextUrl, err := url.Parse(u.String())
+	if err != nil {
+		return nil, handler.Errorf(
+			"invalid url string: %s", err).Status(http.StatusBadRequest)
+	}
+	v := url.Values{}
+	if limit := u.Query().Get(limitQ); limit != "" {
+		v.Add(limitQ, limit)
+	}
+	// ContinuationToken cannot be empty here.
+	v.Add(offsetQ, continuationToken)
+	nextUrl.RawQuery = v.Encode()
+
+	resp := ListResponse{
+		Size:   len(result),
+		Result: result,
+	}
+	resp.Links.Next = nextUrl.String()
+	resp.Links.Self = u.String()
+
+	return &resp, nil
+}

lib/backend/gcsbackend/client.go

@@ -175,8 +175,6 @@ func (c *Client) List(prefix string, opts ...backend.ListOption) (*backend.ListR
 		opt(options)
 	}
 
-	var names []string
-
 	absPrefix := path.Join(c.pather.BasePath(), prefix)
 	pageIterator := c.gcs.GetObjectIterator(absPrefix)
 
@@ -186,19 +184,17 @@ func (c *Client) List(prefix string, opts ...backend.ListOption) (*backend.ListR
 		maxKeys = options.MaxKeys
 		paginationToken = options.ContinuationToken
 	}
-	objectsPage := iterator.NewPager(pageIterator, maxKeys, paginationToken)
-	continuationToken, err := objectsPage.NextPage(&names)
+
+	pager := iterator.NewPager(pageIterator, maxKeys, paginationToken)
+	result, err := c.gcs.NextPage(pager)
 	if err != nil {
 		return nil, err
 	}
 	if !options.Paginated {
-		continuationToken = ""
+		result.ContinuationToken = ""
 	}
 
-	return &backend.ListResult{
-		Names:             names,
-		ContinuationToken: continuationToken,
-	}, nil
+	return result, nil
 }
 
 // isObjectNotFound is helper function for identify non-existing object error.
@@ -264,3 +260,22 @@ func (g *GCSImpl) GetObjectIterator(prefix string) iterator.Pageable {
 	query.Prefix = prefix
 	return g.bucket.Objects(g.ctx, &query)
 }
+
+func (g *GCSImpl) NextPage(pager *iterator.Pager) (*backend.ListResult,
+	error) {
+
+	var objectAttrs []*storage.ObjectAttrs
+	continuationToken, err := pager.NextPage(&objectAttrs)
+	if err != nil {
+		return nil, err
+	}
+
+	names := make([]string, len(objectAttrs))
+	for idx, objectAttr := range objectAttrs {
+		names[idx] = objectAttr.Name
+	}
+	return &backend.ListResult{
+		Names:             names,
+		ContinuationToken: continuationToken,
+	}, nil
+}

lib/backend/gcsbackend/client_test.go

@@ -193,20 +193,32 @@ func TestClientList(t *testing.T) {
 	client := mocks.new()
 
 	contToken := ""
+	mocks.gcs.EXPECT().GetObjectIterator(
+		"/root/test",
+	).AnyTimes().Return(Alphabets(t, maxIterate))
 	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))
 		}
+
+		continuationToken := ""
+		if (i + count) < maxIterate {
+			strconv.Itoa(i + count)
+		}
+		result := &backend.ListResult{
+			Names:             expected,
+			ContinuationToken: continuationToken,
+		}
+		mocks.gcs.EXPECT().NextPage(
+			gomock.Any(),
+		).Return(result, nil)
+
+		result, err := client.List("test", backend.ListWithPagination(),
+			backend.ListWithMaxKeys(count),
+			backend.ListWithContinuationToken(contToken))
+		require.NoError(err)
 		require.Equal(expected, result.Names)
 		contToken = result.ContinuationToken
 		i += count

lib/backend/gcsbackend/gcs.go

@@ -17,6 +17,7 @@ import (
 	"io"
 
 	"cloud.google.com/go/storage"
+	"github.com/uber/kraken/lib/backend"
 	"google.golang.org/api/iterator"
 )
 
@@ -26,4 +27,5 @@ type GCS interface {
 	Download(objectName string, w io.Writer) (int64, error)
 	Upload(objectName string, r io.Reader) (int64, error)
 	GetObjectIterator(prefix string) iterator.Pageable
+	NextPage(pager *iterator.Pager) (*backend.ListResult, error)
 }