Prometheus: Restore FromAlert header

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) {

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{

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)
+		})
+	})
+}