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 { if config.PrepareStmt {
db.ConnPool = &PreparedStmtDB{ db.ConnPool = &PreparedStmtDB{
ConnPool: db.ConnPool, 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 { if config.PrepareStmt {
tx.Statement.ConnPool = &PreparedStmtDB{ tx.Statement.ConnPool = &PreparedStmtDB{
ConnPool: db.Config.ConnPool, ConnPool: db.Config.ConnPool,
stmts: map[string]*sql.Stmt{}, Stmts: map[string]*sql.Stmt{},
} }
} }

View File

@ -7,23 +7,39 @@ import (
) )
type PreparedStmtDB struct { type PreparedStmtDB struct {
stmts map[string]*sql.Stmt Stmts map[string]*sql.Stmt
mux sync.RWMutex mux sync.RWMutex
ConnPool 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) { func (db *PreparedStmtDB) prepare(query string) (*sql.Stmt, error) {
db.mux.RLock() db.mux.RLock()
if stmt, ok := db.stmts[query]; ok { if stmt, ok := db.Stmts[query]; ok {
db.mux.RUnlock() db.mux.RUnlock()
return stmt, nil return stmt, nil
} }
db.mux.RUnlock() db.mux.RUnlock()
db.mux.Lock() 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) stmt, err := db.ConnPool.PrepareContext(context.Background(), query)
if err == nil { if err == nil {
db.stmts[query] = stmt db.Stmts[query] = stmt
} }
db.mux.Unlock() db.mux.Unlock()