diff --git a/pkg/tsdb/tempo/metrics_stream.go b/pkg/tsdb/tempo/metrics_stream.go index 7836395cb99..b2ec7e30a95 100644 --- a/pkg/tsdb/tempo/metrics_stream.go +++ b/pkg/tsdb/tempo/metrics_stream.go @@ -8,7 +8,7 @@ import ( "io" "github.com/grafana/grafana/pkg/tsdb/tempo/traceql" - "google.golang.org/grpc/metadata" + "github.com/grafana/grafana/pkg/tsdb/tempo/util" "github.com/grafana/grafana-plugin-sdk-go/backend" "github.com/grafana/grafana-plugin-sdk-go/backend/tracing" @@ -27,7 +27,6 @@ type PartialTempoQuery struct { func (s *Service) runMetricsStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender, datasource *DatasourceInfo) error { ctx, span := tracing.DefaultTracer().Start(ctx, "datasource.tempo.runMetricsStream") defer span.End() - backend.Logger.Warn("runMetricsStream called heeer") response := &backend.DataResponse{} @@ -59,19 +58,14 @@ func (s *Service) runMetricsStream(ctx context.Context, req *backend.RunStreamRe } if qrr.GetQuery() == "" { - return backend.DownstreamErrorf("tempo metrics stream search query cannot be empty") + return backend.DownstreamErrorf("tempo search query cannot be empty") } qrr.Start = uint64(backendQuery.TimeRange.From.UnixNano()) qrr.End = uint64(backendQuery.TimeRange.To.UnixNano()) - // Setting the user agent for the gRPC call. When DS is decoupled we don't recreate instance when grafana config - // changes or updates, so we have to get it from context. - // Ideally this would be pushed higher, so it's set once for all rpc calls, but we have only one now. - ctx = metadata.AppendToOutgoingContext(ctx, "User-Agent", backend.UserAgentFromContext(ctx).String()) - for key, value := range req.Headers { - ctx = metadata.AppendToOutgoingContext(ctx, key, value) - } + ctx = util.AppendHeadersToOutgoingContext(ctx, req) + if isInstantQuery(tempoQuery.MetricsQueryType) { instantQuery := &tempopb.QueryInstantRequest{ Query: qrr.Query, diff --git a/pkg/tsdb/tempo/search_stream.go b/pkg/tsdb/tempo/search_stream.go index ac5700f56ed..5496c0a7a3e 100644 --- a/pkg/tsdb/tempo/search_stream.go +++ b/pkg/tsdb/tempo/search_stream.go @@ -7,12 +7,11 @@ import ( "fmt" "io" - "google.golang.org/grpc/metadata" - "github.com/grafana/grafana-plugin-sdk-go/backend" "github.com/grafana/grafana-plugin-sdk-go/backend/tracing" "github.com/grafana/grafana-plugin-sdk-go/data" "github.com/grafana/grafana/pkg/tsdb/tempo/kinds/dataquery" + "github.com/grafana/grafana/pkg/tsdb/tempo/util" "github.com/grafana/tempo/pkg/tempopb" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -34,7 +33,6 @@ type StreamSender interface { func (s *Service) runSearchStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender, datasource *DatasourceInfo) error { ctx, span := tracing.DefaultTracer().Start(ctx, "datasource.tempo.runSearchStream") defer span.End() - backend.Logger.Warn("runSearchStream called heeer") response := &backend.DataResponse{} var backendQuery *backend.DataQuery @@ -56,21 +54,14 @@ func (s *Service) runSearchStream(ctx context.Context, req *backend.RunStreamReq } if sr.GetQuery() == "" { - return backend.DownstreamErrorf("tempo run search query cannot be empty") + return backend.DownstreamErrorf("tempo search query cannot be empty") } sr.Start = uint32(backendQuery.TimeRange.From.Unix()) sr.End = uint32(backendQuery.TimeRange.To.Unix()) - // Setting the user agent for the gRPC call. When DS is decoupled we don't recreate instance when grafana config - // changes or updates, so we have to get it from context. - // Ideally this would be pushed higher, so it's set once for all rpc calls, but we have only one now. - ctx = metadata.AppendToOutgoingContext(ctx, "User-Agent", backend.UserAgentFromContext(ctx).String()) - // append the rest of the headers - backend.Logger.Warn("Headers:", "headers", req.Headers) - for key, value := range req.Headers { - ctx = metadata.AppendToOutgoingContext(ctx, key, value) - } + ctx = util.AppendHeadersToOutgoingContext(ctx, req) + stream, err := datasource.StreamingClient.Search(ctx, sr) if err != nil { span.RecordError(err) diff --git a/pkg/tsdb/tempo/stream_handler.go b/pkg/tsdb/tempo/stream_handler.go index f1c4db3ec02..c9fbb942f94 100644 --- a/pkg/tsdb/tempo/stream_handler.go +++ b/pkg/tsdb/tempo/stream_handler.go @@ -6,8 +6,7 @@ import ( "strings" "github.com/grafana/grafana-plugin-sdk-go/backend" - "github.com/grafana/grafana/pkg/components/simplejson" - "github.com/grafana/grafana/pkg/services/datasources" + "github.com/grafana/grafana/pkg/tsdb/tempo/util" ) func (s *Service) SubscribeStream(_ context.Context, req *backend.SubscribeStreamRequest) (*backend.SubscribeStreamResponse, error) { @@ -43,26 +42,13 @@ func (s *Service) RunStream(ctx context.Context, request *backend.RunStreamReque s.logger.Debug("New stream call", "path", request.Path) tempoDatasource, dsInfoErr := s.getDSInfo(ctx, request.PluginContext) - // get team http headers. - plugin := backend.PluginConfigFromContext(ctx) - headers := map[string]string{} - b := plugin.DataSourceInstanceSettings.JSONData - js, err := simplejson.NewJson(b) + // get incoming and team http headers and append to stream request. + headers, err := util.SetHeadersFromIncomingContext(ctx) if err != nil { return err } - teamHttpHeaders, err := datasources.GetTeamHTTPHeaders(js) - if err != nil { - return err - } - - for _, ruleValue := range teamHttpHeaders.Headers { - for _, accessRule := range ruleValue { - headers[accessRule.Header] = accessRule.LBACRule - } - } - backend.Logger.Warn("Team HTTP Headers: %v", headers) request.Headers = headers + if strings.HasPrefix(request.Path, SearchPathPrefix) { if dsInfoErr != nil { return backend.DownstreamErrorf("failed to get datasource information: %w", dsInfoErr) diff --git a/pkg/tsdb/tempo/util/stream_util.go b/pkg/tsdb/tempo/util/stream_util.go new file mode 100644 index 00000000000..72eea926981 --- /dev/null +++ b/pkg/tsdb/tempo/util/stream_util.go @@ -0,0 +1,90 @@ +package util + +import ( + "context" + "fmt" + + "github.com/grafana/grafana-plugin-sdk-go/backend" + "github.com/grafana/grafana/pkg/components/simplejson" + "github.com/grafana/grafana/pkg/services/datasources" + "google.golang.org/grpc/metadata" +) + +// Appends incoming request headers to the outgoing context to make sure none are lost when we make the request to tempo. +func AppendHeadersToOutgoingContext(ctx context.Context, req *backend.RunStreamRequest) context.Context { + // append all incoming headers + for key, value := range req.Headers { + ctx = metadata.AppendToOutgoingContext(ctx, key, value) + } + // Setting the user agent for the gRPC call. When DS is decoupled we don't recreate instance when grafana config + // changes or updates, so we have to get it from context. + // Ideally this would be pushed higher, so it's set once for all rpc calls, but we have only one now. + ctx = metadata.AppendToOutgoingContext(ctx, "User-Agent", backend.UserAgentFromContext(ctx).String()) + return ctx +} + +// When we receive a new query request we should make sure that all incoming HTTP headers are being forwarding to the grpc stream request +// this is to make sure that no headers are lost when we make the actual call to Tempo later on. +func SetHeadersFromIncomingContext(ctx context.Context) (map[string]string, error) { + headers := map[string]string{} + // get the plugin from context + plugin := backend.PluginConfigFromContext(ctx) + + // get the HTTP headers + teamHeaders, error := getTeamHTTPHeaders(plugin) + if error != nil { + return nil, error + } + + // get the rest of the incoming headers + headers, err := getClientOptionsHeaders(ctx, plugin) + if err != nil { + return nil, err + } + + for key, value := range teamHeaders { + headers[key] = value + } + return headers, nil +} + +func getTeamHTTPHeaders(plugin backend.PluginContext) (map[string]string, error) { + headers := map[string]string{} + // Grab the JSON data from the datasource instance settings + jsonData := plugin.DataSourceInstanceSettings.JSONData + js, err := simplejson.NewJson(jsonData) + if err != nil { + return nil, err + } + + // fetch team http headers + teamHttpHeaders, err := datasources.GetTeamHTTPHeaders(js) + if err != nil { + return nil, err + } + + // if present, set the Team HTTP Headers + if teamHttpHeaders != nil { + for _, ruleValue := range teamHttpHeaders.Headers { + for _, accessRule := range ruleValue { + headers[accessRule.Header] = accessRule.LBACRule + } + } + } + return headers, nil +} + +func getClientOptionsHeaders(ctx context.Context, plugin backend.PluginContext) (map[string]string, error) { + headers := map[string]string{} + opts, err := plugin.DataSourceInstanceSettings.HTTPClientOptions(ctx) + if err != nil { + return nil, fmt.Errorf("failed to get HTTP client options: %w", err) + } + + for name, values := range opts.Header { + for _, value := range values { + headers[name] = value + } + } + return headers, nil +}