From c9f594ee5e4e2f58a8ed600e8023d15a4442c209 Mon Sep 17 00:00:00 2001
From: cdimitroulas <c.dimitroulas@gmail.com>
Date: Tue, 19 Nov 2024 17:05:01 +0000
Subject: [PATCH] Added SpanFromWorkflowContext function

---
 contrib/datadog/tracing/interceptor.go      | 14 +++++++
 contrib/datadog/tracing/interceptor_test.go | 43 ++++++++++++++++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/contrib/datadog/tracing/interceptor.go b/contrib/datadog/tracing/interceptor.go
index 954c79598..4faadde58 100644
--- a/contrib/datadog/tracing/interceptor.go
+++ b/contrib/datadog/tracing/interceptor.go
@@ -143,6 +143,20 @@ func (t *tracerImpl) ContextWithSpan(ctx context.Context, span interceptor.Trace
 	return tracer.ContextWithSpan(ctx, span.(*tracerSpan).Span)
 }
 
+// SpanFromWorkflowContext extracts the DataDog Span object from the workflow context.
+func SpanFromWorkflowContext(ctx workflow.Context) (ddtrace.Span, bool) {
+	val := ctx.Value(activeSpanContextKey)
+	if val == nil {
+		return tracer.SpanFromContext(nil)
+	}
+
+	if span, ok := val.(*tracerSpan); ok {
+		return span.Span, true
+	}
+
+	return tracer.SpanFromContext(nil)
+}
+
 func genSpanID(idempotencyKey string) uint64 {
 	h := fnv.New64()
 	// Write() always writes all bytes and never fails; the count and error result are for implementing io.Writer.
diff --git a/contrib/datadog/tracing/interceptor_test.go b/contrib/datadog/tracing/interceptor_test.go
index caeff15a5..1d7943a01 100644
--- a/contrib/datadog/tracing/interceptor_test.go
+++ b/contrib/datadog/tracing/interceptor_test.go
@@ -22,6 +22,7 @@
 package tracing
 
 import (
+	"errors"
 	"strings"
 	"testing"
 
@@ -31,6 +32,9 @@ import (
 
 	"go.temporal.io/sdk/interceptor"
 	"go.temporal.io/sdk/internal/interceptortest"
+	"go.temporal.io/sdk/testsuite"
+	"go.temporal.io/sdk/worker"
+	"go.temporal.io/sdk/workflow"
 )
 
 type testTracer struct {
@@ -133,7 +137,6 @@ func Test_OnFinishOption(t *testing.T) {
 		Tracer: impl,
 		mt:     mt,
 	}
-
 	interceptortest.RunTestWorkflowWithError(t, trc)
 
 	spans := trc.FinishedSpans()
@@ -141,3 +144,41 @@ func Test_OnFinishOption(t *testing.T) {
 	require.Len(t, spans, 1)
 	require.Equal(t, "temporal.RunWorkflow", spans[0].Name)
 }
+
+func setCustomSpanTagWorkflow(ctx workflow.Context) error {
+	span, ok := SpanFromWorkflowContext(ctx)
+
+	if !ok {
+		return errors.New("Did not find span in workflow context")
+	}
+
+	span.SetTag("testTag", "testValue")
+	return nil
+}
+
+func Test_SpanFromWorkflowContext(t *testing.T) {
+	// Start the mock tracer.
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	var suite testsuite.WorkflowTestSuite
+	env := suite.NewTestWorkflowEnvironment()
+	env.RegisterWorkflow(setCustomSpanTagWorkflow)
+
+	impl := NewTracer(TracerOptions{})
+	testTracer := testTracer{
+		Tracer: impl,
+		mt:     mt,
+	}
+
+	// Set tracer interceptor
+	env.SetWorkerOptions(worker.Options{
+		Interceptors: []interceptor.WorkerInterceptor{interceptor.NewTracingInterceptor(testTracer)},
+	})
+
+	env.ExecuteWorkflow(setCustomSpanTagWorkflow)
+
+	require.True(t, env.IsWorkflowCompleted())
+	testSpan := mt.FinishedSpans()[0]
+	require.Equal(t, "testValue", testSpan.Tag("testTag"))
+}