feat(PreparedStmtDB): support reset

This commit is contained in:
a631807682 2022-10-19 13:45:43 +08:00
parent 3f20a543fa
commit 638805ed62
No known key found for this signature in database
GPG Key ID: 137D1D75522168AB
2 changed files with 32 additions and 1 deletions

View File

@ -44,6 +44,12 @@ func (db *PreparedStmtDB) Close() {
}
}
func (db *PreparedStmtDB) Reset() {
db.Close()
db.PreparedSQL = make([]string, 0, 100)
db.Stmts = map[string](*Stmt){}
}
func (db *PreparedStmtDB) prepare(ctx context.Context, conn ConnPool, isTransaction bool, query string) (Stmt, error) {
db.Mux.RLock()
if stmt, ok := db.Stmts[query]; ok && (!stmt.Transaction || isTransaction) {

View File

@ -2,8 +2,8 @@ package tests_test
import (
"context"
"sync"
"errors"
"sync"
"testing"
"time"
@ -168,3 +168,28 @@ func TestPreparedStmtInTransaction(t *testing.T) {
t.Errorf("Failed, got error: %v", err)
}
}
func TestPreparedStmtReset(t *testing.T) {
tx := DB.Session(&gorm.Session{PrepareStmt: true})
pdb, ok := tx.ConnPool.(*gorm.PreparedStmtDB)
if !ok {
t.Fatalf("should assign PreparedStatement Manager back to database when using PrepareStmt mode")
}
user := *GetUser("prepared_stmt_reset", Config{})
tx.Create(&user)
pdb.Mux.Lock()
if len(pdb.PreparedSQL) == 0 || len(pdb.Stmts) == 0 {
pdb.Mux.Unlock()
t.Fatalf("prepared stmt can not be empty")
}
pdb.Mux.Unlock()
pdb.Reset()
pdb.Mux.Lock()
defer pdb.Mux.Unlock()
if len(pdb.PreparedSQL) != 0 || len(pdb.Stmts) != 0 {
t.Fatalf("prepared stmt should be empty")
}
}