contrib/database/sql: Add WithErrorCheck options

contrib/database/sql/conn.go

@@ -239,5 +239,8 @@ func (tp *traceParams) tryTrace(ctx context.Context, qtype queryType, query stri
 			span.SetTag(k, v)
 		}
 	}
-	span.Finish(tracer.WithError(err))
+	if err != nil && (tp.cfg.errCheck == nil || tp.cfg.errCheck(err)) {
+		span.SetTag(ext.Error, err)
+	}
+	span.Finish()
 }

contrib/database/sql/conn_test.go

@@ -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"
@@ -169,6 +171,38 @@ func TestWithChildSpansOnly(t *testing.T) {
 	}
 }
 
+func TestWithErrorCheck(t *testing.T) {
+	testOpts := func(errExist bool, opts ...Option) func(t *testing.T) {
+		return func(t *testing.T) {
+			mt := mocktracer.Start()
+			defer mt.Stop()
+
+			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", testOpts(true))
+	t.Run("errcheck", testOpts(false, WithErrorCheck(func(err error) bool {
+		return !strings.Contains(err.Error(), `Error 1054: Unknown column 'a' in 'field list'`)
+	})))
+
+}
+
 func TestWithCustomTag(t *testing.T) {
 	type sqlRegister struct {
 		name   string

contrib/database/sql/option.go

@@ -18,6 +18,7 @@ type config struct {
 	analyticsRate        float64
 	dsn                  string
 	childSpansOnly       bool
+	errCheck             func(err error) bool
 	tags                 map[string]interface{}
 	commentInjectionMode tracer.SQLCommentInjectionMode
 }
@@ -89,6 +90,15 @@ func WithChildSpansOnly() Option {
 	}
 }
 
+// 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
+	}
+}
+
 // WithCustomTag will attach the value to the span tagged by the key
 func WithCustomTag(key string, value interface{}) Option {
 	return func(cfg *config) {

contrib/database/sql/sql.go

@@ -187,6 +187,9 @@ func OpenDB(c driver.Connector, opts ...Option) *sql.DB {
 	if math.IsNaN(cfg.analyticsRate) {
 		cfg.analyticsRate = rc.analyticsRate
 	}
+	if cfg.errCheck == nil {
+		cfg.errCheck = rc.errCheck
+	}
 	if cfg.commentInjectionMode == tracer.SQLInjectionUndefined {
 		cfg.commentInjectionMode = rc.commentInjectionMode
 	}