diff --git a/pkg/registry/apis/query/client.go b/pkg/registry/apis/query/client.go index 50f7a16434a1..7f621512c09c 100644 --- a/pkg/registry/apis/query/client.go +++ b/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 } diff --git a/pkg/registry/apis/query/query.go b/pkg/registry/apis/query/query.go index ec50d0462159..f17943bee50c 100644 --- a/pkg/registry/apis/query/query.go +++ b/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 } diff --git a/pkg/registry/apis/query/query_test.go b/pkg/registry/apis/query/query_test.go new file mode 100644 index 000000000000..f1bdde43fe31 --- /dev/null +++ b/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 +}