From bd2dc5a03f368252878eef89ed7fb5e486292e54 Mon Sep 17 00:00:00 2001
From: a631807682 <631807682@qq.com>
Date: Tue, 18 Oct 2022 08:07:35 +0800
Subject: [PATCH] fix(GetRows): skip prepare stmt

---
 migrator.go | 28 +++++++++++++++++++++-------
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/migrator.go b/migrator.go
index 26573da..1080ed2 100644
--- a/migrator.go
+++ b/migrator.go
@@ -528,14 +528,28 @@ func (m Migrator) GetRows(currentSchema interface{}, table interface{}) (*sql.Ro
 		name = fmt.Sprintf("%v.%v", currentSchema, table)
 	}
 
-	return m.DB.Session(&gorm.Session{}).Table(name).Limit(1).Scopes(func(d *gorm.DB) *gorm.DB {
-		dialector, _ := m.Dialector.(Dialector)
-		// use simple protocol
-		if !m.DB.PrepareStmt && (dialector.Config != nil && (dialector.Config.DriverName == "" || dialector.Config.DriverName == "pgx")) {
-			d.Statement.Vars = append(d.Statement.Vars, pgx.QuerySimpleProtocol(true))
+	var res []interface{}
+	sql := m.DB.ToSQL(func(tx *gorm.DB) *gorm.DB {
+		return tx.Table(name).Limit(1).Find(&res)
+	})
+	dialector, _ := m.Dialector.(Dialector)
+
+	// skip the prepared stmt
+	if m.DB.PrepareStmt {
+		if pdb, ok := m.DB.ConnPool.(*gorm.PreparedStmtDB); ok {
+			pdb.Mux.Lock()
+			if stmt, ok := pdb.Stmts[sql]; ok {
+				go stmt.Close()
+				delete(pdb.Stmts, sql)
+			}
+			pdb.Mux.Unlock()
 		}
-		return d
-	}).Rows()
+	} else if dialector.Config != nil && (dialector.Config.DriverName == "" || dialector.Config.DriverName == "pgx") {
+		// use simple protocol for pgx
+		return m.DB.Raw(sql, pgx.QuerySimpleProtocol(true)).Rows()
+	}
+
+	return m.DB.Raw(sql).Rows()
 }
 
 func (m Migrator) CurrentSchema(stmt *gorm.Statement, table string) (interface{}, interface{}) {