tablet/gorm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Squashed commit of the following:

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

    Add PreparedStmt tests

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

    prepare_stmt add ctx
This commit is contained in:
Jinzhu 2020-08-31 12:08:33 +08:00
parent 53f8c9fc1c
commit 9b0ad4730f
3 changed files with 60 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
gorm.go
View File

@ -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,

22
prepare_stmt.go
View File

@ -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{}
}

48
tests/prepared_stmt_test.go Normal file
View File

@ -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)
}
}