prepare_stmt add ctx

This commit is contained in:
王岚 2020-08-28 18:40:59 +08:00
parent 0d96f99499
commit 066d54db1f
1 changed files with 11 additions and 11 deletions

View File

@ -25,7 +25,7 @@ func (db *PreparedStmtDB) Close() {
db.Mux.Unlock() db.Mux.Unlock()
} }
func (db *PreparedStmtDB) prepare(query string) (*sql.Stmt, error) { func (db *PreparedStmtDB) prepare(ctx context.Context, query string) (*sql.Stmt, error) {
db.Mux.RLock() db.Mux.RLock()
if stmt, ok := db.Stmts[query]; ok { if stmt, ok := db.Stmts[query]; ok {
db.Mux.RUnlock() db.Mux.RUnlock()
@ -40,7 +40,7 @@ func (db *PreparedStmtDB) prepare(query string) (*sql.Stmt, error) {
return stmt, nil return stmt, nil
} }
stmt, err := db.ConnPool.PrepareContext(context.Background(), query) stmt, err := db.ConnPool.PrepareContext(ctx, query)
if err == nil { if err == nil {
db.Stmts[query] = stmt db.Stmts[query] = stmt
db.PreparedSQL = append(db.PreparedSQL, query) db.PreparedSQL = append(db.PreparedSQL, query)
@ -59,7 +59,7 @@ func (db *PreparedStmtDB) BeginTx(ctx context.Context, opt *sql.TxOptions) (Conn
} }
func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) { func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
stmt, err := db.prepare(query) stmt, err := db.prepare(ctx, query)
if err == nil { if err == nil {
result, err = stmt.ExecContext(ctx, args...) result, err = stmt.ExecContext(ctx, args...)
if err != nil { if err != nil {
@ -73,7 +73,7 @@ func (db *PreparedStmtDB) ExecContext(ctx context.Context, query string, args ..
} }
func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) { func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
stmt, err := db.prepare(query) stmt, err := db.prepare(ctx, query)
if err == nil { if err == nil {
rows, err = stmt.QueryContext(ctx, args...) rows, err = stmt.QueryContext(ctx, args...)
if err != nil { if err != nil {
@ -87,7 +87,7 @@ func (db *PreparedStmtDB) QueryContext(ctx context.Context, query string, args .
} }
func (db *PreparedStmtDB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row { func (db *PreparedStmtDB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
stmt, err := db.prepare(query) stmt, err := db.prepare(ctx, query)
if err == nil { if err == nil {
return stmt.QueryRowContext(ctx, args...) return stmt.QueryRowContext(ctx, args...)
} }
@ -100,9 +100,9 @@ type PreparedStmtTX struct {
} }
func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) { func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ...interface{}) (result sql.Result, err error) {
stmt, err := tx.PreparedStmtDB.prepare(query) stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
if err == nil { if err == nil {
result, err = tx.Tx.Stmt(stmt).ExecContext(ctx, args...) result, err = tx.Tx.StmtContext(ctx, stmt).ExecContext(ctx, args...)
if err != nil { if err != nil {
tx.PreparedStmtDB.Mux.Lock() tx.PreparedStmtDB.Mux.Lock()
stmt.Close() stmt.Close()
@ -114,9 +114,9 @@ func (tx *PreparedStmtTX) ExecContext(ctx context.Context, query string, args ..
} }
func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) { func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
stmt, err := tx.PreparedStmtDB.prepare(query) stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
if err == nil { if err == nil {
rows, err = tx.Tx.Stmt(stmt).QueryContext(ctx, args...) rows, err = tx.Tx.StmtContext(ctx, stmt).QueryContext(ctx, args...)
if err != nil { if err != nil {
tx.PreparedStmtDB.Mux.Lock() tx.PreparedStmtDB.Mux.Lock()
stmt.Close() stmt.Close()
@ -128,9 +128,9 @@ func (tx *PreparedStmtTX) QueryContext(ctx context.Context, query string, args .
} }
func (tx *PreparedStmtTX) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row { func (tx *PreparedStmtTX) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
stmt, err := tx.PreparedStmtDB.prepare(query) stmt, err := tx.PreparedStmtDB.prepare(ctx, query)
if err == nil { if err == nil {
return tx.Tx.Stmt(stmt).QueryRowContext(ctx, args...) return tx.Tx.StmtContext(ctx, stmt).QueryRowContext(ctx, args...)
} }
return &sql.Row{} return &sql.Row{}
} }