public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ds-querier: handle response-headers explicitly (#107304)

Browse Source
This commit is contained in:
Gábor Farkas
2025-07-15 08:31:57 +02:00
committed by GitHub
parent 2c0cfb0dc7
commit 868791d751
4 changed files with 207 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/registry/apis/query/client/plugin.go
+28 -1
View File
@@ -18,6 +18,7 @@ import (
"github.com/grafana/grafana/pkg/plugins"
"github.com/grafana/grafana/pkg/registry/apis/query/clientapi"
"github.com/grafana/grafana/pkg/services/accesscontrol"
"github.com/grafana/grafana/pkg/services/contexthandler"
"github.com/grafana/grafana/pkg/services/datasources"
"github.com/grafana/grafana/pkg/services/pluginsintegration/adapters"
"github.com/grafana/grafana/pkg/services/pluginsintegration/plugincontext"
@@ -109,8 +110,34 @@ func getGrafanaDataSourceSettings(ctx context.Context) (*backend.DataSourceInsta
return adapters.ModelToInstanceSettings(ds, decryptFunc)
}
func (d *pluginClient) QueryData(ctx context.Context, req data.QueryDataRequest) (*clientapi.Response, error) {
// middlewares may set response-http-headers through context, so we need to do extra steps
// first we create an isolated context to be used by query-data
isolatedCtx := contexthandler.CopyWithReqContext(ctx)
qdr, err := d.innerQueryData(isolatedCtx, req)
if err != nil {
return nil, err
}
rsp := &clientapi.Response{
QDR: qdr,
Headers: nil,
}
// we extract the response-headers from the isolated context, and return them explicitly
// in the clientapi.Response structure
reqCtx := contexthandler.FromContext(isolatedCtx)
if reqCtx != nil {
rsp.Headers = reqCtx.Resp.Header()
}
return rsp, nil
}
// ExecuteQueryData implements QueryHelper.
func (d *pluginClient) QueryData(ctx context.Context, req data.QueryDataRequest) (*backend.QueryDataResponse, error) {
func (d *pluginClient) innerQueryData(ctx context.Context, req data.QueryDataRequest) (*backend.QueryDataResponse, error) {
queries, dsRef, err := data.ToDataSourceQueries(req)
if err != nil {
return nil, err
pkg/registry/apis/query/clientapi/clientapi.go
+7 -1
View File
@@ -2,14 +2,20 @@ package clientapi
import (
"context"
"net/http"
"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/services/featuremgmt"
)
type Response struct {
QDR *backend.QueryDataResponse
Headers http.Header
}
type QueryDataClient interface {
QueryData(ctx context.Context, req data.QueryDataRequest) (*backend.QueryDataResponse, error)
QueryData(ctx context.Context, req data.QueryDataRequest) (*Response, error)
}
type InstanceConfigurationSettings struct {
pkg/registry/apis/query/query.go
+64 -30
View File
@@ -5,12 +5,14 @@ import (
"errors"
"fmt"
"net/http"
"slices"
"strconv"
"time"
"github.com/grafana/grafana-plugin-sdk-go/backend"
"github.com/grafana/grafana-plugin-sdk-go/experimental/apis/data/v0alpha1"
"github.com/grafana/grafana/pkg/registry/apis/query/clientapi"
"github.com/grafana/grafana/pkg/services/contexthandler"
"github.com/grafana/grafana/pkg/services/datasources"
"github.com/grafana/grafana/pkg/services/ngalert/models"
"go.opentelemetry.io/otel/attribute"
@@ -197,10 +199,16 @@ func (r *queryREST) Connect(connectCtx context.Context, name string, _ runtime.O
// Actually run the query (includes expressions)
rsp, err := b.execute(ctx, req, instanceConfig)
if err != nil {
b.log.Error("execute error", "http code", query.GetResponseCode(rsp), "err", err)
if rsp != nil { // if we have a response, we assume the err is set in the response
responder.Object(query.GetResponseCode(rsp), &query.QueryDataResponse{
QueryDataResponse: *rsp,
// we extract the QDR, but the response may be nil
var qdr *backend.QueryDataResponse
if rsp != nil {
qdr = rsp.QDR
}
b.log.Error("execute error", "http code", query.GetResponseCode(qdr), "err", err)
if qdr != nil { // if we have a response, we assume the err is set in the response
responder.Object(query.GetResponseCode(qdr), &query.QueryDataResponse{
QueryDataResponse: *qdr,
})
return
} else {
@@ -211,8 +219,14 @@ func (r *queryREST) Connect(connectCtx context.Context, name string, _ runtime.O
}
}
responder.Object(query.GetResponseCode(rsp), &query.QueryDataResponse{
QueryDataResponse: *rsp, // wrap the backend response as a QueryDataResponse
// response headers are communicated back using the context for some reason
reqCtx := contexthandler.FromContext(ctx)
if reqCtx != nil {
mergeHeaders(reqCtx.Resp.Header(), rsp.Headers, b.log)
}
responder.Object(query.GetResponseCode(rsp.QDR), &query.QueryDataResponse{
QueryDataResponse: *rsp.QDR, // wrap the backend response as a QueryDataResponse
})
}), nil
}
@@ -231,39 +245,54 @@ func logEmptyRefids(queries []v0alpha1.DataQuery, logger log.Logger) {
}
}
func (b *QueryAPIBuilder) execute(ctx context.Context, req parsedRequestInfo, instanceConfig clientapi.InstanceConfigurationSettings) (qdr *backend.QueryDataResponse, err error) {
func mergeHeaders(main http.Header, extra http.Header, l log.Logger) {
for headerName, extraValues := range extra {
mainValues := main.Values(headerName)
for _, extraV := range extraValues {
if !slices.Contains(mainValues, extraV) {
main.Add(headerName, extraV)
} else {
l.Warn("skipped duplicate response header", "header", headerName, "value", extraV)
}
}
}
}
func (b *QueryAPIBuilder) execute(ctx context.Context, req parsedRequestInfo, instanceConfig clientapi.InstanceConfigurationSettings) (*clientapi.Response, error) {
var rsp *clientapi.Response
var err error
switch len(req.Requests) {
case 0:
b.log.Debug("executing empty query")
qdr = &backend.QueryDataResponse{}
rsp = &clientapi.Response{QDR: &backend.QueryDataResponse{}}
case 1:
b.log.Debug("executing single query")
qdr, err = b.handleQuerySingleDatasource(ctx, req.Requests[0], instanceConfig)
rsp, err = b.handleQuerySingleDatasource(ctx, req.Requests[0], instanceConfig)
if err != nil {
b.log.Debug("handleQuerySingleDatasource failed", err)
}
if err == nil && isSingleAlertQuery(req) {
b.log.Debug("handling alert query with single query")
qdr, err = b.convertQueryFromAlerting(ctx, req.Requests[0], qdr)
rsp.QDR, err = b.convertQueryFromAlerting(ctx, req.Requests[0], rsp.QDR)
if err != nil {
b.log.Debug("convertQueryFromAlerting failed", "err", err)
}
}
default:
b.log.Debug("executing concurrent queries")
qdr, err = b.executeConcurrentQueries(ctx, req.Requests, instanceConfig)
rsp, err = b.executeConcurrentQueries(ctx, req.Requests, instanceConfig)
if err != nil {
b.log.Debug("error in executeConcurrentQueries", "err", err)
}
}
if err != nil {
b.log.Debug("error in query phase, skipping expressions", "error", err)
return qdr, err //return early here to prevent expressions from being executed if we got an error during the query phase
return rsp, err //return early here to prevent expressions from being executed if we got an error during the query phase
}
if len(req.Expressions) > 0 {
b.log.Debug("executing expressions")
qdr, err = b.handleExpressions(ctx, req, qdr)
rsp.QDR, err = b.handleExpressions(ctx, req, rsp.QDR)
if err != nil {
b.log.Debug("handleExpressions failed", "err", err)
}
@@ -271,18 +300,18 @@ func (b *QueryAPIBuilder) execute(ctx context.Context, req parsedRequestInfo, in
// Remove hidden results
for _, refId := range req.HideBeforeReturn {
r, ok := qdr.Responses[refId]
r, ok := rsp.QDR.Responses[refId]
if ok && r.Error == nil {
delete(qdr.Responses, refId)
delete(rsp.QDR.Responses, refId)
}
}
return qdr, err
return rsp, err
}
// Process a single request
// See: https://github.com/grafana/grafana/blob/v10.2.3/pkg/services/query/query.go#L242
func (b *QueryAPIBuilder) handleQuerySingleDatasource(ctx context.Context, req datasourceRequest, instanceConfig clientapi.InstanceConfigurationSettings) (*backend.QueryDataResponse, error) {
func (b *QueryAPIBuilder) handleQuerySingleDatasource(ctx context.Context, req datasourceRequest, instanceConfig clientapi.InstanceConfigurationSettings) (*clientapi.Response, error) {
ctx, span := b.tracer.Start(ctx, "Query.handleQuerySingleDatasource")
defer span.End()
span.SetAttributes(
@@ -298,7 +327,7 @@ func (b *QueryAPIBuilder) handleQuerySingleDatasource(ctx context.Context, req d
}
}
if allHidden {
return &backend.QueryDataResponse{}, nil
return &clientapi.Response{}, nil
}
client, err := b.clientSupplier.GetDataSourceClient(
@@ -321,14 +350,14 @@ func (b *QueryAPIBuilder) handleQuerySingleDatasource(ctx context.Context, req d
if err == nil && rsp != nil {
for _, q := range req.Request.Queries {
if q.ResultAssertions != nil {
result, ok := rsp.Responses[q.RefID]
result, ok := rsp.QDR.Responses[q.RefID]
if ok && result.Error == nil {
err = q.ResultAssertions.Validate(result.Frames)
if err != nil {
b.log.Error("Validate failed", "err", err)
result.Error = err
result.ErrorSource = backend.ErrorSourceDownstream
rsp.Responses[q.RefID] = result
rsp.QDR.Responses[q.RefID] = result
}
}
}
@@ -344,24 +373,24 @@ func (b *QueryAPIBuilder) handleQuerySingleDatasource(ctx context.Context, req d
}
// buildErrorResponses applies the provided error to each query response in the list. These queries should all belong to the same datasource.
func buildErrorResponse(err error, req datasourceRequest) *backend.QueryDataResponse {
func buildErrorResponse(err error, req datasourceRequest) *clientapi.Response {
rsp := backend.NewQueryDataResponse()
for _, query := range req.Request.Queries {
rsp.Responses[query.RefID] = backend.DataResponse{
Error: err,
}
}
return rsp
return &clientapi.Response{QDR: rsp, Headers: nil}
}
// executeConcurrentQueries executes queries to multiple datasources concurrently and returns the aggregate result.
func (b *QueryAPIBuilder) executeConcurrentQueries(ctx context.Context, requests []datasourceRequest, instanceConfig clientapi.InstanceConfigurationSettings) (*backend.QueryDataResponse, error) {
func (b *QueryAPIBuilder) executeConcurrentQueries(ctx context.Context, requests []datasourceRequest, instanceConfig clientapi.InstanceConfigurationSettings) (*clientapi.Response, error) {
ctx, span := b.tracer.Start(ctx, "Query.executeConcurrentQueries")
defer span.End()
g, ctx := errgroup.WithContext(ctx)
g.SetLimit(b.concurrentQueryLimit) // prevent too many concurrent requests
rchan := make(chan *backend.QueryDataResponse, len(requests))
rchan := make(chan *clientapi.Response, len(requests))
// Create panic recovery function for loop below
recoveryFn := func(req datasourceRequest) {
@@ -386,9 +415,9 @@ func (b *QueryAPIBuilder) executeConcurrentQueries(ctx context.Context, requests
g.Go(func() error {
defer recoveryFn(req)
dqr, err := b.handleQuerySingleDatasource(ctx, req, instanceConfig)
rsp, err := b.handleQuerySingleDatasource(ctx, req, instanceConfig)
if err == nil {
rchan <- dqr
rchan <- rsp
} else {
rchan <- buildErrorResponse(err, req)
}
@@ -402,14 +431,19 @@ func (b *QueryAPIBuilder) executeConcurrentQueries(ctx context.Context, requests
close(rchan)
// Merge the results from each response
resp := backend.NewQueryDataResponse()
rsp := &clientapi.Response{
QDR: backend.NewQueryDataResponse(),
Headers: http.Header{},
}
for result := range rchan {
for refId, dataResponse := range result.Responses {
resp.Responses[refId] = dataResponse
for refId, dataResponse := range result.QDR.Responses {
rsp.QDR.Responses[refId] = dataResponse
}
mergeHeaders(rsp.Headers, result.Headers, b.log)
}
return resp, nil
return rsp, nil
}
// Unlike the implementation in expr/node.go, all datasource queries have been processed first
pkg/registry/apis/query/query_test.go
+108
View File
@@ -173,3 +173,111 @@ func (m mockClient) CheckHealth(ctx context.Context, req *backend.CheckHealthReq
func (m mockClient) GetInstanceConfigurationSettings(_ context.Context) (clientapi.InstanceConfigurationSettings, error) {
return clientapi.InstanceConfigurationSettings{}, nil
}
func TestMergeHeaders(t *testing.T) {
tests := []struct {
name string
h1 http.Header
h2 http.Header
expected http.Header
}{
{
name: "into empty",
h1: http.Header{},
h2: http.Header{
"A": {"1", "2"},
"B": {"3"},
},
expected: http.Header{
"A": {"1", "2"},
"B": {"3"},
},
},
{
name: "from empty",
h1: http.Header{
"A": {"1", "2"},
"B": {"3"},
},
h2: http.Header{},
expected: http.Header{
"A": {"1", "2"},
"B": {"3"},
},
},
{
name: "from nil",
h1: http.Header{
"A": {"1", "2"},
"B": {"3"},
},
h2: nil,
expected: http.Header{
"A": {"1", "2"},
"B": {"3"},
},
},
{
name: "no merging",
h1: http.Header{
"A": {"1", "2"},
},
h2: http.Header{
"B": {"3", "4"},
},
expected: http.Header{
"A": {"1", "2"},
"B": {"3", "4"},
},
},
{
name: "with merging",
h1: http.Header{
"A": {"1", "2"},
},
h2: http.Header{
"A": {"3", "4"},
},
expected: http.Header{
"A": {"1", "2", "3", "4"},
},
},
{
name: "with duplicates",
h1: http.Header{
"A": {"1", "2"},
},
h2: http.Header{
"A": {"2", "3"},
},
expected: http.Header{
"A": {"1", "2", "3"},
},
},
{
name: "with all",
h1: http.Header{
"A": {"1", "2", "3"},
"B": {"4"},
},
h2: http.Header{
"A": {"3", "4"},
"B": {"5"},
"C": {"6"},
},
expected: http.Header{
"A": {"1", "2", "3", "4"},
"B": {"4", "5"},
"C": {"6"},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
h1 := tt.h1.Clone() // don't mutate the test-data
mergeHeaders(h1, tt.h2, log.New("test.logger"))
require.Equal(t, tt.expected, h1)
})
}
}