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:
co-authored by
ismail simsek
parent
a8eb29f1d7
commit
d20afa2a39
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user