Add WithQueryStats option

api/prometheus/v1/api.go

@@ -421,10 +421,36 @@ type Metadata struct {
 	Unit string     `json:"unit"`
 }
 
+type stepStat struct {
+	T int64
+	V float64
+}
+
+type queryTimings struct {
+	EvalTotalTime        float64 `json:"evalTotalTime"`
+	ResultSortTime       float64 `json:"resultSortTime"`
+	QueryPreparationTime float64 `json:"queryPreparationTime"`
+	InnerEvalTime        float64 `json:"innerEvalTime"`
+	ExecQueueTime        float64 `json:"execQueueTime"`
+	ExecTotalTime        float64 `json:"execTotalTime"`
+}
+
+type querySamples struct {
+	TotalQueryableSamplesPerStep []stepStat `json:"totalQueryableSamplesPerStep,omitempty"`
+	TotalQueryableSamples        int        `json:"totalQueryableSamples"`
+	PeakSamples                  int        `json:"peakSamples"`
+}
+
+type QueryStats struct {
+	Timings queryTimings `json:"timings"`
+	Samples querySamples `json:"samples"`
+}
+
 // queryResult contains result data for a query.
 type queryResult struct {
 	Type   model.ValueType `json:"resultType"`
 	Result interface{}     `json:"result"`
+	Stats  *QueryStats     `json:"stats"`
 
 	// The decoded value.
 	v model.Value
@@ -580,7 +606,10 @@ func (qr *queryResult) UnmarshalJSON(b []byte) error {
 	v := struct {
 		Type   model.ValueType `json:"resultType"`
 		Result json.RawMessage `json:"result"`
-	}{}
+		Stats  *QueryStats     `json:"stats"`
+	}{
+		Stats: qr.Stats,
+	}
 
 	err := json.Unmarshal(b, &v)
 	if err != nil {
@@ -819,7 +848,9 @@ func (h *httpAPI) LabelValues(ctx context.Context, label string, matches []strin
 }
 
 type apiOptions struct {
-	timeout time.Duration
+	timeout      time.Duration
+	stats        *QueryStats
+	perStepStats bool
 }
 
 type Option func(c *apiOptions)
@@ -830,6 +861,13 @@ func WithTimeout(timeout time.Duration) Option {
 	}
 }
 
+func WithQueryStats(s *QueryStats, perStep bool) Option {
+	return func(o *apiOptions) {
+		o.stats = s
+		o.perStepStats = perStep
+	}
+}
+
 func (h *httpAPI) Query(ctx context.Context, query string, ts time.Time, opts ...Option) (model.Value, Warnings, error) {
 
 	u := h.client.URL(epQuery, nil)
@@ -845,6 +883,10 @@ func (h *httpAPI) Query(ctx context.Context, query string, ts time.Time, opts ..
 		q.Set("timeout", d.String())
 	}
 
+	if opt.perStepStats {
+		q.Set("stats", "all")
+	}
+
 	q.Set("query", query)
 	if !ts.IsZero() {
 		q.Set("time", formatTime(ts))
@@ -856,6 +898,10 @@ func (h *httpAPI) Query(ctx context.Context, query string, ts time.Time, opts ..
 	}
 
 	var qres queryResult
+	if opt.stats != nil {
+		qres.Stats = opt.stats
+	}
+
 	return model.Value(qres.v), warnings, json.Unmarshal(body, &qres)
 }
 
@@ -878,12 +924,19 @@ func (h *httpAPI) QueryRange(ctx context.Context, query string, r Range, opts ..
 		q.Set("timeout", d.String())
 	}
 
+	if opt.perStepStats {
+		q.Set("stats", "all")
+	}
+
 	_, body, warnings, err := h.client.DoGetFallback(ctx, u, q)
 	if err != nil {
 		return nil, warnings, err
 	}
 
 	var qres queryResult
+	if opt.stats != nil {
+		qres.Stats = opt.stats
+	}
 
 	return model.Value(qres.v), warnings, json.Unmarshal(body, &qres)
 }

api/prometheus/v1/api_test.go

@@ -1300,6 +1300,28 @@ func TestAPIs(t *testing.T) {
 				},
 			},
 		},
+		{
+			do: doQuery("2", testTime, WithQueryStats(nil, true)),
+			inRes: &queryResult{
+				Type: model.ValScalar,
+				Result: &model.Scalar{
+					Value:     2,
+					Timestamp: model.TimeFromUnix(testTime.Unix()),
+				},
+			},
+
+			reqMethod: "POST",
+			reqPath:   "/api/v1/query",
+			reqParam: url.Values{
+				"query": []string{"2"},
+				"time":  []string{testTime.Format(time.RFC3339Nano)},
+				"stats": []string{"all"},
+			},
+			res: &model.Scalar{
+				Value:     2,
+				Timestamp: model.TimeFromUnix(testTime.Unix()),
+			},
+		},
 	}
 
 	var tests []apiTest

api/prometheus/v1/example_test.go

@@ -218,3 +218,30 @@ func ExampleAPI_series() {
 		fmt.Println(lbl)
 	}
 }
+
+func ExampleAPI_queryWithStats() {
+	client, err := api.NewClient(api.Config{
+		Address: "http://demo.robustperception.io:9090",
+	})
+	if err != nil {
+		fmt.Printf("Error creating client: %v\n", err)
+		os.Exit(1)
+	}
+
+	v1api := v1.NewAPI(client)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	stats := v1.QueryStats{}
+	result, warnings, err := v1api.Query(ctx, "up", time.Now(),
+		v1.WithTimeout(5*time.Second),
+		v1.WithQueryStats(&stats, true))
+	if err != nil {
+		fmt.Printf("Error querying Prometheus: %v\n", err)
+		os.Exit(1)
+	}
+	if len(warnings) > 0 {
+		fmt.Printf("Warnings: %v\n", warnings)
+	}
+	fmt.Printf("Result:\n%v\n", result)
+	fmt.Printf("Stats:\n%v\n", stats)
+}