forked from Masterminds/squirrel
-
Notifications
You must be signed in to change notification settings - Fork 38
/
stmtcacher.go
114 lines (97 loc) · 2.94 KB
/
stmtcacher.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
package sqrl
import (
"context"
"database/sql"
"sync"
)
// Preparer is the interface that wraps the Prepare method.
//
// Prepare executes the given query as implemented by database/sql.Prepare.
// Prepare executes the given query as implemented by database/sql.PrepareContext.
type Preparer interface {
Prepare(query string) (*sql.Stmt, error)
PrepareContext(ctx context.Context, query string) (*sql.Stmt, error)
}
// DBProxy groups the Execer, Queryer, QueryRower, and Preparer interfaces.
type DBProxy interface {
Execer
Queryer
QueryRower
Preparer
ExecerContext
QueryerContext
QueryRowerContext
}
type stmtCacher struct {
prep Preparer
cache map[string]*sql.Stmt
mu sync.Mutex
}
// NewStmtCacher returns a DBProxy wrapping prep that caches Prepared Stmts.
//
// Stmts are cached based on the string value of their queries.
func NewStmtCacher(prep Preparer) DBProxy {
return &stmtCacher{prep: prep, cache: make(map[string]*sql.Stmt)}
}
func (sc *stmtCacher) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) {
sc.mu.Lock()
defer sc.mu.Unlock()
stmt, ok := sc.cache[query]
if ok {
return stmt, nil
}
stmt, err := sc.prep.PrepareContext(ctx, query)
if err == nil {
sc.cache[query] = stmt
}
return stmt, err
}
func (sc *stmtCacher) ExecContext(ctx context.Context, query string, args ...interface{}) (res sql.Result, err error) {
stmt, err := sc.PrepareContext(ctx, query)
if err != nil {
return
}
return stmt.ExecContext(ctx, args...)
}
func (sc *stmtCacher) QueryContext(ctx context.Context, query string, args ...interface{}) (rows *sql.Rows, err error) {
stmt, err := sc.PrepareContext(ctx, query)
if err != nil {
return
}
return stmt.QueryContext(ctx, args...)
}
func (sc *stmtCacher) QueryRowContext(ctx context.Context, query string, args ...interface{}) RowScanner {
stmt, err := sc.PrepareContext(ctx, query)
if err != nil {
return &Row{err: err}
}
return stmt.QueryRowContext(ctx, args...)
}
func (sc *stmtCacher) Prepare(query string) (*sql.Stmt, error) {
return sc.PrepareContext(context.Background(), query)
}
func (sc *stmtCacher) Exec(query string, args ...interface{}) (res sql.Result, err error) {
return sc.ExecContext(context.Background(), query, args...)
}
func (sc *stmtCacher) Query(query string, args ...interface{}) (rows *sql.Rows, err error) {
return sc.QueryContext(context.Background(), query, args...)
}
func (sc *stmtCacher) QueryRow(query string, args ...interface{}) RowScanner {
return sc.QueryRowContext(context.Background(), query, args...)
}
// DBProxyBeginner describes a DBProxy that can start transactions
type DBProxyBeginner interface {
DBProxy
Begin() (*sql.Tx, error)
}
type stmtCacheProxy struct {
DBProxy
db *sql.DB
}
// NewStmtCacheProxy creates new cache proxy for statements
func NewStmtCacheProxy(db *sql.DB) DBProxyBeginner {
return &stmtCacheProxy{DBProxy: NewStmtCacher(db), db: db}
}
func (sp *stmtCacheProxy) Begin() (*sql.Tx, error) {
return sp.db.Begin()
}