Squashed commit of the following:

commit 759038a126122d5b3323979fdd7d867a4ab85585
Author: Jinzhu <wosmvp@gmail.com>
Date:   Mon Aug 31 12:06:31 2020 +0800

    Add PreparedStmt tests

commit 066d54db1f
Author: 王岚 <wanglan.backend@bytedance.com>
Date:   Fri Aug 28 18:40:59 2020 +0800

    prepare_stmt add ctx

gorm.go

@@ -169,6 +169,7 @@ func (db *DB) Session(config *Session) *DB {
 
 	if config.PrepareStmt {
 		if v, ok := db.cacheStore.Load("preparedStmt"); ok {
+			tx.Statement = tx.Statement.clone()
 			preparedStmt := v.(*PreparedStmtDB)
 			tx.Statement.ConnPool = &PreparedStmtDB{
 				ConnPool: db.Config.ConnPool,

prepare_stmt.go

@@ -25,7 +25,7 @@ func (db *PreparedStmtDB) Close() {
 	db.Mux.Unlock()
 }
 
-func (db *PreparedStmtDB) prepare(query string) (*sql.Stmt, error) {
+func (db *PreparedStmtDB) prepare(ctx context.Context, query string) (*sql.Stmt, error) {
 	db.Mux.RLock()
 	if stmt, ok := db.Stmts[query]; ok {
 		db.Mux.RUnlock()
@@ -40,7 +40,7 @@ func (db *PreparedStmtDB) prepare(query string) (*sql.Stmt, error) {
 		return stmt, nil
 	}
 
-	stmt, err := db.ConnPool.PrepareContext(context.Background(), query)
+	stmt, err := db.ConnPool.PrepareContext(ctx, query)
 	if err == nil {
 		db.Stmts[query] = stmt
 		db.PreparedSQL = append(db.PreparedSQL, query)
@@ -59,7 +59,7 @@ func (db *PreparedStmtDB) BeginTx(ctx context.Context, opt *sql.TxOptions) (Conn
 }
 
 func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
-	stmt, err := db.prepare(query)
+	stmt, err := db.prepare(ctx, query)
 	if err == nil {
 		result, err = stmt.ExecContext(ctx, args...)
 		if err != nil {
@@ -73,7 +73,7 @@ func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ..
 }
 
 func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
-	stmt, err := db.prepare(query)
+	stmt, err := db.prepare(ctx, query)
 	if err == nil {
 		rows, err = stmt.QueryContext(ctx, args...)
 		if err != nil {
@@ -87,7 +87,7 @@ func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args .
 }
 
 func (db *PreparedStmtDB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
-	stmt, err := db.prepare(query)
+	stmt, err := db.prepare(ctx, query)
 	if err == nil {
 		return stmt.QueryRowContext(ctx, args...)
 	}
@@ -100,9 +100,9 @@ type PreparedStmtTX struct {
 }
 
 func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
-	stmt, err := tx.PreparedStmtDB.prepare(query)
+	stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
 	if err == nil {
-		result, err = tx.Tx.Stmt(stmt).ExecContext(ctx, args...)
+		result, err = tx.Tx.StmtContext(ctx, stmt).ExecContext(ctx, args...)
 		if err != nil {
 			tx.PreparedStmtDB.Mux.Lock()
 			stmt.Close()
@@ -114,9 +114,9 @@ func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ..
 }
 
 func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
-	stmt, err := tx.PreparedStmtDB.prepare(query)
+	stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
 	if err == nil {
-		rows, err = tx.Tx.Stmt(stmt).QueryContext(ctx, args...)
+		rows, err = tx.Tx.StmtContext(ctx, stmt).QueryContext(ctx, args...)
 		if err != nil {
 			tx.PreparedStmtDB.Mux.Lock()
 			stmt.Close()
@@ -128,9 +128,9 @@ func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args .
 }
 
 func (tx *PreparedStmtTX) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
-	stmt, err := tx.PreparedStmtDB.prepare(query)
+	stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
 	if err == nil {
-		return tx.Tx.Stmt(stmt).QueryRowContext(ctx, args...)
+		return tx.Tx.StmtContext(ctx, stmt).QueryRowContext(ctx, args...)
 	}
 	return &sql.Row{}
 }

tests/prepared_stmt_test.go

@@ -0,0 +1,48 @@
+package tests_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"gorm.io/gorm"
+	. "gorm.io/gorm/utils/tests"
+)
+
+func TestPreparedStmt(t *testing.T) {
+	tx := DB.Session(&gorm.Session{PrepareStmt: true})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	txCtx := tx.WithContext(ctx)
+
+	user := *GetUser("prepared_stmt", Config{})
+
+	txCtx.Create(&user)
+
+	var result1 User
+	if err := txCtx.Find(&result1, user.ID).Error; err != nil {
+		t.Fatalf("no error should happen but got %v", err)
+	}
+
+	time.Sleep(time.Second)
+
+	var result2 User
+	if err := tx.Find(&result2, user.ID).Error; err != nil {
+		t.Fatalf("no error should happen but got %v", err)
+	}
+
+	user2 := *GetUser("prepared_stmt2", Config{})
+	if err := txCtx.Create(&user2).Error; err == nil {
+		t.Fatalf("should failed to create with timeout context")
+	}
+
+	if err := tx.Create(&user2).Error; err != nil {
+		t.Fatalf("failed to create, got error %v", err)
+	}
+
+	var result3 User
+	if err := tx.Find(&result3, user2.ID).Error; err != nil {
+		t.Fatalf("no error should happen but got %v", err)
+	}
+}