diff --git a/Makefile b/Makefile index 79fa85ffb..137c08a79 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ SOURCE ?= file go_bindata github github_ee bitbucket aws_s3 google_cloud_storage godoc_vfs gitlab -DATABASE ?= postgres mysql redshift cassandra spanner cockroachdb clickhouse mongodb sqlserver firebird neo4j pgx +DATABASE ?= postgres mysql redshift cassandra spanner cockroachdb yugabytedb clickhouse mongodb sqlserver firebird neo4j pgx DATABASE_TEST ?= $(DATABASE) sqlite sqlite3 sqlcipher VERSION ?= $(shell git describe --tags 2>/dev/null | cut -c 2-) TEST_FLAGS ?= diff --git a/README.md b/README.md index be2a83116..7309cb17c 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,7 @@ Database drivers run migrations. [Add a new database?](database/driver.go) * [Shell](database/shell) ([todo #171](https://github.com/mattes/migrate/issues/171)) * [Google Cloud Spanner](database/spanner) * [CockroachDB](database/cockroachdb) +* [YugabyteDB](database/yugabytedb) * [ClickHouse](database/clickhouse) * [Firebird](database/firebird) * [MS SQL Server](database/sqlserver) diff --git a/database/yugabytedb/README.md b/database/yugabytedb/README.md new file mode 100644 index 000000000..4157cbc4e --- /dev/null +++ b/database/yugabytedb/README.md @@ -0,0 +1,19 @@ +# yugabytedb + +`yugabytedb://user:password@host:port/dbname?query` (`yugabyte://`, and `ysql://` work, too) + +| URL Query | WithInstance Config | Description | +|------------|---------------------|-------------| +| `x-migrations-table` | `MigrationsTable` | Name of the migrations table | +| `x-lock-table` | `LockTable` | Name of the table which maintains the migration lock | +| `x-force-lock` | `ForceLock` | Force lock acquisition to fix faulty migrations which may not have released the schema lock (Boolean, default is `false`) | +| `dbname` | `DatabaseName` | The name of the database to connect to | +| `user` | | The user to sign in as | +| `password` | | The user's password | +| `host` | | The host to connect to. Values that start with / are for unix domain sockets. (default is localhost) | +| `port` | | The port to bind to. (default is 5432) | +| `connect_timeout` | | Maximum wait for connection, in seconds. Zero or not specified means wait indefinitely. | +| `sslcert` | | Cert file location. The file must contain PEM encoded data. | +| `sslkey` | | Key file location. The file must contain PEM encoded data. | +| `sslrootcert` | | The location of the root certificate file. The file must contain PEM encoded data. | +| `sslmode` | | Whether or not to use SSL (disable\|require\|verify-ca\|verify-full) | diff --git a/database/yugabytedb/examples/migrations/1085649617_create_users_table.down.sql b/database/yugabytedb/examples/migrations/1085649617_create_users_table.down.sql new file mode 100644 index 000000000..c99ddcdc8 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1085649617_create_users_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS users; diff --git a/database/yugabytedb/examples/migrations/1085649617_create_users_table.up.sql b/database/yugabytedb/examples/migrations/1085649617_create_users_table.up.sql new file mode 100644 index 000000000..92897dcab --- /dev/null +++ b/database/yugabytedb/examples/migrations/1085649617_create_users_table.up.sql @@ -0,0 +1,5 @@ +CREATE TABLE users ( + user_id integer unique, + name varchar(40), + email varchar(40) +); diff --git a/database/yugabytedb/examples/migrations/1185749658_add_city_to_users.down.sql b/database/yugabytedb/examples/migrations/1185749658_add_city_to_users.down.sql new file mode 100644 index 000000000..940c60712 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1185749658_add_city_to_users.down.sql @@ -0,0 +1 @@ +ALTER TABLE users DROP COLUMN IF EXISTS city; diff --git a/database/yugabytedb/examples/migrations/1185749658_add_city_to_users.up.sql b/database/yugabytedb/examples/migrations/1185749658_add_city_to_users.up.sql new file mode 100644 index 000000000..67823edc9 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1185749658_add_city_to_users.up.sql @@ -0,0 +1,3 @@ +ALTER TABLE users ADD COLUMN city varchar(100); + + diff --git a/database/yugabytedb/examples/migrations/1285849751_add_index_on_user_emails.down.sql b/database/yugabytedb/examples/migrations/1285849751_add_index_on_user_emails.down.sql new file mode 100644 index 000000000..3e87dd229 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1285849751_add_index_on_user_emails.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS users_email_index; diff --git a/database/yugabytedb/examples/migrations/1285849751_add_index_on_user_emails.up.sql b/database/yugabytedb/examples/migrations/1285849751_add_index_on_user_emails.up.sql new file mode 100644 index 000000000..03a04639c --- /dev/null +++ b/database/yugabytedb/examples/migrations/1285849751_add_index_on_user_emails.up.sql @@ -0,0 +1,3 @@ +CREATE UNIQUE INDEX users_email_index ON users (email); + +-- 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. diff --git a/database/yugabytedb/examples/migrations/1385949617_create_books_table.down.sql b/database/yugabytedb/examples/migrations/1385949617_create_books_table.down.sql new file mode 100644 index 000000000..1a0b1a214 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1385949617_create_books_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS books; diff --git a/database/yugabytedb/examples/migrations/1385949617_create_books_table.up.sql b/database/yugabytedb/examples/migrations/1385949617_create_books_table.up.sql new file mode 100644 index 000000000..f1503b518 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1385949617_create_books_table.up.sql @@ -0,0 +1,5 @@ +CREATE TABLE books ( + user_id integer, + name varchar(40), + author varchar(40) +); diff --git a/database/yugabytedb/examples/migrations/1485949617_create_movies_table.down.sql b/database/yugabytedb/examples/migrations/1485949617_create_movies_table.down.sql new file mode 100644 index 000000000..3a5187689 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1485949617_create_movies_table.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS movies; diff --git a/database/yugabytedb/examples/migrations/1485949617_create_movies_table.up.sql b/database/yugabytedb/examples/migrations/1485949617_create_movies_table.up.sql new file mode 100644 index 000000000..f0ef5943b --- /dev/null +++ b/database/yugabytedb/examples/migrations/1485949617_create_movies_table.up.sql @@ -0,0 +1,5 @@ +CREATE TABLE movies ( + user_id integer, + name varchar(40), + director varchar(40) +); diff --git a/database/yugabytedb/examples/migrations/1585849751_just_a_comment.up.sql b/database/yugabytedb/examples/migrations/1585849751_just_a_comment.up.sql new file mode 100644 index 000000000..9b6b57a61 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1585849751_just_a_comment.up.sql @@ -0,0 +1 @@ +-- 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. diff --git a/database/yugabytedb/examples/migrations/1685849751_another_comment.up.sql b/database/yugabytedb/examples/migrations/1685849751_another_comment.up.sql new file mode 100644 index 000000000..9b6b57a61 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1685849751_another_comment.up.sql @@ -0,0 +1 @@ +-- 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. diff --git a/database/yugabytedb/examples/migrations/1785849751_another_comment.up.sql b/database/yugabytedb/examples/migrations/1785849751_another_comment.up.sql new file mode 100644 index 000000000..9b6b57a61 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1785849751_another_comment.up.sql @@ -0,0 +1 @@ +-- 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. diff --git a/database/yugabytedb/examples/migrations/1885849751_another_comment.up.sql b/database/yugabytedb/examples/migrations/1885849751_another_comment.up.sql new file mode 100644 index 000000000..9b6b57a61 --- /dev/null +++ b/database/yugabytedb/examples/migrations/1885849751_another_comment.up.sql @@ -0,0 +1 @@ +-- 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. diff --git a/database/yugabytedb/yugabytedb.go b/database/yugabytedb/yugabytedb.go new file mode 100644 index 000000000..a947f5044 --- /dev/null +++ b/database/yugabytedb/yugabytedb.go @@ -0,0 +1,427 @@ +package yugabytedb + +import ( + "context" + "database/sql" + "errors" + "fmt" + "io" + "io/ioutil" + "net/url" + "regexp" + "strconv" + + "github.com/golang-migrate/migrate/v4" + "github.com/golang-migrate/migrate/v4/database" + "github.com/hashicorp/go-multierror" + "github.com/lib/pq" + "go.uber.org/atomic" +) + +func init() { + db := YugabyteDB{} + database.Register("yugabyte", &db) + database.Register("yugabytedb", &db) + database.Register("ysql", &db) +} + +var DefaultMigrationsTable = "migrations" +var DefaultLockTable = "migrations_locks" + +var ( + ErrNilConfig = errors.New("no config") + ErrNoDatabaseName = errors.New("no database name") + ErrMaxRetriesExceeded = errors.New("max retries exceeded") +) + +type Config struct { + MigrationsTable string + LockTable string + ForceLock bool + DatabaseName string +} + +type YugabyteDB struct { + db *sql.DB + isLocked atomic.Bool + + // Open and WithInstance need to guarantee that config is never nil + config *Config +} + +func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) { + if config == nil { + return nil, ErrNilConfig + } + + if err := instance.Ping(); err != nil { + return nil, err + } + + if config.DatabaseName == "" { + query := `SELECT current_database()` + var databaseName string + if err := instance.QueryRow(query).Scan(&databaseName); err != nil { + return nil, &database.Error{OrigErr: err, Query: []byte(query)} + } + + if len(databaseName) == 0 { + return nil, ErrNoDatabaseName + } + + config.DatabaseName = databaseName + } + + if len(config.MigrationsTable) == 0 { + config.MigrationsTable = DefaultMigrationsTable + } + + if len(config.LockTable) == 0 { + config.LockTable = DefaultLockTable + } + + px := &YugabyteDB{ + db: instance, + config: config, + } + + // ensureVersionTable is a locking operation, so we need to ensureLockTable before we ensureVersionTable. + if err := px.ensureLockTable(); err != nil { + return nil, err + } + + if err := px.ensureVersionTable(); err != nil { + return nil, err + } + + return px, nil +} + +func (c *YugabyteDB) Open(dbURL string) (database.Driver, error) { + purl, err := url.Parse(dbURL) + if err != nil { + return nil, err + } + + // As YugabyteDB uses the postgres protocol, and 'postgres' is already a registered database, we need to replace the + // connect prefix, with the actual protocol, so that the library can differentiate between the implementations + re := regexp.MustCompile("^(yugabyte(db)?|ysql)") + connectString := re.ReplaceAllString(migrate.FilterCustomQuery(purl).String(), "postgres") + + db, err := sql.Open("postgres", connectString) + if err != nil { + return nil, err + } + + migrationsTable := purl.Query().Get("x-migrations-table") + if len(migrationsTable) == 0 { + migrationsTable = DefaultMigrationsTable + } + + lockTable := purl.Query().Get("x-lock-table") + if len(lockTable) == 0 { + lockTable = DefaultLockTable + } + + forceLockQuery := purl.Query().Get("x-force-lock") + forceLock, err := strconv.ParseBool(forceLockQuery) + if err != nil { + forceLock = false + } + + px, err := WithInstance(db, &Config{ + DatabaseName: purl.Path, + MigrationsTable: migrationsTable, + LockTable: lockTable, + ForceLock: forceLock, + }) + if err != nil { + return nil, err + } + + return px, nil +} + +func (c *YugabyteDB) Close() error { + return c.db.Close() +} + +// Locking is done manually with a separate lock table. Implementing advisory locks in YugabyteDB is being discussed +// See: https://github.com/yugabyte/yugabyte-db/issues/3642 +func (c *YugabyteDB) Lock() error { + return database.CasRestoreOnErr(&c.isLocked, false, true, database.ErrLocked, func() (err error) { + return ExecuteInTx(context.Background(), c.db, &sql.TxOptions{Isolation: sql.LevelSerializable}, func(tx *sql.Tx) (err error) { + aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName) + if err != nil { + return err + } + + query := "SELECT * FROM " + c.config.LockTable + " WHERE lock_id = $1" + rows, err := tx.Query(query, aid) + if err != nil { + return database.Error{OrigErr: err, Err: "failed to fetch migration lock", Query: []byte(query)} + } + defer func() { + if errClose := rows.Close(); errClose != nil { + err = multierror.Append(err, errClose) + } + }() + + // If row exists at all, lock is present + locked := rows.Next() + if locked && !c.config.ForceLock { + return database.ErrLocked + } + + query = "INSERT INTO " + c.config.LockTable + " (lock_id) VALUES ($1)" + if _, err := tx.Exec(query, aid); err != nil { + return database.Error{OrigErr: err, Err: "failed to set migration lock", Query: []byte(query)} + } + + return nil + }) + }) +} + +// Locking is done manually with a separate lock table. Implementing advisory locks in YugabyteDB is being discussed +// See: https://github.com/yugabyte/yugabyte-db/issues/3642 +func (c *YugabyteDB) Unlock() error { + return database.CasRestoreOnErr(&c.isLocked, true, false, database.ErrNotLocked, func() (err error) { + aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName) + if err != nil { + return err + } + + // In the event of an implementation (non-migration) error, it is possible for the lock to not be released. Until + // a better locking mechanism is added, a manual purging of the lock table may be required in such circumstances + query := "DELETE FROM " + c.config.LockTable + " WHERE lock_id = $1" + if _, err := c.db.Exec(query, aid); err != nil { + if e, ok := err.(*pq.Error); ok { + // 42P01 is "UndefinedTableError" in YugabyteDB + // https://github.com/yugabyte/yugabyte-db/blob/9c6b8e6beb56eed8eeb357178c0c6b837eb49896/src/postgres/src/backend/utils/errcodes.txt#L366 + if e.Code == "42P01" { + // On drops, the lock table is fully removed; This is fine, and is a valid "unlocked" state for the schema + return nil + } + } + + return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(query)} + } + + return nil + }) +} + +func (c *YugabyteDB) Run(migration io.Reader) error { + migr, err := ioutil.ReadAll(migration) + if err != nil { + return err + } + + // run migration + query := string(migr[:]) + if _, err := c.db.Exec(query); err != nil { + return database.Error{OrigErr: err, Err: "migration failed", Query: migr} + } + + return nil +} + +func (c *YugabyteDB) SetVersion(version int, dirty bool) error { + return ExecuteInTx(context.Background(), c.db, &sql.TxOptions{Isolation: sql.LevelSerializable}, func(tx *sql.Tx) error { + if _, err := tx.Exec(`DELETE FROM "` + c.config.MigrationsTable + `"`); err != nil { + return err + } + + // Also re-write the schema version for nil dirty versions to prevent + // empty schema version for failed down migration on the first migration + // See: https://github.com/golang-migrate/migrate/issues/330 + if version >= 0 || (version == database.NilVersion && dirty) { + if _, err := tx.Exec(`INSERT INTO "`+c.config.MigrationsTable+`" (version, dirty) VALUES ($1, $2)`, version, dirty); err != nil { + return err + } + } + + return nil + }) +} + +func (c *YugabyteDB) Version() (version int, dirty bool, err error) { + query := `SELECT version, dirty FROM "` + c.config.MigrationsTable + `" LIMIT 1` + err = c.db.QueryRow(query).Scan(&version, &dirty) + + switch { + case err == sql.ErrNoRows: + return database.NilVersion, false, nil + + case err != nil: + if e, ok := err.(*pq.Error); ok { + // 42P01 is "UndefinedTableError" in YugabyteDB + // https://github.com/yugabyte/yugabyte-db/blob/9c6b8e6beb56eed8eeb357178c0c6b837eb49896/src/postgres/src/backend/utils/errcodes.txt#L366 + if e.Code == "42P01" { + return database.NilVersion, false, nil + } + } + return 0, false, &database.Error{OrigErr: err, Query: []byte(query)} + + default: + return version, dirty, nil + } +} + +func (c *YugabyteDB) Drop() (err error) { + query := `SELECT table_name FROM information_schema.tables WHERE table_schema=(SELECT current_schema()) AND table_type='BASE TABLE'` + tables, err := c.db.Query(query) + if err != nil { + return &database.Error{OrigErr: err, Query: []byte(query)} + } + defer func() { + if errClose := tables.Close(); errClose != nil { + err = multierror.Append(err, errClose) + } + }() + + // delete one table after another + tableNames := make([]string, 0) + for tables.Next() { + var tableName string + if err := tables.Scan(&tableName); err != nil { + return err + } + if len(tableName) > 0 { + tableNames = append(tableNames, tableName) + } + } + if err := tables.Err(); err != nil { + return &database.Error{OrigErr: err, Query: []byte(query)} + } + + if len(tableNames) > 0 { + for _, t := range tableNames { + query = `DROP TABLE IF EXISTS ` + t + ` CASCADE` + if _, err := c.db.Exec(query); err != nil { + return &database.Error{OrigErr: err, Query: []byte(query)} + } + } + } + + return nil +} + +// ensureVersionTable checks if versions table exists and, if not, creates it. +// Note that this function locks the database +func (c *YugabyteDB) ensureVersionTable() (err error) { + if err = c.Lock(); err != nil { + return err + } + + defer func() { + if e := c.Unlock(); e != nil { + if err == nil { + err = e + } else { + err = multierror.Append(err, e) + } + } + }() + + // check if migration table exists + var count int + query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1` + if err := c.db.QueryRow(query, c.config.MigrationsTable).Scan(&count); err != nil { + return &database.Error{OrigErr: err, Query: []byte(query)} + } + if count == 1 { + return nil + } + + // if not, create the empty migration table + query = `CREATE TABLE "` + c.config.MigrationsTable + `" (version INT NOT NULL PRIMARY KEY, dirty BOOL NOT NULL)` + if _, err := c.db.Exec(query); err != nil { + return &database.Error{OrigErr: err, Query: []byte(query)} + } + return nil +} + +func (c *YugabyteDB) ensureLockTable() error { + // check if lock table exists + var count int + query := `SELECT COUNT(1) FROM information_schema.tables WHERE table_name = $1 AND table_schema = (SELECT current_schema()) LIMIT 1` + if err := c.db.QueryRow(query, c.config.LockTable).Scan(&count); err != nil { + return &database.Error{OrigErr: err, Query: []byte(query)} + } + if count == 1 { + return nil + } + + // if not, create the empty lock table + query = `CREATE TABLE "` + c.config.LockTable + `" (lock_id TEXT NOT NULL PRIMARY KEY)` + if _, err := c.db.Exec(query); err != nil { + return &database.Error{OrigErr: err, Query: []byte(query)} + } + + return nil +} + +func ExecuteInTx(ctx context.Context, db *sql.DB, opts *sql.TxOptions, fn func(tx *sql.Tx) error) (err error) { + // It looks like currently, YugabyteDB doesn't return an error on savepoint release. + // So, we have to retry the whole transaction. + // TODO: do more research on savepoints support, see https://github.com/yugabyte/yugabyte-db/issues/9219 + + const maxRetries = 50 + retryCount := 0 + + for { + if retryCount-1 >= maxRetries { + return fmt.Errorf("%w: reached %d attempts", ErrMaxRetriesExceeded, retryCount-1) + } + + tx, err := db.BeginTx(ctx, opts) + if err != nil { + return err + } + + if err := fn(tx); err != nil { + _ = tx.Rollback() + + if !errIsRetryable(err) { + return err + } + + retryCount++ + + continue + } + + if err := tx.Commit(); err != nil { + _ = tx.Rollback() + + if !errIsRetryable(err) { + return err + } + + retryCount++ + + continue + } + + return nil + } +} + +func errIsRetryable(err error) bool { + pqErr := pq.Error{} + if !errors.As(err, &pqErr) { + return false + } + + code := string(pqErr.Code) + + // Assume that it's safe to retry 08006 and XX000 because we check for lock existence + // before creating and lock ID is primary key. Version field in migrations table is primary key too + // and delete all versions is an idempotend operation. + return code == "40001" || // Serialization error (optimistic locking conflict) + code == "40P01" || // Deadlock + code == "08006" || // Connection failure (node down, need to reconnect) + code == "XX000" // Internal error (may happen during HA) +} diff --git a/database/yugabytedb/yugabytedb_test.go b/database/yugabytedb/yugabytedb_test.go new file mode 100644 index 000000000..0c53ec13e --- /dev/null +++ b/database/yugabytedb/yugabytedb_test.go @@ -0,0 +1,185 @@ +package yugabytedb + +// error codes https://github.com/lib/pq/blob/master/error.go + +import ( + "context" + "database/sql" + "fmt" + "log" + "strings" + "testing" + "time" + + "github.com/dhui/dktest" + "github.com/golang-migrate/migrate/v4" + + _ "github.com/lib/pq" + + dt "github.com/golang-migrate/migrate/v4/database/testing" + "github.com/golang-migrate/migrate/v4/dktesting" + + _ "github.com/golang-migrate/migrate/v4/source/file" +) + +const defaultPort = 5433 + +var ( + opts = dktest.Options{ + Cmd: []string{"bin/yugabyted", "start", "--daemon=false"}, + PortRequired: true, + ReadyFunc: isReady, + Timeout: time.Duration(60) * time.Second, + } + // Released versions: https://docs.yugabyte.com/latest/releases/#current-supported-releases + specs = []dktesting.ContainerSpec{ + {ImageName: "yugabytedb/yugabyte:2.6.16.0-b14", Options: opts}, + {ImageName: "yugabytedb/yugabyte:2.8.4.0-b30", Options: opts}, + {ImageName: "yugabytedb/yugabyte:2.12.2.0-b58", Options: opts}, + {ImageName: "yugabytedb/yugabyte:2.13.0.1-b2", Options: opts}, + } +) + +func isReady(ctx context.Context, c dktest.ContainerInfo) bool { + ip, port, err := c.Port(defaultPort) + if err != nil { + log.Println("port error:", err) + return false + } + + db, err := sql.Open("postgres", fmt.Sprintf("postgres://yugabyte:yugabyte@%v:%v?sslmode=disable", ip, port)) + if err != nil { + log.Println("open error:", err) + return false + } + if err := db.PingContext(ctx); err != nil { + log.Println("ping error:", err) + return false + } + if err := db.Close(); err != nil { + log.Println("close error:", err) + } + return true +} + +func createDB(t *testing.T, c dktest.ContainerInfo) { + ip, port, err := c.Port(defaultPort) + if err != nil { + t.Fatal(err) + } + + db, err := sql.Open("postgres", fmt.Sprintf("postgres://yugabyte:yugabyte@%v:%v?sslmode=disable", ip, port)) + if err != nil { + t.Fatal(err) + } + if err = db.Ping(); err != nil { + t.Fatal(err) + } + defer func() { + if err := db.Close(); err != nil { + t.Error(err) + } + }() + + if _, err = db.Exec("CREATE DATABASE migrate"); err != nil { + t.Fatal(err) + } +} + +func getConnectionString(ip, port string, options ...string) string { + options = append(options, "sslmode=disable") + + return fmt.Sprintf("yugabyte://yugabyte:yugabyte@%v:%v/migrate?%s", ip, port, strings.Join(options, "&")) +} + +func Test(t *testing.T) { + dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) { + createDB(t, ci) + + ip, port, err := ci.Port(defaultPort) + if err != nil { + t.Fatal(err) + } + + addr := getConnectionString(ip, port) + c := &YugabyteDB{} + d, err := c.Open(addr) + if err != nil { + t.Fatal(err) + } + dt.Test(t, d, []byte("SELECT 1")) + }) +} + +func TestMigrate(t *testing.T) { + dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) { + createDB(t, ci) + + ip, port, err := ci.Port(defaultPort) + if err != nil { + t.Fatal(err) + } + + addr := getConnectionString(ip, port) + c := &YugabyteDB{} + d, err := c.Open(addr) + if err != nil { + t.Fatal(err) + } + + m, err := migrate.NewWithDatabaseInstance("file://./examples/migrations", "migrate", d) + if err != nil { + t.Fatal(err) + } + dt.TestMigrate(t, m) + }) +} + +func TestMultiStatement(t *testing.T) { + dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) { + createDB(t, ci) + + ip, port, err := ci.Port(defaultPort) + if err != nil { + t.Fatal(err) + } + + addr := getConnectionString(ip, port) + c := &YugabyteDB{} + d, err := c.Open(addr) + if err != nil { + t.Fatal(err) + } + if err := d.Run(strings.NewReader("CREATE TABLE foo (foo text); CREATE TABLE bar (bar text);")); err != nil { + t.Fatalf("expected err to be nil, got %v", err) + } + + // make sure second table exists + var exists bool + if err := d.(*YugabyteDB).db.QueryRow("SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'bar' AND table_schema = (SELECT current_schema()))").Scan(&exists); err != nil { + t.Fatal(err) + } + if !exists { + t.Fatal("expected table bar to exist") + } + }) +} + +func TestFilterCustomQuery(t *testing.T) { + dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) { + createDB(t, ci) + + ip, port, err := ci.Port(defaultPort) + if err != nil { + t.Fatal(err) + } + + addr := getConnectionString(ip, port, "x-custom=foobar") + c := &YugabyteDB{} + d, err := c.Open(addr) + if err != nil { + t.Fatal(err) + } + dt.Test(t, d, []byte("SELECT 1")) + }) +} diff --git a/internal/cli/build_yugabytedb.go b/internal/cli/build_yugabytedb.go new file mode 100644 index 000000000..9cda6067f --- /dev/null +++ b/internal/cli/build_yugabytedb.go @@ -0,0 +1,8 @@ +//go:build yugabytedb +// +build yugabytedb + +package cli + +import ( + _ "github.com/golang-migrate/migrate/v4/database/yugabytedb" +)