Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 1d8a881

Browse files
taywrobeltaywrobel
committed
Add CockroachDB Support
Adds support for CockroachDB. Cockroach uses the postges wire protocol and has a large amount of common SQL functionality shared with Postgres, so much of the postgres code was able to be copied and modified. Since the protocol is used in determining the driver, and the Postgres protocol is also used by Cockroach, new connect string prefixes were added: cockroach:// cockroachdb:// and crdb-postgres://. These fake protocol strings are replaced in the connect function with the correct `postgres://` protocol. TODO: Tests needed (Cockroach has a docker image, so this shouldn't be too hard)
1 parent 8439d71 commit 1d8a881

18 files changed

+356
-1
lines changed

‎Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
SOURCE ?= file go-bindata github aws-s3 google-cloud-storage
2-
DATABASE ?= postgres mysql redshift cassandra sqlite3 spanner
2+
DATABASE ?= postgres mysql redshift cassandra sqlite3 spanner cockroachdb
33
VERSION ?= $(shell git describe --tags 2>/dev/null | cut -c 2-)
44
TEST_FLAGS ?=
55
REPO_OWNER ?= $(shell cd .. && basename "$$(pwd)")

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ Database drivers run migrations. [Add a new database?](database/driver.go)
3232
* [CrateDB](database/crate) ([todo #170](https://github.com/mattes/migrate/issues/170))
3333
* [Shell](database/shell) ([todo #171](https://github.com/mattes/migrate/issues/171))
3434
* [Google Cloud Spanner](database/spanner)
35+
* [CockroachDB](database/cockroachdb)
3536

3637

3738

‎cli/build_cockroachdb.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
// +build cockroachdb
2+
3+
package main
4+
5+
import (
6+
_ "github.com/mattes/migrate/database/cockroachdb"
7+
)

‎database/cockroachdb/cockroachdb.go‎

Lines changed: 319 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,319 @@
1+
package cockroachdb
2+
3+
import (
4+
"database/sql"
5+
"fmt"
6+
"io"
7+
"io/ioutil"
8+
nurl "net/url"
9+
10+
"github.com/cockroachdb/cockroach-go/crdb"
11+
"github.com/lib/pq"
12+
"github.com/mattes/migrate"
13+
"github.com/mattes/migrate/database"
14+
"regexp"
15+
)
16+
17+
func init() {
18+
db := CockroachDb{}
19+
database.Register("cockroach", &db)
20+
database.Register("cockroachdb", &db)
21+
database.Register("crdb-postgres", &db)
22+
}
23+
24+
var DefaultMigrationsTable = "schema_migrations"
25+
var DefaultLockTable = "schema_lock"
26+
27+
var (
28+
ErrNilConfig = fmt.Errorf("no config")
29+
ErrNoDatabaseName = fmt.Errorf("no database name")
30+
)
31+
32+
type Config struct {
33+
MigrationsTable string
34+
LockTable string
35+
DatabaseName string
36+
}
37+
38+
type CockroachDb struct {
39+
db *sql.DB
40+
isLocked bool
41+
42+
// Open and WithInstance need to guarantee that config is never nil
43+
config *Config
44+
}
45+
46+
func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
47+
if config == nil {
48+
return nil, ErrNilConfig
49+
}
50+
51+
if err := instance.Ping(); err != nil {
52+
return nil, err
53+
}
54+
55+
query := `SELECT current_database()`
56+
var databaseName string
57+
if err := instance.QueryRow(query).Scan(&databaseName); err != nil {
58+
return nil, &database.Error{OrigErr: err, Query: []byte(query)}
59+
}
60+
61+
if len(databaseName) == 0 {
62+
return nil, ErrNoDatabaseName
63+
}
64+
65+
config.DatabaseName = databaseName
66+
67+
if len(config.MigrationsTable) == 0 {
68+
config.MigrationsTable = DefaultMigrationsTable
69+
}
70+
71+
if len(config.LockTable) == 0 {
72+
config.LockTable = DefaultLockTable
73+
}
74+
75+
px := &CockroachDb{
76+
db: instance,
77+
config: config,
78+
}
79+
80+
if err := px.ensureVersionTable(); err != nil {
81+
return nil, err
82+
}
83+
84+
if err := px.ensureLockTable(); err != nil {
85+
return nil, err
86+
}
87+
88+
return px, nil
89+
}
90+
91+
func (c *CockroachDb) Open(url string) (database.Driver, error) {
92+
purl, err := nurl.Parse(url)
93+
if err != nil {
94+
return nil, err
95+
}
96+
97+
// As Cockroach uses the postgres protocol, and 'postgres' is already a registered database, we need to replace the
98+
// connect prefix, with the actual protocol, so that the library can differentiate between the implementations
99+
re := regexp.MustCompile("^(cockroach(db)?|crdb-postgres)")
100+
connectString := re.ReplaceAllString(migrate.FilterCustomQuery(purl).String(), "postgres")
101+
102+
db, err := sql.Open("postgres", connectString)
103+
if err != nil {
104+
return nil, err
105+
}
106+
107+
migrationsTable := purl.Query().Get("x-migrations-table")
108+
if len(migrationsTable) == 0 {
109+
migrationsTable = DefaultMigrationsTable
110+
}
111+
112+
lockTable := purl.Query().Get("x-lock-table")
113+
if len(lockTable) == 0 {
114+
lockTable = DefaultLockTable
115+
}
116+
117+
px, err := WithInstance(db, &Config{
118+
DatabaseName: purl.Path,
119+
MigrationsTable: migrationsTable,
120+
LockTable: lockTable,
121+
})
122+
if err != nil {
123+
return nil, err
124+
}
125+
126+
return px, nil
127+
}
128+
129+
func (c *CockroachDb) Close() error {
130+
return c.db.Close()
131+
}
132+
133+
// Locking is done manually with a separate lock table. Implementing advisory locks in CRDB is being discussed
134+
// See: https://github.com/cockroachdb/cockroach/issues/13546
135+
func (c *CockroachDb) Lock() error {
136+
err := crdb.ExecuteTx(c.db, func(tx *sql.Tx) error {
137+
aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
138+
if err != nil {
139+
return err
140+
}
141+
142+
query := "SELECT * FROM " + c.config.LockTable + " WHERE lock_id = $1"
143+
rows, err := tx.Query(query, aid)
144+
if err != nil {
145+
return database.Error{OrigErr: err, Err: "failed to fetch migration lock", Query: []byte(query)}
146+
}
147+
defer rows.Close()
148+
149+
// If row exists at all, lock is present
150+
locked := rows.Next()
151+
if locked {
152+
return database.Error{Err: "lock could not be acquired; already locked", Query: []byte(query)}
153+
}
154+
155+
query = "INSERT INTO " + c.config.LockTable + " (lock_id) VALUES ($1)"
156+
if _, err := tx.Exec(query, aid) ; err != nil {
157+
return database.Error{OrigErr: err, Err: "failed to set migration lock", Query: []byte(query)}
158+
}
159+
160+
return nil
161+
})
162+
163+
if err != nil {
164+
return err
165+
} else {
166+
c.isLocked = true
167+
return nil
168+
}
169+
}
170+
171+
// Locking is done manually with a separate lock table. Implementing advisory locks in CRDB is being discussed
172+
// See: https://github.com/cockroachdb/cockroach/issues/13546
173+
func (c *CockroachDb) Unlock() error {
174+
aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
175+
if err != nil {
176+
return err
177+
}
178+
179+
// In the event of an implementation (non-migration) error, it is possible for the lock to not be released. Until
180+
// a better locking mechanism is added, a manual purging of the lock table may be required in such circumstances
181+
query := "DELETE FROM " + c.config.LockTable + " WHERE lock_id = $1"
182+
if _, err := c.db.Exec(query, aid); err != nil {
183+
return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(query)}
184+
}
185+
186+
c.isLocked = false
187+
return nil
188+
}
189+
190+
func (c *CockroachDb) Run(migration io.Reader) error {
191+
migr, err := ioutil.ReadAll(migration)
192+
if err != nil {
193+
return err
194+
}
195+
196+
// run migration
197+
query := string(migr[:])
198+
if _, err := c.db.Exec(query); err != nil {
199+
return database.Error{OrigErr: err, Err: "migration failed", Query: migr}
200+
}
201+
202+
return nil
203+
}
204+
205+
func (c *CockroachDb) SetVersion(version int, dirty bool) error {
206+
return crdb.ExecuteTx(c.db, func(tx *sql.Tx) error {
207+
if _, err := tx.Exec( `TRUNCATE "` + c.config.MigrationsTable + `"`); err != nil {
208+
return err
209+
}
210+
211+
if version >= 0 {
212+
if _, err := tx.Exec(`INSERT INTO "` + c.config.MigrationsTable + `" (version, dirty) VALUES ($1, $2)`, version, dirty); err != nil {
213+
return err
214+
}
215+
}
216+
217+
return nil
218+
})
219+
}
220+
221+
func (c *CockroachDb) Version() (version int, dirty bool, err error) {
222+
query := `SELECT version, dirty FROM "` + c.config.MigrationsTable + `" LIMIT 1`
223+
err = c.db.QueryRow(query).Scan(&version, &dirty)
224+
225+
switch {
226+
case err == sql.ErrNoRows:
227+
return database.NilVersion, false, nil
228+
229+
case err != nil:
230+
if e, ok := err.(*pq.Error); ok {
231+
// 42P01 is "UndefinedTableError" in CockroachDB
232+
// https://github.com/cockroachdb/cockroach/blob/master/pkg/sql/pgwire/pgerror/codes.go
233+
if e.Code == "42P01" {
234+
return database.NilVersion, false, nil
235+
}
236+
}
237+
return 0, false, &database.Error{OrigErr: err, Query: []byte(query)}
238+
239+
default:
240+
return version, dirty, nil
241+
}
242+
}
243+
244+
func (c *CockroachDb) Drop() error {
245+
// select all tables in current schema
246+
query := `SELECT table_name FROM information_schema.tables WHERE table_schema=(SELECT current_schema())`
247+
tables, err := c.db.Query(query)
248+
if err != nil {
249+
return &database.Error{OrigErr: err, Query: []byte(query)}
250+
}
251+
defer tables.Close()
252+
253+
// delete one table after another
254+
tableNames := make([]string, 0)
255+
for tables.Next() {
256+
var tableName string
257+
if err := tables.Scan(&tableName); err != nil {
258+
return err
259+
}
260+
if len(tableName) > 0 {
261+
tableNames = append(tableNames, tableName)
262+
}
263+
}
264+
265+
if len(tableNames) > 0 {
266+
// delete one by one ...
267+
for _, t := range tableNames {
268+
query = `DROP TABLE IF EXISTS ` + t + ` CASCADE`
269+
if _, err := c.db.Exec(query); err != nil {
270+
return &database.Error{OrigErr: err, Query: []byte(query)}
271+
}
272+
}
273+
if err := c.ensureVersionTable(); err != nil {
274+
return err
275+
}
276+
}
277+
278+
return nil
279+
}
280+
281+
func (c *CockroachDb) ensureVersionTable() error {
282+
// check if migration table exists
283+
var count int
284+
query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1`
285+
if err := c.db.QueryRow(query, c.config.MigrationsTable).Scan(&count); err != nil {
286+
return &database.Error{OrigErr: err, Query: []byte(query)}
287+
}
288+
if count == 1 {
289+
return nil
290+
}
291+
292+
// if not, create the empty migration table
293+
query = `CREATE TABLE "` + c.config.MigrationsTable + `" (version INT NOT NULL PRIMARY KEY, dirty BOOL NOT NULL)`
294+
if _, err := c.db.Exec(query); err != nil {
295+
return &database.Error{OrigErr: err, Query: []byte(query)}
296+
}
297+
return nil
298+
}
299+
300+
301+
func (c *CockroachDb) ensureLockTable() error {
302+
// check if lock table exists
303+
var count int
304+
query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1`
305+
if err := c.db.QueryRow(query, c.config.LockTable).Scan(&count); err != nil {
306+
return &database.Error{OrigErr: err, Query: []byte(query)}
307+
}
308+
if count == 1 {
309+
return nil
310+
}
311+
312+
// if not, create the empty lock table
313+
query = `CREATE TABLE "` + c.config.LockTable + `" (lock_id INT NOT NULL PRIMARY KEY)`
314+
if _, err := c.db.Exec(query); err != nil {
315+
return &database.Error{OrigErr: err, Query: []byte(query)}
316+
}
317+
318+
return nil
319+
}

‎database/cockroachdb/examples/migrations/1085649617_create_users_table.down.sql‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
DROP TABLE IF EXISTS users;

‎database/cockroachdb/examples/migrations/1085649617_create_users_table.up.sql‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
CREATE TABLE users (
2+
user_id INT UNIQUE,
3+
name STRING(40),
4+
email STRING(40)
5+
);

‎database/cockroachdb/examples/migrations/1185749658_add_city_to_users.down.sql‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
ALTER TABLE users DROP COLUMN IF EXISTS city;

‎database/cockroachdb/examples/migrations/1185749658_add_city_to_users.up.sql‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
ALTER TABLE users ADD COLUMN city TEXT;

‎database/cockroachdb/examples/migrations/1285849751_add_index_on_user_emails.down.sql‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
DROP INDEX IF EXISTS users_email_index;

‎database/cockroachdb/examples/migrations/1285849751_add_index_on_user_emails.up.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
CREATE UNIQUE INDEX IF NOT EXISTS users_email_index ON users (email);
2+
3+
-- Lorem ipsum dolor sit amet, consectetur adipiscing elit. Aenean sed interdum velit, tristique iaculis justo. Pellentesque ut porttitor dolor. Donec sit amet pharetra elit. Cras vel ligula ex. Phasellus posuere.

0 commit comments

Comments
 (0)