Repository: go-gorm/gorm
Pull request #6916 (2024-03-19 14:17:02)

fix: remove `callback` from `callbacks` if `Remove()` called
<!--
Make sure these boxes checked before submitting your pull request.

For significant changes, please open an issue to make an agreement on an implementation design/plan first before starting it.
-->

- [x] Do only one thing
- [x] Non breaking API changes
- [x] Tested

### What did this pull request do?

<!--
provide a general description of the code changes in your pull request
-->

ref: https://github.com/go-gorm/gorm/pull/6910#pullrequestreview-1945116356

Modified the `Remove()` method so that it not only sets the removed flag but also removes the `callback` from the `callbacks` within `compile()`. 
This resolves the following issues:

- The problem where the `Get()` method could retrieve callbacks that have already been removed.
- The problem where removed callbacks could still operate due to the `Remove()` method.

In addition, I have made some modifications to the existing tests.
Previously, the specified `before` and `after` for a `callback` continued to affect the overall order even after being `Remove()`d. I believe that continuing to have an effect after being removed is an unintended behavior.
What do you think? I would appreciate any advice.

### User Case Description

<!-- Your use case -->

#### The problem where the `Get()` method

If a callback named `cb1` is registered, it is expected that calling `Get("cb1")` after `Remove("cb1")` should return `nil`. However, a callback function registered as `cb1` is being returned instead.

```golang
db, _ := gorm.Open(nil, nil)
createCallback := db.Callback().Create()

createCallback.Register("cb1", func(*gorm.DB){})
createCallback.Remove("cb1")

cb := createCallback.Get("cb1") // not nil
```

#### The problem where the `Remove()` method

The expectation is that callbacks marked as removed via `Remove()` should not be executed. However, we encountered an issue where callbacks registered through methods like `Before("*").Register()` continue to operate even after being `Remove()`d. 
This typically occurs in scenarios such as the following:

```golang
db, _ := gorm.Open(nil, nil)
rawCallback := db.Callback().Raw()

rawCallback.Before("*").Register("cb1", func(*gorm.DB) { panic("Remove does not work") })
rawCallback.Remove("cb1")

db.Exec("SELECT 1") // panic
```

#### workaround

For your reference, using `Before("*").Remove()` can be mentioned as a workaround for the above issue.

snackmgmg: Thanks! I have improved the part you mentioned in your comment.
https://github.com/go-gorm/gorm/pull/6916/commits/f255f362be9ef06b868a8cb3f6a99779b509bf1e

Patch:
diff --git a/callbacks.go b/callbacks.go
index 195d172..50b5b0e 100644
--- a/callbacks.go
+++ b/callbacks.go
@@ -187,10 +187,18 @@ func (p *processor) Replace(name string, fn func(*DB)) error {
 
 func (p *processor) compile() (err error) {
 	var callbacks []*callback
+	removedMap := map[string]bool{}
 	for _, callback := range p.callbacks {
 		if callback.match == nil || callback.match(p.db) {
 			callbacks = append(callbacks, callback)
 		}
+		if callback.remove {
+			removedMap[callback.name] = true
+		}
+	}
+
+	if len(removedMap) > 0 {
+		callbacks = removeCallbacks(callbacks, removedMap)
 	}
 	p.callbacks = callbacks
 
@@ -339,3 +347,14 @@ func sortCallbacks(cs []*callback) (fns []func(*DB), err error) {
 
 	return
 }
+
+func removeCallbacks(cs []*callback, nameMap map[string]bool) []*callback {
+	callbacks := make([]*callback, 0, len(cs))
+	for _, callback := range cs {
+		if nameMap[callback.name] {
+			continue
+		}
+		callbacks = append(callbacks, callback)
+	}
+	return callbacks
+}
diff --git a/logger/sql.go b/logger/sql.go
index 8ce8d8b..ad47879 100644
--- a/logger/sql.go
+++ b/logger/sql.go
@@ -34,6 +34,19 @@ var convertibleTypes = []reflect.Type{reflect.TypeOf(time.Time{}), reflect.TypeO
 // RegEx matches only numeric values
 var numericPlaceholderRe = regexp.MustCompile(`\$\d+\$`)
 
+func isNumeric(k reflect.Kind) bool {
+	switch k {
+	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+		return true
+	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+		return true
+	case reflect.Float32, reflect.Float64:
+		return true
+	default:
+		return false
+	}
+}
+
 // ExplainSQL generate SQL string with given parameters, the generated SQL is expected to be used in logger, execute it might introduce a SQL injection vulnerability
 func ExplainSQL(sql string, numericPlaceholder *regexp.Regexp, escaper string, avars ...interface{}) string {
 	var (
@@ -110,6 +123,12 @@ func ExplainSQL(sql string, numericPlaceholder *regexp.Regexp, escaper string, a
 				convertParams(v, idx)
 			} else if rv.Kind() == reflect.Ptr && !rv.IsZero() {
 				convertParams(reflect.Indirect(rv).Interface(), idx)
+			} else if isNumeric(rv.Kind()) {
+				if rv.CanInt() || rv.CanUint() {
+					vars[idx] = fmt.Sprintf("%d", rv.Interface())
+				} else {
+					vars[idx] = fmt.Sprintf("%.6f", rv.Interface())
+				}
 			} else {
 				for _, t := range convertibleTypes {
 					if rv.Type().ConvertibleTo(t) {
diff --git a/prepare_stmt.go b/prepare_stmt.go
index aa94462..c60b5db 100644
--- a/prepare_stmt.go
+++ b/prepare_stmt.go
@@ -3,6 +3,8 @@ package gorm
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
+	"errors"
 	"reflect"
 	"sync"
 )
@@ -147,7 +149,7 @@ func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ..
 	stmt, err := db.prepare(ctx, db.ConnPool, false, query)
 	if err == nil {
 		result, err = stmt.ExecContext(ctx, args...)
-		if err != nil {
+		if errors.Is(err, driver.ErrBadConn) {
 			db.Mux.Lock()
 			defer db.Mux.Unlock()
 			go stmt.Close()
@@ -161,7 +163,7 @@ func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args .
 	stmt, err := db.prepare(ctx, db.ConnPool, false, query)
 	if err == nil {
 		rows, err = stmt.QueryContext(ctx, args...)
-		if err != nil {
+		if errors.Is(err, driver.ErrBadConn) {
 			db.Mux.Lock()
 			defer db.Mux.Unlock()
 
@@ -207,7 +209,7 @@ func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ..
 	stmt, err := tx.PreparedStmtDB.prepare(ctx, tx.Tx, true, query)
 	if err == nil {
 		result, err = tx.Tx.StmtContext(ctx, stmt.Stmt).ExecContext(ctx, args...)
-		if err != nil {
+		if errors.Is(err, driver.ErrBadConn) {
 			tx.PreparedStmtDB.Mux.Lock()
 			defer tx.PreparedStmtDB.Mux.Unlock()
 
@@ -222,7 +224,7 @@ func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args .
 	stmt, err := tx.PreparedStmtDB.prepare(ctx, tx.Tx, true, query)
 	if err == nil {
 		rows, err = tx.Tx.StmtContext(ctx, stmt.Stmt).QueryContext(ctx, args...)
-		if err != nil {
+		if errors.Is(err, driver.ErrBadConn) {
 			tx.PreparedStmtDB.Mux.Lock()
 			defer tx.PreparedStmtDB.Mux.Unlock()
 


Test patch:
diff --git a/logger/sql_test.go b/logger/sql_test.go
index 036ef3a..9002a7e 100644
--- a/logger/sql_test.go
+++ b/logger/sql_test.go
@@ -37,14 +37,18 @@ func format(v []byte, escaper string) string {
 func TestExplainSQL(t *testing.T) {
 	type role string
 	type password []byte
+	type intType int
+	type floatType float64
 	var (
-		tt     = now.MustParse("2020-02-23 11:10:10")
-		myrole = role("admin")
-		pwd    = password("pass")
-		jsVal  = []byte(`{"Name":"test","Val":"test"}`)
-		js     = JSON(jsVal)
-		esVal  = []byte(`{"Name":"test","Val":"test"}`)
-		es     = ExampleStruct{Name: "test", Val: "test"}
+		tt                 = now.MustParse("2020-02-23 11:10:10")
+		myrole             = role("admin")
+		pwd                = password("pass")
+		jsVal              = []byte(`{"Name":"test","Val":"test"}`)
+		js                 = JSON(jsVal)
+		esVal              = []byte(`{"Name":"test","Val":"test"}`)
+		es                 = ExampleStruct{Name: "test", Val: "test"}
+		intVal   intType   = 1
+		floatVal floatType = 1.23
 	)
 
 	results := []struct {
@@ -107,6 +111,18 @@ func TestExplainSQL(t *testing.T) {
 			Vars:          []interface{}{"jinzhu", 1, float32(999.99), true, []byte("12345"), tt, &tt, nil, "w@g.\"com", myrole, pwd, &js, &es},
 			Result:        fmt.Sprintf(`create table users (name, age, height, actived, bytes, create_at, update_at, deleted_at, email, role, pass, json_struct, example_struct) values ("jinzhu", 1, 999.99, true, "12345", "2020-02-23 11:10:10", "2020-02-23 11:10:10", NULL, "w@g.""com", "admin", "pass", %v, %v)`, format(jsVal, `"`), format(esVal, `"`)),
 		},
+		{
+			SQL:           "create table users (name, age, height, actived, bytes, create_at, update_at, deleted_at, email, role, pass, int_val) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+			NumericRegexp: nil,
+			Vars:          []interface{}{"jinzhu?", 1, 999.99, true, []byte("12345"), tt, &tt, nil, "w@g.\"com", myrole, pwd, intVal},
+			Result:        `create table users (name, age, height, actived, bytes, create_at, update_at, deleted_at, email, role, pass, int_val) values ("jinzhu?", 1, 999.99, true, "12345", "2020-02-23 11:10:10", "2020-02-23 11:10:10", NULL, "w@g.""com", "admin", "pass", 1)`,
+		},
+		{
+			SQL:           "create table users (name, age, height, actived, bytes, create_at, update_at, deleted_at, email, role, pass, float_val) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+			NumericRegexp: nil,
+			Vars:          []interface{}{"jinzhu?", 1, 999.99, true, []byte("12345"), tt, &tt, nil, "w@g.\"com", myrole, pwd, floatVal},
+			Result:        `create table users (name, age, height, actived, bytes, create_at, update_at, deleted_at, email, role, pass, float_val) values ("jinzhu?", 1, 999.99, true, "12345", "2020-02-23 11:10:10", "2020-02-23 11:10:10", NULL, "w@g.""com", "admin", "pass", 1.230000)`,
+		},
 	}
 
 	for idx, r := range results {
diff --git a/tests/callbacks_test.go b/tests/callbacks_test.go
index 4479da4..f77209f 100644
--- a/tests/callbacks_test.go
+++ b/tests/callbacks_test.go
@@ -91,7 +91,7 @@ func TestCallbacks(t *testing.T) {
 		},
 		{
 			callbacks: []callback{{h: c1}, {h: c2, before: "c4", after: "c5"}, {h: c3}, {h: c4}, {h: c5}, {h: c2, remove: true}},
-			results:   []string{"c1", "c5", "c3", "c4"},
+			results:   []string{"c1", "c3", "c4", "c5"},
 		},
 		{
 			callbacks: []callback{{h: c1}, {name: "c", h: c2}, {h: c3}, {name: "c", h: c4, replace: true}},
@@ -206,3 +206,49 @@ func TestPluginCallbacks(t *testing.T) {
 		t.Errorf("callbacks tests failed, got %v", msg)
 	}
 }
+
+func TestCallbacksGet(t *testing.T) {
+	db, _ := gorm.Open(nil, nil)
+	createCallback := db.Callback().Create()
+
+	createCallback.Before("*").Register("c1", c1)
+	if cb := createCallback.Get("c1"); reflect.DeepEqual(cb, c1) {
+		t.Errorf("callbacks tests failed, got: %p, want: %p", cb, c1)
+	}
+
+	createCallback.Remove("c1")
+	if cb := createCallback.Get("c2"); cb != nil {
+		t.Errorf("callbacks test failed. got: %p, want: nil", cb)
+	}
+}
+
+func TestCallbacksRemove(t *testing.T) {
+	db, _ := gorm.Open(nil, nil)
+	createCallback := db.Callback().Create()
+
+	createCallback.Before("*").Register("c1", c1)
+	createCallback.After("*").Register("c2", c2)
+	createCallback.Before("c4").Register("c3", c3)
+	createCallback.After("c2").Register("c4", c4)
+
+	// callbacks: []string{"c1", "c3", "c4", "c2"}
+	createCallback.Remove("c1")
+	if ok, msg := assertCallbacks(createCallback, []string{"c3", "c4", "c2"}); !ok {
+		t.Errorf("callbacks tests failed, got %v", msg)
+	}
+
+	createCallback.Remove("c4")
+	if ok, msg := assertCallbacks(createCallback, []string{"c3", "c2"}); !ok {
+		t.Errorf("callbacks tests failed, got %v", msg)
+	}
+
+	createCallback.Remove("c2")
+	if ok, msg := assertCallbacks(createCallback, []string{"c3"}); !ok {
+		t.Errorf("callbacks tests failed, got %v", msg)
+	}
+
+	createCallback.Remove("c3")
+	if ok, msg := assertCallbacks(createCallback, []string{}); !ok {
+		t.Errorf("callbacks tests failed, got %v", msg)
+	}
+}
diff --git a/tests/go.mod b/tests/go.mod
index 5616ebb..3d3901d 100644
--- a/tests/go.mod
+++ b/tests/go.mod
@@ -7,7 +7,7 @@ require (
 	github.com/jinzhu/now v1.1.5
 	github.com/lib/pq v1.10.9
 	github.com/stretchr/testify v1.9.0
-	gorm.io/driver/mysql v1.5.5
+	gorm.io/driver/mysql v1.5.6
 	gorm.io/driver/postgres v1.5.7
 	gorm.io/driver/sqlite v1.5.5
 	gorm.io/driver/sqlserver v1.5.3
diff --git a/tests/prepared_stmt_test.go b/tests/prepared_stmt_test.go
index b234c8b..b86bc3d 100644
--- a/tests/prepared_stmt_test.go
+++ b/tests/prepared_stmt_test.go
@@ -126,33 +126,6 @@ func TestPreparedStmtDeadlock(t *testing.T) {
 	AssertEqual(t, sqlDB.Stats().InUse, 0)
 }
 
-func TestPreparedStmtError(t *testing.T) {
-	tx, err := OpenTestConnection(&gorm.Config{})
-	AssertEqual(t, err, nil)
-
-	sqlDB, _ := tx.DB()
-	sqlDB.SetMaxOpenConns(1)
-
-	tx = tx.Session(&gorm.Session{PrepareStmt: true})
-
-	wg := sync.WaitGroup{}
-	for i := 0; i < 10; i++ {
-		wg.Add(1)
-		go func() {
-			// err prepare
-			tag := Tag{Locale: "zh"}
-			tx.Table("users").Find(&tag)
-			wg.Done()
-		}()
-	}
-	wg.Wait()
-
-	conn, ok := tx.ConnPool.(*gorm.PreparedStmtDB)
-	AssertEqual(t, ok, true)
-	AssertEqual(t, len(conn.Stmts), 0)
-	AssertEqual(t, sqlDB.Stats().InUse, 0)
-}
-
 func TestPreparedStmtInTransaction(t *testing.T) {
 	user := User{Name: "jinzhu"}