From 9b0ad4730f16d6ac7cf18d1aa42d74714959745b Mon Sep 17 00:00:00 2001 From: Jinzhu Date: Mon, 31 Aug 2020 12:08:33 +0800 Subject: [PATCH] Squashed commit of the following: MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit commit 759038a126122d5b3323979fdd7d867a4ab85585 Author: Jinzhu Date: Mon Aug 31 12:06:31 2020 +0800 Add PreparedStmt tests commit 066d54db1fc93ea58c190195104a2d7086623f69 Author: 王岚 Date: Fri Aug 28 18:40:59 2020 +0800 prepare_stmt add ctx --- gorm.go | 1 + prepare_stmt.go | 22 ++++++++--------- tests/prepared_stmt_test.go | 48 +++++++++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 11 deletions(-) create mode 100644 tests/prepared_stmt_test.go diff --git a/gorm.go b/gorm.go index 3c187f42..fec4310b 100644 --- a/gorm.go +++ b/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, diff --git a/prepare_stmt.go b/prepare_stmt.go index 7e87558d..7c80bafe 100644 --- a/prepare_stmt.go +++ b/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{} } diff --git a/tests/prepared_stmt_test.go b/tests/prepared_stmt_test.go new file mode 100644 index 00000000..b81318d3 --- /dev/null +++ b/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) + } +}