Refactor Prepared Statement

gorm.go

@@ -108,11 +108,15 @@ func Open(dialector Dialector, config *Config) (db *DB, err error) {
 		err = config.Dialector.Initialize(db)
 	}
 
+	preparedStmt := &PreparedStmtDB{
+		ConnPool:    db.ConnPool,
+		Stmts:       map[string]*sql.Stmt{},
+		PreparedSQL: make([]string, 0, 100),
+	}
+	db.cacheStore.Store("preparedStmt", preparedStmt)
+
 	if config.PrepareStmt {
-		db.ConnPool = &PreparedStmtDB{
+		db.ConnPool = preparedStmt
-			ConnPool: db.ConnPool,
-			Stmts:    map[string]*sql.Stmt{},
-		}
 	}
 
 	db.Statement = &Statement{
@@ -157,9 +161,13 @@ func (db *DB) Session(config *Session) *DB {
 	}
 
 	if config.PrepareStmt {
-		tx.Statement.ConnPool = &PreparedStmtDB{
+		if v, ok := db.cacheStore.Load("preparedStmt"); ok {
-			ConnPool: db.Config.ConnPool,
+			preparedStmt := v.(*PreparedStmtDB)
-			Stmts:    map[string]*sql.Stmt{},
+			tx.Statement.ConnPool = &PreparedStmtDB{
+				ConnPool: db.Config.ConnPool,
+				mux:      preparedStmt.mux,
+				Stmts:    preparedStmt.Stmts,
+			}
 		}
 	}
 

prepare_stmt.go

@@ -7,16 +7,19 @@ import (
 )
 
 type PreparedStmtDB struct {
-	Stmts map[string]*sql.Stmt
+	Stmts       map[string]*sql.Stmt
-	mux   sync.RWMutex
+	PreparedSQL []string
+	mux         sync.RWMutex
 	ConnPool
 }
 
 func (db *PreparedStmtDB) Close() {
 	db.mux.Lock()
-	for k, stmt := range db.Stmts {
+	for _, query := range db.PreparedSQL {
-		delete(db.Stmts, k)
+		if stmt, ok := db.Stmts[query]; ok {
-		stmt.Close()
+			delete(db.Stmts, query)
+			stmt.Close()
+		}
 	}
 
 	db.mux.Unlock()
@@ -40,6 +43,7 @@ func (db *PreparedStmtDB) prepare(query string) (*sql.Stmt, error) {
 	stmt, err := db.ConnPool.PrepareContext(context.Background(), query)
 	if err == nil {
 		db.Stmts[query] = stmt
+		db.PreparedSQL = append(db.PreparedSQL, query)
 	}
 	db.mux.Unlock()
 

tests/.gitignore

@@ -0,0 +1 @@
+go.sum