From 27458e97ce7223f28849d86e31620263c3ac7c50 Mon Sep 17 00:00:00 2001 From: Martin Polden Date: Mon, 5 Aug 2024 13:06:16 +0200 Subject: [PATCH 1/4] Add --header to document commands --- client/go/internal/cli/cmd/document.go | 45 +++++++++++++-------- client/go/internal/cli/cmd/document_test.go | 13 ++++-- client/go/internal/vespa/document/http.go | 6 ++- 3 files changed, 42 insertions(+), 22 deletions(-) diff --git a/client/go/internal/cli/cmd/document.go b/client/go/internal/cli/cmd/document.go index c9f0c780be34..9d52db136f12 100644 --- a/client/go/internal/cli/cmd/document.go +++ b/client/go/internal/cli/cmd/document.go @@ -21,13 +21,14 @@ 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 @@ -35,11 +36,16 @@ func documentClient(cli *CLI, timeoutSecs int, waiter *Waiter, printCurl bool) ( 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 @@ -47,8 +53,8 @@ func documentClient(cli *CLI, timeoutSecs int, waiter *Waiter, printCurl bool) ( 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 } @@ -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 } @@ -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", @@ -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 } @@ -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", @@ -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 } @@ -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", @@ -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 } @@ -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", @@ -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 } @@ -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 } @@ -250,6 +260,7 @@ func newDocumentGetCmd(cli *CLI) *cobra.Command { timeoutSecs int waitSecs int fieldSet string + headers []string ) cmd := &cobra.Command{ Use: "get id", @@ -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 } diff --git a/client/go/internal/cli/cmd/document_test.go b/client/go/internal/cli/cmd/document_test.go index 3cfc66fdad47..418581ff6784 100644 --- a/client/go/internal/cli/cmd/document_test.go +++ b/client/go/internal/cli/cmd/document_test.go @@ -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) { @@ -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) diff --git a/client/go/internal/vespa/document/http.go b/client/go/internal/vespa/document/http.go index df4d97e2a829..0237c55b1a34 100644 --- a/client/go/internal/vespa/document/http.go +++ b/client/go/internal/vespa/document/http.go @@ -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") } @@ -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) } From 4865d457e7f2ea0c8efe699e61bf46506384791d Mon Sep 17 00:00:00 2001 From: Martin Polden Date: Mon, 5 Aug 2024 13:10:37 +0200 Subject: [PATCH 2/4] Use min/max builtins --- client/go/internal/vespa/document/document.go | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/client/go/internal/vespa/document/document.go b/client/go/internal/vespa/document/document.go index 9c301cd79905..b9af09cf7cc6 100644 --- a/client/go/internal/vespa/document/document.go +++ b/client/go/internal/vespa/document/document.go @@ -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 -} From fdfe6eb0e08f6c95d903f07d45143fd988ff4aa4 Mon Sep 17 00:00:00 2001 From: Martin Polden Date: Mon, 5 Aug 2024 13:28:27 +0200 Subject: [PATCH 3/4] Add --header to visit command --- client/go/internal/cli/cmd/visit.go | 18 +++++++++++++++--- client/go/internal/cli/cmd/visit_test.go | 3 +++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/client/go/internal/cli/cmd/visit.go b/client/go/internal/cli/cmd/visit.go index 7f99df2038eb..7381039de0f2 100644 --- a/client/go/internal/cli/cmd/visit.go +++ b/client/go/internal/cli/cmd/visit.go @@ -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" ) @@ -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) { @@ -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 { @@ -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) } @@ -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 } @@ -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 { @@ -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) @@ -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) diff --git a/client/go/internal/cli/cmd/visit_test.go b/client/go/internal/cli/cmd/visit_test.go index 85594912da29..340089672a19 100644 --- a/client/go/internal/cli/cmd/visit_test.go +++ b/client/go/internal/cli/cmd/visit_test.go @@ -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) @@ -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{ From 48d58263491ed17260bed2686fd87ee7b32a1ba6 Mon Sep 17 00:00:00 2001 From: Martin Polden Date: Mon, 5 Aug 2024 13:29:12 +0200 Subject: [PATCH 4/4] Simplify condition --- client/go/internal/cli/cmd/visit.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/client/go/internal/cli/cmd/visit.go b/client/go/internal/cli/cmd/visit.go index 7381039de0f2..7439b70a27fc 100644 --- a/client/go/internal/cli/cmd/visit.go +++ b/client/go/internal/cli/cmd/visit.go @@ -170,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 { @@ -308,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