Update QueryRange API to support timeouts

Signed-off-by: Joseph Woodward <joseph.woodward@xeuse.com>
prometheus · Apr 16, 2022 · 93ecf8f · 93ecf8f
1 parent e2f36d2
commit 93ecf8f
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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)