feat(querier): propagate all known alerting headers (#92873)

pkg/registry/apis/datasource/sub_query.go

@@ -4,14 +4,14 @@ import (
 	"context"
 	"fmt"
 	"net/http"
-	"strings"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	data "github.com/grafana/grafana-plugin-sdk-go/experimental/apis/data/v0alpha1"
+	query "github.com/grafana/grafana/pkg/apis/query/v0alpha1"
+	query_headers "github.com/grafana/grafana/pkg/registry/apis/query"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apiserver/pkg/registry/rest"
 
-	query "github.com/grafana/grafana/pkg/apis/query/v0alpha1"
 	"github.com/grafana/grafana/pkg/web"
 )
 
@@ -74,29 +74,10 @@ func (r *subQueryREST) Connect(ctx context.Context, name string, opts runtime.Ob
 		ctx = backend.WithGrafanaConfig(ctx, pluginCtx.GrafanaConfig)
 		ctx = contextualMiddlewares(ctx)
 
-		// only forward expected headers, log unexpected ones
-		headers := make(map[string]string)
-		// headers are case insensitive, however some datasources still check for camel casing so we have to send them camel cased
-		expectedHeaders := map[string]string{
-			"fromalert":      "FromAlert",
-			"content-type":   "Content-Type",
-			"content-length": "Content-Length",
-			"user-agent":     "User-Agent",
-			"accept":         "Accept",
-		}
-		for k, v := range req.Header {
-			headerToSend, ok := expectedHeaders[strings.ToLower(k)]
-			if ok {
-				headers[headerToSend] = v[0]
-			} else {
-				r.builder.log.Warn("datasource received an unexpected header, ignoring it", "header", k)
-			}
-		}
-
 		rsp, err := r.builder.client.QueryData(ctx, &backend.QueryDataRequest{
 			Queries:       queries,
 			PluginContext: pluginCtx,
-			Headers:       headers,
+			Headers:       query_headers.ExtractKnownHeaders(req.Header),
 		})
 		if err != nil {
 			responder.Error(err)

pkg/registry/apis/datasource/sub_query_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana/pkg/apis/datasource/v0alpha1"
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
 	"github.com/stretchr/testify/require"
 	"k8s.io/apimachinery/pkg/runtime"
 )
@@ -32,14 +33,24 @@ func TestSubQueryConnect(t *testing.T) {
 
 	rr := httptest.NewRecorder()
 	req := httptest.NewRequest(http.MethodGet, "/some-path", nil)
-	req.Header.Set("fromAlert", "true")
+	req.Header.Set(models.FromAlertHeaderName, "true")
+	req.Header.Set(models.CacheSkipHeaderName, "true")
+	req.Header.Set("X-Rule-Uid", "abc")
+	req.Header.Set("X-Rule-Folder", "folder-1")
+	req.Header.Set("X-Rule-Source", "grafana-ruler")
+	req.Header.Set("X-Grafana-Org-Id", "1")
 	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("some-unexpected-header", "some-value")
 	handler.ServeHTTP(rr, req)
 
 	// test that headers are forwarded and cased appropriately
 	require.Equal(t, map[string]string{
-		"FromAlert":    "true",
-		"Content-Type": "application/json",
+		models.FromAlertHeaderName: "true",
+		models.CacheSkipHeaderName: "true",
+		"X-Rule-Uid":               "abc",
+		"X-Rule-Folder":            "folder-1",
+		"X-Rule-Source":            "grafana-ruler",
+		"X-Grafana-Org-Id":         "1",
 	}, *sqr.builder.client.(mockClient).lastCalledWithHeaders)
 }
 

pkg/registry/apis/query/header_utils.go

@@ -0,0 +1,35 @@
+package query
+
+import (
+	"net/http"
+	"strings"
+
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
+)
+
+// Set of headers that we want to forward to the datasource api servers. Those are used i.e. for
+// cache control or identifying the source of the request.
+//
+// The headers related to grafana alerting can be found here:
+// https://github.com/grafana/grafana-ruler/blob/96e6d4b25c0d973a7615b92b35739511a6fbd72f/pkg/ruler/rulesmanager/ds_query_rule_evaluator.go#L313-L328
+//
+// The usage of strings.ToLower is because the server would convert `FromAlert` to `Fromalert`. So the make matching
+// easier, we just match all headers in lower case.
+var expectedHeaders = map[string]string{
+	strings.ToLower(models.FromAlertHeaderName): models.FromAlertHeaderName,
+	strings.ToLower(models.CacheSkipHeaderName): models.CacheSkipHeaderName,
+	strings.ToLower("X-Rule-Uid"):               "X-Rule-Uid",
+	strings.ToLower("X-Rule-Folder"):            "X-Rule-Folder",
+	strings.ToLower("X-Rule-Source"):            "X-Rule-Source",
+	strings.ToLower("X-Grafana-Org-Id"):         "X-Grafana-Org-Id",
+}
+
+func ExtractKnownHeaders(header http.Header) map[string]string {
+	extractedHeaders := make(map[string]string)
+	for k, v := range header {
+		if headerName, exists := expectedHeaders[strings.ToLower(k)]; exists {
+			extractedHeaders[headerName] = v[0]
+		}
+	}
+	return extractedHeaders
+}

pkg/registry/apis/query/query.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"net/http"
 	"strconv"
-	"strings"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -135,26 +134,8 @@ func (r *queryREST) Connect(connectCtx context.Context, name string, _ runtime.O
 			return
 		}
 
-		// get headers from the original http req and add them to each sub request
-		// headers are case insensitive, however some datasources still check for camel casing so we have to send them camel cased
-		expectedHeaders := map[string]string{
-			"fromalert":      "FromAlert",
-			"content-type":   "Content-Type",
-			"content-length": "Content-Length",
-			"user-agent":     "User-Agent",
-			"accept":         "Accept",
-		}
-
 		for i := range req.Requests {
-			req.Requests[i].Headers = make(map[string]string)
-			for k, v := range httpreq.Header {
-				headerToSend, ok := expectedHeaders[strings.ToLower(k)]
-				if ok {
-					req.Requests[i].Headers[headerToSend] = v[0]
-				} else {
-					b.log.Warn(fmt.Sprintf("query service received an unexpected header, ignoring it: %s", k))
-				}
-			}
+			req.Requests[i].Headers = ExtractKnownHeaders(httpreq.Header)
 		}
 
 		// Actually run the query

pkg/registry/apis/query/query_test.go

@@ -14,6 +14,7 @@ import (
 	"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/services/ngalert/models"
 	"github.com/stretchr/testify/require"
 	"k8s.io/apimachinery/pkg/runtime"
 )
@@ -53,14 +54,23 @@ func TestQueryRestConnectHandler(t *testing.T) {
 			"to": "now"}`),
 	}
 	req := httptest.NewRequest(http.MethodGet, "/some-path", bytes.NewReader(body.Raw))
-	req.Header.Set("fromAlert", "true")
+	req.Header.Set(models.FromAlertHeaderName, "true")
+	req.Header.Set(models.CacheSkipHeaderName, "true")
+	req.Header.Set("X-Rule-Uid", "abc")
+	req.Header.Set("X-Rule-Folder", "folder-1")
+	req.Header.Set("X-Rule-Source", "grafana-ruler")
+	req.Header.Set("X-Grafana-Org-Id", "1")
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("some-unexpected-header", "some-value")
 	handler.ServeHTTP(rr, req)
 
 	require.Equal(t, map[string]string{
-		"FromAlert":    "true",
-		"Content-Type": "application/json",
+		models.FromAlertHeaderName: "true",
+		models.CacheSkipHeaderName: "true",
+		"X-Rule-Uid":               "abc",
+		"X-Rule-Folder":            "folder-1",
+		"X-Rule-Source":            "grafana-ruler",
+		"X-Grafana-Org-Id":         "1",
 	}, *b.client.(mockClient).lastCalledWithHeaders)
 }