From 85f223d3e4684655011f52bf410870c0ec7f6fef Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Sat, 21 Sep 2024 14:36:22 +0300
Subject: [PATCH] feat: allow to configure baggage span processor

---
 uptrace/config.go  | 19 +++++++++++++++++++
 uptrace/tracing.go |  5 +++++
 2 files changed, 24 insertions(+)

diff --git a/uptrace/config.go b/uptrace/config.go
index 8272c3ef..8e9fb3b1 100644
--- a/uptrace/config.go
+++ b/uptrace/config.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/uptrace/uptrace-go/internal"
 
+	"go.opentelemetry.io/contrib/processors/baggagecopy"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/propagation"
@@ -32,6 +33,7 @@ type config struct {
 	textMapPropagator propagation.TextMapPropagator
 	tracerProvider    *sdktrace.TracerProvider
 	traceSampler      sdktrace.Sampler
+	spanProcessors    []sdktrace.SpanProcessor
 	prettyPrint       bool
 	bspOptions        []sdktrace.BatchSpanProcessorOption
 
@@ -219,6 +221,23 @@ func WithTraceSampler(sampler sdktrace.Sampler) TracingOption {
 	})
 }
 
+// WithSpanProcessor configures an additional span processor.
+func WithSpanProcessor(sp sdktrace.SpanProcessor) TracingOption {
+	return tracingOption(func(conf *config) {
+		conf.spanProcessors = append(conf.spanProcessors, sp)
+	})
+}
+
+// WithBaggageSpanProcessor configures the The Baggage span processor.
+//
+// The Baggage span processor duplicates onto a span the attributes found in Baggage in the parent context at the moment the span is started. The passed filter determines which baggage members are added to the span.
+//
+// If filter is nil, all baggage members will be added.
+func WithBaggageSpanProcessor(filter baggagecopy.Filter) TracingOption {
+	sp := baggagecopy.NewSpanProcessor(filter)
+	return WithSpanProcessor(sp)
+}
+
 // WithPropagator sets the global TextMapPropagator used by OpenTelemetry.
 // The default is propagation.TraceContext and propagation.Baggage.
 func WithPropagator(propagator propagation.TextMapPropagator) TracingOption {
diff --git a/uptrace/tracing.go b/uptrace/tracing.go
index 85533f2c..be228ddd 100644
--- a/uptrace/tracing.go
+++ b/uptrace/tracing.go
@@ -54,6 +54,11 @@ func configureTracing(ctx context.Context, client *client, conf *config) {
 	bsp := sdktrace.NewBatchSpanProcessor(exp, bspOptions...)
 	provider.RegisterSpanProcessor(bsp)
 
+	// Register additional span processors.
+	for _, sp := range conf.spanProcessors {
+		provider.RegisterSpanProcessor(sp)
+	}
+
 	if conf.prettyPrint {
 		exporter, err := stdouttrace.New(stdouttrace.WithPrettyPrint())
 		if err != nil {