Set a default query timeout

2 files changed, 19 insertions, 8 deletions

diff --git a/client/go/cmd/query.go b/client/go/cmd/query.go
index 6638c275330..7aa24cf4768 100644
--- a/client/go/cmd/query.go
+++ b/client/go/cmd/query.go
@@ -5,6 +5,7 @@
 package cmd
 
 import (
+	"fmt"
 	"log"
 	"net/http"
 	"net/url"
@@ -19,7 +20,7 @@ var queryTimeoutSecs int
 
 func init() {
 	rootCmd.AddCommand(queryCmd)
-	queryCmd.Flags().IntVarP(&queryTimeoutSecs, "timeout", "T", 10, "Timeout for the query request in seconds")
+	queryCmd.Flags().IntVarP(&queryTimeoutSecs, "timeout", "T", 10, "Timeout for the query in seconds")
 }
 
 var queryCmd = &cobra.Command{
@@ -46,9 +47,19 @@ func query(arguments []string) {
 		key, value := splitArg(arguments[i])
 		urlQuery.Set(key, value)
 	}
+	queryTimeout := urlQuery.Get("timeout")
+	if queryTimeout == "" {
+		// No timeout set by user, use the timeout option
+		queryTimeout = fmt.Sprintf("%ds", queryTimeoutSecs)
+		urlQuery.Set("timeout", queryTimeout)
+	}
 	url.RawQuery = urlQuery.Encode()
-
-	response, err := service.Do(&http.Request{URL: url}, time.Second*time.Duration(queryTimeoutSecs))
+	deadline, err := time.ParseDuration(queryTimeout)
+	if err != nil {
+		fatalErr(err, "Invalid query timeout")
+		return
+	}
+	response, err := service.Do(&http.Request{URL: url}, deadline+time.Second) // Slightly longer than query timeout
 	if err != nil {
 		fatalErr(nil, "Request failed: ", err)
 		return
diff --git a/client/go/cmd/query_test.go b/client/go/cmd/query_test.go
index 55046ae49ba..57c309c8cb3 100644
--- a/client/go/cmd/query_test.go
+++ b/client/go/cmd/query_test.go
@@ -13,25 +13,25 @@ import (
 
 func TestQuery(t *testing.T) {
 	assertQuery(t,
-		"?yql=select+from+sources+%2A+where+title+contains+%27foo%27",
+		"?timeout=10s&yql=select+from+sources+%2A+where+title+contains+%27foo%27",
 		"select from sources * where title contains 'foo'")
 }
 
 func TestQueryNonJsonResult(t *testing.T) {
 	assertQuery(t,
-		"?yql=select+from+sources+%2A+where+title+contains+%27foo%27",
+		"?timeout=10s&yql=select+from+sources+%2A+where+title+contains+%27foo%27",
 		"select from sources * where title contains 'foo'")
 }
 
 func TestQueryWithMultipleParameters(t *testing.T) {
 	assertQuery(t,
-		"?hits=5&yql=select+from+sources+%2A+where+title+contains+%27foo%27",
-		"select from sources * where title contains 'foo'", "hits=5")
+		"?hits=5&timeout=20s&yql=select+from+sources+%2A+where+title+contains+%27foo%27",
+		"select from sources * where title contains 'foo'", "hits=5", "timeout=20s")
 }
 
 func TestQueryWithExplicitYqlParameter(t *testing.T) {
 	assertQuery(t,
-		"?yql=select+from+sources+%2A+where+title+contains+%27foo%27",
+		"?timeout=10s&yql=select+from+sources+%2A+where+title+contains+%27foo%27",
 		"yql=select from sources * where title contains 'foo'")
 }