Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

contrib/gorm.io/gorm.v1: Add err check #940

Merged
merged 11 commits into from
Oct 20, 2021
6 changes: 5 additions & 1 deletion contrib/gorm.io/gorm.v1/gorm.go
Original file line number Diff line number Diff line change
Expand Up @@ -121,5 +121,9 @@ func after(db *gorm.DB, operationName string, cfg *config) {
}

span, _ := tracer.StartSpanFromContext(ctx, operationName, opts...)
span.Finish(tracer.WithError(db.Error))
var dbErr error
if cfg.errCheck(db.Error) {
dbErr = db.Error
}
span.Finish(tracer.WithError(dbErr))
}
39 changes: 39 additions & 0 deletions contrib/gorm.io/gorm.v1/gorm_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -343,3 +343,42 @@ func TestContext(t *testing.T) {
assert.Equal(t, testCtx.Value(key(contextKey)), ctx.Value(key(contextKey)))
})
}

func TestError(t *testing.T) {
vmlellis marked this conversation as resolved.
Show resolved Hide resolved
mt := mocktracer.Start()
defer mt.Stop()

assertErrCheck := func(t *testing.T, mt mocktracer.Tracer, errExist bool, opts ...Option) {
sqltrace.Register("pgx", &stdlib.Driver{})
sqlDb, err := sqltrace.Open("pgx", pgConnString)
assert.Nil(t, err)

db, err := Open(postgres.New(postgres.Config{Conn: sqlDb}), &gorm.Config{}, opts...)
assert.Nil(t, err)
db.AutoMigrate(&Product{})
db.First(&Product{}, Product{Code: "L1210", Price: 2000})

spans := mt.FinishedSpans()
assert.True(t, len(spans) > 1)

// Get last span (gorm.db)
s := spans[len(spans)-1]

assert.Equal(t, errExist, s.Tag(ext.Error) != nil)
}

t.Run("defaults", func(t *testing.T) {
mt := mocktracer.Start()
defer mt.Stop()
assertErrCheck(t, mt, true)
})

t.Run("errcheck", func(t *testing.T) {
mt := mocktracer.Start()
defer mt.Stop()
errFn := func(err error) bool {
return err != gorm.ErrRecordNotFound
}
assertErrCheck(t, mt, false, WithErrorCheck(errFn))
})
}
11 changes: 11 additions & 0 deletions contrib/gorm.io/gorm.v1/option.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ type config struct {
serviceName string
analyticsRate float64
dsn string
errCheck func(err error) bool
}

// Option represents an option that can be passed to Register, Open or OpenDB.
Expand All @@ -28,6 +29,7 @@ func defaults(cfg *config) {
} else {
cfg.analyticsRate = math.NaN()
}
cfg.errCheck = func(error) bool { return true }
}

// WithServiceName sets the given service name when registering a driver,
Expand Down Expand Up @@ -60,3 +62,12 @@ func WithAnalyticsRate(rate float64) Option {
}
}
}

// WithErrorCheck specifies a function fn which determines whether the passed
// error should be marked as an error. The fn is called whenever a gorm operation
// finishes with an error
vmlellis marked this conversation as resolved.
Show resolved Hide resolved
func WithErrorCheck(fn func(err error) bool) Option {
return func(cfg *config) {
cfg.errCheck = fn
}
}