contrib/uptrace/bun: remove trace analytics options and test against …

…all supported databases
DataDog · Jul 5, 2024 · 33f5b50 · 33f5b50
1 parent aedd371
commit 33f5b50
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 137 deletions.
diff --git a/contrib/uptrace/bun/bun.go b/contrib/uptrace/bun/bun.go
@@ -8,7 +8,6 @@ package bun
 
 import (
 	"context"
-	"math"
 
 	"github.com/uptrace/bun"
 	"github.com/uptrace/bun/dialect"
@@ -19,7 +18,10 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/telemetry"
 )
 
-const componentName = "uptrace/bun"
+const (
+	componentName      = "uptrace/bun"
+	defaultServiceName = "bun.db"
+)
 
 func init() {
 	telemetry.LoadIntegration(componentName)
@@ -66,9 +68,6 @@ func (qh *queryHook) BeforeQuery(ctx context.Context, qe *bun.QueryEvent) contex
 			tracer.Tag(ext.DBSystem, dbSystem),
 		}
 	)
-	if !math.IsNaN(qh.cfg.analyticsRate) {
-		opts = append(opts, tracer.Tag(ext.EventSampleRate, qh.cfg.analyticsRate))
-	}
 	_, ctx = tracer.StartSpanFromContext(ctx, "bun.query", opts...)
 	return ctx
 }

diff --git a/contrib/uptrace/bun/bun_test.go b/contrib/uptrace/bun/bun_test.go
@@ -17,6 +17,9 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
 
+	_ "github.com/go-sql-driver/mysql"
+	_ "github.com/lib/pq"
+	_ "github.com/microsoft/go-mssqldb"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/uptrace/bun"
@@ -33,10 +36,11 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
-func setupDB(opts ...Option) *bun.DB {
-	sqlite, err := sql.Open("sqlite", "file::memory:?cache=shared")
+func setupDB(t *testing.T, driverName, dataSourceName string, opts ...Option) *bun.DB {
+	t.Helper()
+	sqlite, err := sql.Open(driverName, dataSourceName)
 	if err != nil {
-		panic(err)
+		t.Fatal(err)
 	}
 
 	db := bun.NewDB(sqlite, sqlitedialect.New())
@@ -54,29 +58,64 @@ func TestSelect(t *testing.T) {
 	mt := mocktracer.Start()
 	defer mt.Stop()
 
-	db := setupDB()
-	parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
-		tracer.ServiceName("fake-http-server"),
-		tracer.SpanType(ext.SpanTypeWeb),
-	)
-
-	var n, rows int64
-	// Using WithContext will make the postgres span a child of
-	// the span inside ctx (parentSpan)
-	res, err := db.NewSelect().ColumnExpr("1").Exec(ctx, &n)
-	parentSpan.Finish()
-	spans := mt.FinishedSpans()
-
-	require.NoError(t, err)
-	rows, _ = res.RowsAffected()
-	assert.Equal(int64(1), rows)
-	assert.Equal(2, len(spans))
-	assert.Equal(nil, err)
-	assert.Equal(int64(1), n)
-	assert.Equal("bun.query", spans[0].OperationName())
-	assert.Equal("http.request", spans[1].OperationName())
-	assert.Equal("uptrace/bun", spans[0].Tag(ext.Component))
-	assert.Equal(ext.DBSystemOtherSQL, spans[0].Tag(ext.DBSystem))
+	tC := []struct {
+		name       string
+		driver     string
+		dataSource string
+		expected   string
+	}{
+		{
+			name:       "SQLite",
+			driver:     "sqlite",
+			dataSource: "file::memory:?cache=shared",
+			expected:   ext.DBSystemOtherSQL,
+		},
+		{
+			name:       "Postgres",
+			driver:     "postgres",
+			dataSource: "postgres://postgres:postgres@127.0.0.1:5432/postgres?sslmode=disable",
+			expected:   ext.DBSystemPostgreSQL,
+		},
+		{
+			name:       "MySQL",
+			driver:     "mysql",
+			dataSource: "test:test@tcp(127.0.0.1:3306)/test",
+			expected:   ext.DBSystemMySQL,
+		},
+		{
+			name:       "MSSQL",
+			driver:     "sqlserver",
+			dataSource: "sqlserver://sa:myPassw0rd@127.0.0.1:1433?database=master",
+			expected:   ext.DBSystemMicrosoftSQLServer,
+		},
+	}
+	for _, tt := range tC {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			db := setupDB(t, tt.driver, tt.dataSource)
+			parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
+				tracer.ServiceName("fake-http-server"),
+				tracer.SpanType(ext.SpanTypeWeb),
+			)
+
+			var n, rows int64
+			res, err := db.NewSelect().ColumnExpr("1").Exec(ctx, &n)
+			parentSpan.Finish()
+			spans := mt.FinishedSpans()
+
+			require.NoError(t, err)
+			rows, _ = res.RowsAffected()
+			assert.Equal(int64(1), rows)
+			assert.Equal(2, len(spans))
+			assert.Equal(nil, err)
+			assert.Equal(int64(1), n)
+			assert.Equal("bun.query", spans[0].OperationName())
+			assert.Equal("http.request", spans[1].OperationName())
+			assert.Equal("uptrace/bun", spans[0].Tag(ext.Component))
+			assert.Equal(ext.DBSystemOtherSQL, spans[0].Tag(ext.DBSystem))
+			mt.Reset()
+		})
+	}
 }
 
 func TestServiceName(t *testing.T) {
@@ -85,7 +124,7 @@ func TestServiceName(t *testing.T) {
 		mt := mocktracer.Start()
 		defer mt.Stop()
 
-		db := setupDB()
+		db := setupDB(t, "sqlite", "file::memory:?cache=shared")
 		parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
 			tracer.ServiceName("fake-http-server"),
 			tracer.SpanType(ext.SpanTypeWeb),
@@ -120,15 +159,13 @@ func TestServiceName(t *testing.T) {
 		mt := mocktracer.Start()
 		defer mt.Stop()
 
-		db := setupDB()
+		db := setupDB(t, "sqlite", "file::memory:?cache=shared")
 		parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
 			tracer.ServiceName("fake-http-server"),
 			tracer.SpanType(ext.SpanTypeWeb),
 		)
 
 		var n int
-		// Using WithContext will make the postgres span a child of
-		// the span inside ctx (parentSpan)
 		res, err := db.NewSelect().ColumnExpr("1").Exec(ctx, &n)
 		parentSpan.Finish()
 		spans := mt.FinishedSpans()
@@ -152,15 +189,13 @@ func TestServiceName(t *testing.T) {
 		mt := mocktracer.Start()
 		defer mt.Stop()
 
-		db := setupDB(WithService("my-service-name"))
+		db := setupDB(t, "sqlite", "file::memory:?cache=shared", WithService("my-service-name"))
 		parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
 			tracer.ServiceName("fake-http-server"),
 			tracer.SpanType(ext.SpanTypeWeb),
 		)
 
 		var n int
-		// Using WithContext will make the postgres span a child of
-		// the span inside ctx (parentSpan)
 		res, err := db.NewSelect().ColumnExpr("1").Exec(ctx, &n)
 		parentSpan.Finish()
 		spans := mt.FinishedSpans()
@@ -179,68 +214,3 @@ func TestServiceName(t *testing.T) {
 		assert.Equal(ext.DBSystemOtherSQL, spans[0].Tag(ext.DBSystem))
 	})
 }
-
-func TestAnalyticsSettings(t *testing.T) {
-	assertRate := func(t *testing.T, mt mocktracer.Tracer, rate interface{}, opts ...Option) {
-		db := setupDB(opts...)
-		parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
-			tracer.ServiceName("fake-http-server"),
-			tracer.SpanType(ext.SpanTypeWeb),
-		)
-
-		var n int
-		_, err := db.NewSelect().ColumnExpr("1").Exec(ctx, &n)
-		parentSpan.Finish()
-
-		require.NoError(t, err)
-
-		spans := mt.FinishedSpans()
-		assert.Len(t, spans, 2)
-		s := spans[0]
-		assert.Equal(t, rate, s.Tag(ext.EventSampleRate))
-	}
-
-	t.Run("defaults", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
-
-		assertRate(t, mt, nil)
-	})
-
-	t.Run("global", func(t *testing.T) {
-		t.Skip("global flag disabled")
-		mt := mocktracer.Start()
-		defer mt.Stop()
-
-		rate := globalconfig.AnalyticsRate()
-		defer globalconfig.SetAnalyticsRate(rate)
-		globalconfig.SetAnalyticsRate(0.4)
-
-		assertRate(t, mt, 0.4)
-	})
-
-	t.Run("enabled", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
-
-		assertRate(t, mt, 1.0, WithAnalytics(true))
-	})
-
-	t.Run("disabled", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
-
-		assertRate(t, mt, nil, WithAnalytics(false))
-	})
-
-	t.Run("override", func(t *testing.T) {
-		mt := mocktracer.Start()
-		defer mt.Stop()
-
-		rate := globalconfig.AnalyticsRate()
-		defer globalconfig.SetAnalyticsRate(rate)
-		globalconfig.SetAnalyticsRate(0.4)
-
-		assertRate(t, mt, 0.23, WithAnalyticsRate(0.23))
-	})
-}
diff --git a/contrib/uptrace/bun/option.go b/contrib/uptrace/bun/option.go
@@ -6,31 +6,22 @@
 package bun
 
 import (
-	"math"
-
-	"gopkg.in/DataDog/dd-trace-go.v1/internal"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
 )
 
 type config struct {
-	serviceName   string
-	analyticsRate float64
+	serviceName string
 }
 
 // Option represents an option that can be used to create or wrap a client.
 type Option func(*config)
 
 func defaults(cfg *config) {
-	service := "bun.db"
+	service := defaultServiceName
 	if svc := globalconfig.ServiceName(); svc != "" {
 		service = svc
 	}
 	cfg.serviceName = service
-	if internal.BoolEnv("DD_TRACE_BUN_ANALYTICS_ENABLED", false) {
-		cfg.analyticsRate = 1.0
-	} else {
-		cfg.analyticsRate = math.NaN()
-	}
 }
 
 // WithService sets the given service name for the client.
@@ -39,26 +30,3 @@ func WithService(name string) Option {
 		cfg.serviceName = name
 	}
 }
-
-// WithAnalytics enables Trace Analytics for all started spans.
-func WithAnalytics(on bool) Option {
-	return func(cfg *config) {
-		if on {
-			cfg.analyticsRate = 1.0
-		} else {
-			cfg.analyticsRate = math.NaN()
-		}
-	}
-}
-
-// WithAnalyticsRate sets the sampling rate for Trace Analytics events
-// correlated to started spans.
-func WithAnalyticsRate(rate float64) Option {
-	return func(cfg *config) {
-		if rate >= 0.0 && rate <= 1.0 {
-			cfg.analyticsRate = rate
-		} else {
-			cfg.analyticsRate = math.NaN()
-		}
-	}
-}