contrib/gorm.io/gorm.v1: propagate span ctx

- Start span in before callback instead of after to allow child context to be given to underlying DB driver
- Fix row and raw callbacks not being ordered correctly
- Don't generate spans for DryRun statements (e.g. when using a statement or subquery as var)

fixes DataDog#849

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"
@@ -76,86 +74,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, operationName, cfg)
+			after(db, cfg)
+		}
+	}
+
+	beforeFunc := func(operationName string) func(*gorm.DB) {
+		return func(db *gorm.DB) {
+			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) {
@@ -167,10 +159,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))
 }