From b098318340a3b652243c915faf422a43330b9f4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dario=20Casta=C3=B1=C3=A9?= Date: Thu, 4 Jul 2024 15:18:40 +0200 Subject: [PATCH] contrib/uptrace/bun: initial implementation --- contrib/uptrace/bun/bun.go | 83 ++++++++++ contrib/uptrace/bun/bun_test.go | 246 ++++++++++++++++++++++++++++ contrib/uptrace/bun/example_test.go | 31 ++++ contrib/uptrace/bun/option.go | 64 ++++++++ go.mod | 4 +- go.sum | 8 +- 6 files changed, 433 insertions(+), 3 deletions(-) create mode 100644 contrib/uptrace/bun/bun.go create mode 100644 contrib/uptrace/bun/bun_test.go create mode 100644 contrib/uptrace/bun/example_test.go create mode 100644 contrib/uptrace/bun/option.go diff --git a/contrib/uptrace/bun/bun.go b/contrib/uptrace/bun/bun.go new file mode 100644 index 0000000000..fcec303952 --- /dev/null +++ b/contrib/uptrace/bun/bun.go @@ -0,0 +1,83 @@ +// 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" + "math" + + "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" + +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), + } + ) + if !math.IsNaN(qh.cfg.analyticsRate) { + opts = append(opts, tracer.Tag(ext.EventSampleRate, qh.cfg.analyticsRate)) + } + _, ctx = tracer.StartSpanFromContext(ctx, "bun", 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)) +} diff --git a/contrib/uptrace/bun/bun_test.go b/contrib/uptrace/bun/bun_test.go new file mode 100644 index 0000000000..d1d038ffc0 --- /dev/null +++ b/contrib/uptrace/bun/bun_test.go @@ -0,0 +1,246 @@ +// 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/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(opts ...Option) *bun.DB { + sqlite, err := sql.Open("sqlite", "file::memory:?cache=shared") + if err != nil { + panic(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() + + 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", 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)) +} + +func TestServiceName(t *testing.T) { + t.Run("default", func(t *testing.T) { + assert := assert.New(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 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() + + 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", 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)) + }) + + t.Run("global", func(t *testing.T) { + globalconfig.SetServiceName("global-service") + defer globalconfig.SetServiceName("") + + assert := assert.New(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 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() + + 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", 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(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() + + 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", 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)) + }) +} + +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/example_test.go b/contrib/uptrace/bun/example_test.go new file mode 100644 index 0000000000..4c29f6c71b --- /dev/null +++ b/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) +} diff --git a/contrib/uptrace/bun/option.go b/contrib/uptrace/bun/option.go new file mode 100644 index 0000000000..8432ea52c0 --- /dev/null +++ b/contrib/uptrace/bun/option.go @@ -0,0 +1,64 @@ +// 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 ( + "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 +} + +// 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" + 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. +func WithService(name string) Option { + return func(cfg *config) { + 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() + } + } +} diff --git a/go.mod b/go.mod index c879e8440f..02331de13b 100644 --- a/go.mod +++ b/go.mod @@ -85,6 +85,8 @@ require ( github.com/tidwall/buntdb v1.3.0 github.com/tinylib/msgp v1.1.8 github.com/twitchtv/twirp v8.1.3+incompatible + github.com/uptrace/bun v1.1.17 + github.com/uptrace/bun/dialect/sqlitedialect v1.1.17 github.com/urfave/negroni v1.0.0 github.com/valyala/fasthttp v1.51.0 github.com/vektah/gqlparser/v2 v2.5.16 @@ -244,7 +246,7 @@ require ( github.com/valyala/fasttemplate v1.2.2 // indirect github.com/valyala/tcplisten v1.0.0 // indirect github.com/vmihailenco/bufpool v0.1.11 // indirect - github.com/vmihailenco/msgpack/v5 v5.3.5 // indirect + github.com/vmihailenco/msgpack/v5 v5.4.1 // indirect github.com/vmihailenco/tagparser v0.1.2 // indirect github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect github.com/xdg-go/pbkdf2 v1.0.0 // indirect diff --git a/go.sum b/go.sum index dde9e28268..bd07095c44 100644 --- a/go.sum +++ b/go.sum @@ -2048,6 +2048,10 @@ github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2 github.com/ugorji/go v1.1.4/go.mod h1:uQMGLiO92mf5W77hV/PUCpI3pbzQx3CRekS0kk+RGrc= github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU= github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= +github.com/uptrace/bun v1.1.17 h1:qxBaEIo0hC/8O3O6GrMDKxqyT+mw5/s0Pn/n6xjyGIk= +github.com/uptrace/bun v1.1.17/go.mod h1:hATAzivtTIRsSJR4B8AXR+uABqnQxr3myKDKEf5iQ9U= +github.com/uptrace/bun/dialect/sqlitedialect v1.1.17 h1:i8NFU9r8YuavNFaYlNqi4ppn+MgoHtqLgpWQDrVTjm0= +github.com/uptrace/bun/dialect/sqlitedialect v1.1.17/go.mod h1:YF0FO4VVnY9GHNH6rM4r3STlVEBxkOc6L88Bm5X5mzA= github.com/urfave/cli v0.0.0-20171014202726-7bc6a0acffa5/go.mod h1:70zkFmudgCuE/ngEzBv17Jvp/497gISqfk5gWijbERA= github.com/urfave/cli v1.20.0/go.mod h1:70zkFmudgCuE/ngEzBv17Jvp/497gISqfk5gWijbERA= github.com/urfave/cli v1.22.1/go.mod h1:Gos4lmkARVdJ6EkW0WaNv/tZAAMe9V7XWyB60NtXRu0= @@ -2076,8 +2080,8 @@ github.com/vishvananda/netns v0.0.0-20210104183010-2eb08e3e575f/go.mod h1:DD4vA1 github.com/vmihailenco/bufpool v0.1.11 h1:gOq2WmBrq0i2yW5QJ16ykccQ4wH9UyEsgLm6czKAd94= github.com/vmihailenco/bufpool v0.1.11/go.mod h1:AFf/MOy3l2CFTKbxwt0mp2MwnqjNEs5H/UxrkA5jxTQ= github.com/vmihailenco/msgpack/v5 v5.3.4/go.mod h1:7xyJ9e+0+9SaZT0Wt1RGleJXzli6Q/V5KbhBonMG9jc= -github.com/vmihailenco/msgpack/v5 v5.3.5 h1:5gO0H1iULLWGhs2H5tbAHIZTV8/cYafcFOr9znI5mJU= -github.com/vmihailenco/msgpack/v5 v5.3.5/go.mod h1:7xyJ9e+0+9SaZT0Wt1RGleJXzli6Q/V5KbhBonMG9jc= +github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8= +github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok= github.com/vmihailenco/tagparser v0.1.2 h1:gnjoVuB/kljJ5wICEEOpx98oXMWPLj22G67Vbd1qPqc= github.com/vmihailenco/tagparser v0.1.2/go.mod h1:OeAg3pn3UbLjkWt+rN9oFYB6u/cQgqMEUPoW2WPyhdI= github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=