Prometheus: Use contextual middleware for req headers and simplify client creation (#51061)

* Use contextual middleware and simplify client creation

* Fix tests

* Add test for the header propagation

* Fix tests and lint

* Update pkg/tsdb/prometheus/prometheus.go

Co-authored-by: ismail simsek <ismailsimsek09@gmail.com>

Co-authored-by: ismail simsek <ismailsimsek09@gmail.com>
This commit is contained in:
Andrej Ocenas
2022-06-23 14:48:16 +02:00
committed by GitHub
co-authored by ismail simsek
parent a8eb29f1d7
commit d20afa2a39
17 changed files with 394 additions and 697 deletions
@@ -5,6 +5,7 @@ import (
"encoding/json"
"fmt"
"math"
"net/http"
"regexp"
"sort"
"strconv"
@@ -12,14 +13,13 @@ import (
"time"
"github.com/grafana/grafana-plugin-sdk-go/backend"
sdkHTTPClient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
"github.com/grafana/grafana-plugin-sdk-go/data"
"github.com/grafana/grafana/pkg/infra/httpclient"
"github.com/grafana/grafana/pkg/infra/log"
"github.com/grafana/grafana/pkg/infra/tracing"
"github.com/grafana/grafana/pkg/services/featuremgmt"
"github.com/grafana/grafana/pkg/setting"
"github.com/grafana/grafana/pkg/tsdb/intervalv2"
"github.com/grafana/grafana/pkg/tsdb/prometheus/buffered/promclient"
"github.com/grafana/grafana/pkg/tsdb/prometheus/middleware"
"github.com/grafana/grafana/pkg/tsdb/prometheus/utils"
"github.com/grafana/grafana/pkg/util/maputil"
apiv1 "github.com/prometheus/client_golang/api/prometheus/v1"
"github.com/prometheus/common/model"
@@ -57,17 +57,24 @@ var (
type Buffered struct {
intervalCalculator intervalv2.Calculator
tracer tracing.Tracer
getClient clientGetter
client apiv1.API
log log.Logger
ID int64
URL string
TimeInterval string
}
func New(httpClientProvider httpclient.Provider, cfg *setting.Cfg, features featuremgmt.FeatureToggles, tracer tracing.Tracer, settings backend.DataSourceInstanceSettings, plog log.Logger) (*Buffered, error) {
var jsonData map[string]interface{}
if err := json.Unmarshal(settings.JSONData, &jsonData); err != nil {
return nil, fmt.Errorf("error reading settings: %w", err)
// New creates and object capable of executing and parsing a Prometheus queries. It's "buffered" because there is
// another implementation capable of streaming parse the response.
func New(roundTripper http.RoundTripper, tracer tracing.Tracer, settings backend.DataSourceInstanceSettings, plog log.Logger) (*Buffered, error) {
promClient, err := CreateClient(roundTripper, settings.URL)
if err != nil {
return nil, fmt.Errorf("error creating prom client: %v", err)
}
jsonData, err := utils.GetJsonData(settings)
if err != nil {
return nil, fmt.Errorf("error getting jsonData: %w", err)
}
timeInterval, err := maputil.GetStringOptional(jsonData, "timeInterval")
@@ -75,23 +82,34 @@ func New(httpClientProvider httpclient.Provider, cfg *setting.Cfg, features feat
return nil, err
}
p := promclient.NewProvider(settings, jsonData, httpClientProvider, cfg, features, plog)
pc, err := promclient.NewProviderCache(p)
if err != nil {
return nil, err
}
return &Buffered{
intervalCalculator: intervalv2.NewCalculator(),
tracer: tracer,
log: plog,
getClient: pc.GetClient,
client: promClient,
TimeInterval: timeInterval,
ID: settings.ID,
URL: settings.URL,
}, nil
}
func (b *Buffered) runQueries(ctx context.Context, client apiv1.API, queries []*PrometheusQuery) (*backend.QueryDataResponse, error) {
func (b *Buffered) ExecuteTimeSeriesQuery(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
// Add headers from the request to context so they are added later on by a context middleware. This is because
// prom client does not allow us to do this directly.
ctxWithHeaders := sdkHTTPClient.WithContextualMiddleware(ctx, middleware.ReqHeadersMiddleware(req.Headers))
queries, err := b.parseTimeSeriesQuery(req)
if err != nil {
result := backend.QueryDataResponse{
Responses: backend.Responses{},
}
return &result, fmt.Errorf("error parsing time series query: %v", err)
}
return b.runQueries(ctxWithHeaders, queries)
}
func (b *Buffered) runQueries(ctx context.Context, queries []*PrometheusQuery) (*backend.QueryDataResponse, error) {
result := backend.QueryDataResponse{
Responses: backend.Responses{},
}
@@ -99,11 +117,12 @@ func (b *Buffered) runQueries(ctx context.Context, client apiv1.API, queries []*
for _, query := range queries {
b.log.Debug("Sending query", "start", query.Start, "end", query.End, "step", query.Step, "query", query.Expr)
ctx, span := b.tracer.Start(ctx, "datasource.prometheus")
span.SetAttributes("expr", query.Expr, attribute.Key("expr").String(query.Expr))
span.SetAttributes("start_unixnano", query.Start, attribute.Key("start_unixnano").Int64(query.Start.UnixNano()))
span.SetAttributes("stop_unixnano", query.End, attribute.Key("stop_unixnano").Int64(query.End.UnixNano()))
defer span.End()
ctx, endSpan := utils.StartTrace(ctx, b.tracer, "datasource.prometheus", []utils.Attribute{
{Key: "expr", Value: query.Expr, Kv: attribute.Key("expr").String(query.Expr)},
{Key: "start_unixnano", Value: query.Start, Kv: attribute.Key("start_unixnano").Int64(query.Start.UnixNano())},
{Key: "stop_unixnano", Value: query.End, Kv: attribute.Key("stop_unixnano").Int64(query.End.UnixNano())},
})
defer endSpan()
response := make(map[TimeSeriesQueryType]interface{})
@@ -115,7 +134,7 @@ func (b *Buffered) runQueries(ctx context.Context, client apiv1.API, queries []*
}
if query.RangeQuery {
rangeResponse, _, err := client.QueryRange(ctx, query.Expr, timeRange)
rangeResponse, _, err := b.client.QueryRange(ctx, query.Expr, timeRange)
if err != nil {
b.log.Error("Range query failed", "query", query.Expr, "err", err)
result.Responses[query.RefId] = backend.DataResponse{Error: err}
@@ -125,7 +144,7 @@ func (b *Buffered) runQueries(ctx context.Context, client apiv1.API, queries []*
}
if query.InstantQuery {
instantResponse, _, err := client.Query(ctx, query.Expr, query.End)
instantResponse, _, err := b.client.Query(ctx, query.Expr, query.End)
if err != nil {
b.log.Error("Instant query failed", "query", query.Expr, "err", err)
result.Responses[query.RefId] = backend.DataResponse{Error: err}
@@ -137,7 +156,7 @@ func (b *Buffered) runQueries(ctx context.Context, client apiv1.API, queries []*
// This is a special case
// If exemplar query returns error, we want to only log it and continue with other results processing
if query.ExemplarQuery {
exemplarResponse, err := client.QueryExemplars(ctx, query.Expr, timeRange.Start, timeRange.End)
exemplarResponse, err := b.client.QueryExemplars(ctx, query.Expr, timeRange.Start, timeRange.End)
if err != nil {
b.log.Error("Exemplar query failed", "query", query.Expr, "err", err)
} else {
@@ -163,23 +182,6 @@ func (b *Buffered) runQueries(ctx context.Context, client apiv1.API, queries []*
return &result, nil
}
func (b *Buffered) ExecuteTimeSeriesQuery(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
client, err := b.getClient(req.Headers)
if err != nil {
return nil, err
}
queries, err := b.parseTimeSeriesQuery(req)
if err != nil {
result := backend.QueryDataResponse{
Responses: backend.Responses{},
}
return &result, err
}
return b.runQueries(ctx, client, queries)
}
func formatLegend(metric model.Metric, query *PrometheusQuery) string {
var legend = metric.String()
@@ -209,18 +211,18 @@ func formatLegend(metric model.Metric, query *PrometheusQuery) string {
return legend
}
func (b *Buffered) parseTimeSeriesQuery(queryContext *backend.QueryDataRequest) ([]*PrometheusQuery, error) {
func (b *Buffered) parseTimeSeriesQuery(req *backend.QueryDataRequest) ([]*PrometheusQuery, error) {
qs := []*PrometheusQuery{}
for _, query := range queryContext.Queries {
for _, query := range req.Queries {
model := &QueryModel{}
err := json.Unmarshal(query.JSON, model)
if err != nil {
return nil, err
return nil, fmt.Errorf("error unmarshaling query model: %v", err)
}
//Final interval value
interval, err := calculatePrometheusInterval(model, b.TimeInterval, query, b.intervalCalculator)
if err != nil {
return nil, err
return nil, fmt.Errorf("error calculating interval: %v", err)
}
// Interpolate variables in expr
@@ -234,7 +236,7 @@ func (b *Buffered) parseTimeSeriesQuery(queryContext *backend.QueryDataRequest)
// We never want to run exemplar query for alerting
exemplarQuery := model.ExemplarQuery
if queryContext.Headers["FromAlert"] == "true" {
if req.Headers["FromAlert"] == "true" {
exemplarQuery = false
}