From bf5259e0b24eba6b7b2e78993f0cc28c5dc74725 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Pereira?= <adrapereira@gmail.com>
Date: Wed, 11 Oct 2023 13:52:49 +0100
Subject: [PATCH] Added spans to trace.go

---
 pkg/tsdb/tempo/tempo.go |  5 ++++-
 pkg/tsdb/tempo/trace.go | 18 ++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/pkg/tsdb/tempo/tempo.go b/pkg/tsdb/tempo/tempo.go
index 7a75c70b689..b8b6600b1f4 100644
--- a/pkg/tsdb/tempo/tempo.go
+++ b/pkg/tsdb/tempo/tempo.go
@@ -10,6 +10,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
 	"github.com/grafana/grafana/pkg/infra/httpclient"
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/infra/tracing"
 	"github.com/grafana/grafana/pkg/tsdb/tempo/kinds/dataquery"
 	"github.com/grafana/tempo/pkg/tempopb"
 )
@@ -17,11 +18,13 @@ import (
 type Service struct {
 	im     instancemgmt.InstanceManager
 	logger log.Logger
+	tracer tracing.Tracer
 }
 
-func ProvideService(httpClientProvider httpclient.Provider) *Service {
+func ProvideService(httpClientProvider httpclient.Provider, tracer tracing.Tracer) *Service {
 	return &Service{
 		logger: log.New("tsdb.tempo"),
+		tracer: tracer,
 		im:     datasource.NewInstanceManager(newInstanceSettings(httpClientProvider)),
 	}
 }
diff --git a/pkg/tsdb/tempo/trace.go b/pkg/tsdb/tempo/trace.go
index 9bd4ff33f6e..4fb6718c7f7 100644
--- a/pkg/tsdb/tempo/trace.go
+++ b/pkg/tsdb/tempo/trace.go
@@ -11,12 +11,20 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/data"
 	"github.com/grafana/grafana/pkg/tsdb/tempo/kinds/dataquery"
 	"go.opentelemetry.io/collector/pdata/ptrace"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/trace"
 )
 
 func (s *Service) getTrace(ctx context.Context, pCtx backend.PluginContext, query backend.DataQuery) (*backend.DataResponse, error) {
 	result := &backend.DataResponse{}
 	refID := query.RefID
 
+	ctx, span := s.tracer.Start(ctx, "datasource.tempo.getTrace", trace.WithAttributes(
+		attribute.String("queryType", query.QueryType),
+	))
+	defer span.End()
+
 	model := &dataquery.TempoQuery{}
 	err := json.Unmarshal(query.JSON, model)
 	if err != nil {
@@ -34,11 +42,15 @@ func (s *Service) getTrace(ctx context.Context, pCtx backend.PluginContext, quer
 
 	request, err := s.createRequest(ctx, dsInfo, *model.Query, query.TimeRange.From.Unix(), query.TimeRange.To.Unix())
 	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
 		return result, err
 	}
 
 	resp, err := dsInfo.HTTPClient.Do(request)
 	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
 		return result, fmt.Errorf("failed get to tempo: %w", err)
 	}
 
@@ -55,6 +67,8 @@ func (s *Service) getTrace(ctx context.Context, pCtx backend.PluginContext, quer
 
 	if resp.StatusCode != http.StatusOK {
 		result.Error = fmt.Errorf("failed to get trace with id: %v Status: %s Body: %s", model.Query, resp.Status, string(body))
+		span.RecordError(result.Error)
+		span.SetStatus(codes.Error, result.Error.Error())
 		return result, nil
 	}
 
@@ -62,11 +76,15 @@ func (s *Service) getTrace(ctx context.Context, pCtx backend.PluginContext, quer
 	otTrace, err := pbUnmarshaler.UnmarshalTraces(body)
 
 	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
 		return &backend.DataResponse{}, fmt.Errorf("failed to convert tempo response to Otlp: %w", err)
 	}
 
 	frame, err := TraceToFrame(otTrace)
 	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
 		return &backend.DataResponse{}, fmt.Errorf("failed to transform trace %v to data frame: %w", model.Query, err)
 	}
 	frame.RefID = refID