Skip to content

Commit 53256fb

Browse files
austin-denobleaustin-denoble
committed
rename bodyutil package to argsio, rename functions in package to be a bit clearer, clean up flags.go to leverage argsio rather that duplicating a bunch of logic for reading out of flag args, add unit tests for flags_test.go
1 parent 32afd74 commit 53256fb

9 files changed

Lines changed: 269 additions & 173 deletions

File tree

‎internal/pkg/cli/command/index/vector/fetch.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ package vector
33
import (
44
"context"
55

6-
"github.com/pinecone-io/cli/internal/pkg/utils/bodyutil"
6+
"github.com/pinecone-io/cli/internal/pkg/utils/argio"
77
"github.com/pinecone-io/cli/internal/pkg/utils/exit"
88
"github.com/pinecone-io/cli/internal/pkg/utils/flags"
99
"github.com/pinecone-io/cli/internal/pkg/utils/help"
@@ -75,7 +75,7 @@ func runFetchCmd(ctx context.Context, options fetchCmdOptions) {
7575

7676
// Apply body overlay if provided
7777
if options.body != "" {
78-
if b, src, err := bodyutil.DecodeBodyArgs[fetchBody](options.body); err != nil {
78+
if b, src, err := argio.DecodeBodyArgs[fetchBody](options.body); err != nil {
7979
msg.FailMsg("Failed to parse fetch body (%s): %s", style.Emphasis(src.Label), err)
8080
exit.Errorf(err, "Failed to parse fetch body (%s): %v", src.Label, err)
8181
} else if b != nil {

‎internal/pkg/cli/command/index/vector/query.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ package vector
33
import (
44
"context"
55

6-
"github.com/pinecone-io/cli/internal/pkg/utils/bodyutil"
6+
"github.com/pinecone-io/cli/internal/pkg/utils/argio"
77
"github.com/pinecone-io/cli/internal/pkg/utils/exit"
88
"github.com/pinecone-io/cli/internal/pkg/utils/flags"
99
"github.com/pinecone-io/cli/internal/pkg/utils/help"
@@ -90,7 +90,7 @@ func runQueryCmd(ctx context.Context, options queryCmdOptions) {
9090

9191
// Apply body overlay if provided
9292
if options.body != "" {
93-
if b, src, err := bodyutil.DecodeBodyArgs[queryBody](options.body); err != nil {
93+
if b, src, err := argio.DecodeBodyArgs[queryBody](options.body); err != nil {
9494
msg.FailMsg("Failed to parse query body (%s): %s", style.Emphasis(src.Label), err)
9595
exit.Errorf(err, "Failed to parse query body (%s): %v", src.Label, err)
9696
} else if b != nil {

‎internal/pkg/cli/command/index/vector/update.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ package vector
33
import (
44
"context"
55

6-
"github.com/pinecone-io/cli/internal/pkg/utils/bodyutil"
6+
"github.com/pinecone-io/cli/internal/pkg/utils/argio"
77
"github.com/pinecone-io/cli/internal/pkg/utils/exit"
88
"github.com/pinecone-io/cli/internal/pkg/utils/flags"
99
"github.com/pinecone-io/cli/internal/pkg/utils/help"
@@ -82,7 +82,7 @@ func runUpdateCmd(ctx context.Context, options updateCmdOptions) {
8282

8383
// Apply body overlay if provided
8484
if options.body != "" {
85-
if b, src, err := bodyutil.DecodeBodyArgs[updateBody](options.body); err != nil {
85+
if b, src, err := argio.DecodeBodyArgs[updateBody](options.body); err != nil {
8686
msg.FailMsg("Failed to parse update body (%s): %s", style.Emphasis(src.Label), err)
8787
exit.Errorf(err, "Failed to parse update body (%s): %v", src.Label, err)
8888
} else if b != nil {

‎internal/pkg/cli/command/index/vector/upsert.go‎

Lines changed: 46 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,12 @@
11
package vector
22

33
import (
4+
"bytes"
45
"context"
6+
"encoding/json"
7+
"io"
58

6-
"github.com/pinecone-io/cli/internal/pkg/utils/bodyutil"
9+
"github.com/pinecone-io/cli/internal/pkg/utils/argio"
710
"github.com/pinecone-io/cli/internal/pkg/utils/exit"
811
"github.com/pinecone-io/cli/internal/pkg/utils/help"
912
"github.com/pinecone-io/cli/internal/pkg/utils/msg"
@@ -38,7 +41,10 @@ func NewUpsertCmd() *cobra.Command {
3841
pc index vector upsert --index-name my-index --namespace my-namespace ./vectors.json
3942
pc index vector upsert --index-name my-index --namespace my-namespace --file - < ./vectors.json
4043
pc index vector upsert --index-name my-index --body @./payload.json
44+
pc index vector upsert --index-name my-index --namespace my-namespace @./vectors.jsonl
4145
cat payload.json | pc index vector upsert --index-name my-index --body @-
46+
47+
Body may be a JSON object with "vectors": [...] or JSONL of Vector objects.
4248
`),
4349
Run: func(cmd *cobra.Command, args []string) {
4450
runUpsertCmd(cmd.Context(), options)
@@ -47,8 +53,8 @@ func NewUpsertCmd() *cobra.Command {
4753

4854
cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of index to upsert into")
4955
cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "namespace to upsert into")
50-
cmd.Flags().StringVar(&options.body, "body", "", "request body JSON (inline, @path.json, or @- for stdin; only one argument may use stdin; max size: see PC_CLI_MAX_JSON_BYTES)")
51-
cmd.Flags().IntVarP(&options.batchSize, "batch-size", "b", 1000, "size of batches to upsert (default: 1000)")
56+
cmd.Flags().StringVar(&options.body, "body", "", "request body JSON or JSONL (inline, @path.json[l], or @- for stdin; only one argument may use stdin; max size: see PC_CLI_MAX_JSON_BYTES)")
57+
cmd.Flags().IntVarP(&options.batchSize, "batch-size", "b", 500, "size of batches to upsert (default: 500)")
5258
cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON")
5359
_ = cmd.MarkFlagRequired("index-name")
5460
_ = cmd.MarkFlagRequired("body")
@@ -57,8 +63,13 @@ func NewUpsertCmd() *cobra.Command {
5763
}
5864

5965
func runUpsertCmd(ctx context.Context, options upsertCmdOptions) {
60-
var payload *upsertBody
61-
payload, src, err := bodyutil.DecodeBodyArgs[upsertBody](options.body)
66+
b, src, err := argio.ReadAll(options.body, true)
67+
if err != nil {
68+
msg.FailMsg("Failed to read upsert body (%s): %s", style.Emphasis(src.Label), err)
69+
exit.Error(err, "Failed to read upsert body")
70+
}
71+
72+
payload, err := parseUpsertBody(b)
6273
if err != nil {
6374
msg.FailMsg("Failed to parse upsert body (%s): %s", style.Emphasis(src.Label), err)
6475
exit.Error(err, "Failed to parse upsert body")
@@ -124,3 +135,33 @@ func runUpsertCmd(ctx context.Context, options upsertCmdOptions) {
124135
}
125136
}
126137
}
138+
139+
func parseUpsertBody(b []byte) (*upsertBody, error) {
140+
var payload upsertBody
141+
// First try and decode as JSON: {"vectors":[...]}
142+
{
143+
dec := json.NewDecoder(bytes.NewReader(b))
144+
dec.DisallowUnknownFields()
145+
if err := dec.Decode(&payload); err == nil && len(payload.Vectors) > 0 {
146+
return &payload, nil
147+
}
148+
}
149+
150+
// Fallback: JSONL/stream of pinecone.Vector values
151+
var vectors []pinecone.Vector
152+
dec := json.NewDecoder(bytes.NewReader(b))
153+
dec.DisallowUnknownFields()
154+
for {
155+
var v pinecone.Vector
156+
if err := dec.Decode(&v); err == io.EOF {
157+
break
158+
} else if err != nil {
159+
return nil, err
160+
}
161+
vectors = append(vectors, v)
162+
}
163+
if len(vectors) == 0 {
164+
return nil, io.EOF
165+
}
166+
return &upsertBody{Vectors: vectors}, nil
167+
}

‎internal/pkg/cli/command/index/vector/upsert_test.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package vector
2+
3+
import (
4+
"testing"
5+
)
6+
7+
func TestParseUpsertBody_JSONL(t *testing.T) {
8+
jsonl := `{"id":"a","values":[1,2,3]}
9+
{"id":"b","values":[4,5,6]}
10+
`
11+
payload, err := parseUpsertBody([]byte(jsonl))
12+
if err != nil {
13+
t.Fatalf("unexpected error: %v", err)
14+
}
15+
if payload == nil || len(payload.Vectors) != 2 {
16+
t.Fatalf("expected 2 vectors, got %+v", payload)
17+
}
18+
if payload.Vectors[0].Id != "a" || payload.Vectors[1].Id != "b" {
19+
t.Fatalf("unexpected ids: %v, %v", payload.Vectors[0].Id, payload.Vectors[1].Id)
20+
}
21+
}

internal/pkg/utils/bodyutil/body.go renamed to internal/pkg/utils/argio/argio.go

Lines changed: 37 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package bodyutil
1+
package argio
22

33
import (
44
"encoding/json"
@@ -24,55 +24,78 @@ type SourceInfo struct {
2424
Label string `json:"label"` // "inline", file path, "stdin"
2525
}
2626

27-
// OpenArgReader returns an io.ReadCloser for inline/@file/@- with size limits applied.
28-
// When isBody is true, the body size limit is used, otherwise the flag size limit.
29-
func OpenArgReader(spec string, isBody bool) (io.ReadCloser, SourceInfo, error) {
27+
// OpenReader returns an io.ReadCloser for inline/@file/@- values with inputpolicy.MaxBodyJSONBytes limit applied.
28+
func OpenReader(value string) (io.ReadCloser, SourceInfo, error) {
3029
limit := inputpolicy.MaxBodyJSONBytes
3130
switch {
32-
case spec == "":
31+
case value == "": // empty value is inline
3332
return nil, SourceInfo{Kind: SourceInline, Label: "inline"}, nil
34-
case spec == "@-":
33+
case value == "@-": // @- is stdin
3534
r, err := stdin.ReaderOnce(limit)
3635
if err != nil {
3736
return nil, SourceInfo{Kind: SourceStdin, Label: "stdin"}, fmt.Errorf("stdin already consumed; only one argument may use stdin")
3837
}
38+
3939
return r, SourceInfo{Kind: SourceStdin, Label: "stdin"}, nil
40-
case len(spec) > 0 && spec[0] == '@':
41-
path := spec[1:]
40+
case len(value) > 0 && value[0] == '@': // @file is a file
41+
path := value[1:]
4242
if err := inputpolicy.ValidatePath(path); err != nil {
4343
return nil, SourceInfo{Kind: SourceFile, Label: path}, err
4444
}
45+
4546
f, err := os.Open(path)
4647
if err != nil {
4748
return nil, SourceInfo{Kind: SourceFile, Label: path}, err
4849
}
50+
4951
return struct {
5052
io.Reader
5153
io.Closer
5254
}{Reader: io.LimitReader(f, limit), Closer: f}, SourceInfo{Kind: SourceFile, Label: path}, nil
53-
default:
54-
return io.NopCloser(strings.NewReader(spec)), SourceInfo{Kind: SourceInline, Label: "inline"}, nil
55+
default: // if no stdin and no file, it's inline
56+
return io.NopCloser(strings.NewReader(value)), SourceInfo{Kind: SourceInline, Label: "inline"}, nil
57+
}
58+
}
59+
60+
// ReadAll reads the entire argument from the inline/@file/@- value into memory using the same limits
61+
// as OpenReader. The returned bytes are bounded by inputpolicy.MaxBodyJSONBytes.
62+
func ReadAll(value string) ([]byte, SourceInfo, error) {
63+
rc, src, err := OpenReader(value)
64+
if err != nil {
65+
return nil, src, err
66+
}
67+
if rc == nil {
68+
return nil, src, fmt.Errorf("empty input from %s", src.Label)
5569
}
70+
defer rc.Close()
71+
72+
b, err := io.ReadAll(rc)
73+
if err != nil {
74+
return nil, src, err
75+
}
76+
77+
return b, src, nil
5678
}
5779

5880
// DecodeBodyArgs unmarshals a JSON body argument (inline/@file/@-) using a bounded reader.
59-
func DecodeBodyArgs[T any](spec string) (*T, SourceInfo, error) {
60-
rc, src, err := OpenArgReader(spec, true)
81+
func DecodeBodyArgs[T any](value string) (*T, SourceInfo, error) {
82+
rc, src, err := OpenReader(value)
6183
if err != nil {
6284
return nil, src, err
6385
}
86+
6487
var closer io.Closer
6588
if rc != nil {
6689
closer = rc
6790
defer closer.Close()
6891
}
92+
6993
dec := json.NewDecoder(rc)
7094
dec.DisallowUnknownFields()
7195
var out T
7296
if err := dec.Decode(&out); err != nil {
7397
return nil, src, fmt.Errorf("invalid JSON from %s: %w", src.Label, err)
7498
}
99+
75100
return &out, src, nil
76101
}
77-
78-
// NOTE: we intentionally import strings and use strings.NewReader for clarity.

internal/pkg/utils/bodyutil/bodyutil_test.go renamed to internal/pkg/utils/argio/argio_test.go

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package bodyutil
1+
package argio
22

33
import (
44
"io"
@@ -7,7 +7,7 @@ import (
77
)
88

99
func TestOpenArgReader_Inline(t *testing.T) {
10-
rc, src, err := OpenArgReader(`{"a":1}`, true)
10+
rc, src, err := OpenReader(`{"a":1}`)
1111
if err != nil {
1212
t.Fatalf("unexpected error: %v", err)
1313
}
@@ -20,3 +20,16 @@ func TestOpenArgReader_Inline(t *testing.T) {
2020
t.Fatalf("unexpected body %q", string(b))
2121
}
2222
}
23+
24+
func TestReadAll_Inline(t *testing.T) {
25+
b, src, err := ReadAll(`{"a":1}`)
26+
if err != nil {
27+
t.Fatalf("unexpected error: %v", err)
28+
}
29+
if src.Kind != SourceInline || src.Label != "inline" {
30+
t.Fatalf("unexpected source: %+v", src)
31+
}
32+
if !strings.Contains(string(b), `"a":1`) {
33+
t.Fatalf("unexpected body %q", string(b))
34+
}
35+
}

0 commit comments

Comments
 (0)