-
Notifications
You must be signed in to change notification settings - Fork 0
/
store.go
187 lines (156 loc) · 4.1 KB
/
store.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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
package postgres
import (
"context"
"database/sql"
"encoding/json"
"github.com/gokv/store"
)
// Store holds the SQL statements prepared against a Postgresql table.
// Initialise with New.
type Store struct {
getStmt *sql.Stmt
getAllStmt *sql.Stmt
addStmt *sql.Stmt
setStmt *sql.Stmt
updateStmt *sql.Stmt
deleteStmt *sql.Stmt
ping func(context.Context) error
}
// New creates a table of name tablename if it does not exist, and prepares
// statements against it.
// The table has two columms: "key" is the TEXT primary key and "v" is a JSONb column holding the values.
func New(db *sql.DB, tablename string) (s Store, err error) {
if _, err = db.Exec(`CREATE TABLE IF NOT EXISTS "` + tablename + `" (key TEXT NOT NULL PRIMARY KEY, v jsonb NOT NULL)`); err != nil {
return s, err
}
if s.getStmt, err = db.Prepare(`SELECT v FROM "` + tablename + `" WHERE key=$1`); err != nil {
_ = s.Close()
return s, err
}
if s.getAllStmt, err = db.Prepare(`SELECT v FROM "` + tablename + `"`); err != nil {
_ = s.Close()
return s, err
}
if s.addStmt, err = db.Prepare(`INSERT INTO "` + tablename + `" (key, v) VALUES ($1, $2)`); err != nil {
_ = s.Close()
return s, err
}
if s.setStmt, err = db.Prepare(`INSERT INTO "` + tablename + `" (key, v) VALUES ($1, $2) ON CONFLICT (key) DO UPDATE SET v=$2`); err != nil {
_ = s.Close()
return s, err
}
if s.updateStmt, err = db.Prepare(`UPDATE "` + tablename + `" SET v=$2 WHERE key=$1`); err != nil {
_ = s.Close()
return s, err
}
if s.deleteStmt, err = db.Prepare(`DELETE FROM "` + tablename + `" WHERE key=$1`); err != nil {
_ = s.Close()
return s, err
}
s.ping = db.PingContext
return s, err
}
// Close releases the resources associated with the Store.
// Returns the first error encountered while closing the prepared statements.
func (s Store) Close() (err error) {
for _, stmt := range []*sql.Stmt{
s.getStmt,
s.getAllStmt,
s.addStmt,
s.setStmt,
s.updateStmt,
s.deleteStmt,
} {
if stmt != nil {
if e := stmt.Close(); err == nil {
err = e
}
}
}
return err
}
// Get retrieves a new item by key and unmarshals it into v, or returns false if
// not found.
// Err is non-nil in case of failure.
func (s Store) Get(ctx context.Context, key interface{}, v json.Unmarshaler) (bool, error) {
var b []byte
if err := s.getStmt.QueryRowContext(ctx, key).Scan(&b); err != nil {
if err == sql.ErrNoRows {
return false, nil
}
return false, err
}
return true, v.UnmarshalJSON(b)
}
// GetAll appends to c every item in the store.
// Err is non-nil in case of failure.
func (s Store) GetAll(ctx context.Context, c store.Collection) error {
rows, err := s.getAllStmt.QueryContext(ctx)
if err != nil {
return err
}
defer rows.Close()
var b []byte
for rows.Next() {
if err = rows.Scan(&b); err != nil {
return err
}
if err = c.New().UnmarshalJSON(b); err != nil {
return err
}
}
return rows.Err()
}
func (s Store) Add(ctx context.Context, key interface{}, v json.Marshaler) error {
b, err := v.MarshalJSON()
if err != nil {
return err
}
_, err = s.addStmt.ExecContext(ctx, key, b)
return err
}
func (s Store) Set(ctx context.Context, key interface{}, v json.Marshaler) error {
b, err := v.MarshalJSON()
if err != nil {
return err
}
_, err = s.setStmt.ExecContext(ctx, key, b)
return err
}
// Update assigns the given value to the given key, if it exists.
// Err is non-nil if the key was not already present, or in case of failure.
func (s Store) Update(ctx context.Context, key interface{}, v json.Marshaler) error {
b, err := v.MarshalJSON()
if err != nil {
return err
}
res, err := s.updateStmt.ExecContext(ctx, key, b)
if err != nil {
return err
}
n, err := res.RowsAffected()
if err != nil {
return err
}
if n < 1 {
return store.ErrNoRows
}
return nil
}
func (s Store) Delete(ctx context.Context, key interface{}) error {
res, err := s.deleteStmt.ExecContext(ctx, key)
if err != nil {
return err
}
n, err := res.RowsAffected()
if err != nil {
return err
}
if n < 1 {
return store.ErrNoRows
}
return nil
}
func (s Store) Ping(ctx context.Context) error {
return s.ping(ctx)
}