POC: command attributes hook

instrumentation/go.mongodb.org/mongo-driver/mongo/otelmongo/config.go

@@ -4,7 +4,9 @@
 package otelmongo // import "go.opentelemetry.io/contrib/instrumentation/go.mongodb.org/mongo-driver/mongo/otelmongo"
 
 import (
+	"go.mongodb.org/mongo-driver/event"
 	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -13,11 +15,13 @@
 
 // config is used to configure the mongo tracer.
 type config struct {
-	TracerProvider trace.TracerProvider
-
-	Tracer trace.Tracer
-
+	TracerProvider           trace.TracerProvider
+	Tracer                   trace.Tracer
 	CommandAttributeDisabled bool
+
+	// CommandAttributesFn is a function that returns a list of attributes using
+	// the CommandStartedEvent. This is used to add custom attributes to the span.
+	CommandAttributesFn func(event *event.CommandStartedEvent) []attribute.KeyValue
 }
 
 // newConfig returns a config with all Options set.
@@ -66,3 +70,12 @@
 		cfg.CommandAttributeDisabled = disabled
 	})
 }
+
+// WithCommandAttributesFn specifies a function that returns a list of
+// attributes using the CommandStartedEvent. This is used to add custom
+// attributes to the span.
+func WithCommandAttributesFn(fn func(event *event.CommandStartedEvent) []attribute.KeyValue) Option {
+	return optionFunc(func(cfg *config) {
+		cfg.CommandAttributesFn = fn
+	})
+}

instrumentation/go.mongodb.org/mongo-driver/mongo/otelmongo/internal/semconv/event_monitor.go

@@ -44,6 +44,7 @@
 type AttributeOptions struct {
 	collectionName           string
 	commandAttributeDisabled bool
+	commandAttributes        []attribute.KeyValue
 }
 
 // AttributeOption is a function type that modifies AttributeOptions.
@@ -65,6 +66,13 @@
 	}
 }
 
+// WithCommandAttributes is a functional option to set command attributes.
+func WithCommandAttributes(attrs ...attribute.KeyValue) AttributeOption {
+	return func(opts *AttributeOptions) {
+		opts.commandAttributes = attrs
+	}
+}
+
 // CommandStartedTraceAttrs generates trace attributes for a CommandStartedEvent
 // based on the EventMonitor version.
 func (m EventMonitor) CommandStartedTraceAttrs(
@@ -128,7 +136,13 @@
 	attrs = append(attrs, semconv1260.NetworkPeerAddress(net.JoinHostPort(hostname, strconv.Itoa(port))))
 
 	if !opts.commandAttributeDisabled {
-		attrs = append(attrs, semconv1260.DBQueryText(sanitizeCommand(evt.Command)))
+		if opts.commandAttributes != nil {
+			for _, attr := range opts.commandAttributes {
+				attrs = append(attrs, attr)
+			}
+		} else {
+			attrs = append(attrs, semconv1260.DBQueryText(sanitizeCommand(evt.Command)))
+		}
 	}
 
 	if opts.collectionName != "" {
@@ -157,7 +171,13 @@
 	attrs = append(attrs, semconv1210.NetPeerName(hostname))
 
 	if !opts.commandAttributeDisabled {
-		attrs = append(attrs, semconv1210.DBStatement(sanitizeCommand(evt.Command)))
+		if opts.commandAttributes != nil {
+			for _, attr := range opts.commandAttributes {
+				attrs = append(attrs, attr)
+			}
+		} else {
+			attrs = append(attrs, semconv1210.DBStatement(sanitizeCommand(evt.Command)))
+		}
 	}
 
 	if opts.collectionName != "" {

instrumentation/go.mongodb.org/mongo-driver/mongo/otelmongo/mongo.go

@@ -35,6 +35,10 @@
 		semconv.WithCommandAttributeDisabled(m.cfg.CommandAttributeDisabled),
 	}
 
+	if m.cfg.CommandAttributesFn != nil {
+		attrOptions = append(attrOptions, semconv.WithCommandAttributes(m.cfg.CommandAttributesFn(evt)...))
+	}
+
 	var spanName string
 	if collection, err := extractCollection(evt); err == nil && collection != "" {
 		spanName = collection + "."