Skip to content

Commit fb5587a

Browse files
committed
client: add pagination with --all flag (PRINFRA-125)
1 parent 6df9ce8 commit fb5587a

16 files changed

Lines changed: 717 additions & 14 deletions

File tree

cmd/heygen/builder.go

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@ package main
22

33
import (
44
"encoding/json"
5+
"errors"
56
"fmt"
67
"io"
78
"os"
89
"strconv"
910
"strings"
1011

12+
"github.com/heygen-com/heygen-cli/internal/client"
1113
"github.com/heygen-com/heygen-cli/internal/command"
1214
clierrors "github.com/heygen-com/heygen-cli/internal/errors"
1315
"github.com/spf13/cobra"
@@ -41,6 +43,33 @@ func buildCobraCommand(spec *command.Spec, ctx *cmdContext) *cobra.Command {
4143
return err
4244
}
4345

46+
if spec.Paginated {
47+
allPages, _ := cmd.Flags().GetBool("all")
48+
cursorFlagName := cursorFlagForSpec(spec)
49+
cursorSet := cursorFlagName != "" && cmd.Flags().Changed(cursorFlagName)
50+
51+
if allPages && cursorSet {
52+
return clierrors.NewUsage(fmt.Sprintf("--all and --%s are mutually exclusive", cursorFlagName))
53+
}
54+
55+
if allPages {
56+
result, err := ctx.client.ExecuteAll(spec, inv)
57+
if err != nil {
58+
var truncErr *client.ErrPaginationTruncated
59+
if errors.As(err, &truncErr) {
60+
fmt.Fprintf(cmd.ErrOrStderr(), "Warning: %s\n", truncErr.Error())
61+
if fmtErr := ctx.formatter.Data(truncErr.Data); fmtErr != nil {
62+
return fmtErr
63+
}
64+
return clierrors.New(truncErr.Error())
65+
}
66+
return err
67+
}
68+
69+
return ctx.formatter.Data(result)
70+
}
71+
}
72+
4473
result, err := ctx.client.Execute(spec, inv)
4574
if err != nil {
4675
return err
@@ -55,6 +84,10 @@ func buildCobraCommand(spec *command.Spec, ctx *cmdContext) *cobra.Command {
5584
registerFlag(cmd, flag)
5685
}
5786

87+
if spec.Paginated {
88+
cmd.Flags().Bool("all", false, "Fetch all pages (returns flat JSON array instead of API envelope)")
89+
}
90+
5891
// Add -d/--data for commands with JSON request bodies
5992
if spec.BodyEncoding == "json" {
6093
cmd.Flags().StringVarP(&rawData, "data", "d", "",
@@ -86,6 +119,15 @@ func buildUseLine(spec *command.Spec) string {
86119
return strings.Join(parts, " ")
87120
}
88121

122+
func cursorFlagForSpec(spec *command.Spec) string {
123+
for _, flag := range spec.Flags {
124+
if flag.JSONName == spec.TokenParam {
125+
return flag.Name
126+
}
127+
}
128+
return ""
129+
}
130+
89131
// registerFlag adds a typed flag to the Cobra command based on the FlagSpec.
90132
func registerFlag(cmd *cobra.Command, flag command.FlagSpec) {
91133
helpText := flag.Help

cmd/heygen/builder_test.go

Lines changed: 193 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"bytes"
55
"encoding/json"
66
"errors"
7+
"fmt"
78
"io"
89
"net/http"
910
"os"
@@ -17,13 +18,14 @@ import (
1718
// videoListSpec mirrors the hand-written video list command as a Spec,
1819
// proving the generic builder produces identical behavior.
1920
var videoListSpec = &command.Spec{
20-
Group: "video",
21-
Name: "list",
22-
Summary: "List videos",
23-
Endpoint: "/v3/videos",
24-
Method: "GET",
25-
// TokenField/DataField for pagination (used by paginator, not tested here)
21+
Group: "video",
22+
Name: "list",
23+
Summary: "List videos",
24+
Endpoint: "/v3/videos",
25+
Method: "GET",
26+
Paginated: true,
2627
TokenField: "next_token",
28+
TokenParam: "token",
2729
DataField: "data",
2830
Flags: []command.FlagSpec{
2931
{Name: "limit", Type: "int", Source: "query", JSONName: "limit"},
@@ -92,6 +94,174 @@ func TestGenBuilder_VideoList_Flags(t *testing.T) {
9294
}
9395
}
9496

97+
func TestGenBuilder_VideoList_AllPages(t *testing.T) {
98+
var calls int
99+
srv := setupTestServer(t, map[string]testHandler{
100+
"GET /v3/videos": {
101+
StatusCode: 200,
102+
ValidateRequest: func(t *testing.T, r *http.Request) {
103+
t.Helper()
104+
calls++
105+
switch calls {
106+
case 1:
107+
if got := r.URL.Query().Get("token"); got != "" {
108+
t.Fatalf("first page token = %q, want empty", got)
109+
}
110+
case 2:
111+
if got := r.URL.Query().Get("token"); got != "cursor_2" {
112+
t.Fatalf("second page token = %q, want %q", got, "cursor_2")
113+
}
114+
default:
115+
t.Fatalf("unexpected request count %d", calls)
116+
}
117+
},
118+
Body: `{"data":[{"id":"v1"},{"id":"v2"}],"next_token":"cursor_2"}`,
119+
},
120+
})
121+
defer srv.Close()
122+
123+
originalHandler := srv.Config.Handler
124+
srv.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
125+
if r.URL.Query().Get("token") == "cursor_2" {
126+
w.WriteHeader(http.StatusOK)
127+
_, _ = w.Write([]byte(`{"data":[{"id":"v3"}],"next_token":null}`))
128+
return
129+
}
130+
originalHandler.ServeHTTP(w, r)
131+
})
132+
133+
res := runGenCommand(t, srv.URL, "test-key", videoListSpec, "list", "--all")
134+
135+
if res.ExitCode != 0 {
136+
t.Fatalf("ExitCode = %d, want 0\nstderr: %s", res.ExitCode, res.Stderr)
137+
}
138+
var parsed []map[string]any
139+
if err := json.Unmarshal([]byte(res.Stdout), &parsed); err != nil {
140+
t.Fatalf("stdout is not valid JSON array: %v\nstdout: %s", err, res.Stdout)
141+
}
142+
if len(parsed) != 3 {
143+
t.Fatalf("len(parsed) = %d, want 3", len(parsed))
144+
}
145+
}
146+
147+
func TestGenBuilder_VideoList_AllPages_SinglePage(t *testing.T) {
148+
srv := setupTestServer(t, map[string]testHandler{
149+
"GET /v3/videos": {
150+
StatusCode: 200,
151+
Body: `{"data":[{"id":"v1"}],"next_token":null}`,
152+
},
153+
})
154+
defer srv.Close()
155+
156+
res := runGenCommand(t, srv.URL, "test-key", videoListSpec, "list", "--all")
157+
158+
if res.ExitCode != 0 {
159+
t.Fatalf("ExitCode = %d, want 0\nstderr: %s", res.ExitCode, res.Stderr)
160+
}
161+
var parsed []map[string]any
162+
if err := json.Unmarshal([]byte(res.Stdout), &parsed); err != nil {
163+
t.Fatalf("stdout is not valid JSON array: %v\nstdout: %s", err, res.Stdout)
164+
}
165+
if len(parsed) != 1 {
166+
t.Fatalf("len(parsed) = %d, want 1", len(parsed))
167+
}
168+
}
169+
170+
func TestGenBuilder_VideoList_AllAndTokenConflict(t *testing.T) {
171+
srv := setupTestServer(t, map[string]testHandler{})
172+
defer srv.Close()
173+
174+
res := runGenCommand(t, srv.URL, "test-key", videoListSpec, "list", "--all", "--token", "cursor_abc")
175+
176+
if res.ExitCode != 2 {
177+
t.Fatalf("ExitCode = %d, want 2\nstderr: %s", res.ExitCode, res.Stderr)
178+
}
179+
if !strings.Contains(res.Stderr, "--all and --token are mutually exclusive") {
180+
t.Fatalf("stderr = %s, want conflict message", res.Stderr)
181+
}
182+
}
183+
184+
func TestGenBuilder_VideoList_NoAllFlag_NonPaginated(t *testing.T) {
185+
nonPaginated := &command.Spec{
186+
Group: "video",
187+
Name: "get",
188+
Summary: "Get video",
189+
Endpoint: "/v3/videos/{video_id}",
190+
Method: "GET",
191+
Args: []command.ArgSpec{
192+
{Name: "video-id", Param: "video_id"},
193+
},
194+
Examples: []string{"heygen video get <video-id>"},
195+
}
196+
197+
srv := setupTestServer(t, map[string]testHandler{})
198+
defer srv.Close()
199+
200+
res := runGenCommand(t, srv.URL, "test-key", nonPaginated, "get", "vid_123", "--all")
201+
202+
if res.ExitCode != 2 {
203+
t.Fatalf("ExitCode = %d, want 2\nstderr: %s", res.ExitCode, res.Stderr)
204+
}
205+
if !strings.Contains(res.Stderr, "unknown flag: --all") {
206+
t.Fatalf("stderr = %s, want unknown flag error", res.Stderr)
207+
}
208+
}
209+
210+
func TestGenBuilder_VideoList_AllPages_Truncated(t *testing.T) {
211+
var calls int
212+
srv := setupTestServer(t, map[string]testHandler{
213+
"GET /v3/videos": {
214+
StatusCode: 200,
215+
ValidateRequest: func(t *testing.T, r *http.Request) {
216+
t.Helper()
217+
calls++
218+
},
219+
Body: truncatedPageBody(0, 2500, "cursor_1"),
220+
},
221+
})
222+
defer srv.Close()
223+
224+
originalHandler := srv.Config.Handler
225+
srv.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
226+
token := r.URL.Query().Get("token")
227+
switch token {
228+
case "":
229+
originalHandler.ServeHTTP(w, r)
230+
case "cursor_1":
231+
calls++
232+
w.WriteHeader(http.StatusOK)
233+
_, _ = w.Write([]byte(truncatedPageBody(2500, 2500, "cursor_2")))
234+
case "cursor_2":
235+
calls++
236+
w.WriteHeader(http.StatusOK)
237+
_, _ = w.Write([]byte(truncatedPageBody(5000, 2500, "cursor_3")))
238+
case "cursor_3":
239+
calls++
240+
w.WriteHeader(http.StatusOK)
241+
_, _ = w.Write([]byte(truncatedPageBody(7500, 2500, "cursor_4")))
242+
default:
243+
t.Fatalf("unexpected token %q", token)
244+
}
245+
})
246+
247+
res := runGenCommand(t, srv.URL, "test-key", videoListSpec, "list", "--all")
248+
249+
if res.ExitCode != 1 {
250+
t.Fatalf("ExitCode = %d, want 1\nstderr: %s", res.ExitCode, res.Stderr)
251+
}
252+
if !strings.Contains(res.Stderr, "Warning: pagination stopped at 10000 items") {
253+
t.Fatalf("stderr = %s, want truncation warning", res.Stderr)
254+
}
255+
256+
var parsed []map[string]any
257+
if err := json.Unmarshal([]byte(res.Stdout), &parsed); err != nil {
258+
t.Fatalf("stdout is not valid JSON array: %v\nstdout: %s", err, res.Stdout)
259+
}
260+
if len(parsed) != 10000 {
261+
t.Fatalf("len(parsed) = %d, want 10000", len(parsed))
262+
}
263+
}
264+
95265
func TestGenBuilder_PostWithBodyFlags(t *testing.T) {
96266
var gotBody map[string]any
97267

@@ -436,3 +606,20 @@ func runGeneratedRootCommand(t *testing.T, serverURL, apiKey string, groups map[
436606
ExitCode: exitCode,
437607
}
438608
}
609+
610+
func truncatedPageBody(start, count int, nextToken string) string {
611+
items := make([]map[string]any, 0, count)
612+
for i := 0; i < count; i++ {
613+
items = append(items, map[string]any{"id": fmt.Sprintf("v%d", start+i)})
614+
}
615+
body := map[string]any{
616+
"data": items,
617+
}
618+
if nextToken == "" {
619+
body["next_token"] = nil
620+
} else {
621+
body["next_token"] = nextToken
622+
}
623+
raw, _ := json.Marshal(body)
624+
return string(raw)
625+
}

codegen/grouper.go

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ import (
5353
// → sessions → sub-group, {session_id} → arg, stop → sub-group
5454
// → x-cli-action: true → no terminal verb appended
5555
// → result: heygen video-agent sessions stop <session-id>
56+
//
5657
// GroupDescriptions maps group name → description from the OpenAPI tag.
5758
// Used by the builder for group command help text.
5859
type GroupDescriptions map[string]string
@@ -157,7 +158,7 @@ func buildSpec(
157158
}
158159

159160
// Pagination
160-
_, spec.TokenField, spec.DataField = detectPagination(op)
161+
spec.Paginated, spec.TokenField, spec.TokenParam, spec.DataField = detectPagination(op, pathItem)
161162

162163
// Flags from query params
163164
for _, paramRef := range collectParams(pathItem, op) {
@@ -439,7 +440,7 @@ func isComplexField(s *openapi3.Schema) bool {
439440

440441
// --- Response analysis ---
441442

442-
func detectPagination(op *openapi3.Operation) (hasMore bool, tokenField, dataField string) {
443+
func detectPagination(op *openapi3.Operation, pathItem *openapi3.PathItem) (paginated bool, tokenField, tokenParam, dataField string) {
443444
respSchema := successResponseSchema(op)
444445
if respSchema == nil {
445446
return
@@ -455,11 +456,39 @@ func detectPagination(op *openapi3.Operation) (hasMore bool, tokenField, dataFie
455456
for _, schema := range schemasToCheck {
456457
for _, candidate := range []string{"next_token", "token", "cursor"} {
457458
if _, ok := schema.Properties[candidate]; ok {
458-
return true, candidate, dataField
459+
tokenField = candidate
460+
break
459461
}
460462
}
463+
if tokenField != "" {
464+
break
465+
}
466+
}
467+
if tokenField == "" {
468+
return false, "", "", dataField
469+
}
470+
471+
tokenParam = detectCursorParam(pathItem, op)
472+
if tokenParam == "" {
473+
return false, tokenField, "", dataField
461474
}
462-
return
475+
476+
return true, tokenField, tokenParam, dataField
477+
}
478+
479+
func detectCursorParam(pathItem *openapi3.PathItem, op *openapi3.Operation) string {
480+
params := collectParams(pathItem, op)
481+
for _, paramRef := range params {
482+
param := paramRef.Value
483+
if param == nil || param.In != "query" {
484+
continue
485+
}
486+
switch param.Name {
487+
case "token", "cursor", "page_token":
488+
return param.Name
489+
}
490+
}
491+
return ""
463492
}
464493

465494
func successResponseSchema(op *openapi3.Operation) *openapi3.Schema {

codegen/grouper_test.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,9 +156,15 @@ func TestGroupEndpoints_Pagination(t *testing.T) {
156156
if s.Name != "list" {
157157
continue
158158
}
159+
if !s.Paginated {
160+
t.Error("Paginated = false, want true")
161+
}
159162
if s.TokenField != "token" {
160163
t.Errorf("TokenField = %q, want 'token'", s.TokenField)
161164
}
165+
if s.TokenParam != "token" {
166+
t.Errorf("TokenParam = %q, want 'token'", s.TokenParam)
167+
}
162168
if s.DataField != "data" {
163169
t.Errorf("DataField = %q, want 'data'", s.DataField)
164170
}

0 commit comments

Comments
 (0)