feat(contrib/gomongodb.org/mongo-driver): add optional query tag truncation

contrib/go.mongodb.org/mongo-driver.v2/mongo/mongo.go

@@ -41,13 +41,11 @@ type monitor struct {
 
 func (m *monitor) Started(ctx context.Context, evt *event.CommandStartedEvent) {
 	hostname, port := peerInfo(evt)
-	b, _ := bson.MarshalExtJSON(evt.Command, false, false)
 	opts := []tracer.StartSpanOption{
 		tracer.SpanType(ext.SpanTypeMongoDB),
 		tracer.ServiceName(m.cfg.serviceName),
 		tracer.ResourceName("mongo." + evt.CommandName),
 		tracer.Tag(ext.DBInstance, evt.DatabaseName),
-		tracer.Tag(m.cfg.spanName, string(b)),
 		tracer.Tag(ext.DBType, "mongo"),
 		tracer.Tag(ext.PeerHostname, hostname),
 		tracer.Tag(ext.NetworkDestinationName, hostname),
@@ -56,6 +54,13 @@ func (m *monitor) Started(ctx context.Context, evt *event.CommandStartedEvent) {
 		tracer.Tag(ext.SpanKind, ext.SpanKindClient),
 		tracer.Tag(ext.DBSystem, ext.DBSystemMongoDB),
 	}
+	if m.cfg.maxQuerySize != 0 {
+		b, _ := bson.MarshalExtJSON(evt.Command, false, false)
+		if m.cfg.maxQuerySize > 0 && len(b) > m.cfg.maxQuerySize {
+			b = b[:m.cfg.maxQuerySize]
+		}
+		opts = append(opts, tracer.Tag(m.cfg.spanName, string(b)))
+	}
 	span, _ := tracer.StartSpanFromContext(ctx, m.cfg.spanName, opts...)
 	key := spanKey{
 		ConnectionID: evt.ConnectionID,

contrib/go.mongodb.org/mongo-driver.v2/mongo/mongo_test.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -80,3 +81,68 @@ func Test(t *testing.T) {
 	assert.Equal(t, ext.SpanKindClient, s.Tag(ext.SpanKind))
 	assert.Equal(t, "mongodb", s.Tag(ext.DBSystem))
 }
+
+func TestTruncation(t *testing.T) {
+	getQuery := func(t *testing.T, max int) (string, bool) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
+		defer cancel()
+
+		span, ctx := tracer.StartSpanFromContext(ctx, "mongodb-test")
+
+		addr := "mongodb://localhost:27017/?connect=direct"
+		opts := options.Client()
+		opts.Monitor = NewMonitor(WithMaxQuerySize(max))
+		opts.ApplyURI(addr)
+		client, err := mongo.Connect(opts)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		_, err = client.
+			Database("test-database").
+			Collection("test-collection").
+			UpdateOne(
+				ctx,
+				bson.D{{Key: "_id", Value: "68536ec8d906742797f5705a"}},
+				bson.D{{Key: "$set", Value: map[string]any{"test-item": "test-value"}}},
+			)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		span.Finish()
+
+		spans := mt.FinishedSpans()
+		value, ok := spans[0].Tag("mongodb.query").(string)
+		return value, ok
+	}
+
+	t.Run("zero", func(t *testing.T) {
+		// Should *not* attach the tag.
+		_, ok := getQuery(t, 0)
+		assert.False(t, ok)
+	})
+
+	t.Run("positive", func(t *testing.T) {
+		// Should truncate.
+		actual, _ := getQuery(t, 50)
+		assert.Equal(t, actual, `{"update":"test-collection","ordered":true,"lsid":`)
+	})
+
+	t.Run("negative", func(t *testing.T) {
+		// Should *not* truncate. The actual query contains a random session ID, so we just check the end which is deterministic.
+		actual, _ := getQuery(t, -1)
+		wantSuffix := `"u":{"$set":{"test-item":"test-value"}}}]}`
+		assert.True(t, strings.HasSuffix(actual, `"u":{"$set":{"test-item":"test-value"}}}]}`), "query %q does not end with %q", actual, wantSuffix)
+	})
+
+	t.Run("greater than query size", func(t *testing.T) {
+		// Should *not* truncate. The actual query contains a random session ID, so we just check the end which is deterministic.
+		actual, _ := getQuery(t, 1000) // arbitrary value > the size of the query we will be truncating
+		wantSuffix := `"u":{"$set":{"test-item":"test-value"}}}]}`
+		assert.True(t, strings.HasSuffix(actual, `"u":{"$set":{"test-item":"test-value"}}}]}`), "query %q does not end with %q", actual, wantSuffix)
+	})
+}

contrib/go.mongodb.org/mongo-driver.v2/mongo/option.go

@@ -10,8 +10,9 @@ import (
 )
 
 type config struct {
-	serviceName string
-	spanName    string
+	serviceName  string
+	spanName     string
+	maxQuerySize int
 }
 
 // Option describes options for the Mongo integration.
@@ -29,6 +30,7 @@ func (fn OptionFn) apply(cfg *config) {
 func defaults(cfg *config) {
 	cfg.serviceName = instr.ServiceName(instrumentation.ComponentDefault, nil)
 	cfg.spanName = instr.OperationName(instrumentation.ComponentDefault, nil)
+	cfg.maxQuerySize = -1
 }
 
 // WithService sets the given service name for this integration spans.
@@ -37,3 +39,18 @@ func WithService(name string) OptionFn {
 		cfg.serviceName = name
 	}
 }
+
+// WithMaxQuerySize sets the maximum query size (in bytes) before queries
+// are truncated when attached as a span tag.
+//
+// If negative (the default), query truncation is disabled and the query
+// will always be attached in full.
+//
+// If zero, traces will not include a query tag.
+//
+// Defaults to -1.
+func WithMaxQuerySize(max int) OptionFn {
+	return func(cfg *config) {
+		cfg.maxQuerySize = max
+	}
+}

contrib/go.mongodb.org/mongo-driver/mongo/mongo.go

@@ -45,13 +45,11 @@ type monitor struct {
 
 func (m *monitor) Started(ctx context.Context, evt *event.CommandStartedEvent) {
 	hostname, port := peerInfo(evt)
-	b, _ := bson.MarshalExtJSON(evt.Command, false, false)
 	opts := []tracer.StartSpanOption{
 		tracer.SpanType(ext.SpanTypeMongoDB),
 		tracer.ServiceName(m.cfg.serviceName),
 		tracer.ResourceName("mongo." + evt.CommandName),
 		tracer.Tag(ext.DBInstance, evt.DatabaseName),
-		tracer.Tag(m.cfg.spanName, string(b)),
 		tracer.Tag(ext.DBType, "mongo"),
 		tracer.Tag(ext.PeerHostname, hostname),
 		tracer.Tag(ext.NetworkDestinationName, hostname),
@@ -63,6 +61,13 @@ func (m *monitor) Started(ctx context.Context, evt *event.CommandStartedEvent) {
 	if !math.IsNaN(m.cfg.analyticsRate) {
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, m.cfg.analyticsRate))
 	}
+	if m.cfg.maxQuerySize != 0 {
+		b, _ := bson.MarshalExtJSON(evt.Command, false, false)
+		if m.cfg.maxQuerySize > 0 && len(b) > m.cfg.maxQuerySize {
+			b = b[:m.cfg.maxQuerySize]
+		}
+		opts = append(opts, tracer.Tag(m.cfg.spanName, string(b)))
+	}
 	span, _ := tracer.StartSpanFromContext(ctx, m.cfg.spanName, opts...)
 	key := spanKey{
 		ConnectionID: evt.ConnectionID,

contrib/go.mongodb.org/mongo-driver/mongo/mongo_test.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -136,3 +137,68 @@ func TestAnalyticsSettings(t *testing.T) {
 		assertRate(t, mt, 0.23, WithAnalyticsRate(0.23))
 	})
 }
+
+func TestTruncation(t *testing.T) {
+	getQuery := func(t *testing.T, max int) (string, bool) {
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
+		defer cancel()
+
+		span, ctx := tracer.StartSpanFromContext(ctx, "mongodb-test")
+
+		addr := "mongodb://localhost:27017/?connect=direct"
+		opts := options.Client()
+		opts.Monitor = NewMonitor(WithMaxQuerySize(max))
+		opts.ApplyURI(addr)
+		client, err := mongo.Connect(ctx, opts)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		_, err = client.
+			Database("test-database").
+			Collection("test-collection").
+			UpdateOne(
+				ctx,
+				bson.D{{Key: "_id", Value: "68536ec8d906742797f5705a"}},
+				bson.D{{Key: "$set", Value: map[string]any{"test-item": "test-value"}}},
+			)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		span.Finish()
+
+		spans := mt.FinishedSpans()
+		value, ok := spans[0].Tag("mongodb.query").(string)
+		return value, ok
+	}
+
+	t.Run("zero", func(t *testing.T) {
+		// Should *not* attach the tag.
+		_, ok := getQuery(t, 0)
+		assert.False(t, ok)
+	})
+
+	t.Run("positive", func(t *testing.T) {
+		// Should truncate.
+		actual, _ := getQuery(t, 50)
+		assert.Equal(t, actual, `{"update":"test-collection","ordered":true,"lsid":`)
+	})
+
+	t.Run("negative", func(t *testing.T) {
+		// Should *not* truncate. The actual query contains a random session ID, so we just check the end which is deterministic.
+		actual, _ := getQuery(t, -1)
+		wantSuffix := `"u":{"$set":{"test-item":"test-value"}}}]}`
+		assert.True(t, strings.HasSuffix(actual, `"u":{"$set":{"test-item":"test-value"}}}]}`), "query %q does not end with %q", actual, wantSuffix)
+	})
+
+	t.Run("greater than query size", func(t *testing.T) {
+		// Should *not* truncate. The actual query contains a random session ID, so we just check the end which is deterministic.
+		actual, _ := getQuery(t, 1000) // arbitrary value > the size of the query we will be truncating
+		wantSuffix := `"u":{"$set":{"test-item":"test-value"}}}]}`
+		assert.True(t, strings.HasSuffix(actual, `"u":{"$set":{"test-item":"test-value"}}}]}`), "query %q does not end with %q", actual, wantSuffix)
+	})
+}

contrib/go.mongodb.org/mongo-driver/mongo/option.go

@@ -15,6 +15,7 @@ type config struct {
 	serviceName   string
 	spanName      string
 	analyticsRate float64
+	maxQuerySize  int
 }
 
 // Option describes options for the Mongo integration.
@@ -33,6 +34,7 @@ func defaults(cfg *config) {
 	cfg.serviceName = instr.ServiceName(instrumentation.ComponentDefault, nil)
 	cfg.spanName = instr.OperationName(instrumentation.ComponentDefault, nil)
 	cfg.analyticsRate = instr.AnalyticsRate(false)
+	cfg.maxQuerySize = -1
 }
 
 // WithService sets the given service name for this integration spans.
@@ -64,3 +66,18 @@ func WithAnalyticsRate(rate float64) OptionFn {
 		}
 	}
 }
+
+// WithMaxQuerySize sets the maximum query size (in bytes) before queries
+// are truncated when attached as a span tag.
+//
+// If negative (the default), query truncation is disabled and the query
+// will always be attached in full.
+//
+// If zero, traces will not include a query tag.
+//
+// Defaults to -1.
+func WithMaxQuerySize(max int) OptionFn {
+	return func(cfg *config) {
+		cfg.maxQuerySize = max
+	}
+}