Skip to content

Commit

Permalink
Add --header to document commands
Browse files Browse the repository at this point in the history
  • Loading branch information
@mpolden
mpolden committed Aug 5, 2024
1 parent 687b93a commit 27458e9
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 22 deletions.
45 changes: 28 additions & 17 deletions client/go/internal/cli/cmd/document.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,34 +21,40 @@ import (
"github.com/vespa-engine/vespa/client/go/internal/vespa/document"
)

func addDocumentFlags(cli *CLI, cmd *cobra.Command, printCurl *bool, timeoutSecs, waitSecs *int) {
func addDocumentFlags(cli *CLI, cmd *cobra.Command, printCurl *bool, timeoutSecs, waitSecs *int, headers *[]string) {
cmd.PersistentFlags().BoolVarP(printCurl, "verbose", "v", false, "Print the equivalent curl command for the document operation")
cmd.PersistentFlags().IntVarP(timeoutSecs, "timeout", "T", 60, "Timeout for the document request in seconds")
cmd.PersistentFlags().StringSliceVarP(headers, "header", "", nil, "Add a header to the HTTP request, on the format 'Header: Value'. This can be specified multiple times")
cli.bindWaitFlag(cmd, 0, waitSecs)
}

func documentClient(cli *CLI, timeoutSecs int, waiter *Waiter, printCurl bool) (*document.Client, *vespa.Service, error) {
func documentClient(cli *CLI, timeoutSecs int, waiter *Waiter, printCurl bool, headers []string) (*document.Client, *vespa.Service, error) {
docService, err := documentService(cli, waiter)
if err != nil {
return nil, nil, err
}
if printCurl {
docService.CurlWriter = vespa.CurlWriter{Writer: cli.Stderr}
}
header, err := httputil.ParseHeader(headers)
if err != nil {
return nil, nil, err
}
client, err := document.NewClient(document.ClientOptions{
Compression: document.CompressionAuto,
Timeout: time.Duration(timeoutSecs) * time.Second,
BaseURL: docService.BaseURL,
NowFunc: time.Now,
Header: header,
}, []httputil.Client{docService})
if err != nil {
return nil, nil, err
}
return client, docService, nil
}

func sendOperation(op document.Operation, args []string, timeoutSecs int, waiter *Waiter, printCurl bool, cli *CLI) error {
client, service, err := documentClient(cli, timeoutSecs, waiter, printCurl)
func sendOperation(op document.Operation, args []string, timeoutSecs int, waiter *Waiter, printCurl bool, cli *CLI, headers []string) error {
client, service, err := documentClient(cli, timeoutSecs, waiter, printCurl, headers)
if err != nil {
return err
}
Expand Down Expand Up @@ -91,8 +97,8 @@ func sendOperation(op document.Operation, args []string, timeoutSecs int, waiter
return printResult(cli, operationResult(false, doc, service, result), false)
}

func readDocument(id string, timeoutSecs int, waiter *Waiter, printCurl bool, cli *CLI, fieldSet string) error {
client, service, err := documentClient(cli, timeoutSecs, waiter, printCurl)
func readDocument(id string, timeoutSecs int, waiter *Waiter, printCurl bool, cli *CLI, fieldSet string, headers []string) error {
client, service, err := documentClient(cli, timeoutSecs, waiter, printCurl, headers)
if err != nil {
return err
}
Expand Down Expand Up @@ -127,6 +133,7 @@ func newDocumentCmd(cli *CLI) *cobra.Command {
printCurl bool
timeoutSecs int
waitSecs int
headers []string
)
cmd := &cobra.Command{
Use: "document json-file",
Expand All @@ -147,10 +154,10 @@ should be used instead of this.`,
Args: cobra.ExactArgs(1),
RunE: func(cmd *cobra.Command, args []string) error {
waiter := cli.waiter(time.Duration(waitSecs)*time.Second, cmd)
return sendOperation(-1, args, timeoutSecs, waiter, printCurl, cli)
return sendOperation(-1, args, timeoutSecs, waiter, printCurl, cli, headers)
},
}
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs)
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs, &headers)
return cmd
}

Expand All @@ -159,6 +166,7 @@ func newDocumentPutCmd(cli *CLI) *cobra.Command {
printCurl bool
timeoutSecs int
waitSecs int
headers []string
)
cmd := &cobra.Command{
Use: "put [id] json-file",
Expand All @@ -173,10 +181,10 @@ $ vespa document put id:mynamespace:music::a-head-full-of-dreams src/test/resour
SilenceUsage: true,
RunE: func(cmd *cobra.Command, args []string) error {
waiter := cli.waiter(time.Duration(waitSecs)*time.Second, cmd)
return sendOperation(document.OperationPut, args, timeoutSecs, waiter, printCurl, cli)
return sendOperation(document.OperationPut, args, timeoutSecs, waiter, printCurl, cli, headers)
},
}
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs)
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs, &headers)
return cmd
}

Expand All @@ -185,6 +193,7 @@ func newDocumentUpdateCmd(cli *CLI) *cobra.Command {
printCurl bool
timeoutSecs int
waitSecs int
headers []string
)
cmd := &cobra.Command{
Use: "update [id] json-file",
Expand All @@ -198,10 +207,10 @@ $ vespa document update id:mynamespace:music::a-head-full-of-dreams src/test/res
SilenceUsage: true,
RunE: func(cmd *cobra.Command, args []string) error {
waiter := cli.waiter(time.Duration(waitSecs)*time.Second, cmd)
return sendOperation(document.OperationUpdate, args, timeoutSecs, waiter, printCurl, cli)
return sendOperation(document.OperationUpdate, args, timeoutSecs, waiter, printCurl, cli, headers)
},
}
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs)
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs, &headers)
return cmd
}

Expand All @@ -210,6 +219,7 @@ func newDocumentRemoveCmd(cli *CLI) *cobra.Command {
printCurl bool
timeoutSecs int
waitSecs int
headers []string
)
cmd := &cobra.Command{
Use: "remove id | json-file",
Expand All @@ -224,7 +234,7 @@ $ vespa document remove id:mynamespace:music::a-head-full-of-dreams`,
RunE: func(cmd *cobra.Command, args []string) error {
waiter := cli.waiter(time.Duration(waitSecs)*time.Second, cmd)
if strings.HasPrefix(args[0], "id:") {
client, service, err := documentClient(cli, timeoutSecs, waiter, printCurl)
client, service, err := documentClient(cli, timeoutSecs, waiter, printCurl, headers)
if err != nil {
return err
}
Expand All @@ -236,11 +246,11 @@ $ vespa document remove id:mynamespace:music::a-head-full-of-dreams`,
result := client.Send(doc)
return printResult(cli, operationResult(false, doc, service, result), false)
} else {
return sendOperation(document.OperationRemove, args, timeoutSecs, waiter, printCurl, cli)
return sendOperation(document.OperationRemove, args, timeoutSecs, waiter, printCurl, cli, headers)
}
},
}
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs)
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs, &headers)
return cmd
}

Expand All @@ -250,6 +260,7 @@ func newDocumentGetCmd(cli *CLI) *cobra.Command {
timeoutSecs int
waitSecs int
fieldSet string
headers []string
)
cmd := &cobra.Command{
Use: "get id",
Expand All @@ -260,11 +271,11 @@ func newDocumentGetCmd(cli *CLI) *cobra.Command {
Example: `$ vespa document get id:mynamespace:music::a-head-full-of-dreams`,
RunE: func(cmd *cobra.Command, args []string) error {
waiter := cli.waiter(time.Duration(waitSecs)*time.Second, cmd)
return readDocument(args[0], timeoutSecs, waiter, printCurl, cli, fieldSet)
return readDocument(args[0], timeoutSecs, waiter, printCurl, cli, fieldSet, headers)
},
}
cmd.Flags().StringVar(&fieldSet, "field-set", "", "Fields to include when reading document")
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs)
addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs, &headers)
return cmd
}

Expand Down
13 changes: 10 additions & 3 deletions client/go/internal/cli/cmd/document_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,8 +106,16 @@ func TestDocumentPutTransportError(t *testing.T) {
}

func TestDocumentGet(t *testing.T) {
assertDocumentGet([]string{"document", "get", "id:mynamespace:music::a-head-full-of-dreams"},
client := &mock.HTTPClient{}
assertDocumentGet(client, []string{"document", "get", "id:mynamespace:music::a-head-full-of-dreams"},
"id:mynamespace:music::a-head-full-of-dreams", t)
}

func TestDocumentGetWithHeader(t *testing.T) {
client := &mock.HTTPClient{}
assertDocumentGet(client, []string{"document", "get", "--header", "X-Foo: Bar", "id:mynamespace:music::a-head-full-of-dreams"},
"id:mynamespace:music::a-head-full-of-dreams", t)
assert.Equal(t, "Bar", client.LastRequest.Header.Get("X-Foo"))
}

func assertDocumentSend(args []string, expectedOperation string, expectedMethod string, expectedDocumentId string, expectedPayloadFile string, t *testing.T) {
Expand Down Expand Up @@ -154,8 +162,7 @@ func assertDocumentSend(args []string, expectedOperation string, expectedMethod
}
}

func assertDocumentGet(args []string, documentId string, t *testing.T) {
client := &mock.HTTPClient{}
func assertDocumentGet(client *mock.HTTPClient, args []string, documentId string, t *testing.T) {
documentURL := "http://127.0.0.1:8080"
client.NextResponseString(200, "{\"fields\":{\"foo\":\"bar\"}}")
cli, stdout, _ := newTestCLI(t)
Expand Down
6 changes: 4 additions & 2 deletions client/go/internal/vespa/document/http.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,9 @@ func (c *Client) newRequest(method, url string, body io.Reader, gzipped bool) (*
for k, v := range c.options.Header {
req.Header[k] = v
}
req.Header.Set("Content-Type", "application/json; charset=utf-8")
if method != http.MethodGet {
req.Header.Set("Content-Type", "application/json; charset=utf-8")
}
if gzipped {
req.Header.Set("Content-Encoding", "gzip")
}
Expand Down Expand Up @@ -296,7 +298,7 @@ func (c *Client) Get(id Id, fieldSet string) Result {
}
url := buf.String()
result := Result{Id: id}
req, err := http.NewRequest(http.MethodGet, url, nil)
req, err := c.newRequest(http.MethodGet, url, nil, false)
if err != nil {
return resultWithErr(result, err, 0)
}
Expand Down

0 comments on commit 27458e9

Please sign in to comment.