From 34e02e282dc4a3cb55ca6441b489ec182e654d59 Mon Sep 17 00:00:00 2001 From: Ivan Ryabov Date: Fri, 28 Jun 2024 21:42:56 +1000 Subject: [PATCH] client: Add Option to provide limit query param for APIs that support it (#1544) * client: Add Option to provide limit query param for APIs that support it Signed-off-by: Ivan Ryabov * Renamed formatOptions -> addOptionalURLParams and comment as per review feedback Signed-off-by: Ivan Ryabov --------- Signed-off-by: Ivan Ryabov --- api/prometheus/v1/api.go | 70 +++++++++++++++++++---------------- api/prometheus/v1/api_test.go | 16 ++++---- 2 files changed, 47 insertions(+), 39 deletions(-) diff --git a/api/prometheus/v1/api.go b/api/prometheus/v1/api.go index 1cfe8d8..cddf027 100644 --- a/api/prometheus/v1/api.go +++ b/api/prometheus/v1/api.go @@ -475,9 +475,9 @@ type API interface { // Flags returns the flag values that Prometheus was launched with. Flags(ctx context.Context) (FlagsResult, error) // LabelNames returns the unique label names present in the block in sorted order by given time range and matchers. - LabelNames(ctx context.Context, matches []string, startTime, endTime time.Time) ([]string, Warnings, error) + LabelNames(ctx context.Context, matches []string, startTime, endTime time.Time, opts ...Option) ([]string, Warnings, error) // LabelValues performs a query for the values of the given label, time range and matchers. - LabelValues(ctx context.Context, label string, matches []string, startTime, endTime time.Time) (model.LabelValues, Warnings, error) + LabelValues(ctx context.Context, label string, matches []string, startTime, endTime time.Time, opts ...Option) (model.LabelValues, Warnings, error) // 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. @@ -489,7 +489,7 @@ type API interface { // Runtimeinfo returns the various runtime information properties about the Prometheus server. Runtimeinfo(ctx context.Context) (RuntimeinfoResult, error) // Series finds series by label matchers. - Series(ctx context.Context, matches []string, startTime, endTime time.Time) ([]model.LabelSet, Warnings, error) + Series(ctx context.Context, matches []string, startTime, endTime time.Time, opts ...Option) ([]model.LabelSet, Warnings, error) // Snapshot creates a snapshot of all current data into snapshots/- // under the TSDB's data directory and returns the directory as response. Snapshot(ctx context.Context, skipHead bool) (SnapshotResult, error) @@ -502,7 +502,7 @@ type API interface { // Metadata returns metadata about metrics currently scraped by the metric name. Metadata(ctx context.Context, metric, limit string) (map[string][]Metadata, error) // TSDB returns the cardinality statistics. - TSDB(ctx context.Context) (TSDBResult, error) + TSDB(ctx context.Context, opts ...Option) (TSDBResult, error) // WalReplay returns the current replay status of the wal. WalReplay(ctx context.Context) (WalReplayStatus, error) } @@ -1024,9 +1024,10 @@ func (h *httpAPI) Runtimeinfo(ctx context.Context) (RuntimeinfoResult, error) { return res, err } -func (h *httpAPI) LabelNames(ctx context.Context, matches []string, startTime, endTime time.Time) ([]string, Warnings, error) { +func (h *httpAPI) LabelNames(ctx context.Context, matches []string, startTime, endTime time.Time, opts ...Option) ([]string, Warnings, error) { u := h.client.URL(epLabels, nil) - q := u.Query() + q := addOptionalURLParams(u.Query(), opts) + if !startTime.IsZero() { q.Set("start", formatTime(startTime)) } @@ -1046,9 +1047,10 @@ func (h *httpAPI) LabelNames(ctx context.Context, matches []string, startTime, e return labelNames, w, err } -func (h *httpAPI) LabelValues(ctx context.Context, label string, matches []string, startTime, endTime time.Time) (model.LabelValues, Warnings, error) { +func (h *httpAPI) LabelValues(ctx context.Context, label string, matches []string, startTime, endTime time.Time, opts ...Option) (model.LabelValues, Warnings, error) { u := h.client.URL(epLabelValues, map[string]string{"name": label}) - q := u.Query() + q := addOptionalURLParams(u.Query(), opts) + if !startTime.IsZero() { q.Set("start", formatTime(startTime)) } @@ -1076,6 +1078,7 @@ func (h *httpAPI) LabelValues(ctx context.Context, label string, matches []strin type apiOptions struct { timeout time.Duration + limit uint64 } type Option func(c *apiOptions) @@ -1088,20 +1091,35 @@ func WithTimeout(timeout time.Duration) Option { } } -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() +// WithLimit provides an optional maximum number of returned entries for APIs that support limit parameter +// e.g. https://prometheus.io/docs/prometheus/latest/querying/api/#instant-querie:~:text=%3A%20End%20timestamp.-,limit%3D%3Cnumber%3E,-%3A%20Maximum%20number%20of +func WithLimit(limit uint64) Option { + return func(o *apiOptions) { + o.limit = limit + } +} +func addOptionalURLParams(q url.Values, opts []Option) url.Values { opt := &apiOptions{} for _, o := range opts { o(opt) } - d := opt.timeout - if d > 0 { - q.Set("timeout", d.String()) + if opt.timeout > 0 { + q.Set("timeout", opt.timeout.String()) } + if opt.limit > 0 { + q.Set("limit", strconv.FormatUint(opt.limit, 10)) + } + + return q +} + +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 := addOptionalURLParams(u.Query(), opts) + q.Set("query", query) if !ts.IsZero() { q.Set("time", formatTime(ts)) @@ -1118,36 +1136,25 @@ func (h *httpAPI) Query(ctx context.Context, query string, ts time.Time, opts .. 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() + q := addOptionalURLParams(u.Query(), opts) q.Set("query", query) q.Set("start", formatTime(r.Start)) 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 } var qres queryResult - return qres.v, warnings, json.Unmarshal(body, &qres) } -func (h *httpAPI) Series(ctx context.Context, matches []string, startTime, endTime time.Time) ([]model.LabelSet, Warnings, error) { +func (h *httpAPI) Series(ctx context.Context, matches []string, startTime, endTime time.Time, opts ...Option) ([]model.LabelSet, Warnings, error) { u := h.client.URL(epSeries, nil) - q := u.Query() + q := addOptionalURLParams(u.Query(), opts) for _, m := range matches { q.Add("match[]", m) @@ -1166,8 +1173,7 @@ func (h *httpAPI) Series(ctx context.Context, matches []string, startTime, endTi } var mset []model.LabelSet - err = json.Unmarshal(body, &mset) - return mset, warnings, err + return mset, warnings, json.Unmarshal(body, &mset) } func (h *httpAPI) Snapshot(ctx context.Context, skipHead bool) (SnapshotResult, error) { @@ -1278,8 +1284,10 @@ func (h *httpAPI) Metadata(ctx context.Context, metric, limit string) (map[strin return res, err } -func (h *httpAPI) TSDB(ctx context.Context) (TSDBResult, error) { +func (h *httpAPI) TSDB(ctx context.Context, opts ...Option) (TSDBResult, error) { u := h.client.URL(epTSDB, nil) + q := addOptionalURLParams(u.Query(), opts) + u.RawQuery = q.Encode() req, err := http.NewRequest(http.MethodGet, u.String(), nil) if err != nil { diff --git a/api/prometheus/v1/api_test.go b/api/prometheus/v1/api_test.go index e87e483..01ddf45 100644 --- a/api/prometheus/v1/api_test.go +++ b/api/prometheus/v1/api_test.go @@ -154,15 +154,15 @@ func TestAPIs(t *testing.T) { } } - doLabelNames := func(matches []string, startTime, endTime time.Time) func() (interface{}, Warnings, error) { + doLabelNames := func(matches []string, startTime, endTime time.Time, opts ...Option) func() (interface{}, Warnings, error) { return func() (interface{}, Warnings, error) { - return promAPI.LabelNames(context.Background(), matches, startTime, endTime) + return promAPI.LabelNames(context.Background(), matches, startTime, endTime, opts...) } } - doLabelValues := func(matches []string, label string, startTime, endTime time.Time) func() (interface{}, Warnings, error) { + doLabelValues := func(matches []string, label string, startTime, endTime time.Time, opts ...Option) func() (interface{}, Warnings, error) { return func() (interface{}, Warnings, error) { - return promAPI.LabelValues(context.Background(), label, matches, startTime, endTime) + return promAPI.LabelValues(context.Background(), label, matches, startTime, endTime, opts...) } } @@ -178,9 +178,9 @@ func TestAPIs(t *testing.T) { } } - doSeries := func(matcher string, startTime, endTime time.Time) func() (interface{}, Warnings, error) { + doSeries := func(matcher string, startTime, endTime time.Time, opts ...Option) func() (interface{}, Warnings, error) { return func() (interface{}, Warnings, error) { - return promAPI.Series(context.Background(), []string{matcher}, startTime, endTime) + return promAPI.Series(context.Background(), []string{matcher}, startTime, endTime, opts...) } } @@ -219,9 +219,9 @@ func TestAPIs(t *testing.T) { } } - doTSDB := func() func() (interface{}, Warnings, error) { + doTSDB := func(opts ...Option) func() (interface{}, Warnings, error) { return func() (interface{}, Warnings, error) { - v, err := promAPI.TSDB(context.Background()) + v, err := promAPI.TSDB(context.Background(), opts...) return v, nil, err } }