-
Notifications
You must be signed in to change notification settings - Fork 2
/
dynmgrm.go
318 lines (278 loc) · 8.7 KB
/
dynmgrm.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
//go:generate mockgen -destination=internal/mocks/mock_dynmgrm.go -package=mocks -source=./dynmgrm.go
//go:generate mockgen -destination=internal/mocks/mock_gorm_expression.go -package=mocks gorm.io/gorm/clause Expression
//go:generate mockgen -destination=internal/mocks/mock_gorm_builder.go -package=mocks gorm.io/gorm/clause Builder
package dynmgrm
import (
"database/sql"
"errors"
"fmt"
"github.com/miyamo2/godynamo"
"gorm.io/gorm/migrator"
"strconv"
"strings"
"gorm.io/gorm"
"gorm.io/gorm/callbacks"
"gorm.io/gorm/clause"
"gorm.io/gorm/logger"
"gorm.io/gorm/schema"
)
// compatibility check
var _ gorm.Dialector = (*Dialector)(nil)
// KeySchemaDataType is the data type for the DynamoDB key schema.
type KeySchemaDataType string
// String returns the string representation of the KeySchemaDataType.
func (d KeySchemaDataType) String() string {
return string(d)
}
const (
// DriverName is the driver name for DynamoDB.
DriverName = "godynamo"
DBName = "dynamodb"
)
// Define KeySchemaDataType
const (
// KeySchemaDataTypeString is the data type for string.
KeySchemaDataTypeString KeySchemaDataType = "string"
// KeySchemaDataTypeNumber is the data type for number.
KeySchemaDataTypeNumber KeySchemaDataType = "number"
// KeySchemaDataTypeBinary is the data type for binary.
KeySchemaDataTypeBinary KeySchemaDataType = "binary"
)
var (
queryClauses = []string{"SELECT", "FROM", "WHERE"}
createClauses = []string{"INSERT", "VALUES"}
updateClauses = []string{"UPDATE", "SET", "WHERE"}
deleteClauses = []string{"DELETE", "FROM", "WHERE"}
clauseBuilders = map[string]clause.ClauseBuilder{
"VALUES": toClauseBuilder(buildValuesClause),
"SET": toClauseBuilder(buildSetClause),
}
)
// config is the configuration for the DynamoDB connection.
type config struct {
region string
akId string
secret string
endpoint string
timeout int
conn gorm.ConnPool
}
// DBOpener is the interface for opening a database.
type DBOpener interface {
DSN() string
DriverName() string
Apply() (*sql.DB, error)
}
type CallbacksRegisterer interface {
Register(db *gorm.DB, config *callbacks.Config)
}
// Dialector gorm dialector for DynamoDB
type Dialector struct {
conn gorm.ConnPool
// dbOpener is used for testing
dbOpener DBOpener
// callbacksRegisterer is used for testing
callbacksRegisterer CallbacksRegisterer
}
// DialectorOption is the option for the DynamoDB dialector.
type DialectorOption func(*config)
// WithRegion sets the region for the DynamoDB connection.
//
// Default: https://github.com/miyamo2/godynamo?tab=readme-ov-file#data-source-name-dsn-format-for-aws-dynamodb
func WithRegion(region string) func(*config) {
return func(config *config) {
config.region = region
}
}
// WithAccessKeyID sets the access key ID for the DynamoDB connection.
//
// Default: https://github.com/miyamo2/godynamo?tab=readme-ov-file#data-source-name-dsn-format-for-aws-dynamodb
func WithAccessKeyID(accessKeyId string) func(*config) {
return func(config *config) {
config.akId = accessKeyId
}
}
// WithSecretKey sets the secret key for the DynamoDB connection.
//
// Default: https://github.com/miyamo2/godynamo?tab=readme-ov-file#data-source-name-dsn-format-for-aws-dynamodb
func WithSecretKey(secretKey string) func(*config) {
return func(config *config) {
config.secret = secretKey
}
}
// WithEndpoint sets the endpoint for the DynamoDB connection.
func WithEndpoint(endpoint string) func(*config) {
return func(config *config) {
config.endpoint = endpoint
}
}
// WithTimeout sets the timeout milliseconds for the DynamoDB connection.
//
// Default: https://github.com/miyamo2/godynamo?tab=readme-ov-file#data-source-name-dsn-format-for-aws-dynamodb
func WithTimeout(timeout int) func(*config) {
return func(config *config) {
config.timeout = timeout
}
}
// WithConnection sets the exist connection for the DynamoDB.
func WithConnection(conn gorm.ConnPool) func(*config) {
return func(config *config) {
config.conn = conn
}
}
// Open returns a new DynamoDB dialector based on the DSN.
//
// e.g. "region=ap-northeast-1;AkId=<YOUR_ACCESS_KEY_ID>;SecretKey=<YOUR_SECRET_KEY>"
func Open(dsn string) gorm.Dialector {
return &Dialector{
dbOpener: dbOpener{dsn: dsn, driverName: DriverName},
callbacksRegisterer: &callbacksRegisterer{},
}
}
// New returns a new DynamoDB dialector with options.
func New(option ...DialectorOption) gorm.Dialector {
conf := config{}
buildConfig(&conf, option...)
return &Dialector{
conn: conf.conn,
dbOpener: dbOpener{dsn: parseConnectionString(conf), driverName: DriverName},
callbacksRegisterer: &callbacksRegisterer{},
}
}
func buildConfig(conf *config, option ...DialectorOption) {
for _, opt := range option {
opt(conf)
}
}
func parseConnectionString(config config) string {
dsnbuf := strings.Builder{}
if config.region != "" {
writeConnectionParameter(&dsnbuf, "region", config.region)
}
if config.akId != "" {
writeConnectionParameter(&dsnbuf, "akId", config.akId)
}
if config.secret != "" {
writeConnectionParameter(&dsnbuf, "secretKey", config.secret)
}
if config.endpoint != "" {
writeConnectionParameter(&dsnbuf, "endpoint", config.endpoint)
}
if config.timeout != 0 {
writeConnectionParameter(&dsnbuf, "timeout", strconv.Itoa(config.timeout))
}
return dsnbuf.String()
}
func writeConnectionParameter(dsnbuf *strings.Builder, key, value string) {
if dsnbuf.Len() > 0 {
dsnbuf.WriteString(";")
}
dsnbuf.WriteString(fmt.Sprintf("%s=%s", key, value))
}
// Name returns the name of the db.
func (dialector Dialector) Name() string {
return DBName
}
// Initialize initializes the DynamoDB connection.
func (dialector Dialector) Initialize(db *gorm.DB) (err error) {
if dialector.conn != nil {
db.ConnPool = dialector.conn
} else {
conn, err := dialector.dbOpener.Apply()
if err != nil {
return err
}
db.ConnPool = conn
}
dialector.callbacksRegisterer.Register(
db,
&callbacks.Config{
CreateClauses: createClauses,
QueryClauses: queryClauses,
UpdateClauses: updateClauses,
DeleteClauses: deleteClauses,
})
for k, v := range clauseBuilders {
db.ClauseBuilders[k] = v
}
return
}
// DefaultValueOf returns the default value of the field.
func (dialector Dialector) DefaultValueOf(field *schema.Field) clause.Expression {
return clause.Expr{SQL: ""}
}
// BindVarTo writes the bind variable of [goodynamo] to [clauses.Writer].
//
// [goodynamo]: https://pkg.go.dev/github.com/miyamo2/godynamo
// [clauses.Writer]: https://pkg.go.dev/gorm.io/gorm/clause#Writer
func (dialector Dialector) BindVarTo(writer clause.Writer, stmt *gorm.Statement, v interface{}) {
writer.WriteByte('?')
}
// QuoteTo escapes identifiers in SQL queries
func (dialector Dialector) QuoteTo(writer clause.Writer, str string) {
writer.WriteString(fmt.Sprintf(`"%s"`, str))
}
// Explain returns the SQL string with the variables replaced.
// Explain is typically used only for logging, dry runs, and migration.
func (dialector Dialector) Explain(sql string, vars ...interface{}) string {
return logger.ExplainSQL(sql, nil, `"`, vars...)
}
// DataTypeOf maps GORM's data types to DynamoDB's data types.
// DataTypeOf works only with migration, so it will not return data types that are not allowed in PK, SK.
func (dialector Dialector) DataTypeOf(field *schema.Field) string {
switch field.DataType {
case schema.Bool:
return KeySchemaDataTypeString.String()
case schema.Int, schema.Uint, schema.Float:
return KeySchemaDataTypeNumber.String()
case schema.String:
return KeySchemaDataTypeString.String()
case schema.Time:
return KeySchemaDataTypeString.String()
case schema.Bytes:
return KeySchemaDataTypeBinary.String()
default:
return KeySchemaDataTypeString.String()
}
}
// Migrator returns the migrator for DynamoDB.
func (dialector Dialector) Migrator(db *gorm.DB) gorm.Migrator {
return &Migrator{
db: db,
base: migrator.Migrator{
Config: migrator.Config{
DB: db,
Dialector: dialector,
},
},
}
}
// Translate it will translate the error to native gorm errors.
func (dialector Dialector) Translate(err error) error {
switch {
case errors.Is(err, godynamo.ErrTxCommitting),
errors.Is(err, godynamo.ErrTxRollingBack),
errors.Is(err, godynamo.ErrInTx),
errors.Is(err, godynamo.ErrInvalidTxStage),
errors.Is(err, godynamo.ErrNoTx):
return gorm.ErrInvalidTransaction
}
return err
}
type dbOpener struct {
dsn string
driverName string
}
func (o dbOpener) DSN() string {
return o.dsn
}
func (o dbOpener) DriverName() string {
return o.driverName
}
func (o dbOpener) Apply() (*sql.DB, error) {
return sql.Open(o.DriverName(), o.DSN())
}
type callbacksRegisterer struct{}
func (c *callbacksRegisterer) Register(db *gorm.DB, config *callbacks.Config) {
callbacks.RegisterDefaultCallbacks(db, config)
}