diff options
author | Martin Polden <mpolden@mpolden.no> | 2024-03-15 13:21:57 +0100 |
---|---|---|
committer | Martin Polden <mpolden@mpolden.no> | 2024-03-15 13:25:01 +0100 |
commit | 734d6487482fca2203317f2ef9903ada794bf684 (patch) | |
tree | 8567d281abb77146217e4caa9ce4fb37e8707952 /client/go | |
parent | c44558af0a82be9dda1f1ba8bf84337789a906e5 (diff) |
CLI: Add query header option
Diffstat (limited to 'client/go')
-rw-r--r-- | client/go/internal/cli/cmd/query.go | 38 | ||||
-rw-r--r-- | client/go/internal/cli/cmd/query_test.go | 20 |
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 |