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

clean up code into helper functions and utils file

Browse Source
This commit is contained in:
Jocelyn Collado-Kuri
2026-01-09 12:56:42 -08:00
parent f1e388ff4e
commit 9be719c9aa
4 changed files with 102 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/tsdb/tempo/metrics_stream.go
+4 -10
View File
@@ -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,
pkg/tsdb/tempo/search_stream.go
+4 -13
View File
@@ -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)
pkg/tsdb/tempo/stream_handler.go
+4 -18
View File
@@ -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)
pkg/tsdb/tempo/util/stream_util.go
+90
View File
@@ -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
}