contrib/jinzhu/gorm: allow custom tags. (#723)

This change allows setting custom tags via a callback with access to *gorm.Scope. The new option is called WithCustomTag. Closes #716
DataDog · Sep 8, 2020 · 498a455 · 498a455
1 parent 9135761
commit 498a455
Show file tree

Hide file tree

Showing 6 changed files with 120 additions and 0 deletions.
diff --git a/contrib/gopkg.in/jinzhu/gorm.v1/gorm.go b/contrib/gopkg.in/jinzhu/gorm.v1/gorm.go
@@ -124,6 +124,11 @@ func after(scope *gorm.Scope, operationName string) {
 	if !math.IsNaN(cfg.analyticsRate) {
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
 	}
+	if cfg.tagFns != nil {
+		for key, tagFn := range cfg.tagFns {
+			opts = append(opts, tracer.Tag(key, tagFn(scope)))
+		}
+	}
 
 	span, _ := tracer.StartSpanFromContext(ctx, operationName, opts...)
 	span.Finish(tracer.WithError(scope.DB().Error))

diff --git a/contrib/gopkg.in/jinzhu/gorm.v1/gorm_test.go b/contrib/gopkg.in/jinzhu/gorm.v1/gorm_test.go
@@ -303,3 +303,44 @@ func TestContext(t *testing.T) {
 		assert.Equal(t, context.Background(), ctx)
 	})
 }
+
+func TestCustomTags(t *testing.T) {
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	sqltrace.Register("postgres", &pq.Driver{})
+	db, err := Open("postgres", "postgres://postgres:postgres@127.0.0.1:5432/postgres?sslmode=disable",
+		WithCustomTag("custom_tag", func(scope *gorm.Scope) interface{} {
+			return scope.SQLVars[3]
+		}),
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer db.Close()
+	db.AutoMigrate(&Product{})
+
+	parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
+		tracer.ServiceName("fake-http-server"),
+		tracer.SpanType(ext.SpanTypeWeb),
+	)
+
+	db = WithContext(ctx, db)
+	db.Create(&Product{Code: "L1212", Price: 1000})
+
+	parentSpan.Finish()
+
+	spans := mt.FinishedSpans()
+	assert.True(len(spans) >= 3)
+
+	// We deterministically expect the span to be the third last,
+	// followed by the underlying postgres DB trace and the above http.request span.
+	span := spans[len(spans)-3]
+	assert.Equal("gorm.create", span.OperationName())
+	assert.Equal(ext.SpanTypeSQL, span.Tag(ext.SpanType))
+	assert.Equal("L1212", span.Tag("custom_tag"))
+	assert.Equal(
+		`INSERT INTO "products" ("created_at","updated_at","deleted_at","code","price") VALUES ($1,$2,$3,$4,$5) RETURNING "products"."id"`,
+		span.Tag(ext.ResourceName))
+}
diff --git a/contrib/gopkg.in/jinzhu/gorm.v1/option.go b/contrib/gopkg.in/jinzhu/gorm.v1/option.go
@@ -9,12 +9,15 @@ import (
 	"math"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/internal"
+
+	"gopkg.in/jinzhu/gorm.v1"
 )
 
 type config struct {
 	serviceName   string
 	analyticsRate float64
 	dsn           string
+	tagFns        map[string]func(scope *gorm.Scope) interface{}
 }
 
 // Option represents an option that can be passed to Register, Open or OpenDB.
@@ -60,3 +63,14 @@ func WithAnalyticsRate(rate float64) Option {
 		}
 	}
 }
+
+// WithCustomTag will cause the given tagFn to be evaluated after executing
+// a query and attach the result to the span tagged by the key.
+func WithCustomTag(tag string, tagFn func(scope *gorm.Scope) interface{}) Option {
+	return func(cfg *config) {
+		if cfg.tagFns == nil {
+			cfg.tagFns = make(map[string]func(scope *gorm.Scope) interface{})
+		}
+		cfg.tagFns[tag] = tagFn
+	}
+}
diff --git a/contrib/jinzhu/gorm/gorm.go b/contrib/jinzhu/gorm/gorm.go
@@ -124,6 +124,11 @@ func after(scope *gorm.Scope, operationName string) {
 	if !math.IsNaN(cfg.analyticsRate) {
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
 	}
+	if cfg.tagFns != nil {
+		for key, tagFn := range cfg.tagFns {
+			opts = append(opts, tracer.Tag(key, tagFn(scope)))
+		}
+	}
 
 	span, _ := tracer.StartSpanFromContext(ctx, operationName, opts...)
 	span.Finish(tracer.WithError(scope.DB().Error))

diff --git a/contrib/jinzhu/gorm/gorm_test.go b/contrib/jinzhu/gorm/gorm_test.go
@@ -303,3 +303,44 @@ func TestContext(t *testing.T) {
 		assert.Equal(t, context.Background(), ctx)
 	})
 }
+
+func TestCustomTags(t *testing.T) {
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	sqltrace.Register("postgres", &pq.Driver{})
+	db, err := Open("postgres", "postgres://postgres:postgres@127.0.0.1:5432/postgres?sslmode=disable",
+		WithCustomTag("custom_tag", func(scope *gorm.Scope) interface{} {
+			return scope.SQLVars[3]
+		}),
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer db.Close()
+	db.AutoMigrate(&Product{})
+
+	parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
+		tracer.ServiceName("fake-http-server"),
+		tracer.SpanType(ext.SpanTypeWeb),
+	)
+
+	db = WithContext(ctx, db)
+	db.Create(&Product{Code: "L1212", Price: 1000})
+
+	parentSpan.Finish()
+
+	spans := mt.FinishedSpans()
+	assert.True(len(spans) >= 3)
+
+	// We deterministically expect the span to be the third last,
+	// followed by the underlying postgres DB trace and the above http.request span.
+	span := spans[len(spans)-3]
+	assert.Equal("gorm.create", span.OperationName())
+	assert.Equal(ext.SpanTypeSQL, span.Tag(ext.SpanType))
+	assert.Equal("L1212", span.Tag("custom_tag"))
+	assert.Equal(
+		`INSERT INTO "products" ("created_at","updated_at","deleted_at","code","price") VALUES ($1,$2,$3,$4,$5) RETURNING "products"."id"`,
+		span.Tag(ext.ResourceName))
+}
diff --git a/contrib/jinzhu/gorm/option.go b/contrib/jinzhu/gorm/option.go
@@ -9,12 +9,15 @@ import (
 	"math"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/internal"
+
+	"github.com/jinzhu/gorm"
 )
 
 type config struct {
 	serviceName   string
 	analyticsRate float64
 	dsn           string
+	tagFns        map[string]func(scope *gorm.Scope) interface{}
 }
 
 // Option represents an option that can be passed to Register, Open or OpenDB.
@@ -60,3 +63,14 @@ func WithAnalyticsRate(rate float64) Option {
 		}
 	}
 }
+
+// WithCustomTag will cause the given tagFn to be evaluated after executing
+// a query and attach the result to the span tagged by the key.
+func WithCustomTag(tag string, tagFn func(scope *gorm.Scope) interface{}) Option {
+	return func(cfg *config) {
+		if cfg.tagFns == nil {
+			cfg.tagFns = make(map[string]func(scope *gorm.Scope) interface{})
+		}
+		cfg.tagFns[tag] = tagFn
+	}
+}