-
Notifications
You must be signed in to change notification settings - Fork 466
/
postgresdbaccess.go
321 lines (257 loc) · 7.75 KB
/
postgresdbaccess.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
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
// ------------------------------------------------------------
// Copyright (c) Microsoft Corporation and Dapr Contributors.
// Licensed under the MIT License.
// ------------------------------------------------------------
package postgresql
import (
"database/sql"
"encoding/base64"
"encoding/json"
"fmt"
"strconv"
"github.com/agrea/ptr"
"github.com/dapr/components-contrib/state"
"github.com/dapr/components-contrib/state/utils"
"github.com/dapr/kit/logger"
// Blank import for the underlying PostgreSQL driver.
_ "github.com/jackc/pgx/v4/stdlib"
)
const (
connectionStringKey = "connectionString"
errMissingConnectionString = "missing connection string"
tableName = "state"
)
// postgresDBAccess implements dbaccess.
type postgresDBAccess struct {
logger logger.Logger
metadata state.Metadata
db *sql.DB
connectionString string
}
// newPostgresDBAccess creates a new instance of postgresAccess.
func newPostgresDBAccess(logger logger.Logger) *postgresDBAccess {
logger.Debug("Instantiating new PostgreSQL state store")
return &postgresDBAccess{
logger: logger,
}
}
// Init sets up PostgreSQL connection and ensures that the state table exists.
func (p *postgresDBAccess) Init(metadata state.Metadata) error {
p.logger.Debug("Initializing PostgreSQL state store")
p.metadata = metadata
if val, ok := metadata.Properties[connectionStringKey]; ok && val != "" {
p.connectionString = val
} else {
p.logger.Error("Missing postgreSQL connection string")
return fmt.Errorf(errMissingConnectionString)
}
db, err := sql.Open("pgx", p.connectionString)
if err != nil {
p.logger.Error(err)
return err
}
p.db = db
pingErr := db.Ping()
if pingErr != nil {
return pingErr
}
err = p.ensureStateTable(tableName)
if err != nil {
return err
}
return nil
}
// Set makes an insert or update to the database.
func (p *postgresDBAccess) Set(req *state.SetRequest) error {
return state.SetWithOptions(p.setValue, req)
}
// setValue is an internal implementation of set to enable passing the logic to state.SetWithRetries as a func.
func (p *postgresDBAccess) setValue(req *state.SetRequest) error {
p.logger.Debug("Setting state value in PostgreSQL")
err := state.CheckRequestOptions(req.Options)
if err != nil {
return err
}
if req.Key == "" {
return fmt.Errorf("missing key in set operation")
}
if v, ok := req.Value.(string); ok && v == "" {
return fmt.Errorf("empty string is not allowed in set operation")
}
v := req.Value
byteArray, isBinary := req.Value.([]uint8)
if isBinary {
v = base64.StdEncoding.EncodeToString(byteArray)
}
// Convert to json string
bt, _ := utils.Marshal(v, json.Marshal)
value := string(bt)
var result sql.Result
// Sprintf is required for table name because sql.DB does not substitute parameters for table names.
// Other parameters use sql.DB parameter substitution.
if req.ETag == nil {
result, err = p.db.Exec(fmt.Sprintf(
`INSERT INTO %s (key, value, isbinary) VALUES ($1, $2, $3)
ON CONFLICT (key) DO UPDATE SET value = $2, isbinary = $3, updatedate = NOW();`,
tableName), req.Key, value, isBinary)
} else {
// Convert req.ETag to integer for postgres compatibility
var etag int
etag, err = strconv.Atoi(*req.ETag)
if err != nil {
return state.NewETagError(state.ETagInvalid, err)
}
// When an etag is provided do an update - no insert
result, err = p.db.Exec(fmt.Sprintf(
`UPDATE %s SET value = $1, isbinary = $2, updatedate = NOW()
WHERE key = $3 AND xmin = $4;`,
tableName), value, isBinary, req.Key, etag)
}
if err != nil {
if req.ETag != nil && *req.ETag != "" {
return state.NewETagError(state.ETagMismatch, err)
}
return err
}
rows, err := result.RowsAffected()
if err != nil {
return err
}
if rows != 1 {
return fmt.Errorf("no item was updated")
}
return nil
}
// Get returns data from the database. If data does not exist for the key an empty state.GetResponse will be returned.
func (p *postgresDBAccess) Get(req *state.GetRequest) (*state.GetResponse, error) {
p.logger.Debug("Getting state value from PostgreSQL")
if req.Key == "" {
return nil, fmt.Errorf("missing key in get operation")
}
var value string
var isBinary bool
var etag int
err := p.db.QueryRow(fmt.Sprintf("SELECT value, isbinary, xmin as etag FROM %s WHERE key = $1", tableName), req.Key).Scan(&value, &isBinary, &etag)
if err != nil {
// If no rows exist, return an empty response, otherwise return the error.
if err == sql.ErrNoRows {
return &state.GetResponse{}, nil
}
return nil, err
}
if isBinary {
var s string
var data []byte
if err = json.Unmarshal([]byte(value), &s); err != nil {
return nil, err
}
if data, err = base64.StdEncoding.DecodeString(s); err != nil {
return nil, err
}
return &state.GetResponse{
Data: data,
ETag: ptr.String(strconv.Itoa(etag)),
Metadata: req.Metadata,
}, nil
}
return &state.GetResponse{
Data: []byte(value),
ETag: ptr.String(strconv.Itoa(etag)),
Metadata: req.Metadata,
}, nil
}
// Delete removes an item from the state store.
func (p *postgresDBAccess) Delete(req *state.DeleteRequest) error {
return state.DeleteWithOptions(p.deleteValue, req)
}
// deleteValue is an internal implementation of delete to enable passing the logic to state.DeleteWithRetries as a func.
func (p *postgresDBAccess) deleteValue(req *state.DeleteRequest) error {
p.logger.Debug("Deleting state value from PostgreSQL")
if req.Key == "" {
return fmt.Errorf("missing key in delete operation")
}
var result sql.Result
var err error
if req.ETag == nil {
result, err = p.db.Exec("DELETE FROM state WHERE key = $1", req.Key)
} else {
// Convert req.ETag to integer for postgres compatibility
etag, conversionError := strconv.Atoi(*req.ETag)
if conversionError != nil {
return state.NewETagError(state.ETagInvalid, err)
}
result, err = p.db.Exec("DELETE FROM state WHERE key = $1 and xmin = $2", req.Key, etag)
}
if err != nil {
return err
}
rows, err := result.RowsAffected()
if err != nil {
return err
}
if rows != 1 && req.ETag != nil && *req.ETag != "" {
return state.NewETagError(state.ETagMismatch, nil)
}
return nil
}
func (p *postgresDBAccess) ExecuteMulti(sets []state.SetRequest, deletes []state.DeleteRequest) error {
p.logger.Debug("Executing multiple PostgreSQL operations")
tx, err := p.db.Begin()
if err != nil {
return err
}
if len(deletes) > 0 {
for _, d := range deletes {
da := d // Fix for gosec G601: Implicit memory aliasing in for loop.
err = p.Delete(&da)
if err != nil {
tx.Rollback()
return err
}
}
}
if len(sets) > 0 {
for _, s := range sets {
sa := s // Fix for gosec G601: Implicit memory aliasing in for loop.
err = p.Set(&sa)
if err != nil {
tx.Rollback()
return err
}
}
}
err = tx.Commit()
return err
}
// Close implements io.Close.
func (p *postgresDBAccess) Close() error {
if p.db != nil {
return p.db.Close()
}
return nil
}
func (p *postgresDBAccess) ensureStateTable(stateTableName string) error {
exists, err := tableExists(p.db, stateTableName)
if err != nil {
return err
}
if !exists {
p.logger.Info("Creating PostgreSQL state table")
createTable := fmt.Sprintf(`CREATE TABLE %s (
key text NOT NULL PRIMARY KEY,
value jsonb NOT NULL,
isbinary boolean NOT NULL,
insertdate TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
updatedate TIMESTAMP WITH TIME ZONE NULL);`, stateTableName)
_, err = p.db.Exec(createTable)
if err != nil {
return err
}
}
return nil
}
func tableExists(db *sql.DB, tableName string) (bool, error) {
var exists bool = false
err := db.QueryRow("SELECT EXISTS (SELECT FROM pg_tables where tablename = $1)", tableName).Scan(&exists)
return exists, err
}