From 8f46a7a40bb1f21e945685a55c58f887dfc8a721 Mon Sep 17 00:00:00 2001
From: Sam Sullivan <samsullivan09@gmail.com>
Date: Wed, 13 Mar 2024 03:43:48 -0700
Subject: [PATCH] contrib/99designs/gqlgen: add WithCustomTag option (#2598)

---
 contrib/99designs/gqlgen/option.go      | 28 ++++++++++++++++++-------
 contrib/99designs/gqlgen/tracer.go      | 16 ++++++++++----
 contrib/99designs/gqlgen/tracer_test.go | 10 +++++++++
 3 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/contrib/99designs/gqlgen/option.go b/contrib/99designs/gqlgen/option.go
index bcd68640cd..d098dda82f 100644
--- a/contrib/99designs/gqlgen/option.go
+++ b/contrib/99designs/gqlgen/option.go
@@ -17,14 +17,16 @@ const defaultServiceName = "graphql"
 type config struct {
 	serviceName   string
 	analyticsRate float64
+	tags          map[string]interface{}
 }
 
 // An Option configures the gqlgen integration.
-type Option func(t *config)
+type Option func(cfg *config)
 
-func defaults(t *config) {
-	t.serviceName = namingschema.ServiceNameOverrideV0(defaultServiceName, defaultServiceName)
-	t.analyticsRate = globalconfig.AnalyticsRate()
+func defaults(cfg *config) {
+	cfg.serviceName = namingschema.ServiceNameOverrideV0(defaultServiceName, defaultServiceName)
+	cfg.analyticsRate = globalconfig.AnalyticsRate()
+	cfg.tags = make(map[string]interface{})
 }
 
 // WithAnalytics enables or disables Trace Analytics for all started spans.
@@ -37,14 +39,24 @@ func WithAnalytics(on bool) Option {
 
 // WithAnalyticsRate sets the sampling rate for Trace Analytics events correlated to started spans.
 func WithAnalyticsRate(rate float64) Option {
-	return func(t *config) {
-		t.analyticsRate = rate
+	return func(cfg *config) {
+		cfg.analyticsRate = rate
 	}
 }
 
 // WithServiceName sets the given service name for the gqlgen server.
 func WithServiceName(name string) Option {
-	return func(t *config) {
-		t.serviceName = name
+	return func(cfg *config) {
+		cfg.serviceName = name
+	}
+}
+
+// WithCustomTag will attach the value to the span tagged by the key.
+func WithCustomTag(key string, value interface{}) Option {
+	return func(cfg *config) {
+		if cfg.tags == nil {
+			cfg.tags = make(map[string]interface{})
+		}
+		cfg.tags[key] = value
 	}
 }
diff --git a/contrib/99designs/gqlgen/tracer.go b/contrib/99designs/gqlgen/tracer.go
index ae541a8343..9a7954959b 100644
--- a/contrib/99designs/gqlgen/tracer.go
+++ b/contrib/99designs/gqlgen/tracer.go
@@ -139,13 +139,17 @@ func (t *gqlTracer) InterceptOperation(ctx context.Context, next graphql.Operati
 func (t *gqlTracer) InterceptField(ctx context.Context, next graphql.Resolver) (res any, err error) {
 	opCtx := graphql.GetOperationContext(ctx)
 	fieldCtx := graphql.GetFieldContext(ctx)
-	opts := []tracer.StartSpanOption{
+	opts := make([]tracer.StartSpanOption, 0, 6+len(t.cfg.tags))
+	for k, v := range t.cfg.tags {
+		opts = append(opts, tracer.Tag(k, v))
+	}
+	opts = append(opts,
 		tracer.Tag(tagGraphqlField, fieldCtx.Field.Name),
 		tracer.Tag(tagGraphqlOperationType, opCtx.Operation.Operation),
 		tracer.Tag(ext.Component, componentName),
 		tracer.ResourceName(fmt.Sprintf("%s.%s", fieldCtx.Object, fieldCtx.Field.Name)),
 		tracer.Measured(),
-	}
+	)
 	if !math.IsNaN(t.cfg.analyticsRate) {
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, t.cfg.analyticsRate))
 	}
@@ -172,14 +176,18 @@ func (*gqlTracer) InterceptResponse(ctx context.Context, next graphql.ResponseHa
 // also creates child spans (orphans in the case of a subscription) for the
 // read, parsing and validation phases of the operation.
 func (t *gqlTracer) createRootSpan(ctx context.Context, opCtx *graphql.OperationContext) (ddtrace.Span, context.Context) {
-	opts := []tracer.StartSpanOption{
+	opts := make([]tracer.StartSpanOption, 0, 7+len(t.cfg.tags))
+	for k, v := range t.cfg.tags {
+		opts = append(opts, tracer.Tag(k, v))
+	}
+	opts = append(opts,
 		tracer.SpanType(ext.SpanTypeGraphQL),
 		tracer.Tag(ext.SpanKind, ext.SpanKindServer),
 		tracer.ServiceName(t.cfg.serviceName),
 		tracer.Tag(ext.Component, componentName),
 		tracer.ResourceName(opCtx.RawQuery),
 		tracer.StartTime(opCtx.Stats.OperationStart),
-	}
+	)
 	if !math.IsNaN(t.cfg.analyticsRate) {
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, t.cfg.analyticsRate))
 	}
diff --git a/contrib/99designs/gqlgen/tracer_test.go b/contrib/99designs/gqlgen/tracer_test.go
index eda266e82c..1fa928e2bc 100644
--- a/contrib/99designs/gqlgen/tracer_test.go
+++ b/contrib/99designs/gqlgen/tracer_test.go
@@ -65,6 +65,16 @@ func TestOptions(t *testing.T) {
 				assert.Equal(0.5, root.Tag(ext.EventSampleRate))
 			},
 		},
+		"WithCustomTag": {
+			tracerOpts: []Option{
+				WithCustomTag("customTag1", "customValue1"),
+				WithCustomTag("customTag2", "customValue2"),
+			},
+			test: func(assert *assert.Assertions, root mocktracer.Span) {
+				assert.Equal("customValue1", root.Tag("customTag1"))
+				assert.Equal("customValue2", root.Tag("customTag2"))
+			},
+		},
 	} {
 		t.Run(name, func(t *testing.T) {
 			assert := assert.New(t)