Merge pull request #30648 from vespa-engine/mpolden/query-header

CLI: Add query header option

2 files changed, 50 insertions, 8 deletions

diff --git a/client/go/internal/cli/cmd/query.go b/client/go/internal/cli/cmd/query.go
index bddf3af06f9..879c19c78a7 100644
--- a/client/go/internal/cli/cmd/query.go
+++ b/client/go/internal/cli/cmd/query.go
@@ -28,11 +28,14 @@ func newQueryCmd(cli *CLI) *cobra.Command {
 		queryTimeoutSecs int
 		waitSecs         int
 		format           string
+		headers          []string
 	)
 	cmd := &cobra.Command{
-		Use:     "query query-parameters",
-		Short:   "Issue a query to Vespa",
-		Example: `$ vespa query "yql=select * from music where album contains 'head'" hits=5`,
+		Use:   "query query-parameters",
+		Short: "Issue a query to Vespa",
+		Example: `$ vespa query "yql=select * from music where album contains 'head'" hits=5
+$ vespa query --format=plain "yql=select * from music where album contains 'head'" hits=5
+$ vespa query --header="X-First-Name: Joe" "yql=select * from music where album contains 'head'" hits=5`,
 		Long: `Issue a query to Vespa.
 
 Any parameter from https://docs.vespa.ai/en/reference/query-api-reference.html
@@ -42,11 +45,12 @@ can be set by the syntax [parameter-name]=[value].`,
 		SilenceUsage:      true,
 		Args:              cobra.MinimumNArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return query(cli, args, queryTimeoutSecs, waitSecs, printCurl, format)
+			return query(cli, args, queryTimeoutSecs, waitSecs, printCurl, format, headers)
 		},
 	}
-	cmd.PersistentFlags().BoolVarP(&printCurl, "verbose", "v", false, "Print the equivalent curl command for the query")
-	cmd.PersistentFlags().StringVarP(&format, "format", "", "human", "Output format. Must be 'human' (human-readable) or 'plain' (no formatting)")
+	cmd.Flags().BoolVarP(&printCurl, "verbose", "v", false, "Print the equivalent curl command for the query")
+	cmd.Flags().StringVarP(&format, "format", "", "human", "Output format. Must be 'human' (human-readable) or 'plain' (no formatting)")
+	cmd.Flags().StringSliceVarP(&headers, "header", "", nil, "Add a header to the HTTP request, on the format 'Header: Value'. This can be specified multiple times")
 	cmd.Flags().IntVarP(&queryTimeoutSecs, "timeout", "T", 10, "Timeout for the query in seconds")
 	cli.bindWaitFlag(cmd, 0, &waitSecs)
 	return cmd
@@ -63,7 +67,21 @@ func printCurl(stderr io.Writer, url string, service *vespa.Service) error {
 	return err
 }
 
-func query(cli *CLI, arguments []string, timeoutSecs, waitSecs int, curl bool, format string) error {
+func parseHeaders(headers []string) (http.Header, error) {
+	h := make(http.Header)
+	for _, header := range headers {
+		kv := strings.SplitN(header, ":", 2)
+		if len(kv) < 2 {
+			return nil, fmt.Errorf("invalid header %q: missing colon separator", header)
+		}
+		k := kv[0]
+		v := strings.TrimSpace(kv[1])
+		h.Add(k, v)
+	}
+	return h, nil
+}
+
+func query(cli *CLI, arguments []string, timeoutSecs, waitSecs int, curl bool, format string, headers []string) error {
 	target, err := cli.target(targetOptions{})
 	if err != nil {
 		return err
@@ -100,7 +118,11 @@ func query(cli *CLI, arguments []string, timeoutSecs, waitSecs int, curl bool, f
 			return err
 		}
 	}
-	response, err := service.Do(&http.Request{URL: url}, deadline+time.Second) // Slightly longer than query timeout
+	header, err := parseHeaders(headers)
+	if err != nil {
+		return err
+	}
+	response, err := service.Do(&http.Request{Header: header, URL: url}, deadline+time.Second) // Slightly longer than query timeout
 	if err != nil {
 		return fmt.Errorf("request failed: %w", err)
 	}
diff --git a/client/go/internal/cli/cmd/query_test.go b/client/go/internal/cli/cmd/query_test.go
index 3a2eeba159a..af0ae1763e9 100644
--- a/client/go/internal/cli/cmd/query_test.go
+++ b/client/go/internal/cli/cmd/query_test.go
@@ -80,6 +80,26 @@ func TestServerError(t *testing.T) {
 	assertQueryServiceError(t, 501, "server error message")
 }
 
+func TestQueryHeader(t *testing.T) {
+	client := &mock.HTTPClient{}
+	client.NextResponseString(200, "{\"query\":\"result\"}")
+
+	cli, _, stderr := newTestCLI(t)
+	cli.httpClient = client
+
+	assert.Nil(t, cli.Run("-t", "http://127.0.0.1:8080", "query",
+		"--header", "X-Foo: bar",
+		"--header", "X-Foo: baz",
+		"--header", "X-Bar:   foo bar  ",
+		"select from sources * where title contains 'foo'"))
+	assert.Equal(t, []string{"bar", "baz"}, client.LastRequest.Header.Values("X-Foo"))
+	assert.Equal(t, "foo bar", client.LastRequest.Header.Get("X-Bar"))
+
+	assert.NotNil(t, cli.Run("-t", "http://127.0.0.1:8080", "query",
+		"--header", "X-Foo", "select from sources * where title contains 'foo'"))
+	assert.Equal(t, "Error: invalid header \"X-Foo\": missing colon separator\n", stderr.String())
+}
+
 func TestStreamingQuery(t *testing.T) {
 	body := `
 event: token