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