From e812ec5c46c24a6108e3c41c620cf56d7c8221aa Mon Sep 17 00:00:00 2001
From: Joseph Woodward <joseph.woodward@xeuse.com>
Date: Sat, 16 Apr 2022 02:53:10 +0100
Subject: [PATCH] Update QueryRange API to support timeouts

---
 api/prometheus/v1/api.go          | 14 ++++++++++++--
 api/prometheus/v1/api_test.go     | 11 ++++++-----
 api/prometheus/v1/example_test.go |  2 +-
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/api/prometheus/v1/api.go b/api/prometheus/v1/api.go
index 3416b2b59..83b0efee7 100644
--- a/api/prometheus/v1/api.go
+++ b/api/prometheus/v1/api.go
@@ -240,7 +240,7 @@ type API interface {
 	// Query performs a query for the given time.
 	Query(ctx context.Context, query string, ts time.Time, opts ...Option) (model.Value, Warnings, error)
 	// QueryRange performs a query for the given range.
-	QueryRange(ctx context.Context, query string, r Range) (model.Value, Warnings, error)
+	QueryRange(ctx context.Context, query string, r Range, opts ...Option) (model.Value, Warnings, error)
 	// QueryExemplars performs a query for exemplars by the given query and time range.
 	QueryExemplars(ctx context.Context, query string, startTime time.Time, endTime time.Time) ([]ExemplarQueryResult, error)
 	// Buildinfo returns various build information properties about the Prometheus server
@@ -859,7 +859,7 @@ func (h *httpAPI) Query(ctx context.Context, query string, ts time.Time, opts ..
 	return model.Value(qres.v), warnings, json.Unmarshal(body, &qres)
 }
 
-func (h *httpAPI) QueryRange(ctx context.Context, query string, r Range) (model.Value, Warnings, error) {
+func (h *httpAPI) QueryRange(ctx context.Context, query string, r Range, opts ...Option) (model.Value, Warnings, error) {
 	u := h.client.URL(epQueryRange, nil)
 	q := u.Query()
 
@@ -868,6 +868,16 @@ func (h *httpAPI) QueryRange(ctx context.Context, query string, r Range) (model.
 	q.Set("end", formatTime(r.End))
 	q.Set("step", strconv.FormatFloat(r.Step.Seconds(), 'f', -1, 64))
 
+	opt := &apiOptions{}
+	for _, o := range opts {
+		o(opt)
+	}
+
+	d := opt.timeout
+	if d > 0 {
+		q.Set("timeout", d.String())
+	}
+
 	_, body, warnings, err := h.client.DoGetFallback(ctx, u, q)
 	if err != nil {
 		return nil, warnings, err
diff --git a/api/prometheus/v1/api_test.go b/api/prometheus/v1/api_test.go
index ab88573a9..9c4f0d542 100644
--- a/api/prometheus/v1/api_test.go
+++ b/api/prometheus/v1/api_test.go
@@ -176,9 +176,9 @@ func TestAPIs(t *testing.T) {
 		}
 	}
 
-	doQueryRange := func(q string, rng Range) func() (interface{}, Warnings, error) {
+	doQueryRange := func(q string, rng Range, opts ...Option) func() (interface{}, Warnings, error) {
 		return func() (interface{}, Warnings, error) {
-			return promAPI.QueryRange(context.Background(), q, rng)
+			return promAPI.QueryRange(context.Background(), q, rng, opts...)
 		}
 	}
 
@@ -258,8 +258,9 @@ func TestAPIs(t *testing.T) {
 			reqMethod: "POST",
 			reqPath:   "/api/v1/query",
 			reqParam: url.Values{
-				"query": []string{"2"},
-				"time":  []string{testTime.Format(time.RFC3339Nano)},
+				"query":   []string{"2"},
+				"time":    []string{testTime.Format(time.RFC3339Nano)},
+				"timeout": []string{"s"},
 			},
 			res: &model.Scalar{
 				Value:     2,
@@ -365,7 +366,7 @@ func TestAPIs(t *testing.T) {
 				Start: testTime.Add(-time.Minute),
 				End:   testTime,
 				Step:  time.Minute,
-			}),
+			}, WithTimeout(5*time.Second)),
 			inErr: fmt.Errorf("some error"),
 
 			reqMethod: "POST",
diff --git a/api/prometheus/v1/example_test.go b/api/prometheus/v1/example_test.go
index 01b2a194e..f0ee3b667 100644
--- a/api/prometheus/v1/example_test.go
+++ b/api/prometheus/v1/example_test.go
@@ -67,7 +67,7 @@ func ExampleAPI_queryRange() {
 		End:   time.Now(),
 		Step:  time.Minute,
 	}
-	result, warnings, err := v1api.QueryRange(ctx, "rate(prometheus_tsdb_head_samples_appended_total[5m])", r)
+	result, warnings, err := v1api.QueryRange(ctx, "rate(prometheus_tsdb_head_samples_appended_total[5m])", r, v1.WithTimeout(5*time.Second))
 	if err != nil {
 		fmt.Printf("Error querying Prometheus: %v\n", err)
 		os.Exit(1)