Skip to content

Commit

Permalink
feat(PreparedStmtDB): support reset
Browse files Browse the repository at this point in the history
  • Loading branch information
a631807682 committed Oct 19, 2022
1 parent 3f20a54 commit 638805e
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 1 deletion.
6 changes: 6 additions & 0 deletions prepare_stmt.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
27 changes: 26 additions & 1 deletion tests/prepared_stmt_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@ package tests_test

import (
"context"
"sync"
"errors"
"sync"
"testing"
"time"

Expand Down Expand Up @@ -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")
}
}

0 comments on commit 638805e

Please sign in to comment.