From be59f6beb7c5be7833269fd0e118971f19dbd415 Mon Sep 17 00:00:00 2001
From: soh335 <sugarbabe335@gmail.com>
Date: Wed, 1 Jun 2022 16:59:08 +0900
Subject: [PATCH] contrib/database/sql: Add WithErrorCheck options

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 https://github.com/DataDog/dd-trace-go/issues/1314
---
 contrib/database/sql/conn.go      |  5 +++-
 contrib/database/sql/conn_test.go | 46 +++++++++++++++++++++++++++++++
 contrib/database/sql/option.go    | 10 +++++++
 contrib/database/sql/sql.go       |  7 +++++
 4 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/contrib/database/sql/conn.go b/contrib/database/sql/conn.go
index d382301a5d..d4a56532f6 100644
--- a/contrib/database/sql/conn.go
+++ b/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 tp.cfg.errCheck(err) {
+		span.SetTag(ext.Error, err)
+	}
 }
diff --git a/contrib/database/sql/conn_test.go b/contrib/database/sql/conn_test.go
index 3203bf0e27..94a030a96c 100644
--- a/contrib/database/sql/conn_test.go
+++ b/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,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))
+	})
+}
diff --git a/contrib/database/sql/option.go b/contrib/database/sql/option.go
index 47bd4208bd..a6312e5f5b 100644
--- a/contrib/database/sql/option.go
+++ b/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
+	}
+}
diff --git a/contrib/database/sql/sql.go b/contrib/database/sql/sql.go
index 2383f4a499..383a03422a 100644
--- a/contrib/database/sql/sql.go
+++ b/contrib/database/sql/sql.go
@@ -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,