From 7cc9420eea41c8bf3e0178ea7e85442df9f15381 Mon Sep 17 00:00:00 2001 From: Kyle Brandt Date: Thu, 15 Sep 2022 12:00:29 -0400 Subject: [PATCH 1/3] Revert "Prometheus: Remove middleware for custom headers (#51518)" This reverts commit 23725013687f88bbcdf328d537b47db2d1c8c3ae. --- .../prometheus/buffered/time_series_query.go | 8 ++- .../buffered/time_series_query_test.go | 55 +++++++++++++++++++ pkg/tsdb/prometheus/middleware/req_headers.go | 24 ++++++++ 3 files changed, 86 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 544f568d3058..52bedad8c2fc 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,10 @@ 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. + ctxWithHeaders := sdkHTTPClient.WithContextualMiddleware(ctx, middleware.ReqHeadersMiddleware(req.Headers)) + queries, err := b.parseTimeSeriesQuery(req) if err != nil { result := backend.QueryDataResponse{ @@ -100,7 +106,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 99073e3a0ac6..d6f9c4b67cc9 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 should end up in the outgoing request to prometheus + 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"}, "foo": []string{"bar"}}, 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 000000000000..16bae1f8b5da --- /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) + }) + }) +} From 7cc2dbf54ea1172127fb5d6332d35a99686b980e Mon Sep 17 00:00:00 2001 From: Kyle Brandt Date: Thu, 15 Sep 2022 12:18:08 -0400 Subject: [PATCH 2/3] Prometheus: Add fromAlert header back but only fromAlert header #54403 --- pkg/tsdb/prometheus/buffered/time_series_query.go | 9 ++++++++- pkg/tsdb/prometheus/prometheus.go | 2 ++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/pkg/tsdb/prometheus/buffered/time_series_query.go b/pkg/tsdb/prometheus/buffered/time_series_query.go index 52bedad8c2fc..4fed7d2f3f96 100644 --- a/pkg/tsdb/prometheus/buffered/time_series_query.go +++ b/pkg/tsdb/prometheus/buffered/time_series_query.go @@ -96,7 +96,14 @@ 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. - ctxWithHeaders := sdkHTTPClient.WithContextualMiddleware(ctx, middleware.ReqHeadersMiddleware(req.Headers)) + + 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 { diff --git a/pkg/tsdb/prometheus/prometheus.go b/pkg/tsdb/prometheus/prometheus.go index 48d4b417c5f3..d4d17ad4389e 100644 --- a/pkg/tsdb/prometheus/prometheus.go +++ b/pkg/tsdb/prometheus/prometheus.go @@ -126,6 +126,8 @@ func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest) return data, err } + if req.Headers[] + return i.buffered.ExecuteTimeSeriesQuery(ctx, req) } From 71ad56e97f8791874b843d630b8537e6a02bebe5 Mon Sep 17 00:00:00 2001 From: Kyle Brandt Date: Mon, 19 Sep 2022 07:47:48 -0400 Subject: [PATCH 3/3] cleaup --- pkg/tsdb/prometheus/buffered/time_series_query_test.go | 4 ++-- pkg/tsdb/prometheus/prometheus.go | 2 -- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/pkg/tsdb/prometheus/buffered/time_series_query_test.go b/pkg/tsdb/prometheus/buffered/time_series_query_test.go index d6f9c4b67cc9..f1fc8f2c64d9 100644 --- a/pkg/tsdb/prometheus/buffered/time_series_query_test.go +++ b/pkg/tsdb/prometheus/buffered/time_series_query_test.go @@ -59,7 +59,7 @@ func TestPrometheus_ExecuteTimeSeriesQuery(t *testing.T) { _, err = buffered.ExecuteTimeSeriesQuery(context.Background(), &backend.QueryDataRequest{ PluginContext: backend.PluginContext{}, - // This header should end up in the outgoing request to prometheus + // 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}`), @@ -67,7 +67,7 @@ func TestPrometheus_ExecuteTimeSeriesQuery(t *testing.T) { }) require.NoError(t, err) require.NotNil(t, rt.Req) - require.Equal(t, http.Header{"Content-Type": []string{"application/x-www-form-urlencoded"}, "foo": []string{"bar"}}, rt.Req.Header) + require.Equal(t, http.Header{"Content-Type": []string{"application/x-www-form-urlencoded"}, "Idempotency-Key": []string(nil)}, rt.Req.Header) }) } diff --git a/pkg/tsdb/prometheus/prometheus.go b/pkg/tsdb/prometheus/prometheus.go index d4d17ad4389e..48d4b417c5f3 100644 --- a/pkg/tsdb/prometheus/prometheus.go +++ b/pkg/tsdb/prometheus/prometheus.go @@ -126,8 +126,6 @@ func (s *Service) QueryData(ctx context.Context, req *backend.QueryDataRequest) return data, err } - if req.Headers[] - return i.buffered.ExecuteTimeSeriesQuery(ctx, req) }