From 27288276a260da5373af7db2a776f8bb6a233ec2 Mon Sep 17 00:00:00 2001 From: Kyle Brandt Date: Thu, 22 Sep 2022 11:03:13 -0400 Subject: [PATCH] Prometheus: Restore FromAlert header (#55255) This restores the FromAlert header to prometheus for Grafana managed alert Queries. It does this by reverting "Prometheus: Remove middleware for custom headers (#51518)" , but also changing it so it is only the FromAlert header. This reverts commit 23725013687f88bbcdf328d537b47db2d1c8c3ae. --- .../prometheus/buffered/time_series_query.go | 15 ++++- .../buffered/time_series_query_test.go | 55 +++++++++++++++++++ pkg/tsdb/prometheus/middleware/req_headers.go | 24 ++++++++ 3 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 pkg/tsdb/prometheus/middleware/req_headers.go diff --git a/pkg/tsdb/prometheus/buffered/time_series_query.go b/pkg/tsdb/prometheus/buffered/time_series_query.go index 544f568d305..4fed7d2f3f9 100644 --- a/pkg/tsdb/prometheus/buffered/time_series_query.go +++ b/pkg/tsdb/prometheus/buffered/time_series_query.go @@ -13,10 +13,12 @@ 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/log" "github.com/grafana/grafana/pkg/infra/tracing" "github.com/grafana/grafana/pkg/tsdb/intervalv2" + "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" @@ -92,6 +94,17 @@ func New(roundTripper http.RoundTripper, tracer tracing.Tracer, settings backend } 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. + + addHeaders := make(map[string]string) + + if req.Headers["FromAlert"] == "true" { + addHeaders["FromAlert"] = "true" + } + + ctxWithHeaders := sdkHTTPClient.WithContextualMiddleware(ctx, middleware.ReqHeadersMiddleware(addHeaders)) + queries, err := b.parseTimeSeriesQuery(req) if err != nil { result := backend.QueryDataResponse{ @@ -100,7 +113,7 @@ func (b *Buffered) ExecuteTimeSeriesQuery(ctx context.Context, req *backend.Quer return &result, fmt.Errorf("error parsing time series query: %v", err) } - return b.runQueries(ctx, queries) + return b.runQueries(ctxWithHeaders, queries) } func (b *Buffered) runQueries(ctx context.Context, queries []*PrometheusQuery) (*backend.QueryDataResponse, error) { diff --git a/pkg/tsdb/prometheus/buffered/time_series_query_test.go b/pkg/tsdb/prometheus/buffered/time_series_query_test.go index 99073e3a0ac..f1fc8f2c64d 100644 --- a/pkg/tsdb/prometheus/buffered/time_series_query_test.go +++ b/pkg/tsdb/prometheus/buffered/time_series_query_test.go @@ -1,13 +1,17 @@ package buffered import ( + "context" "math" + "net/http" "testing" "time" "github.com/google/go-cmp/cmp" "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/log/logtest" "github.com/grafana/grafana/pkg/tsdb/intervalv2" apiv1 "github.com/prometheus/client_golang/api/prometheus/v1" p "github.com/prometheus/common/model" @@ -16,6 +20,57 @@ import ( var now = time.Now() +type FakeRoundTripper struct { + Req *http.Request +} + +func (frt *FakeRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + frt.Req = req + return &http.Response{}, nil +} + +func FakeMiddleware(rt *FakeRoundTripper) sdkhttpclient.Middleware { + return sdkhttpclient.NamedMiddlewareFunc("fake", func(opts sdkhttpclient.Options, next http.RoundTripper) http.RoundTripper { + return rt + }) +} + +func TestPrometheus_ExecuteTimeSeriesQuery(t *testing.T) { + t.Run("adding req headers", func(t *testing.T) { + // This makes sure we add req headers from the front end request to the request to prometheus. We do that + // through contextual middleware so this setup is a bit complex and the test itself goes a bit too much into + // internals. + + // This ends the trip and saves the request on the instance so we can inspect it. + rt := &FakeRoundTripper{} + // DefaultMiddlewares also contain contextual middleware which is the one we need to use. + middlewares := sdkhttpclient.DefaultMiddlewares() + middlewares = append(middlewares, FakeMiddleware(rt)) + + // Setup http client in at least similar way to how grafana provides it to the service + provider := sdkhttpclient.NewProvider(sdkhttpclient.ProviderOptions{Middlewares: sdkhttpclient.DefaultMiddlewares()}) + roundTripper, err := provider.GetTransport(sdkhttpclient.Options{ + Middlewares: middlewares, + }) + require.NoError(t, err) + + buffered, err := New(roundTripper, nil, backend.DataSourceInstanceSettings{JSONData: []byte("{}")}, &logtest.Fake{}) + require.NoError(t, err) + + _, err = buffered.ExecuteTimeSeriesQuery(context.Background(), &backend.QueryDataRequest{ + PluginContext: backend.PluginContext{}, + // This header is dropped, as only FromAlert header will be added to outgoing requests + Headers: map[string]string{"foo": "bar"}, + Queries: []backend.DataQuery{{ + JSON: []byte(`{"expr": "metric{label=\"test\"}", "rangeQuery": true}`), + }}, + }) + require.NoError(t, err) + require.NotNil(t, rt.Req) + require.Equal(t, http.Header{"Content-Type": []string{"application/x-www-form-urlencoded"}, "Idempotency-Key": []string(nil)}, rt.Req.Header) + }) +} + func TestPrometheus_timeSeriesQuery_formatLegend(t *testing.T) { t.Run("converting metric name", func(t *testing.T) { metric := map[p.LabelName]p.LabelValue{ diff --git a/pkg/tsdb/prometheus/middleware/req_headers.go b/pkg/tsdb/prometheus/middleware/req_headers.go new file mode 100644 index 00000000000..16bae1f8b5d --- /dev/null +++ b/pkg/tsdb/prometheus/middleware/req_headers.go @@ -0,0 +1,24 @@ +package middleware + +import ( + "net/http" + + sdkHTTPClient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient" +) + +// ReqHeadersMiddleware is used so that we can pass req headers through the prometheus go client as it does not allow +// access to the request directly. Should be used together with WithContextualMiddleware so that it is attached to +// the context of each request with its unique headers. +func ReqHeadersMiddleware(headers map[string]string) sdkHTTPClient.Middleware { + return sdkHTTPClient.NamedMiddlewareFunc("prometheus-req-headers-middleware", func(opts sdkHTTPClient.Options, next http.RoundTripper) http.RoundTripper { + return sdkHTTPClient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range headers { + // As custom headers middleware is before contextual we may overwrite custom headers here with those + // that came with the request which probably makes sense. + req.Header[k] = []string{v} + } + + return next.RoundTrip(req) + }) + }) +}