From 80099dffa14624b2c88911d486aa98b170cec101 Mon Sep 17 00:00:00 2001
From: Joseph Woodward <joseph.woodward@xeuse.com>
Date: Mon, 28 Mar 2022 06:23:29 +0100
Subject: [PATCH] Add timeout parameter for queries

Signed-off-by: Joseph Woodward <joseph.woodward@xeuse.com>
---
 api/prometheus/v1/api.go          | 27 +++++++++++++++++++++++++--
 api/prometheus/v1/api_test.go     |  6 +++---
 api/prometheus/v1/example_test.go |  2 +-
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/api/prometheus/v1/api.go b/api/prometheus/v1/api.go
index 5ed091225..1de5040dd 100644
--- a/api/prometheus/v1/api.go
+++ b/api/prometheus/v1/api.go
@@ -238,7 +238,7 @@ type API interface {
 	// LabelValues performs a query for the values of the given label, time range and matchers.
 	LabelValues(ctx context.Context, label string, matches []string, startTime time.Time, endTime time.Time) (model.LabelValues, Warnings, error)
 	// Query performs a query for the given time.
-	Query(ctx context.Context, query string, ts time.Time) (model.Value, Warnings, error)
+	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)
 	// QueryExemplars performs a query for exemplars by the given query and time range.
@@ -818,10 +818,33 @@ func (h *httpAPI) LabelValues(ctx context.Context, label string, matches []strin
 	return labelValues, w, json.Unmarshal(body, &labelValues)
 }
 
-func (h *httpAPI) Query(ctx context.Context, query string, ts time.Time) (model.Value, Warnings, error) {
+type apiOptions struct {
+	timeout time.Duration
+}
+
+type Option func(c *apiOptions)
+
+func WithTimeout(timeout time.Duration) Option {
+	return func(o *apiOptions) {
+		o.timeout = timeout
+	}
+}
+
+func (h *httpAPI) Query(ctx context.Context, query string, ts time.Time, opts ...Option) (model.Value, Warnings, error) {
+
 	u := h.client.URL(epQuery, nil)
 	q := u.Query()
 
+	opt := &apiOptions{}
+	for _, o := range opts {
+		o(opt)
+	}
+
+	t := opt.timeout.Milliseconds()
+	if t > 0 {
+		q.Set("timeout", fmt.Sprint(t))
+	}
+
 	q.Set("query", query)
 	if !ts.IsZero() {
 		q.Set("time", formatTime(ts))
diff --git a/api/prometheus/v1/api_test.go b/api/prometheus/v1/api_test.go
index e7d0b4c96..ab88573a9 100644
--- a/api/prometheus/v1/api_test.go
+++ b/api/prometheus/v1/api_test.go
@@ -170,9 +170,9 @@ func TestAPIs(t *testing.T) {
 		}
 	}
 
-	doQuery := func(q string, ts time.Time) func() (interface{}, Warnings, error) {
+	doQuery := func(q string, ts time.Time, opts ...Option) func() (interface{}, Warnings, error) {
 		return func() (interface{}, Warnings, error) {
-			return promAPI.Query(context.Background(), q, ts)
+			return promAPI.Query(context.Background(), q, ts, opts...)
 		}
 	}
 
@@ -246,7 +246,7 @@ func TestAPIs(t *testing.T) {
 
 	queryTests := []apiTest{
 		{
-			do: doQuery("2", testTime),
+			do: doQuery("2", testTime, WithTimeout(5*time.Second)),
 			inRes: &queryResult{
 				Type: model.ValScalar,
 				Result: &model.Scalar{
diff --git a/api/prometheus/v1/example_test.go b/api/prometheus/v1/example_test.go
index 818290262..01b2a194e 100644
--- a/api/prometheus/v1/example_test.go
+++ b/api/prometheus/v1/example_test.go
@@ -39,7 +39,7 @@ func ExampleAPI_query() {
 	v1api := v1.NewAPI(client)
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-	result, warnings, err := v1api.Query(ctx, "up", time.Now())
+	result, warnings, err := v1api.Query(ctx, "up", time.Now(), v1.WithTimeout(5*time.Second))
 	if err != nil {
 		fmt.Printf("Error querying Prometheus: %v\n", err)
 		os.Exit(1)