contrib/uptrace/bun: initial implementation (#2771)

Co-authored-by: Rodrigo Argüello <rodrigo.arguello@datadoghq.com>

contrib/uptrace/bun/bun.go

@@ -0,0 +1,82 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+// Package bun provides helper functions for tracing the github.com/uptrace/bun package (https://github.com/uptrace/bun).
+package bun
+
+import (
+	"context"
+
+	"github.com/uptrace/bun"
+	"github.com/uptrace/bun/dialect"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/telemetry"
+)
+
+const (
+	componentName      = "uptrace/bun"
+	defaultServiceName = "bun.db"
+)
+
+func init() {
+	telemetry.LoadIntegration(componentName)
+	tracer.MarkIntegrationImported("github.com/uptrace/bun")
+}
+
+// Wrap augments the given DB with tracing.
+func Wrap(db *bun.DB, opts ...Option) {
+	cfg := new(config)
+	defaults(cfg)
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	log.Debug("contrib/uptrace/bun: Wrapping Database")
+	db.AddQueryHook(&queryHook{cfg: cfg})
+}
+
+type queryHook struct {
+	cfg *config
+}
+
+var _ bun.QueryHook = (*queryHook)(nil)
+
+// BeforeQuery starts a span before a query is executed.
+func (qh *queryHook) BeforeQuery(ctx context.Context, qe *bun.QueryEvent) context.Context {
+	var dbSystem string
+	switch qe.DB.Dialect().Name() {
+	case dialect.PG:
+		dbSystem = ext.DBSystemPostgreSQL
+	case dialect.MySQL:
+		dbSystem = ext.DBSystemMySQL
+	case dialect.MSSQL:
+		dbSystem = ext.DBSystemMicrosoftSQLServer
+	default:
+		dbSystem = ext.DBSystemOtherSQL
+	}
+	var (
+		query = qe.Query
+		opts  = []ddtrace.StartSpanOption{
+			tracer.SpanType(ext.SpanTypeSQL),
+			tracer.ResourceName(string(query)),
+			tracer.ServiceName(qh.cfg.serviceName),
+			tracer.Tag(ext.Component, componentName),
+			tracer.Tag(ext.DBSystem, dbSystem),
+		}
+	)
+	_, ctx = tracer.StartSpanFromContext(ctx, "bun.query", opts...)
+	return ctx
+}
+
+// AfterQuery finishes a span when a query returns.
+func (qh *queryHook) AfterQuery(ctx context.Context, qe *bun.QueryEvent) {
+	span, ok := tracer.SpanFromContext(ctx)
+	if !ok {
+		return
+	}
+	span.Finish(tracer.WithError(qe.Err))
+}

contrib/uptrace/bun/bun_test.go

@@ -0,0 +1,216 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package bun
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"os"
+	"testing"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
+	"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"
+	"github.com/uptrace/bun/dialect/sqlitedialect"
+	_ "modernc.org/sqlite"
+)
+
+func TestMain(m *testing.M) {
+	_, ok := os.LookupEnv("INTEGRATION")
+	if !ok {
+		fmt.Println("--- SKIP: to enable integration test, set the INTEGRATION environment variable")
+		os.Exit(0)
+	}
+	os.Exit(m.Run())
+}
+
+func setupDB(t *testing.T, driverName, dataSourceName string, opts ...Option) *bun.DB {
+	t.Helper()
+	sqlite, err := sql.Open(driverName, dataSourceName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	db := bun.NewDB(sqlite, sqlitedialect.New())
+	Wrap(db, opts...)
+
+	return db
+}
+
+func TestImplementsHook(_ *testing.T) {
+	var _ bun.QueryHook = (*queryHook)(nil)
+}
+
+func TestSelect(t *testing.T) {
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	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) {
+	t.Run("default", func(t *testing.T) {
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		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
+		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.Len(spans, 2)
+		assert.Equal(nil, err)
+		assert.Equal(1, n)
+		assert.Equal("bun.query", spans[0].OperationName())
+		assert.Equal("http.request", spans[1].OperationName())
+		assert.Equal("bun.db", spans[0].Tag(ext.ServiceName))
+		assert.Equal("fake-http-server", spans[1].Tag(ext.ServiceName))
+		assert.Equal("uptrace/bun", spans[0].Tag(ext.Component))
+		assert.Equal(ext.DBSystemOtherSQL, spans[0].Tag(ext.DBSystem))
+		assert.Equal(spans[0].ParentID(), spans[1].SpanID())
+	})
+
+	t.Run("global", func(t *testing.T) {
+		prevName := globalconfig.ServiceName()
+		defer globalconfig.SetServiceName(prevName)
+		globalconfig.SetServiceName("global-service")
+
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		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
+		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(1, n)
+		assert.Equal("bun.query", spans[0].OperationName())
+		assert.Equal("http.request", spans[1].OperationName())
+		assert.Equal("global-service", spans[0].Tag(ext.ServiceName))
+		assert.Equal("fake-http-server", spans[1].Tag(ext.ServiceName))
+		assert.Equal("uptrace/bun", spans[0].Tag(ext.Component))
+		assert.Equal(ext.DBSystemOtherSQL, spans[0].Tag(ext.DBSystem))
+	})
+
+	t.Run("custom", func(t *testing.T) {
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		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
+		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(1, n)
+		assert.Equal("bun.query", spans[0].OperationName())
+		assert.Equal("http.request", spans[1].OperationName())
+		assert.Equal("my-service-name", spans[0].Tag(ext.ServiceName))
+		assert.Equal("fake-http-server", spans[1].Tag(ext.ServiceName))
+		assert.Equal("uptrace/bun", spans[0].Tag(ext.Component))
+		assert.Equal(ext.DBSystemOtherSQL, spans[0].Tag(ext.DBSystem))
+	})
+}

contrib/uptrace/bun/example_test.go

@@ -0,0 +1,31 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package bun_test
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/uptrace/bun"
+	"github.com/uptrace/bun/dialect/sqlitedialect"
+	buntrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/uptrace/bun"
+	_ "modernc.org/sqlite"
+)
+
+func Example() {
+	sqlite, err := sql.Open("sqlite", "file::memory:?cache=shared")
+	if err != nil {
+		panic(err)
+	}
+	db := bun.NewDB(sqlite, sqlitedialect.New())
+
+	// Wrap the connection with the APM hook.
+	buntrace.Wrap(db)
+	var user struct {
+		Name string
+	}
+	_ = db.NewSelect().Column("name").Table("users").Scan(context.Background(), &user)
+}

contrib/uptrace/bun/option.go

@@ -0,0 +1,32 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024 Datadog, Inc.
+
+package bun
+
+import (
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
+)
+
+type config struct {
+	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 := defaultServiceName
+	if svc := globalconfig.ServiceName(); svc != "" {
+		service = svc
+	}
+	cfg.serviceName = service
+}
+
+// WithService sets the given service name for the client.
+func WithService(name string) Option {
+	return func(cfg *config) {
+		cfg.serviceName = name
+	}
+}