From 48a686a6033f5f500259d8acf1618bf76b11ddfe Mon Sep 17 00:00:00 2001 From: Joseph Woodward Date: Thu, 21 Apr 2022 06:23:16 +0100 Subject: [PATCH] Update query API to support timeouts (#1014) * Add timeout parameter for queries Signed-off-by: Joseph Woodward * Update api/prometheus/v1/api.go Co-authored-by: Kemal Akkoyun Signed-off-by: Joseph Woodward * Update api/prometheus/v1/api.go Co-authored-by: Kemal Akkoyun Signed-off-by: Joseph Woodward * Pass timeout as stringified time.Duration instead of millisecond value Signed-off-by: Joseph Woodward * Update QueryRange API to support timeouts Signed-off-by: Joseph Woodward * Add timeout to test request params Signed-off-by: Joseph Woodward Co-authored-by: Kemal Akkoyun --- api/prometheus/v1/api.go | 41 ++++++++++++++++++++++++++++--- api/prometheus/v1/api_test.go | 26 +++++++++++--------- api/prometheus/v1/example_test.go | 4 +-- 3 files changed, 53 insertions(+), 18 deletions(-) diff --git a/api/prometheus/v1/api.go b/api/prometheus/v1/api.go index 5ed091225..83b0efee7 100644 --- a/api/prometheus/v1/api.go +++ b/api/prometheus/v1/api.go @@ -238,9 +238,9 @@ 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) + 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 @@ -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) + } + + d := opt.timeout + if d > 0 { + q.Set("timeout", d.String()) + } + q.Set("query", query) if !ts.IsZero() { q.Set("time", formatTime(ts)) @@ -836,7 +859,7 @@ func (h *httpAPI) Query(ctx context.Context, query string, ts time.Time) (model. 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() @@ -845,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 e7d0b4c96..82c234bbd 100644 --- a/api/prometheus/v1/api_test.go +++ b/api/prometheus/v1/api_test.go @@ -170,15 +170,15 @@ 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...) } } - 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...) } } @@ -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{ @@ -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{(5 * time.Second).String()}, }, res: &model.Scalar{ Value: 2, @@ -365,16 +366,17 @@ 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", reqPath: "/api/v1/query_range", reqParam: url.Values{ - "query": []string{"2"}, - "start": []string{testTime.Add(-time.Minute).Format(time.RFC3339Nano)}, - "end": []string{testTime.Format(time.RFC3339Nano)}, - "step": []string{time.Minute.String()}, + "query": []string{"2"}, + "start": []string{testTime.Add(-time.Minute).Format(time.RFC3339Nano)}, + "end": []string{testTime.Format(time.RFC3339Nano)}, + "step": []string{time.Minute.String()}, + "timeout": []string{(5 * time.Second).String()}, }, err: fmt.Errorf("some error"), }, diff --git a/api/prometheus/v1/example_test.go b/api/prometheus/v1/example_test.go index 818290262..f0ee3b667 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) @@ -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)