contrib/database/sql: Add WithErrorCheck options

contrib/database/sql/conn.go

@@ -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 err != nil && tp.cfg.errCheck != nil && tp.cfg.errCheck(err) {
+		span.SetTag(ext.Error, err)
+	}
 }

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"
@@ -168,3 +170,35 @@ 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'`)
+	})))
+
+}

contrib/database/sql/option.go

@@ -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.
@@ -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
+	}
+}

contrib/database/sql/sql.go

@@ -186,6 +186,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
+	}
 	cfg.childSpansOnly = rc.childSpansOnly
 	tc := &tracedConnector{
 		connector:  c,