Skip to content

Commit

Permalink
contrib/database/sql: Add WithErrorCheck options
Browse files Browse the repository at this point in the history
WithErrorCheck specifies a function fn which determines whether the passed
error should be marked as an error. The fn is called whenever a database/sql operation
finishes with an error.

close #1314
  • Loading branch information
soh335 committed Jun 2, 2022
1 parent 3528a0d commit be59f6b
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 1 deletion.
5 changes: 4 additions & 1 deletion contrib/database/sql/conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -203,5 +203,8 @@ func (tp *traceParams) tryTrace(ctx context.Context, qtype queryType, query stri
span.SetTag(k, v)
}
}
span.Finish(tracer.WithError(err))
defer span.Finish()
if tp.cfg.errCheck(err) {
span.SetTag(ext.Error, err)
}
}
46 changes: 46 additions & 0 deletions contrib/database/sql/conn_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,10 @@ import (
"context"
"database/sql/driver"
"log"
"strings"
"testing"

"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"

"github.com/go-sql-driver/mysql"
Expand Down Expand Up @@ -168,3 +170,47 @@ func TestWithChildSpansOnly(t *testing.T) {
})
}
}

func TestWithErrorCheck(t *testing.T) {
mt := mocktracer.Start()
defer mt.Stop()

assertErrCheck := func(t *testing.T, mt mocktracer.Tracer, errExist bool, opts ...Option) {
Register("mysql", &mysql.MySQLDriver{})
defer unregister("mysql")

db, err := Open("mysql", "test:test@tcp(127.0.0.1:3306)/test", opts...)
if err != nil {
log.Fatal(err)
}
defer db.Close()

db.QueryContext(context.Background(), "SELECT a FROM "+tableName)

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

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 {
if strings.Contains(err.Error(), `Error 1054: Unknown column 'a' in 'field list'`) {
return false
}
return true
}
assertErrCheck(t, mt, false, WithErrorCheck(errFn))
})
}
10 changes: 10 additions & 0 deletions contrib/database/sql/option.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ type config struct {
analyticsRate float64
dsn string
childSpansOnly bool
errCheck func(err error) bool
}

// Option represents an option that can be passed to Register, Open or OpenDB.
Expand Down Expand Up @@ -83,3 +84,12 @@ func WithChildSpansOnly() Option {
cfg.childSpansOnly = true
}
}

// WithErrorCheck specifies a function fn which determines whether the passed
// error should be marked as an error. The fn is called whenever a database/sql operation
// finishes with an error
func WithErrorCheck(fn func(err error) bool) Option {
return func(cfg *config) {
cfg.errCheck = fn
}
}
7 changes: 7 additions & 0 deletions contrib/database/sql/sql.go
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,13 @@ func OpenDB(c driver.Connector, opts ...Option) *sql.DB {
if math.IsNaN(cfg.analyticsRate) {
cfg.analyticsRate = rc.analyticsRate
}
if cfg.errCheck == nil {
if rc.errCheck == nil {
cfg.errCheck = func(err error) bool { return true }
} else {
cfg.errCheck = rc.errCheck
}
}
cfg.childSpansOnly = rc.childSpansOnly
tc := &tracedConnector{
connector: c,
Expand Down

0 comments on commit be59f6b

Please sign in to comment.