Allow close prepared statements, double check before prepare

This commit is contained in:
Jinzhu 2020-06-21 18:01:50 +08:00
parent d0764bead1
commit 7851faa094
2 changed files with 21 additions and 5 deletions

View File

@ -102,7 +102,7 @@ func Open(dialector Dialector, config *Config) (db *DB, err error) {
if config.PrepareStmt {
db.ConnPool = &PreparedStmtDB{
ConnPool: db.ConnPool,
stmts: map[string]*sql.Stmt{},
Stmts: map[string]*sql.Stmt{},
}
}
@ -146,7 +146,7 @@ func (db *DB) Session(config *Session) *DB {
if config.PrepareStmt {
tx.Statement.ConnPool = &PreparedStmtDB{
ConnPool: db.Config.ConnPool,
stmts: map[string]*sql.Stmt{},
Stmts: map[string]*sql.Stmt{},
}
}

View File

@ -7,23 +7,39 @@ import (
)
type PreparedStmtDB struct {
stmts map[string]*sql.Stmt
Stmts map[string]*sql.Stmt
mux sync.RWMutex
ConnPool
}
func (db *PreparedStmtDB) Close() {
db.mux.Lock()
for k, stmt := range db.Stmts {
delete(db.Stmts, k)
stmt.Close()
}
db.mux.Unlock()
}
func (db *PreparedStmtDB) prepare(query string) (*sql.Stmt, error) {
db.mux.RLock()
if stmt, ok := db.stmts[query]; ok {
if stmt, ok := db.Stmts[query]; ok {
db.mux.RUnlock()
return stmt, nil
}
db.mux.RUnlock()
db.mux.Lock()
// double check
if stmt, ok := db.Stmts[query]; ok {
db.mux.Unlock()
return stmt, nil
}
stmt, err := db.ConnPool.PrepareContext(context.Background(), query)
if err == nil {
db.stmts[query] = stmt
db.Stmts[query] = stmt
}
db.mux.Unlock()