diff --git a/contrib/gopkg.in/jinzhu/gorm.v1/gorm.go b/contrib/gopkg.in/jinzhu/gorm.v1/gorm.go index abcc8469ab..6aebb23334 100644 --- 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 index 87140b35bf..89808386a9 100644 --- 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 index 7b799adbfe..0515cf2562 100644 --- 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 index 135cc43ee8..4362d9eda7 100644 --- 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 index ef1a7f38f3..f4baed4353 100644 --- 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 index 7b799adbfe..094410f372 100644 --- 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 + } +}