fix(contrib/gorm.io): allow resource name to be overridden by custom tags (#3909)

Co-authored-by: rodrigo.arguello <[email protected]>

Author: will-bem

contrib/gorm.io/gorm.v1/gorm.go

@@ -65,8 +65,7 @@ func Open(dialector gorm.Dialector, cfg *gorm.Config, opts ...Option) (*gorm.DB,
 }
 
 func withCallbacks(db *gorm.DB, opts ...Option) (*gorm.DB, error) {
-	cfg := new(config)
-	defaults(cfg)
+	cfg := newConfigWithDefaults()
 	for _, fn := range opts {
 		fn.apply(cfg)
 	}
@@ -151,12 +150,6 @@ func before(db *gorm.DB, operationName string, cfg *config) {
 	if !math.IsNaN(cfg.analyticsRate) {
 		opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate))
 	}
-	for key, tagFn := range cfg.tagFns {
-		if tagFn != nil {
-			opts = append(opts, tracer.Tag(key, tagFn(db)))
-		}
-	}
-
 	_, ctx := tracer.StartSpanFromContext(db.Statement.Context, operationName, opts...)
 	db.Statement.Context = ctx
 }
@@ -175,6 +168,15 @@ func after(db *gorm.DB, cfg *config) {
 			dbErr = db.Error
 		}
 		span.SetTag(ext.ResourceName, db.Statement.SQL.String())
+
+		// process tagFns after setting the resource name to allow the resource
+		// name to be overridden
+		for key, tagFn := range cfg.tagFns {
+			if tagFn != nil {
+				span.SetTag(key, tagFn(db))
+			}
+		}
+
 		span.Finish(tracer.WithError(dbErr))
 	}
 }

contrib/gorm.io/gorm.v1/gorm_test.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"strings"
 	"testing"
 
 	sqltrace "github.com/DataDog/dd-trace-go/contrib/database/sql/v2"
@@ -610,3 +611,32 @@ func TestPlugin(t *testing.T) {
 	assert.NotNil(t, db.Callback().Raw().Get("dd-trace-go:before_raw_query"))
 	assert.NotNil(t, db.Callback().Raw().Get("dd-trace-go:before_raw_query"))
 }
+
+func newStubDB() *gorm.DB {
+	sqlbuilder := strings.Builder{}
+	sqlbuilder.WriteString("SELECT * FROM products WHERE id = ?")
+	statement := &gorm.Statement{SQL: sqlbuilder, Context: context.Background()}
+	config := &gorm.Config{DryRun: false}
+	return &gorm.DB{Statement: statement, Config: config}
+}
+
+func TestCustomResourceName(t *testing.T) {
+	// This test is meant to ensure that the resource name can be overridden by a custom tag.
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	cfg := newConfigWithDefaults()
+	customResourceName := "custom resource name"
+	WithCustomTag(ext.ResourceName, func(db *gorm.DB) interface{} {
+		return customResourceName
+	})(cfg)
+
+	db := newStubDB()
+	before(db, "gorm.query", cfg)
+	after(db, cfg)
+
+	spans := mt.FinishedSpans()
+	assert.True(len(spans) > 0)
+	assert.Equal(customResourceName, spans[len(spans)-1].Tag(ext.ResourceName))
+}

contrib/gorm.io/gorm.v1/option.go

@@ -31,11 +31,13 @@ func (fn OptionFn) apply(cfg *config) {
 	fn(cfg)
 }
 
-func defaults(cfg *config) {
+func newConfigWithDefaults() *config {
+	cfg := new(config)
 	cfg.serviceName = "gorm.db"
 	cfg.analyticsRate = instr.AnalyticsRate(false)
 	cfg.errCheck = func(error) bool { return true }
 	cfg.tagFns = make(map[string]func(db *gorm.DB) interface{})
+	return cfg
 }
 
 // WithService sets the given service name when registering a driver,