2020-06-05 05:08:22 +03:00
|
|
|
package gorm
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"database/sql"
|
|
|
|
"sync"
|
|
|
|
)
|
|
|
|
|
2021-01-10 11:46:06 +03:00
|
|
|
type Stmt struct {
|
|
|
|
*sql.Stmt
|
|
|
|
Transaction bool
|
|
|
|
}
|
|
|
|
|
2020-06-05 05:08:22 +03:00
|
|
|
type PreparedStmtDB struct {
|
2021-01-10 11:46:06 +03:00
|
|
|
Stmts map[string]Stmt
|
2020-07-28 09:26:09 +03:00
|
|
|
PreparedSQL []string
|
2020-10-21 10:47:46 +03:00
|
|
|
Mux *sync.RWMutex
|
2020-06-05 05:08:22 +03:00
|
|
|
ConnPool
|
|
|
|
}
|
|
|
|
|
2021-03-19 10:54:32 +03:00
|
|
|
func (db *PreparedStmtDB) GetDB() (*sql.DB, error) {
|
|
|
|
if dbConnector, ok := db.ConnPool.(GetDBConnector); ok && dbConnector != nil {
|
|
|
|
return dbConnector.GetDBConn()
|
|
|
|
}
|
|
|
|
|
|
|
|
if sqldb, ok := db.ConnPool.(*sql.DB); ok {
|
|
|
|
return sqldb, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
return nil, ErrInvaildDB
|
|
|
|
}
|
|
|
|
|
2020-06-21 13:01:50 +03:00
|
|
|
func (db *PreparedStmtDB) Close() {
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Lock()
|
2020-07-28 09:26:09 +03:00
|
|
|
for _, query := range db.PreparedSQL {
|
|
|
|
if stmt, ok := db.Stmts[query]; ok {
|
|
|
|
delete(db.Stmts, query)
|
|
|
|
stmt.Close()
|
|
|
|
}
|
2020-06-21 13:01:50 +03:00
|
|
|
}
|
|
|
|
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Unlock()
|
2020-06-21 13:01:50 +03:00
|
|
|
}
|
|
|
|
|
2021-01-10 11:46:06 +03:00
|
|
|
func (db *PreparedStmtDB) prepare(ctx context.Context, conn ConnPool, isTransaction bool, query string) (Stmt, error) {
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.RLock()
|
2021-01-10 11:46:06 +03:00
|
|
|
if stmt, ok := db.Stmts[query]; ok && (!stmt.Transaction || isTransaction) {
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.RUnlock()
|
2020-06-05 05:08:22 +03:00
|
|
|
return stmt, nil
|
|
|
|
}
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.RUnlock()
|
2020-06-05 05:08:22 +03:00
|
|
|
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Lock()
|
2020-06-21 13:01:50 +03:00
|
|
|
// double check
|
2021-01-10 11:46:06 +03:00
|
|
|
if stmt, ok := db.Stmts[query]; ok && (!stmt.Transaction || isTransaction) {
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Unlock()
|
2020-06-21 13:01:50 +03:00
|
|
|
return stmt, nil
|
2021-01-10 11:46:06 +03:00
|
|
|
} else if ok {
|
|
|
|
stmt.Close()
|
2020-06-21 13:01:50 +03:00
|
|
|
}
|
|
|
|
|
2020-12-14 13:30:43 +03:00
|
|
|
stmt, err := conn.PrepareContext(ctx, query)
|
2020-06-05 05:08:22 +03:00
|
|
|
if err == nil {
|
2021-01-10 11:46:06 +03:00
|
|
|
db.Stmts[query] = Stmt{Stmt: stmt, Transaction: isTransaction}
|
2020-07-28 09:26:09 +03:00
|
|
|
db.PreparedSQL = append(db.PreparedSQL, query)
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Unlock()
|
2020-06-05 05:08:22 +03:00
|
|
|
|
2021-01-10 11:46:06 +03:00
|
|
|
return db.Stmts[query], err
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
func (db *PreparedStmtDB) BeginTx(ctx context.Context, opt *sql.TxOptions) (ConnPool, error) {
|
|
|
|
if beginner, ok := db.ConnPool.(TxBeginner); ok {
|
|
|
|
tx, err := beginner.BeginTx(ctx, opt)
|
|
|
|
return &PreparedStmtTX{PreparedStmtDB: db, Tx: tx}, err
|
|
|
|
}
|
|
|
|
return nil, ErrInvalidTransaction
|
|
|
|
}
|
|
|
|
|
2020-07-03 05:26:18 +03:00
|
|
|
func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
|
2021-01-10 11:46:06 +03:00
|
|
|
stmt, err := db.prepare(ctx, db.ConnPool, false, query)
|
2020-06-05 05:08:22 +03:00
|
|
|
if err == nil {
|
2020-07-03 05:26:18 +03:00
|
|
|
result, err = stmt.ExecContext(ctx, args...)
|
|
|
|
if err != nil {
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Lock()
|
2020-07-03 05:26:18 +03:00
|
|
|
stmt.Close()
|
|
|
|
delete(db.Stmts, query)
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Unlock()
|
2020-07-03 05:26:18 +03:00
|
|
|
}
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
2020-07-03 05:26:18 +03:00
|
|
|
return result, err
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
|
|
|
|
2020-07-03 05:26:18 +03:00
|
|
|
func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
|
2021-01-10 11:46:06 +03:00
|
|
|
stmt, err := db.prepare(ctx, db.ConnPool, false, query)
|
2020-06-05 05:08:22 +03:00
|
|
|
if err == nil {
|
2020-07-03 05:26:18 +03:00
|
|
|
rows, err = stmt.QueryContext(ctx, args...)
|
|
|
|
if err != nil {
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Lock()
|
2020-07-03 05:26:18 +03:00
|
|
|
stmt.Close()
|
|
|
|
delete(db.Stmts, query)
|
2020-08-03 16:48:36 +03:00
|
|
|
db.Mux.Unlock()
|
2020-07-03 05:26:18 +03:00
|
|
|
}
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
2020-07-03 05:26:18 +03:00
|
|
|
return rows, err
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
func (db *PreparedStmtDB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
2021-01-10 11:46:06 +03:00
|
|
|
stmt, err := db.prepare(ctx, db.ConnPool, false, query)
|
2020-06-05 05:08:22 +03:00
|
|
|
if err == nil {
|
|
|
|
return stmt.QueryRowContext(ctx, args...)
|
|
|
|
}
|
|
|
|
return &sql.Row{}
|
|
|
|
}
|
|
|
|
|
|
|
|
type PreparedStmtTX struct {
|
|
|
|
*sql.Tx
|
|
|
|
PreparedStmtDB *PreparedStmtDB
|
|
|
|
}
|
|
|
|
|
2020-09-03 14:16:55 +03:00
|
|
|
func (tx *PreparedStmtTX) Commit() error {
|
|
|
|
if tx.Tx != nil {
|
|
|
|
return tx.Tx.Commit()
|
|
|
|
}
|
|
|
|
return ErrInvalidTransaction
|
|
|
|
}
|
|
|
|
|
|
|
|
func (tx *PreparedStmtTX) Rollback() error {
|
|
|
|
if tx.Tx != nil {
|
|
|
|
return tx.Tx.Rollback()
|
|
|
|
}
|
|
|
|
return ErrInvalidTransaction
|
|
|
|
}
|
|
|
|
|
2020-07-03 05:26:18 +03:00
|
|
|
func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
|
2021-01-10 11:46:06 +03:00
|
|
|
stmt, err := tx.PreparedStmtDB.prepare(ctx, tx.Tx, true, query)
|
2020-06-05 05:08:22 +03:00
|
|
|
if err == nil {
|
2021-01-10 11:46:06 +03:00
|
|
|
result, err = tx.Tx.StmtContext(ctx, stmt.Stmt).ExecContext(ctx, args...)
|
2020-07-03 05:26:18 +03:00
|
|
|
if err != nil {
|
2020-08-03 16:48:36 +03:00
|
|
|
tx.PreparedStmtDB.Mux.Lock()
|
2020-07-03 05:26:18 +03:00
|
|
|
stmt.Close()
|
|
|
|
delete(tx.PreparedStmtDB.Stmts, query)
|
2020-08-03 16:48:36 +03:00
|
|
|
tx.PreparedStmtDB.Mux.Unlock()
|
2020-07-03 05:26:18 +03:00
|
|
|
}
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
2020-07-03 05:26:18 +03:00
|
|
|
return result, err
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
|
|
|
|
2020-07-03 05:26:18 +03:00
|
|
|
func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
|
2021-01-10 11:46:06 +03:00
|
|
|
stmt, err := tx.PreparedStmtDB.prepare(ctx, tx.Tx, true, query)
|
2020-06-05 05:08:22 +03:00
|
|
|
if err == nil {
|
2021-01-10 11:46:06 +03:00
|
|
|
rows, err = tx.Tx.Stmt(stmt.Stmt).QueryContext(ctx, args...)
|
2020-07-03 05:26:18 +03:00
|
|
|
if err != nil {
|
2020-08-03 16:48:36 +03:00
|
|
|
tx.PreparedStmtDB.Mux.Lock()
|
2020-07-03 05:26:18 +03:00
|
|
|
stmt.Close()
|
|
|
|
delete(tx.PreparedStmtDB.Stmts, query)
|
2020-08-03 16:48:36 +03:00
|
|
|
tx.PreparedStmtDB.Mux.Unlock()
|
2020-07-03 05:26:18 +03:00
|
|
|
}
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
2020-07-03 05:26:18 +03:00
|
|
|
return rows, err
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
|
|
|
|
|
|
|
func (tx *PreparedStmtTX) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
|
2021-01-10 11:46:06 +03:00
|
|
|
stmt, err := tx.PreparedStmtDB.prepare(ctx, tx.Tx, true, query)
|
2020-06-05 05:08:22 +03:00
|
|
|
if err == nil {
|
2021-01-10 11:46:06 +03:00
|
|
|
return tx.Tx.StmtContext(ctx, stmt.Stmt).QueryRowContext(ctx, args...)
|
2020-06-05 05:08:22 +03:00
|
|
|
}
|
|
|
|
return &sql.Row{}
|
|
|
|
}
|