feat(cdn) add limit flag to distribution list

Author: cgoetz-inovex

internal/cmd/beta/cdn/distribution/list/list.go

@@ -3,6 +3,7 @@ package list
 import (
 	"context"
 	"fmt"
+	"math"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -23,10 +24,13 @@ import (
 type inputModel struct {
 	*globalflags.GlobalFlagModel
 	SortBy string
+	Limit  *int32
 }
 
 const (
-	sortByFlag = "sort-by"
+	sortByFlag  = "sort-by"
+	limitFlag   = ""
+	maxPageSize = int32(100)
 )
 
 func NewCmd(params *params.CmdParams) *cobra.Command {
@@ -76,6 +80,7 @@ var sortByFlagOptions = []string{"id", "createdAt", "updatedAt", "originUrl", "s
 func configureFlags(cmd *cobra.Command) {
 	// same default as apiClient
 	cmd.Flags().Var(flags.EnumFlag(false, "createdAt", sortByFlagOptions...), sortByFlag, fmt.Sprintf("Sort entries by a specific field, one of %q", sortByFlagOptions))
+	cmd.Flags().Int64(limitFlag, 0, "Limit the output to the first n elements")
 }
 
 func parseInput(p *print.Printer, cmd *cobra.Command, _ []string) (*inputModel, error) {
@@ -84,6 +89,14 @@ func parseInput(p *print.Printer, cmd *cobra.Command, _ []string) (*inputModel,
 		return nil, &errors.ProjectIdError{}
 	}
 
+	limit := flags.FlagToInt32Pointer(p, cmd, limitFlag)
+	if limit != nil && *limit < 1 {
+		return nil, &errors.FlagValidationError{
+			Flag:    limitFlag,
+			Details: "must be greater than 0",
+		}
+	}
+
 	model := inputModel{
 		GlobalFlagModel: globalFlags,
 		SortBy:          flags.FlagWithDefaultToStringValue(p, cmd, sortByFlag),
@@ -93,10 +106,10 @@ func parseInput(p *print.Printer, cmd *cobra.Command, _ []string) (*inputModel,
 	return &model, nil
 }
 
-func buildRequest(ctx context.Context, model *inputModel, apiClient *cdn.APIClient, nextPageID cdn.ListDistributionsResponseGetNextPageIdentifierAttributeType) cdn.ApiListDistributionsRequest {
+func buildRequest(ctx context.Context, model *inputModel, apiClient *cdn.APIClient, nextPageID cdn.ListDistributionsResponseGetNextPageIdentifierAttributeType, pageLimit int32) cdn.ApiListDistributionsRequest {
 	req := apiClient.ListDistributions(ctx, model.GlobalFlagModel.ProjectId)
 	req = req.SortBy(model.SortBy)
-	req = req.PageSize(100)
+	req = req.PageSize(pageLimit)
 	if nextPageID != nil {
 		req = req.PageIdentifier(*nextPageID)
 	}
@@ -135,17 +148,24 @@ func outputResult(p *print.Printer, outputFormat string, distributions []cdn.Dis
 func fetchDistributions(ctx context.Context, model *inputModel, apiClient *cdn.APIClient) ([]cdn.Distribution, error) {
 	var nextPageID cdn.ListDistributionsResponseGetNextPageIdentifierAttributeType
 	var distributions []cdn.Distribution
+	received := int32(0)
+	limit := int32(math.MaxInt32)
+	if model.Limit != nil {
+		limit = utils.Min(limit, *model.Limit)
+	}
 	for {
-		request := buildRequest(ctx, model, apiClient, nextPageID)
+		want := utils.Min(maxPageSize, limit-received)
+		request := buildRequest(ctx, model, apiClient, nextPageID, want)
 		response, err := request.Execute()
 		if err != nil {
 			return nil, fmt.Errorf("list distributions: %w", err)
 		}
-		nextPageID = response.NextPageIdentifier
 		if response.Distributions != nil {
 			distributions = append(distributions, *response.Distributions...)
 		}
-		if nextPageID == nil {
+		nextPageID = response.NextPageIdentifier
+		received += want
+		if nextPageID == nil || received >= limit {
 			break
 		}
 	}

internal/cmd/beta/cdn/distribution/list/list_test.go

@@ -4,8 +4,10 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"slices"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -205,7 +207,7 @@ func TestBuildRequest(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.description, func(t *testing.T) {
-			req := buildRequest(testCtx, tt.inputModel, testClient, tt.nextPageID)
+			req := buildRequest(testCtx, tt.inputModel, testClient, tt.nextPageID, maxPageSize)
 			diff := cmp.Diff(req, tt.expected,
 				cmp.AllowUnexported(tt.expected),
 				cmpopts.EquateComparable(testCtx),
@@ -256,9 +258,21 @@ func fixtureDistribution(id string) cdn.Distribution {
 	}
 }
 
+func fixtureDistributions(count int) []cdn.Distribution {
+	distributions := make([]cdn.Distribution, count)
+	for i := 0; i < count; i++ {
+		id := fmt.Sprintf("dist-%d", i+1)
+		distributions[i] = cdn.Distribution{
+			Id: &id,
+		}
+	}
+	return distributions
+}
+
 func TestFetchDistributions(t *testing.T) {
 	tests := []struct {
 		description string
+		limit       int
 		responses   []testResponse
 		expected    []cdn.Distribution
 		fails       bool
@@ -323,6 +337,20 @@ func TestFetchDistributions(t *testing.T) {
 			},
 			fails: true,
 		},
+		{
+			description: "limit across 2 pages",
+			limit:       110,
+			responses: []testResponse{
+				fixtureTestResponse(
+					responseNextPageID(utils.Ptr(testNextPageID)),
+					responseDistributions(fixtureDistributions(100)...),
+				),
+				fixtureTestResponse(
+					responseDistributions(fixtureDistributions(10)...),
+				),
+			},
+			expected: slices.Concat(fixtureDistributions(100), fixtureDistributions(10)),
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.description, func(t *testing.T) {
@@ -350,7 +378,14 @@ func TestFetchDistributions(t *testing.T) {
 			if err != nil {
 				t.Fatalf("failed to create test client: %v", err)
 			}
-			got, err := fetchDistributions(testCtx, fixtureInputModel(), client)
+			var mods []func(m *inputModel)
+			if tt.limit > 0 {
+				mods = append(mods, func(m *inputModel) {
+					m.Limit = utils.Ptr(int32(tt.limit))
+				})
+			}
+			model := fixtureInputModel(mods...)
+			got, err := fetchDistributions(testCtx, model, client)
 			if err != nil {
 				if !tt.fails {
 					t.Fatalf("fetchDistributions() unexpected error: %v", err)