QueryService: Forward headers to datasource clients (#92329)

Query Service: pass expected headers

pkg/registry/apis/query/client.go

@@ -9,14 +9,13 @@ import (
 // The query runner interface
 type DataSourceClientSupplier interface {
 	// Get a client for a given datasource
-	// NOTE: authorization headers are not yet added and the client may be shared across multiple users
-	GetDataSourceClient(ctx context.Context, ref data.DataSourceRef) (data.QueryDataClient, error)
+	GetDataSourceClient(ctx context.Context, ref data.DataSourceRef, headers map[string]string) (data.QueryDataClient, error)
 }
 
 type CommonDataSourceClientSupplier struct {
 	Client data.QueryDataClient
 }
 
-func (s *CommonDataSourceClientSupplier) GetDataSourceClient(ctx context.Context, ref data.DataSourceRef) (data.QueryDataClient, error) {
+func (s *CommonDataSourceClientSupplier) GetDataSourceClient(_ context.Context, _ data.DataSourceRef, _ map[string]string) (data.QueryDataClient, error) {
 	return s.Client, nil
 }

pkg/registry/apis/query/query.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -48,7 +49,7 @@ func newQueryREST(builder *QueryAPIBuilder) *queryREST {
 }
 
 func (r *queryREST) New() runtime.Object {
-	// This is added as the "ResponseType" regarless what ProducesObject() says :)
+	// This is added as the "ResponseType" regardless what ProducesObject() says :)
 	return &query.QueryDataResponse{}
 }
 
@@ -134,6 +135,28 @@ 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))
+				}
+			}
+		}
+
 		// Actually run the query
 		rsp, err := b.execute(ctx, req)
 		if err != nil {
@@ -192,11 +215,14 @@ func (b *QueryAPIBuilder) handleQuerySingleDatasource(ctx context.Context, req d
 		return &backend.QueryDataResponse{}, nil
 	}
 
-	// Add user headers... here or in client.QueryData
-	client, err := b.client.GetDataSourceClient(ctx, v0alpha1.DataSourceRef{
-		Type: req.PluginId,
-		UID:  req.UID,
-	})
+	client, err := b.client.GetDataSourceClient(
+		ctx,
+		v0alpha1.DataSourceRef{
+			Type: req.PluginId,
+			UID:  req.UID,
+		},
+		req.Headers,
+	)
 	if err != nil {
 		return nil, err
 	}

pkg/registry/apis/query/query_test.go

@@ -0,0 +1,98 @@
+package query
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	data "github.com/grafana/grafana-plugin-sdk-go/experimental/apis/data/v0alpha1"
+	"github.com/grafana/grafana/pkg/expr"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/runtime"
+)
+
+func TestQueryRestConnectHandler(t *testing.T) {
+	b := &QueryAPIBuilder{
+		client: mockClient{
+			lastCalledWithHeaders: &map[string]string{},
+		},
+		tracer: tracing.InitializeTracerForTest(),
+		parser: newQueryParser(expr.NewExpressionQueryReader(featuremgmt.WithFeatures()),
+			&legacyDataSourceRetriever{}, tracing.InitializeTracerForTest()),
+		log: log.New("test"),
+	}
+	qr := newQueryREST(b)
+	ctx := context.Background()
+	mr := mockResponder{}
+
+	handler, err := qr.Connect(ctx, "name", nil, mr)
+	require.NoError(t, err)
+
+	rr := httptest.NewRecorder()
+	body := runtime.RawExtension{
+		Raw: []byte(`{ 
+			"queries": [
+				{
+					"datasource": {
+					"type": "prometheus",
+					"uid": "demo-prometheus"
+					},
+					"expr": "sum(go_gc_duration_seconds)",
+					"range": false,
+					"instant": true
+				}
+			],
+			"from": "now-1h",
+			"to": "now"}`),
+	}
+	req := httptest.NewRequest(http.MethodGet, "/some-path", bytes.NewReader(body.Raw))
+	req.Header.Set("fromAlert", "true")
+	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",
+	}, *b.client.(mockClient).lastCalledWithHeaders)
+}
+
+type mockResponder struct {
+}
+
+// Object writes the provided object to the response. Invoking this method multiple times is undefined.
+func (m mockResponder) Object(statusCode int, obj runtime.Object) {
+}
+
+// Error writes the provided error to the response. This method may only be invoked once.
+func (m mockResponder) Error(err error) {
+}
+
+type mockClient struct {
+	lastCalledWithHeaders *map[string]string
+}
+
+func (m mockClient) GetDataSourceClient(ctx context.Context, ref data.DataSourceRef, headers map[string]string) (data.QueryDataClient, error) {
+	*m.lastCalledWithHeaders = headers
+
+	return nil, fmt.Errorf("mock error")
+}
+
+func (m mockClient) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	return nil, fmt.Errorf("mock error")
+}
+
+func (m mockClient) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+	return nil
+}
+
+func (m mockClient) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	return nil, nil
+}