Skip to content

Commit

Permalink
Merge pull request #32085 from vespa-engine/mpolden/more-header
Browse files Browse the repository at this point in the history 
Add --header option to document and visit commands
  • Loading branch information
@bratseth
bratseth authored Aug 6, 2024
2 parents 91a68f3 + 48d5826 commit ccaf08c
Show file tree
Hide file tree
Showing 6 changed files with 62 additions and 43 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
22 changes: 17 additions & 5 deletions client/go/internal/cli/cmd/visit.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
"time"

"github.com/spf13/cobra"
"github.com/vespa-engine/vespa/client/go/internal/httputil"
"github.com/vespa-engine/vespa/client/go/internal/ioutil"
"github.com/vespa-engine/vespa/client/go/internal/vespa"
)
Expand All @@ -37,7 +38,10 @@ type visitArgs struct {
bucketSpaces []string
waitSecs int
verbose bool
cli *CLI
headers []string

cli *CLI
header http.Header
}

func (v *visitArgs) writeBytes(b []byte) {
Expand Down Expand Up @@ -109,6 +113,11 @@ $ vespa visit --field-set "[id]" # list document IDs
if !result.Success {
return fmt.Errorf("argument error: %s", result.Message)
}
header, err := httputil.ParseHeader(vArgs.headers)
if err != nil {
return err
}
vArgs.header = header
waiter := cli.waiter(time.Duration(vArgs.waitSecs)*time.Second, cmd)
service, err := documentService(cli, waiter)
if err != nil {
Expand All @@ -117,7 +126,7 @@ $ vespa visit --field-set "[id]" # list document IDs
if vArgs.verbose {
service.CurlWriter = vespa.CurlWriter{Writer: cli.Stderr}
}
result = probeHandler(service, cli)
result = probeHandler(&vArgs, service, cli)
if result.Success {
result = visitClusters(&vArgs, service)
}
Expand All @@ -142,6 +151,7 @@ $ vespa visit --field-set "[id]" # list document IDs
cmd.Flags().IntVar(&vArgs.slices, "slices", -1, `Split the document corpus into this number of independent slices`)
cmd.Flags().StringSliceVar(&vArgs.bucketSpaces, "bucket-space", []string{"global", "default"}, `The "default" or "global" bucket space`)
cmd.Flags().BoolVarP(&vArgs.verbose, "verbose", "v", false, `Print the equivalent curl command for the visit operation`)
cmd.Flags().StringSliceVarP(&vArgs.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, &vArgs.waitSecs)
return cmd
}
Expand All @@ -160,7 +170,7 @@ func getEpoch(timeStamp string) (int64, error) {

func checkArguments(vArgs visitArgs) (res OperationResult) {
if vArgs.slices > 0 || vArgs.sliceId > -1 {
if !(vArgs.slices > 0 && vArgs.sliceId > -1) {
if vArgs.slices <= 0 || vArgs.sliceId <= -1 {
return Failure("Both 'slices' and 'slice-id' must be set")
}
if vArgs.sliceId >= vArgs.slices {
Expand Down Expand Up @@ -209,7 +219,7 @@ func parseHandlersOutput(r io.Reader) (*HandlersInfo, error) {
return &handlersInfo, err
}

func probeHandler(service *vespa.Service, cli *CLI) (res OperationResult) {
func probeHandler(vArgs *visitArgs, service *vespa.Service, cli *CLI) (res OperationResult) {
urlPath := service.BaseURL + "/"
url, urlParseError := url.Parse(urlPath)
if urlParseError != nil {
Expand All @@ -218,6 +228,7 @@ func probeHandler(service *vespa.Service, cli *CLI) (res OperationResult) {
request := &http.Request{
URL: url,
Method: "GET",
Header: vArgs.header,
}
timeout := time.Duration(90) * time.Second
response, err := service.Do(request, timeout)
Expand Down Expand Up @@ -297,7 +308,7 @@ func probeVisit(vArgs *visitArgs, service *vespa.Service) []string {

func runVisit(vArgs *visitArgs, service *vespa.Service) (res OperationResult) {
vArgs.debugPrint(fmt.Sprintf("trying to visit: '%s'", vArgs.contentCluster))
var totalDocuments int = 0
var totalDocuments = 0
var continuationToken string
for {
var vvo *VespaVisitOutput
Expand Down Expand Up @@ -376,6 +387,7 @@ func runOneVisit(vArgs *visitArgs, service *vespa.Service, contToken string) (*V
request := &http.Request{
URL: url,
Method: "GET",
Header: vArgs.header,
}
timeout := time.Duration(900) * time.Second
response, err := service.Do(request, timeout)
Expand Down
3 changes: 3 additions & 0 deletions client/go/internal/cli/cmd/visit_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,9 @@ func TestRunOneVisit(t *testing.T) {
op := func(service *vespa.Service) {
vArgs := visitArgs{
contentCluster: "fooCC",
header: make(http.Header),
}
vArgs.header.Set("X-Foo", "Bar")
vvo, res := runOneVisit(&vArgs, service, "BBBB")
assert.Equal(t, true, res.Success)
assert.Equal(t, "visited fooCC", res.Message)
Expand All @@ -75,6 +77,7 @@ func TestRunOneVisit(t *testing.T) {
}
req := withMockClient(t, withResponse, op)
assert.Equal(t, "cluster=fooCC&continuation=BBBB", req.URL.RawQuery)
assert.Equal(t, "Bar", req.Header.Get("X-Foo"))

op = func(service *vespa.Service) {
vArgs := visitArgs{
Expand Down
16 changes: 0 additions & 16 deletions client/go/internal/vespa/document/document.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -401,19 +401,3 @@ func (g *Generator) Read(p []byte) (int, error) {
}
return g.buf.Read(p)
}

type number interface{ float64 | int64 | int }

func min[T number](x, y T) T {
if x < y {
return x
}
return y
}

func max[T number](x, y T) T {
if x > y {
return x
}
return y
}
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 ccaf08c

Please sign in to comment.