-
-
Notifications
You must be signed in to change notification settings - Fork 242
/
dialect_cockroach.go
375 lines (319 loc) · 10.8 KB
/
dialect_cockroach.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
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
package pop
import (
"bytes"
"database/sql"
"fmt"
"io"
"net/url"
"os"
"os/exec"
"path/filepath"
"regexp"
"strings"
"sync"
"github.com/gobuffalo/fizz"
"github.com/gobuffalo/fizz/translators"
"github.com/gobuffalo/pop/v6/columns"
"github.com/gobuffalo/pop/v6/internal/defaults"
"github.com/gobuffalo/pop/v6/logging"
"github.com/gofrs/uuid"
_ "github.com/jackc/pgx/v4/stdlib" // Import PostgreSQL driver
"github.com/jmoiron/sqlx"
)
const nameCockroach = "cockroach"
const portCockroach = "26257"
const selectTablesQueryCockroach = "select table_name from information_schema.tables where table_schema = 'public' and table_type = 'BASE TABLE' and table_name <> ? and table_catalog = ?"
const selectTablesQueryCockroachV1 = "select table_name from information_schema.tables where table_name <> ? and table_schema = ?"
func init() {
AvailableDialects = append(AvailableDialects, nameCockroach)
dialectSynonyms["cockroachdb"] = nameCockroach
dialectSynonyms["crdb"] = nameCockroach
finalizer[nameCockroach] = finalizerCockroach
newConnection[nameCockroach] = newCockroach
}
var _ dialect = &cockroach{}
// ServerInfo holds informational data about connected database server.
type cockroachInfo struct {
VersionString string `db:"version"`
product string `db:"-"`
license string `db:"-"`
version string `db:"-"`
buildInfo string `db:"-"`
client string `db:"-"`
}
type cockroach struct {
commonDialect
translateCache map[string]string
mu sync.Mutex
info cockroachInfo
}
func (p *cockroach) Name() string {
return nameCockroach
}
func (p *cockroach) DefaultDriver() string {
return "pgx"
}
func (p *cockroach) Details() *ConnectionDetails {
return p.ConnectionDetails
}
func (p *cockroach) Create(c *Connection, model *Model, cols columns.Columns) error {
keyType, err := model.PrimaryKeyType()
if err != nil {
return err
}
switch keyType {
case "int", "int64":
cols.Remove(model.IDField())
w := cols.Writeable()
var query string
if len(w.Cols) > 0 {
query = fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s) RETURNING %s", p.Quote(model.TableName()), w.QuotedString(p), w.SymbolizedString(), model.IDField())
} else {
query = fmt.Sprintf("INSERT INTO %s DEFAULT VALUES RETURNING %s", p.Quote(model.TableName()), model.IDField())
}
txlog(logging.SQL, c, query, model.Value)
rows, err := c.Store.NamedQueryContext(model.ctx, query, model.Value)
if err != nil {
return fmt.Errorf("named insert: %w", err)
}
defer rows.Close()
if !rows.Next() {
if err := rows.Err(); err != nil {
return fmt.Errorf("named insert: next: %w", err)
}
return fmt.Errorf("named insert: %w", sql.ErrNoRows)
}
var id interface{}
if err := rows.Scan(&id); err != nil {
return fmt.Errorf("named insert: scan: %w", err)
}
if err := rows.Close(); err != nil {
return fmt.Errorf("named insert: close: %w", err)
}
model.setID(id)
return nil
case "UUID":
var query string
if model.ID() == emptyUUID {
cols.Remove(model.IDField())
w := cols.Writeable()
if len(w.Cols) > 0 {
query = fmt.Sprintf("INSERT INTO %s (%s, %s) VALUES (gen_random_uuid(), %s) RETURNING %s", p.Quote(model.TableName()), model.IDField(), w.QuotedString(p), w.SymbolizedString(), model.IDField())
} else {
query = fmt.Sprintf("INSERT INTO %s (%s) VALUES (gen_random_uuid()) RETURNING %s", p.Quote(model.TableName()), model.IDField(), model.IDField())
}
} else {
w := cols.Writeable()
w.Add(model.IDField())
query = fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s) RETURNING %s", p.Quote(model.TableName()), w.QuotedString(p), w.SymbolizedString(), model.IDField())
}
txlog(logging.SQL, c, query, model.Value)
rows, err := c.Store.NamedQueryContext(model.ctx, query, model.Value)
if err != nil {
return fmt.Errorf("named insert: %w", err)
}
defer rows.Close()
if !rows.Next() {
if err := rows.Err(); err != nil {
return fmt.Errorf("named insert: next: %w", err)
}
return fmt.Errorf("named insert: %w", sql.ErrNoRows)
}
var id uuid.UUID
if err := rows.Scan(&id); err != nil {
return fmt.Errorf("named insert: scan: %w", err)
}
if err := rows.Close(); err != nil {
return fmt.Errorf("named insert: close: %w", err)
}
model.setID(id)
return nil
}
return genericCreate(c, model, cols, p)
}
func (p *cockroach) Update(c *Connection, model *Model, cols columns.Columns) error {
return genericUpdate(c, model, cols, p)
}
func (p *cockroach) UpdateQuery(c *Connection, model *Model, cols columns.Columns, query Query) (int64, error) {
return genericUpdateQuery(c, model, cols, p, query, sqlx.DOLLAR)
}
func (p *cockroach) Destroy(c *Connection, model *Model) error {
stmt := p.TranslateSQL(fmt.Sprintf("DELETE FROM %s AS %s WHERE %s", p.Quote(model.TableName()), model.Alias(), model.WhereID()))
_, err := genericExec(c, stmt, model.ID())
return err
}
func (p *cockroach) Delete(c *Connection, model *Model, query Query) error {
return genericDelete(c, model, query)
}
func (p *cockroach) SelectOne(c *Connection, model *Model, query Query) error {
return genericSelectOne(c, model, query)
}
func (p *cockroach) SelectMany(c *Connection, models *Model, query Query) error {
return genericSelectMany(c, models, query)
}
func (p *cockroach) CreateDB() error {
// createdb -h db -p 5432 -U cockroach enterprise_development
deets := p.ConnectionDetails
db, err := openPotentiallyInstrumentedConnection(p, p.urlWithoutDb())
if err != nil {
return fmt.Errorf("error creating Cockroach database %s: %w", deets.Database, err)
}
defer db.Close()
query := fmt.Sprintf("CREATE DATABASE %s", p.Quote(deets.Database))
log(logging.SQL, query)
_, err = db.Exec(query)
if err != nil {
return fmt.Errorf("error creating Cockroach database %s: %w", deets.Database, err)
}
log(logging.Info, "created database %s", deets.Database)
return nil
}
func (p *cockroach) DropDB() error {
deets := p.ConnectionDetails
db, err := openPotentiallyInstrumentedConnection(p, p.urlWithoutDb())
if err != nil {
return fmt.Errorf("error dropping Cockroach database %s: %w", deets.Database, err)
}
defer db.Close()
query := fmt.Sprintf("DROP DATABASE %s CASCADE;", p.Quote(deets.Database))
log(logging.SQL, query)
_, err = db.Exec(query)
if err != nil {
return fmt.Errorf("error dropping Cockroach database %s: %w", deets.Database, err)
}
log(logging.Info, "dropped database %s", deets.Database)
return nil
}
func (p *cockroach) URL() string {
c := p.ConnectionDetails
if c.URL != "" {
return c.URL
}
s := "postgres://%s:%s@%s:%s/%s?%s"
return fmt.Sprintf(s, c.User, url.QueryEscape(c.Password), c.Host, c.Port, c.Database, c.OptionsString(""))
}
func (p *cockroach) urlWithoutDb() string {
c := p.ConnectionDetails
s := "postgres://%s:%s@%s:%s/?%s"
return fmt.Sprintf(s, c.User, url.QueryEscape(c.Password), c.Host, c.Port, c.OptionsString(""))
}
func (p *cockroach) MigrationURL() string {
return p.URL()
}
func (p *cockroach) TranslateSQL(sql string) string {
defer p.mu.Unlock()
p.mu.Lock()
if csql, ok := p.translateCache[sql]; ok {
return csql
}
csql := sqlx.Rebind(sqlx.DOLLAR, sql)
p.translateCache[sql] = csql
return csql
}
func (p *cockroach) FizzTranslator() fizz.Translator {
return translators.NewCockroach(p.URL(), p.Details().Database)
}
func (p *cockroach) DumpSchema(w io.Writer) error {
cmd := exec.Command("cockroach", "sql", "-e", "SHOW CREATE ALL TABLES", "-d", p.Details().Database, "--format", "raw")
c := p.ConnectionDetails
if defaults.String(c.option("sslmode"), "disable") == "disable" || strings.Contains(c.RawOptions, "sslmode=disable") {
cmd.Args = append(cmd.Args, "--insecure")
}
return cockroachDumpSchema(p.Details(), cmd, w)
}
func cockroachDumpSchema(deets *ConnectionDetails, cmd *exec.Cmd, w io.Writer) error {
log(logging.SQL, strings.Join(cmd.Args, " "))
var bb bytes.Buffer
cmd.Stdout = &bb
cmd.Stderr = os.Stderr
err := cmd.Run()
if err != nil {
return err
}
// --format raw returns comments prefixed with # which is invalid, so we make it a valid SQL comment.
result := regexp.MustCompile("(?m)^#").ReplaceAll(bb.Bytes(), []byte("-- #"))
if _, err := w.Write(result); err != nil {
return err
}
x := bytes.TrimSpace(result)
if len(x) == 0 {
return fmt.Errorf("unable to dump schema for %s", deets.Database)
}
log(logging.Info, "dumped schema for %s", deets.Database)
return nil
}
func (p *cockroach) LoadSchema(r io.Reader) error {
return genericLoadSchema(p, r)
}
func (p *cockroach) TruncateAll(tx *Connection) error {
type table struct {
TableName string `db:"table_name"`
}
tableQuery := p.tablesQuery()
var tables []table
if err := tx.RawQuery(tableQuery, tx.MigrationTableName(), tx.Dialect.Details().Database).All(&tables); err != nil {
return err
}
if len(tables) == 0 {
return nil
}
tableNames := make([]string, len(tables))
for i, t := range tables {
tableNames[i] = t.TableName
//! work around for current limitation of DDL and DML at the same transaction.
// it should be fixed when cockroach support it or with other approach.
// https://www.cockroachlabs.com/docs/stable/known-limitations.html#schema-changes-within-transactions
if err := tx.RawQuery(fmt.Sprintf("delete from %s", p.Quote(t.TableName))).Exec(); err != nil {
return err
}
}
return nil
// TODO!
// return tx3.RawQuery(fmt.Sprintf("truncate %s cascade;", strings.Join(tableNames, ", "))).Exec()
}
func (p *cockroach) AfterOpen(c *Connection) error {
if err := c.RawQuery(`select version() AS "version"`).First(&p.info); err != nil {
return err
}
if s := strings.Split(p.info.VersionString, " "); len(s) > 3 {
p.info.product = s[0]
p.info.license = s[1]
p.info.version = s[2]
p.info.buildInfo = s[3]
}
log(logging.Debug, "server: %v %v %v", p.info.product, p.info.license, p.info.version)
return nil
}
func newCockroach(deets *ConnectionDetails) (dialect, error) {
deets.Dialect = "postgres"
d := &cockroach{
commonDialect: commonDialect{ConnectionDetails: deets},
translateCache: map[string]string{},
mu: sync.Mutex{},
}
d.info.client = deets.option("application_name")
return d, nil
}
func finalizerCockroach(cd *ConnectionDetails) {
appName := filepath.Base(os.Args[0])
cd.setOptionWithDefault("application_name", cd.option("application_name"), appName)
cd.Port = defaults.String(cd.Port, portCockroach)
if cd.URL != "" {
cd.URL = "postgres://" + trimCockroachPrefix(cd.URL)
}
}
func trimCockroachPrefix(u string) string {
parts := strings.Split(u, "://")
if len(parts) != 2 {
return u
}
return parts[1]
}
func (p *cockroach) tablesQuery() string {
// See https://www.cockroachlabs.com/docs/stable/information-schema.html for more info about information schema changes
tableQuery := selectTablesQueryCockroach
if strings.HasPrefix(p.info.version, "v1.") {
tableQuery = selectTablesQueryCockroachV1
}
return tableQuery
}