added context with span

trace/span.go

@@ -18,6 +18,11 @@ func SpanFromContext(ctx context.Context) Span {
 	return trace.SpanFromContext(ctx)
 }
 
+// ContextWithSpan returns a child context from ctx with span embedded
+func ContextWithSpan(ctx context.Context, span Span) context.Context {
+	return trace.ContextWithSpan(ctx, span)
+}
+
 // NewSpanFromContext creates a new span from the given context.
 // If the context already has a span attached to it, the new span will be a child of the existing span.
 // If the context does not have a span attached to it, the new span will be a root span.

trace/span_test.go

@@ -51,6 +51,20 @@ func TestSpanFromContext(t *testing.T) {
 	}
 }
 
+func TestContextWithSpan(t *testing.T) {
+	provider, err := NewProvider(
+		WithContext(context.Background()),
+		WithConfig(&config.OpenTelemetry{Enabled: true}),
+	)
+	assert.NoError(t, err)
+	_, sampleSpan := provider.Tracer().Start(context.Background(), "sample span")
+
+	ctx := ContextWithSpan(context.Background(), sampleSpan)
+	if got := SpanFromContext(ctx); got != sampleSpan {
+		t.Errorf("ContextWithSpan() = %v, want %v", got, sampleSpan)
+	}
+}
+
 func TestNewSpanFromContext(t *testing.T) {
 	assert := assert.New(t)