Merge pull request #19378 from vespa-engine/mpolden/timeout-flag

Add timeout flag to document and query commands

6 files changed, 51 insertions, 32 deletions

diff --git a/client/go/cmd/document.go b/client/go/cmd/document.go
index 1d27a475172..cd0170684cf 100644
--- a/client/go/cmd/document.go
+++ b/client/go/cmd/document.go
@@ -9,13 +9,17 @@ import (
 	"io"
 	"io/ioutil"
 	"strings"
+	"time"
 
 	"github.com/spf13/cobra"
 	"github.com/vespa-engine/vespa/client/go/util"
 	"github.com/vespa-engine/vespa/client/go/vespa"
 )
 
-var printCurl bool
+var (
+	printCurl      bool
+	docTimeoutSecs int
+)
 
 func init() {
 	rootCmd.AddCommand(documentCmd)
@@ -24,6 +28,7 @@ func init() {
 	documentCmd.AddCommand(documentRemoveCmd)
 	documentCmd.AddCommand(documentGetCmd)
 	documentCmd.PersistentFlags().BoolVarP(&printCurl, "verbose", "v", false, "Print the equivalent curl command for the document operation")
+	documentCmd.PersistentFlags().IntVarP(&docTimeoutSecs, "timeout", "T", 60, "Timeout for the document request in seconds")
 }
 
 var documentCmd = &cobra.Command{
@@ -43,7 +48,7 @@ should be used instead of this.`,
 	DisableAutoGenTag: true,
 	Args:              cobra.ExactArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
-		printResult(vespa.Send(args[0], documentService(), curlOutput()), false)
+		printResult(vespa.Send(args[0], documentService(), operationOptions()), false)
 	},
 }
 
@@ -59,9 +64,9 @@ $ vespa document put id:mynamespace:music::a-head-full-of-dreams src/test/resour
 	DisableAutoGenTag: true,
 	Run: func(cmd *cobra.Command, args []string) {
 		if len(args) == 1 {
-			printResult(vespa.Put("", args[0], documentService(), curlOutput()), false)
+			printResult(vespa.Put("", args[0], documentService(), operationOptions()), false)
 		} else {
-			printResult(vespa.Put(args[0], args[1], documentService(), curlOutput()), false)
+			printResult(vespa.Put(args[0], args[1], documentService(), operationOptions()), false)
 		}
 	},
 }
@@ -77,9 +82,9 @@ $ vespa document update id:mynamespace:music::a-head-full-of-dreams src/test/res
 	DisableAutoGenTag: true,
 	Run: func(cmd *cobra.Command, args []string) {
 		if len(args) == 1 {
-			printResult(vespa.Update("", args[0], documentService(), curlOutput()), false)
+			printResult(vespa.Update("", args[0], documentService(), operationOptions()), false)
 		} else {
-			printResult(vespa.Update(args[0], args[1], documentService(), curlOutput()), false)
+			printResult(vespa.Update(args[0], args[1], documentService(), operationOptions()), false)
 		}
 	},
 }
@@ -95,9 +100,9 @@ $ vespa document remove id:mynamespace:music::a-head-full-of-dreams`,
 	DisableAutoGenTag: true,
 	Run: func(cmd *cobra.Command, args []string) {
 		if strings.HasPrefix(args[0], "id:") {
-			printResult(vespa.RemoveId(args[0], documentService(), curlOutput()), false)
+			printResult(vespa.RemoveId(args[0], documentService(), operationOptions()), false)
 		} else {
-			printResult(vespa.RemoveOperation(args[0], documentService(), curlOutput()), false)
+			printResult(vespa.RemoveOperation(args[0], documentService(), operationOptions()), false)
 		}
 	},
 }
@@ -109,12 +114,19 @@ var documentGetCmd = &cobra.Command{
 	DisableAutoGenTag: true,
 	Example:           `$ vespa document get id:mynamespace:music::a-head-full-of-dreams`,
 	Run: func(cmd *cobra.Command, args []string) {
-		printResult(vespa.Get(args[0], documentService(), curlOutput()), true)
+		printResult(vespa.Get(args[0], documentService(), operationOptions()), true)
 	},
 }
 
 func documentService() *vespa.Service { return getService("document", 0) }
 
+func operationOptions() vespa.OperationOptions {
+	return vespa.OperationOptions{
+		CurlOutput: curlOutput(),
+		Timeout:    time.Second * time.Duration(docTimeoutSecs),
+	}
+}
+
 func curlOutput() io.Writer {
 	if printCurl {
 		return stderr
diff --git a/client/go/cmd/query.go b/client/go/cmd/query.go
index 5e2b268865d..66e85763899 100644
--- a/client/go/cmd/query.go
+++ b/client/go/cmd/query.go
@@ -15,8 +15,11 @@ import (
 	"github.com/vespa-engine/vespa/client/go/util"
 )
 
+var queryTimeoutSecs int
+
 func init() {
 	rootCmd.AddCommand(queryCmd)
+	queryCmd.Flags().IntVarP(&queryTimeoutSecs, "timeout", "T", 10, "Timeout for the query request in seconds")
 }
 
 var queryCmd = &cobra.Command{
@@ -45,7 +48,7 @@ func query(arguments []string) {
 	}
 	url.RawQuery = urlQuery.Encode()
 
-	response, err := service.Do(&http.Request{URL: url}, time.Second*10)
+	response, err := service.Do(&http.Request{URL: url}, time.Second*time.Duration(queryTimeoutSecs))
 	if err != nil {
 		printErr(nil, "Request failed: ", err)
 		return
diff --git a/client/go/cmd/root.go b/client/go/cmd/root.go
index 7fe704e4918..ab6c19a665b 100644
--- a/client/go/cmd/root.go
+++ b/client/go/cmd/root.go
@@ -10,14 +10,13 @@ import (
 	"log"
 	"os"
 
-	"github.com/logrusorgru/aurora"
+	"github.com/logrusorgru/aurora/v3"
 	"github.com/mattn/go-colorable"
 	"github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
 )
 
 var (
-	// TODO: add timeout flag
 	rootCmd = &cobra.Command{
 		Use:   "vespa command-name",
 		Short: "The command-line tool for Vespa.ai",
diff --git a/client/go/go.mod b/client/go/go.mod
index 509eb273c6c..27faff3fd0b 100644
--- a/client/go/go.mod
+++ b/client/go/go.mod
@@ -4,7 +4,7 @@ go 1.15
 
 require (
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
-	github.com/logrusorgru/aurora v2.0.3+incompatible
+	github.com/logrusorgru/aurora/v3 v3.0.0
 	github.com/mattn/go-colorable v0.0.9
 	github.com/mattn/go-isatty v0.0.3
 	github.com/spf13/cobra v1.2.1
diff --git a/client/go/go.sum b/client/go/go.sum
index 97328690ee5..0462d0575a1 100644
--- a/client/go/go.sum
+++ b/client/go/go.sum
@@ -180,8 +180,8 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
-github.com/logrusorgru/aurora v2.0.3+incompatible h1:tOpm7WcpBTn4fjmVfgpQq0EfczGlG91VSDkswnjF5A8=
-github.com/logrusorgru/aurora v2.0.3+incompatible/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=
+github.com/logrusorgru/aurora/v3 v3.0.0 h1:R6zcoZZbvVcGMvDCKo45A9U/lzYyzl5NfYIvznmDfE4=
+github.com/logrusorgru/aurora/v3 v3.0.0/go.mod h1:vsR12bk5grlLvLXAYrBsb5Oc/N+LxAlxggSjiwMnCUc=
 github.com/magiconair/properties v1.8.5 h1:b6kJs+EmPFMYGkow9GiUyCyOvIwYetYJ3fSaWak/Gls=
 github.com/magiconair/properties v1.8.5/go.mod h1:y3VJvCyxH9uVvJTWEGAELF3aiYNyPKd5NZ3oSwXrF60=
 github.com/mattn/go-colorable v0.0.9 h1:UVL0vNpWh04HeJXV0KLcaT7r06gOH2l4OW6ddYRUIY4=
diff --git a/client/go/vespa/document.go b/client/go/vespa/document.go
index cfac1930199..5e01d180b5f 100644
--- a/client/go/vespa/document.go
+++ b/client/go/vespa/document.go
@@ -19,24 +19,24 @@ import (
 )
 
 // Sends the operation given in the file
-func Send(jsonFile string, service *Service, curlOutput io.Writer) util.OperationResult {
-	return sendOperation("", jsonFile, service, anyOperation, curlOutput)
+func Send(jsonFile string, service *Service, options OperationOptions) util.OperationResult {
+	return sendOperation("", jsonFile, service, anyOperation, options)
 }
 
-func Put(documentId string, jsonFile string, service *Service, curlOutput io.Writer) util.OperationResult {
-	return sendOperation(documentId, jsonFile, service, putOperation, curlOutput)
+func Put(documentId string, jsonFile string, service *Service, options OperationOptions) util.OperationResult {
+	return sendOperation(documentId, jsonFile, service, putOperation, options)
 }
 
-func Update(documentId string, jsonFile string, service *Service, curlOutput io.Writer) util.OperationResult {
-	return sendOperation(documentId, jsonFile, service, updateOperation, curlOutput)
+func Update(documentId string, jsonFile string, service *Service, options OperationOptions) util.OperationResult {
+	return sendOperation(documentId, jsonFile, service, updateOperation, options)
 }
 
-func RemoveId(documentId string, service *Service, curlOutput io.Writer) util.OperationResult {
-	return sendOperation(documentId, "", service, removeOperation, curlOutput)
+func RemoveId(documentId string, service *Service, options OperationOptions) util.OperationResult {
+	return sendOperation(documentId, "", service, removeOperation, options)
 }
 
-func RemoveOperation(jsonFile string, service *Service, curlOutput io.Writer) util.OperationResult {
-	return sendOperation("", jsonFile, service, removeOperation, curlOutput)
+func RemoveOperation(jsonFile string, service *Service, options OperationOptions) util.OperationResult {
+	return sendOperation("", jsonFile, service, removeOperation, options)
 }
 
 const (
@@ -46,7 +46,12 @@ const (
 	removeOperation string = "remove"
 )
 
-func sendOperation(documentId string, jsonFile string, service *Service, operation string, curlOutput io.Writer) util.OperationResult {
+type OperationOptions struct {
+	CurlOutput io.Writer
+	Timeout    time.Duration
+}
+
+func sendOperation(documentId string, jsonFile string, service *Service, operation string, options OperationOptions) util.OperationResult {
 	header := http.Header{}
 	header.Add("Content-Type", "application/json")
 
@@ -95,7 +100,7 @@ func sendOperation(documentId string, jsonFile string, service *Service, operati
 		Header: header,
 		Body:   ioutil.NopCloser(bytes.NewReader(documentData)),
 	}
-	response, err := serviceDo(service, request, jsonFile, curlOutput)
+	response, err := serviceDo(service, request, jsonFile, options)
 	if response == nil {
 		return util.Failure("Request failed: " + err.Error())
 	}
@@ -134,7 +139,7 @@ func operationToHTTPMethod(operation string) string {
 	panic("Unexpected document operation ''" + operation + "'")
 }
 
-func serviceDo(service *Service, request *http.Request, filename string, curlOutput io.Writer) (*http.Response, error) {
+func serviceDo(service *Service, request *http.Request, filename string, options OperationOptions) (*http.Response, error) {
 	cmd, err := curl.RawArgs(request.URL.String())
 	if err != nil {
 		return nil, err
@@ -149,13 +154,13 @@ func serviceDo(service *Service, request *http.Request, filename string, curlOut
 	cmd.Certificate = service.TLSOptions.CertificateFile
 	cmd.PrivateKey = service.TLSOptions.PrivateKeyFile
 	out := cmd.String() + "\n"
-	if _, err := io.WriteString(curlOutput, out); err != nil {
+	if _, err := io.WriteString(options.CurlOutput, out); err != nil {
 		return nil, err
 	}
-	return service.Do(request, time.Second*60)
+	return service.Do(request, options.Timeout)
 }
 
-func Get(documentId string, service *Service, curlOutput io.Writer) util.OperationResult {
+func Get(documentId string, service *Service, options OperationOptions) util.OperationResult {
 	documentPath, documentPathError := IdToURLPath(documentId)
 	if documentPathError != nil {
 		return util.Failure("Invalid document id '" + documentId + "': " + documentPathError.Error())
@@ -170,7 +175,7 @@ func Get(documentId string, service *Service, curlOutput io.Writer) util.Operati
 		URL:    url,
 		Method: "GET",
 	}
-	response, err := serviceDo(service, request, "", curlOutput)
+	response, err := serviceDo(service, request, "", options)
 	if response == nil {
 		return util.Failure("Request failed: " + err.Error())
 	}