contrib/gorm.io/gorm.v1: propagate parent span (#2759)

Co-authored-by: Michel Chretien <[email protected]> Co-authored-by: Michel Chrétien <[email protected]>
DataDog · Jul 4, 2024 · f00b4a7 · f00b4a7
1 parent fcda398
commit f00b4a7
Show file tree

Hide file tree

Showing 3 changed files with 201 additions and 45 deletions.
diff --git a/contrib/gorm.io/gorm.v1/example_test.go b/contrib/gorm.io/gorm.v1/example_test.go
@@ -7,6 +7,7 @@ package gorm_test
 
 import (
 	"context"
+	"errors"
 	"log"
 
 	sqltrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql"
@@ -41,6 +42,31 @@ func ExampleOpen() {
 	db.Where("name = ?", "jinzhu").First(&user)
 }
 
+// ExampleNewTracePlugin illustrates how to trace gorm using the gorm.Plugin api.
+func ExampleNewTracePlugin() {
+	// Register augments the provided driver with tracing, enabling it to be loaded by gorm.Open and the gormtrace.TracePlugin.
+	sqltrace.Register("pgx", &stdlib.Driver{}, sqltrace.WithServiceName("my-service"))
+	sqlDb, err := sqltrace.Open("pgx", "postgres://pqgotest:password@localhost/pqgotest?sslmode=disable")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db, err := gorm.Open(postgres.New(postgres.Config{Conn: sqlDb}), &gorm.Config{})
+	if err != nil {
+		log.Fatal(err)
+	}
+	var user User
+
+	errCheck := gormtrace.WithErrorCheck(func(err error) bool {
+		return !errors.Is(err, gorm.ErrRecordNotFound)
+	})
+	if err := db.Use(gormtrace.NewTracePlugin(errCheck)); err != nil {
+		log.Fatal(err)
+	}
+
+	// All calls through gorm.DB are now traced.
+	db.Where("name = ?", "jinzhu").First(&user)
+}
+
 func ExampleContext() {
 	// Register augments the provided driver with tracing, enabling it to be loaded by gormtrace.Open.
 	sqltrace.Register("pgx", &stdlib.Driver{}, sqltrace.WithServiceName("my-service"))

diff --git a/contrib/gorm.io/gorm.v1/gorm.go b/contrib/gorm.io/gorm.v1/gorm.go
@@ -7,9 +7,7 @@
 package gorm
 
 import (
-	"context"
 	"math"
-	"time"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
@@ -33,6 +31,26 @@ const (
 	gormSpanStartTimeKey = key("dd-trace-go:span")
 )
 
+type tracePlugin struct {
+	options []Option
+}
+
+// NewTracePlugin returns a new gorm.Plugin that enhances the underlying *gorm.DB with tracing.
+func NewTracePlugin(opts ...Option) gorm.Plugin {
+	return tracePlugin{
+		options: opts,
+	}
+}
+
+func (tracePlugin) Name() string {
+	return "DDTracePlugin"
+}
+
+func (g tracePlugin) Initialize(db *gorm.DB) error {
+	_, err := withCallbacks(db, g.options...)
+	return err
+}
+
 // Open opens a new (traced) database connection. The used driver must be formerly registered
 // using (gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql).Register.
 func Open(dialector gorm.Dialector, cfg *gorm.Config, opts ...Option) (*gorm.DB, error) {
@@ -57,86 +75,80 @@ func withCallbacks(db *gorm.DB, opts ...Option) (*gorm.DB, error) {
 	}
 	log.Debug("Registering Callbacks: %#v", cfg)
 
-	afterFunc := func(operationName string) func(*gorm.DB) {
+	afterFunc := func() func(*gorm.DB) {
+		return func(db *gorm.DB) {
+			after(db, cfg)
+		}
+	}
+
+	beforeFunc := func(operationName string) func(*gorm.DB) {
 		return func(db *gorm.DB) {
-			after(db, operationName, cfg)
+			before(db, operationName, cfg)
 		}
 	}
 
 	cb := db.Callback()
-	err := cb.Create().Before("gorm:create").Register("dd-trace-go:before_create", before)
+	err := cb.Create().Before("gorm:create").Register("dd-trace-go:before_create", beforeFunc("gorm.create"))
 	if err != nil {
 		return db, err
 	}
-	err = cb.Create().After("gorm:create").Register("dd-trace-go:after_create", afterFunc("gorm.create"))
+	err = cb.Create().After("gorm:create").Register("dd-trace-go:after_create", afterFunc())
 	if err != nil {
 		return db, err
 	}
-	err = cb.Update().Before("gorm:update").Register("dd-trace-go:before_update", before)
+	err = cb.Update().Before("gorm:update").Register("dd-trace-go:before_update", beforeFunc("gorm.update"))
 	if err != nil {
 		return db, err
 	}
-	err = cb.Update().After("gorm:update").Register("dd-trace-go:after_update", afterFunc("gorm.update"))
+	err = cb.Update().After("gorm:update").Register("dd-trace-go:after_update", afterFunc())
 	if err != nil {
 		return db, err
 	}
-	err = cb.Delete().Before("gorm:delete").Register("dd-trace-go:before_delete", before)
+	err = cb.Delete().Before("gorm:delete").Register("dd-trace-go:before_delete", beforeFunc("gorm.delete"))
 	if err != nil {
 		return db, err
 	}
-	err = cb.Delete().After("gorm:delete").Register("dd-trace-go:after_delete", afterFunc("gorm.delete"))
+	err = cb.Delete().After("gorm:delete").Register("dd-trace-go:after_delete", afterFunc())
 	if err != nil {
 		return db, err
 	}
-	err = cb.Query().Before("gorm:query").Register("dd-trace-go:before_query", before)
+	err = cb.Query().Before("gorm:query").Register("dd-trace-go:before_query", beforeFunc("gorm.query"))
 	if err != nil {
 		return db, err
 	}
-	err = cb.Query().After("gorm:query").Register("dd-trace-go:after_query", afterFunc("gorm.query"))
+	err = cb.Query().After("gorm:query").Register("dd-trace-go:after_query", afterFunc())
 	if err != nil {
 		return db, err
 	}
-	err = cb.Row().Before("gorm:query").Register("dd-trace-go:before_row_query", before)
+	err = cb.Row().Before("gorm:row").Register("dd-trace-go:before_row_query", beforeFunc("gorm.row_query"))
 	if err != nil {
 		return db, err
 	}
-	err = cb.Row().After("gorm:query").Register("dd-trace-go:after_row_query", afterFunc("gorm.row_query"))
+	err = cb.Row().After("gorm:row").Register("dd-trace-go:after_row_query", afterFunc())
 	if err != nil {
 		return db, err
 	}
-	err = cb.Raw().Before("gorm:query").Register("dd-trace-go:before_raw_query", before)
+	err = cb.Raw().Before("gorm:raw").Register("dd-trace-go:before_raw_query", beforeFunc("gorm.raw_query"))
 	if err != nil {
 		return db, err
 	}
-	err = cb.Raw().After("gorm:query").Register("dd-trace-go:after_raw_query", afterFunc("gorm.raw_query"))
+	err = cb.Raw().After("gorm:raw").Register("dd-trace-go:after_raw_query", afterFunc())
 	if err != nil {
 		return db, err
 	}
 	return db, nil
 }
 
-func before(scope *gorm.DB) {
-	if scope.Statement != nil && scope.Statement.Context != nil {
-		scope.Statement.Context = context.WithValue(scope.Statement.Context, gormSpanStartTimeKey, time.Now())
-	}
-}
-
-func after(db *gorm.DB, operationName string, cfg *config) {
+func before(db *gorm.DB, operationName string, cfg *config) {
 	if db.Statement == nil || db.Statement.Context == nil {
 		return
 	}
-
-	ctx := db.Statement.Context
-	t, ok := ctx.Value(gormSpanStartTimeKey).(time.Time)
-	if !ok {
+	if db.Config == nil || db.Config.DryRun {
 		return
 	}
-
 	opts := []ddtrace.StartSpanOption{
-		tracer.StartTime(t),
 		tracer.ServiceName(cfg.serviceName),
 		tracer.SpanType(ext.SpanTypeSQL),
-		tracer.ResourceName(db.Statement.SQL.String()),
 		tracer.Tag(ext.Component, componentName),
 	}
 	if !math.IsNaN(cfg.analyticsRate) {
@@ -148,10 +160,24 @@ func after(db *gorm.DB, operationName string, cfg *config) {
 		}
 	}
 
-	span, _ := tracer.StartSpanFromContext(ctx, operationName, opts...)
-	var dbErr error
-	if cfg.errCheck(db.Error) {
-		dbErr = db.Error
+	_, ctx := tracer.StartSpanFromContext(db.Statement.Context, operationName, opts...)
+	db.Statement.Context = ctx
+}
+
+func after(db *gorm.DB, cfg *config) {
+	if db.Statement == nil || db.Statement.Context == nil {
+		return
+	}
+	if db.Config == nil || db.Config.DryRun {
+		return
+	}
+	span, ok := tracer.SpanFromContext(db.Statement.Context)
+	if ok {
+		var dbErr error
+		if cfg.errCheck(db.Error) {
+			dbErr = db.Error
+		}
+		span.SetTag(ext.ResourceName, db.Statement.SQL.String())
+		span.Finish(tracer.WithError(dbErr))
 	}
-	span.Finish(tracer.WithError(dbErr))
 }