Refactor Prepared Statement

gorm.go

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

prepare_stmt.go

@@ -8,16 +8,19 @@ import (
 
 type PreparedStmtDB struct {
 	Stmts       map[string]*sql.Stmt
+	PreparedSQL []string
 	mux         sync.RWMutex
 	ConnPool
 }
 
 func (db *PreparedStmtDB) Close() {
 	db.mux.Lock()
-	for k, stmt := range db.Stmts {
-		delete(db.Stmts, k)
+	for _, query := range db.PreparedSQL {
+		if stmt, ok := db.Stmts[query]; ok {
+			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