Add --header to visit command

vespa-engine · Aug 5, 2024 · fdfe6eb · fdfe6eb
1 parent 4865d45
commit fdfe6eb
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git 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
@@ -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{