From beedba2ddf378bb27d3a57299c9f39911ff35cad Mon Sep 17 00:00:00 2001
From: ttoad <530901331qq@gmail.com>
Date: Tue, 10 Dec 2024 19:47:56 +0800
Subject: [PATCH] feat: Injection trace into req.head (#1088)

---
 pkg/infinity/headers.go      | 16 +++++++-
 pkg/infinity/request.go      |  1 +
 pkg/infinity/request_test.go | 72 ++++++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/pkg/infinity/headers.go b/pkg/infinity/headers.go
index 766394f5..dade63ca 100644
--- a/pkg/infinity/headers.go
+++ b/pkg/infinity/headers.go
@@ -2,6 +2,7 @@ package infinity
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"fmt"
 	"mime/multipart"
@@ -9,6 +10,8 @@ import (
 	"strings"
 
 	"github.com/grafana/grafana-infinity-datasource/pkg/models"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
 )
 
 const dummyHeader = "xxxxxxxx"
@@ -148,6 +151,16 @@ func ApplyForwardedOAuthIdentity(requestHeaders map[string]string, settings mode
 	return req
 }
 
+// ApplyTraceHead injecting trace
+// https://opentelemetry.io/docs/specs/otel/context/api-propagators/#textmap-propagator
+func ApplyTraceHead(ctx context.Context, req *http.Request) *http.Request {
+	prop := otel.GetTextMapPropagator()
+	if prop != nil {
+		prop.Inject(ctx, propagation.HeaderCarrier(req.Header))
+	}
+	return req
+}
+
 func getQueryReqHeader(requestHeaders map[string]string, headerName string) string {
 	for name, value := range requestHeaders {
 		if strings.EqualFold(headerName, name) {
@@ -156,4 +169,5 @@ func getQueryReqHeader(requestHeaders map[string]string, headerName string) stri
 	}
 
 	return ""
-}
\ No newline at end of file
+}
+
diff --git a/pkg/infinity/request.go b/pkg/infinity/request.go
index 0b9e2b99..43b72353 100644
--- a/pkg/infinity/request.go
+++ b/pkg/infinity/request.go
@@ -34,6 +34,7 @@ func GetRequest(ctx context.Context, settings models.InfinitySettings, body io.R
 	req = ApplyBearerToken(settings, req, includeSect)
 	req = ApplyApiKeyAuth(settings, req, includeSect)
 	req = ApplyForwardedOAuthIdentity(requestHeaders, settings, req, includeSect)
+	req = ApplyTraceHead(ctx, req)
 	return req, err
 }
 
diff --git a/pkg/infinity/request_test.go b/pkg/infinity/request_test.go
index 39393b47..c185b5ae 100644
--- a/pkg/infinity/request_test.go
+++ b/pkg/infinity/request_test.go
@@ -3,12 +3,16 @@ package infinity_test
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"testing"
 
 	"github.com/grafana/grafana-infinity-datasource/pkg/infinity"
 	"github.com/grafana/grafana-infinity-datasource/pkg/models"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
 )
 
 func Test_getQueryURL(t *testing.T) {
@@ -193,6 +197,7 @@ func Test_getQueryURL(t *testing.T) {
 		})
 	}
 }
+
 func TestClient_GetExecutedURL(t *testing.T) {
 	tests := []struct {
 		name     string
@@ -266,3 +271,70 @@ func TestNormalizeURL(t *testing.T) {
 		})
 	}
 }
+
+const (
+	traceIDStr = "4bf92f3577b34da6a3ce929d0e0e4736"
+	spanIDStr  = "00f067aa0ba902b7"
+)
+
+var (
+	traceID = mustTraceIDFromHex(traceIDStr)
+	spanID  = mustSpanIDFromHex(spanIDStr)
+)
+
+func mustTraceIDFromHex(s string) (t trace.TraceID) {
+	var err error
+	t, err = trace.TraceIDFromHex(s)
+	if err != nil {
+		panic(err)
+	}
+	return
+}
+
+func mustSpanIDFromHex(s string) (t trace.SpanID) {
+	var err error
+	t, err = trace.SpanIDFromHex(s)
+	if err != nil {
+		panic(err)
+	}
+	return
+}
+
+func TestInjectTrace(t *testing.T) {
+	tests := []struct {
+		name   string
+		sc     trace.SpanContext
+		before func()
+	}{
+		{
+			name:   "empty propagation",
+			before: func() {},
+			sc:     trace.NewSpanContext(trace.SpanContextConfig{}),
+		},
+
+		{
+			name: "inject trace",
+			before: func() {
+				otel.SetTextMapPropagator(propagation.TraceContext{})
+			},
+			sc: trace.NewSpanContext(trace.SpanContextConfig{
+				TraceID: traceID,
+				SpanID:  spanID,
+				Remote:  true,
+			}),
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			test.before()
+			ctx := trace.ContextWithRemoteSpanContext(context.Background(), test.sc)
+			req, _ := http.NewRequest(http.MethodGet, "https://github.com/grafana", nil)
+			req = infinity.ApplyTraceHead(ctx, req)
+			if test.sc.HasTraceID() {
+				newCtx := otel.GetTextMapPropagator().Extract(context.Background(), propagation.HeaderCarrier(req.Header))
+				require.Equal(t, test.sc.TraceID(), trace.SpanFromContext(newCtx).SpanContext().TraceID())
+			}
+		})
+	}
+}