From 61f94d8f2182e60f39ea49ed8103c3efe3a2ba0b Mon Sep 17 00:00:00 2001 From: Arne Juul Date: Tue, 12 Nov 2024 13:37:21 +0100 Subject: [PATCH 1/3] add the --postFile option to "vespa query" --- client/go/internal/cli/cmd/query.go | 89 +++++++++++++++++++++++++---- 1 file changed, 77 insertions(+), 12 deletions(-) diff --git a/client/go/internal/cli/cmd/query.go b/client/go/internal/cli/cmd/query.go index 5fa225777f02..1c1d85fccef1 100644 --- a/client/go/internal/cli/cmd/query.go +++ b/client/go/internal/cli/cmd/query.go @@ -5,11 +5,13 @@ package cmd import ( + "bytes" "encoding/json" "fmt" "io" "net/http" "net/url" + "os" "strings" "time" @@ -28,6 +30,7 @@ func newQueryCmd(cli *CLI) *cobra.Command { queryTimeoutSecs int waitSecs int format string + postFile string headers []string ) cmd := &cobra.Command{ @@ -43,13 +46,17 @@ can be set by the syntax [parameter-name]=[value].`, // TODO: Support referencing a query json file DisableAutoGenTag: true, SilenceUsage: true, - Args: cobra.MinimumNArgs(1), + Args: cobra.MinimumNArgs(0), RunE: func(cmd *cobra.Command, args []string) error { + if len(args) == 0 && postFile == "" { + return fmt.Errorf("requires at least 1 arg") + } waiter := cli.waiter(time.Duration(waitSecs)*time.Second, cmd) - return query(cli, args, queryTimeoutSecs, printCurl, format, headers, waiter) + return query(cli, args, queryTimeoutSecs, printCurl, format, postFile, headers, waiter) }, } cmd.Flags().BoolVarP(&printCurl, "verbose", "v", false, "Print the equivalent curl command for the query") + cmd.Flags().StringVarP(&postFile, "postFile", "", "", "Use the given JSON file, POST to the query API") 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") @@ -57,18 +64,31 @@ can be set by the syntax [parameter-name]=[value].`, return cmd } -func printCurl(stderr io.Writer, url string, service *vespa.Service) error { - cmd, err := curl.RawArgs(url) +func printCurl(stderr io.Writer, req *http.Request, postFile string, service *vespa.Service) error { + cmd, err := curl.RawArgs(req.URL.String()) if err != nil { return err } + cmd.Method = req.Method + if postFile != "" { + cmd.WithBodyFile(postFile) + } + for k, vl := range req.Header { + for _, v := range vl { + cmd.Header(k, v) + } + } cmd.Certificate = service.TLSOptions.CertificateFile cmd.PrivateKey = service.TLSOptions.PrivateKeyFile _, err = io.WriteString(stderr, cmd.String()+"\n") return err } -func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format string, headers []string, waiter *Waiter) error { +type nopCloser struct{ io.Reader } + +func (nopCloser) Close() error { return nil } + +func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format string, postFile string, headers []string, waiter *Waiter) error { target, err := cli.target(targetOptions{}) if err != nil { return err @@ -83,6 +103,9 @@ func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format stri return fmt.Errorf("invalid format: %s", format) } url, _ := url.Parse(service.BaseURL + "/search/") + if strings.HasSuffix(service.BaseURL, "/") { + url, _ = url.Parse(service.BaseURL + "search/") + } urlQuery := url.Query() for i := range len(arguments) { key, value := splitArg(arguments[i]) @@ -94,21 +117,34 @@ func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format stri queryTimeout = fmt.Sprintf("%ds", timeoutSecs) urlQuery.Set("timeout", queryTimeout) } - url.RawQuery = urlQuery.Encode() deadline, err := time.ParseDuration(queryTimeout) if err != nil { return fmt.Errorf("invalid query timeout: %w", err) } - if curl { - if err := printCurl(cli.Stderr, url.String(), service); err != nil { - return err - } - } header, err := httputil.ParseHeader(headers) if err != nil { return err } - response, err := service.Do(&http.Request{Header: header, URL: url}, deadline+time.Second) // Slightly longer than query timeout + hReq := &http.Request{Header: header, URL: url} + if postFile != "" { + json, err := getJsonFrom(postFile, urlQuery) + if err != nil { + return fmt.Errorf("bad JSON in postFile '%s': %w", postFile, err) + } + header.Set("Content-Type", "application/json") + hReq.Method = "POST" + hReq.Body = nopCloser{bytes.NewBuffer(bytes.Clone(json))} + if err != nil { + return fmt.Errorf("bad postFile '%s': %w", postFile, err) + } + } + url.RawQuery = urlQuery.Encode() + if curl { + if err := printCurl(cli.Stderr, hReq, postFile, service); err != nil { + return err + } + } + response, err := service.Do(hReq, deadline+time.Second) // Slightly longer than query timeout if err != nil { return fmt.Errorf("request failed: %w", err) } @@ -207,3 +243,32 @@ func splitArg(argument string) (string, string) { } return parts[0], parts[1] } + +func getJsonFrom(fn string, query url.Values) ([]byte, error) { + parsed := make(map[string]any) + f, err := os.Open(fn) + if err != nil { + return nil, err + } + body, err := io.ReadAll(f) + if err != nil { + return nil, err + } + err = json.Unmarshal(body, &parsed) + if err != nil { + return nil, err + } + for k, vl := range query { + if len(vl) == 1 { + parsed[k] = vl[0] + } else { + parsed[k] = vl + } + query.Del(k) + } + b, err := json.Marshal(parsed) + if err != nil { + return nil, err + } + return b, nil +} From 1edbc890153d2664116f715cafeaceb0eb827ed5 Mon Sep 17 00:00:00 2001 From: Arne Juul Date: Wed, 13 Nov 2024 09:21:51 +0100 Subject: [PATCH 2/3] udpates after review: * rename "--postFile" to just "--file" * add unit tests for "--file" * reword help description --- client/go/internal/cli/cmd/query.go | 13 ++----- client/go/internal/cli/cmd/query_test.go | 48 ++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 10 deletions(-) diff --git a/client/go/internal/cli/cmd/query.go b/client/go/internal/cli/cmd/query.go index 1c1d85fccef1..47c1635f3218 100644 --- a/client/go/internal/cli/cmd/query.go +++ b/client/go/internal/cli/cmd/query.go @@ -56,7 +56,7 @@ can be set by the syntax [parameter-name]=[value].`, }, } cmd.Flags().BoolVarP(&printCurl, "verbose", "v", false, "Print the equivalent curl command for the query") - cmd.Flags().StringVarP(&postFile, "postFile", "", "", "Use the given JSON file, POST to the query API") + cmd.Flags().StringVarP(&postFile, "file", "", "", "Read query parameters from the given JSON file and send a POST request, with overrides from arguments") 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") @@ -84,10 +84,6 @@ func printCurl(stderr io.Writer, req *http.Request, postFile string, service *ve return err } -type nopCloser struct{ io.Reader } - -func (nopCloser) Close() error { return nil } - func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format string, postFile string, headers []string, waiter *Waiter) error { target, err := cli.target(targetOptions{}) if err != nil { @@ -102,10 +98,7 @@ func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format stri default: return fmt.Errorf("invalid format: %s", format) } - url, _ := url.Parse(service.BaseURL + "/search/") - if strings.HasSuffix(service.BaseURL, "/") { - url, _ = url.Parse(service.BaseURL + "search/") - } + url, _ := url.Parse(strings.TrimSuffix(service.BaseURL, "/") + "/search/") urlQuery := url.Query() for i := range len(arguments) { key, value := splitArg(arguments[i]) @@ -133,7 +126,7 @@ func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format stri } header.Set("Content-Type", "application/json") hReq.Method = "POST" - hReq.Body = nopCloser{bytes.NewBuffer(bytes.Clone(json))} + hReq.Body = io.NopCloser(bytes.NewBuffer(bytes.Clone(json))) if err != nil { return fmt.Errorf("bad postFile '%s': %w", postFile, err) } diff --git a/client/go/internal/cli/cmd/query_test.go b/client/go/internal/cli/cmd/query_test.go index f5b113b6acb7..ef78335f726f 100644 --- a/client/go/internal/cli/cmd/query_test.go +++ b/client/go/internal/cli/cmd/query_test.go @@ -6,10 +6,13 @@ package cmd import ( "net/http" + "os" + "path/filepath" "strconv" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/vespa-engine/vespa/client/go/internal/mock" ) @@ -134,6 +137,51 @@ data: { assertStreamingQuery(t, bodyWithError, bodyWithError, "--format=plain") } +func TestQueryPostFile(t *testing.T) { + mockResponse := `{"query":"result"}"` + client := &mock.HTTPClient{ReadBody: true} + client.NextResponseString(200, mockResponse) + cli, stdout, _ := newTestCLI(t) + cli.httpClient = client + + tmpFileName := filepath.Join(t.TempDir(), "tq1.json") + jsonQuery := []byte(`{"yql": "some yql here"}`) + require.Nil(t, os.WriteFile(tmpFileName, jsonQuery, 0644)) + + assert.Nil(t, cli.Run("-t", "http://127.0.0.1:8080", "query", "--file", tmpFileName)) + assert.Equal(t, mockResponse+"\n", stdout.String()) + assert.Equal(t, `{"timeout":"10s","yql":"some yql here"}`, string(client.LastBody)) + assert.Equal(t, []string{"application/json"}, client.LastRequest.Header.Values("Content-Type")) + assert.Equal(t, "POST", client.LastRequest.Method) + assert.Equal(t, "http://127.0.0.1:8080/search/", client.LastRequest.URL.String()) +} + +func TestQueryPostFileWithArgs(t *testing.T) { + mockResponse := `{"query":"result"}"` + client := &mock.HTTPClient{ReadBody: true} + client.NextResponseString(200, mockResponse) + cli, _, _ := newTestCLI(t) + cli.httpClient = client + + tmpFileName := filepath.Join(t.TempDir(), "tq2.json") + jsonQuery := []byte(`{"yql": "some yql here"}`) + require.Nil(t, os.WriteFile(tmpFileName, jsonQuery, 0644)) + + assert.Nil(t, cli.Run( + "-t", "http://foo.bar:1234/", + "query", + "--file", tmpFileName, + "yql=foo bar", + "tracelevel=3", + "dispatch.docsumRetryLimit=42")) + assert.Equal(t, + `{"dispatch.docsumRetryLimit":"42","timeout":"10s","tracelevel":"3","yql":"foo bar"}`, + string(client.LastBody)) + assert.Equal(t, []string{"application/json"}, client.LastRequest.Header.Values("Content-Type")) + assert.Equal(t, "POST", client.LastRequest.Method) + assert.Equal(t, "http://foo.bar:1234/search/", client.LastRequest.URL.String()) +} + func assertStreamingQuery(t *testing.T, expectedOutput, body string, args ...string) { t.Helper() client := &mock.HTTPClient{} From e207eddd30c078326a7284a5748438ce36636d04 Mon Sep 17 00:00:00 2001 From: Arne Juul Date: Wed, 13 Nov 2024 09:31:04 +0100 Subject: [PATCH 3/3] refactor option arguments into a struct type --- client/go/internal/cli/cmd/query.go | 60 +++++++++++++++-------------- 1 file changed, 31 insertions(+), 29 deletions(-) diff --git a/client/go/internal/cli/cmd/query.go b/client/go/internal/cli/cmd/query.go index 47c1635f3218..95d8799c9b66 100644 --- a/client/go/internal/cli/cmd/query.go +++ b/client/go/internal/cli/cmd/query.go @@ -24,15 +24,17 @@ import ( "github.com/vespa-engine/vespa/client/go/internal/vespa" ) +type queryOptions struct { + printCurl bool + queryTimeoutSecs int + waitSecs int + format string + postFile string + headers []string +} + func newQueryCmd(cli *CLI) *cobra.Command { - var ( - printCurl bool - queryTimeoutSecs int - waitSecs int - format string - postFile string - headers []string - ) + opts := queryOptions{} cmd := &cobra.Command{ Use: "query query-parameters", Short: "Issue a query to Vespa", @@ -48,19 +50,19 @@ can be set by the syntax [parameter-name]=[value].`, SilenceUsage: true, Args: cobra.MinimumNArgs(0), RunE: func(cmd *cobra.Command, args []string) error { - if len(args) == 0 && postFile == "" { + if len(args) == 0 && opts.postFile == "" { return fmt.Errorf("requires at least 1 arg") } - waiter := cli.waiter(time.Duration(waitSecs)*time.Second, cmd) - return query(cli, args, queryTimeoutSecs, printCurl, format, postFile, headers, waiter) + waiter := cli.waiter(time.Duration(opts.waitSecs)*time.Second, cmd) + return query(cli, args, &opts, waiter) }, } - cmd.Flags().BoolVarP(&printCurl, "verbose", "v", false, "Print the equivalent curl command for the query") - cmd.Flags().StringVarP(&postFile, "file", "", "", "Read query parameters from the given JSON file and send a POST request, with overrides from arguments") - 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) + cmd.Flags().BoolVarP(&opts.printCurl, "verbose", "v", false, "Print the equivalent curl command for the query") + cmd.Flags().StringVarP(&opts.postFile, "file", "", "", "Read query parameters from the given JSON file and send a POST request, with overrides from arguments") + cmd.Flags().StringVarP(&opts.format, "format", "", "human", "Output format. Must be 'human' (human-readable) or 'plain' (no formatting)") + cmd.Flags().StringSliceVarP(&opts.headers, "header", "", nil, "Add a header to the HTTP request, on the format 'Header: Value'. This can be specified multiple times") + cmd.Flags().IntVarP(&opts.queryTimeoutSecs, "timeout", "T", 10, "Timeout for the query in seconds") + cli.bindWaitFlag(cmd, 0, &opts.waitSecs) return cmd } @@ -84,7 +86,7 @@ func printCurl(stderr io.Writer, req *http.Request, postFile string, service *ve return err } -func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format string, postFile string, headers []string, waiter *Waiter) error { +func query(cli *CLI, arguments []string, opts *queryOptions, waiter *Waiter) error { target, err := cli.target(targetOptions{}) if err != nil { return err @@ -93,10 +95,10 @@ func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format stri if err != nil { return err } - switch format { + switch opts.format { case "plain", "human": default: - return fmt.Errorf("invalid format: %s", format) + return fmt.Errorf("invalid format: %s", opts.format) } url, _ := url.Parse(strings.TrimSuffix(service.BaseURL, "/") + "/search/") urlQuery := url.Query() @@ -107,33 +109,33 @@ func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format stri queryTimeout := urlQuery.Get("timeout") if queryTimeout == "" { // No timeout set by user, use the timeout option - queryTimeout = fmt.Sprintf("%ds", timeoutSecs) + queryTimeout = fmt.Sprintf("%ds", opts.queryTimeoutSecs) urlQuery.Set("timeout", queryTimeout) } deadline, err := time.ParseDuration(queryTimeout) if err != nil { return fmt.Errorf("invalid query timeout: %w", err) } - header, err := httputil.ParseHeader(headers) + header, err := httputil.ParseHeader(opts.headers) if err != nil { return err } hReq := &http.Request{Header: header, URL: url} - if postFile != "" { - json, err := getJsonFrom(postFile, urlQuery) + if opts.postFile != "" { + json, err := getJsonFrom(opts.postFile, urlQuery) if err != nil { - return fmt.Errorf("bad JSON in postFile '%s': %w", postFile, err) + return fmt.Errorf("bad JSON in postFile '%s': %w", opts.postFile, err) } header.Set("Content-Type", "application/json") hReq.Method = "POST" hReq.Body = io.NopCloser(bytes.NewBuffer(bytes.Clone(json))) if err != nil { - return fmt.Errorf("bad postFile '%s': %w", postFile, err) + return fmt.Errorf("bad postFile '%s': %w", opts.postFile, err) } } url.RawQuery = urlQuery.Encode() - if curl { - if err := printCurl(cli.Stderr, hReq, postFile, service); err != nil { + if opts.printCurl { + if err := printCurl(cli.Stderr, hReq, opts.postFile, service); err != nil { return err } } @@ -144,7 +146,7 @@ func query(cli *CLI, arguments []string, timeoutSecs int, curl bool, format stri defer response.Body.Close() if response.StatusCode == 200 { - if err := printResponse(response.Body, response.Header.Get("Content-Type"), format, cli); err != nil { + if err := printResponse(response.Body, response.Header.Get("Content-Type"), opts.format, cli); err != nil { return err } } else if response.StatusCode/100 == 4 {