Add wait command (#35)

README.md

@@ -90,6 +90,7 @@ dbmate migrate   # run any pending migrations
 dbmate rollback  # roll back the most recent migration
 dbmate down      # alias for rollback
 dbmate dump      # write the database schema.sql file
+dbmate wait      # wait for the database server to become available
 ```
 
 ## Usage
@@ -222,6 +223,43 @@ On Ubuntu or Debian systems, you can fix this by installing `postgresql-client`,
 
 > Note: The `schema.sql` file will contain a complete schema for your database, even if some tables or columns were created outside of dbmate migrations.
 
+### Waiting For The Database
+
+If you use a Docker development environment for your project, you may encounter issues with the database not being immediately ready when running migrations or unit tests. This can be due to the database server having only just started.
+
+In general, your application should be resilient to not having a working database connection on startup. However, for the purpose of running migrations or unit tests, this is not practical. The `wait` command avoids this situation by allowing you to pause a script or other application until the database is available. Dbmate will attempt a connection to the database server every second, up to a maximum of 60 seconds.
+
+If the database is available, `wait` will return no output:
+
+```sh
+$ dbmate wait
+```
+
+If the database is unavailable, `wait` will block until the database becomes available:
+
+```sh
+$ dbmate wait
+Waiting for database....
+```
+
+You can chain `wait` together with other commands if you sometimes see failures caused by the database not yet being ready:
+
+```sh
+$ dbmate wait && dbmate up
+Waiting for database....
+Creating: myapp_development
+```
+
+If the database is still not available after 60 seconds, the command will return an error:
+
+```sh
+$ dbmate wait
+Waiting for database............................................................
+Error: unable to connect to database: pq: role "foobar" does not exist
+```
+
+Please note that the `wait` command does not verify whether your specified database exists, only that the server is available and ready (so it will return success if the database server is available, but your database has not yet been created).
+
 ### Options
 
 The following command line options are available with all commands. You must use command line arguments in the order `dbmate [global options] command [command options]`.
@@ -265,6 +303,7 @@ Why another database schema migration tool? Dbmate was inspired by many other to
 |Support for creating and dropping databases||||:white_check_mark:||:white_check_mark:|
 |Support for saving schema dump files||||:white_check_mark:||:white_check_mark:|
 |Timestamp-versioned migration files|:white_check_mark:|||:white_check_mark:|:white_check_mark:|:white_check_mark:|
+|Ability to wait for database to become ready||||||:white_check_mark:|
 |Database connection string loaded from environment variables||||||:white_check_mark:|
 |Automatically load .env file||||||:white_check_mark:|
 |No separate configuration file||||:white_check_mark:|:white_check_mark:|:white_check_mark:|

main.go

@@ -104,6 +104,13 @@ func NewApp() *cli.App {
 				return db.DumpSchema()
 			}),
 		},
+		{
+			Name:  "wait",
+			Usage: "Wait for the database to become available",
+			Action: action(func(db *dbmate.DB, c *cli.Context) error {
+				return db.Wait()
+			}),
+		},
 	}
 
 	return app

pkg/dbmate/db.go

@@ -13,10 +13,16 @@ import (
 )
 
 // DefaultMigrationsDir specifies default directory to find migration files
-var DefaultMigrationsDir = "./db/migrations"
+const DefaultMigrationsDir = "./db/migrations"
 
 // DefaultSchemaFile specifies default location for schema.sql
-var DefaultSchemaFile = "./db/schema.sql"
+const DefaultSchemaFile = "./db/schema.sql"
+
+// DefaultWaitInterval specifies length of time between connection attempts
+const DefaultWaitInterval = time.Second
+
+// DefaultWaitTimeout specifies maximum time for connection attempts
+const DefaultWaitTimeout = 60 * time.Second
 
 // DB allows dbmate actions to be performed on a specified database
 type DB struct {
@@ -24,6 +30,8 @@ type DB struct {
 	DatabaseURL    *url.URL
 	MigrationsDir  string
 	SchemaFile     string
+	WaitInterval   time.Duration
+	WaitTimeout    time.Duration
 }
 
 // New initializes a new dbmate database
@@ -33,6 +41,8 @@ func New(databaseURL *url.URL) *DB {
 		DatabaseURL:    databaseURL,
 		MigrationsDir:  DefaultMigrationsDir,
 		SchemaFile:     DefaultSchemaFile,
+		WaitInterval:   DefaultWaitInterval,
+		WaitTimeout:    DefaultWaitTimeout,
 	}
 }
 
@@ -41,6 +51,40 @@ func (db *DB) GetDriver() (Driver, error) {
 	return GetDriver(db.DatabaseURL.Scheme)
 }
 
+// Wait blocks until the database server is available. It does not verify that
+// the specified database exists, only that the host is ready to accept connections.
+func (db *DB) Wait() error {
+	drv, err := db.GetDriver()
+	if err != nil {
+		return err
+	}
+
+	// attempt connection to database server
+	err = drv.Ping(db.DatabaseURL)
+	if err == nil {
+		// connection successful
+		return nil
+	}
+
+	fmt.Print("Waiting for database")
+	for i := 0 * time.Second; i < db.WaitTimeout; i += db.WaitInterval {
+		fmt.Print(".")
+		time.Sleep(db.WaitInterval)
+
+		// attempt connection to database server
+		err = drv.Ping(db.DatabaseURL)
+		if err == nil {
+			// connection successful
+			fmt.Print("\n")
+			return nil
+		}
+	}
+
+	// if we find outselves here, we could not connect within the timeout
+	fmt.Print("\n")
+	return fmt.Errorf("unable to connect to database: %s", err)
+}
+
 // CreateAndMigrate creates the database (if necessary) and runs migrations
 func (db *DB) CreateAndMigrate() error {
 	drv, err := db.GetDriver()

pkg/dbmate/db_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -37,6 +38,32 @@ func TestNew(t *testing.T) {
 	require.Equal(t, u.String(), db.DatabaseURL.String())
 	require.Equal(t, "./db/migrations", db.MigrationsDir)
 	require.Equal(t, "./db/schema.sql", db.SchemaFile)
+	require.Equal(t, time.Second, db.WaitInterval)
+	require.Equal(t, 60*time.Second, db.WaitTimeout)
+}
+
+func TestWait(t *testing.T) {
+	u := postgresTestURL(t)
+	db := newTestDB(t, u)
+
+	// speed up our retry loop for testing
+	db.WaitInterval = time.Millisecond
+	db.WaitTimeout = 5 * time.Millisecond
+
+	// drop database
+	err := db.Drop()
+	require.Nil(t, err)
+
+	// test wait
+	err = db.Wait()
+	require.Nil(t, err)
+
+	// test invalid connection
+	u.Host = "postgres:404"
+	err = db.Wait()
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "unable to connect to database: dial tcp")
+	require.Contains(t, err.Error(), "getsockopt: connection refused")
 }
 
 func TestDumpSchema(t *testing.T) {

pkg/dbmate/driver.go

@@ -17,6 +17,7 @@ type Driver interface {
 	SelectMigrations(*sql.DB, int) (map[string]bool, error)
 	InsertMigration(Transaction, string) error
 	DeleteMigration(Transaction, string) error
+	Ping(*url.URL) error
 }
 
 var drivers = map[string]Driver{}

pkg/dbmate/mysql.go

@@ -225,3 +225,15 @@ func (drv MySQLDriver) DeleteMigration(db Transaction, version string) error {
 
 	return err
 }
+
+// Ping verifies a connection to the database server. It does not verify whether the
+// specified database exists.
+func (drv MySQLDriver) Ping(u *url.URL) error {
+	db, err := drv.openRootDB(u)
+	if err != nil {
+		return err
+	}
+	defer mustClose(db)
+
+	return db.Ping()
+}

pkg/dbmate/mysql_test.go

@@ -254,3 +254,22 @@ func TestMySQLDeleteMigration(t *testing.T) {
 	require.Nil(t, err)
 	require.Equal(t, 1, count)
 }
+
+func TestMySQLPing(t *testing.T) {
+	drv := MySQLDriver{}
+	u := mySQLTestURL(t)
+
+	// drop any existing database
+	err := drv.DropDatabase(u)
+	require.Nil(t, err)
+
+	// ping database
+	err = drv.Ping(u)
+	require.Nil(t, err)
+
+	// ping invalid host should return error
+	u.Host = "mysql:404"
+	err = drv.Ping(u)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "getsockopt: connection refused")
+}

pkg/dbmate/postgres.go

@@ -173,3 +173,15 @@ func (drv PostgresDriver) DeleteMigration(db Transaction, version string) error
 
 	return err
 }
+
+// Ping verifies a connection to the database server. It does not verify whether the
+// specified database exists.
+func (drv PostgresDriver) Ping(u *url.URL) error {
+	db, err := drv.openPostgresDB(u)
+	if err != nil {
+		return err
+	}
+	defer mustClose(db)
+
+	return db.Ping()
+}

pkg/dbmate/postgres_test.go

@@ -236,3 +236,22 @@ func TestPostgresDeleteMigration(t *testing.T) {
 	require.Nil(t, err)
 	require.Equal(t, 1, count)
 }
+
+func TestPostgresPing(t *testing.T) {
+	drv := PostgresDriver{}
+	u := postgresTestURL(t)
+
+	// drop any existing database
+	err := drv.DropDatabase(u)
+	require.Nil(t, err)
+
+	// ping database
+	err = drv.Ping(u)
+	require.Nil(t, err)
+
+	// ping invalid host should return error
+	u.Host = "postgres:404"
+	err = drv.Ping(u)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "getsockopt: connection refused")
+}

pkg/dbmate/sqlite.go

@@ -164,3 +164,16 @@ func (drv SQLiteDriver) DeleteMigration(db Transaction, version string) error {
 
 	return err
 }
+
+// Ping verifies a connection to the database. Due to the way SQLite works, by
+// testing whether the database is valid, it will automatically create the database
+// if it does not already exist.
+func (drv SQLiteDriver) Ping(u *url.URL) error {
+	db, err := drv.Open(u)
+	if err != nil {
+		return err
+	}
+	defer mustClose(db)
+
+	return db.Ping()
+}

pkg/dbmate/sqlite_test.go

@@ -40,6 +40,7 @@ func prepTestSQLiteDB(t *testing.T) *sql.DB {
 func TestSQLiteCreateDropDatabase(t *testing.T) {
 	drv := SQLiteDriver{}
 	u := sqliteTestURL(t)
+	path := sqlitePath(u)
 
 	// drop any existing database
 	err := drv.DropDatabase(u)
@@ -50,7 +51,7 @@ func TestSQLiteCreateDropDatabase(t *testing.T) {
 	require.Nil(t, err)
 
 	// check that database exists
-	_, err = os.Stat(sqlitePath(u))
+	_, err = os.Stat(path)
 	require.Nil(t, err)
 
 	// drop the database
@@ -58,7 +59,7 @@ func TestSQLiteCreateDropDatabase(t *testing.T) {
 	require.Nil(t, err)
 
 	// check that database no longer exists
-	_, err = os.Stat(sqlitePath(u))
+	_, err = os.Stat(path)
 	require.NotNil(t, err)
 	require.Equal(t, true, os.IsNotExist(err))
 }
@@ -212,3 +213,37 @@ func TestSQLiteDeleteMigration(t *testing.T) {
 	require.Nil(t, err)
 	require.Equal(t, 1, count)
 }
+
+func TestSQLitePing(t *testing.T) {
+	drv := SQLiteDriver{}
+	u := sqliteTestURL(t)
+	path := sqlitePath(u)
+
+	// drop any existing database
+	err := drv.DropDatabase(u)
+	require.Nil(t, err)
+
+	// ping database
+	err = drv.Ping(u)
+	require.Nil(t, err)
+
+	// check that the database was created (sqlite-only behavior)
+	_, err = os.Stat(path)
+	require.Nil(t, err)
+
+	// drop the database
+	err = drv.DropDatabase(u)
+	require.Nil(t, err)
+
+	// create directory where database file is expected
+	err = os.Mkdir(path, 0755)
+	require.Nil(t, err)
+	defer func() {
+		err = os.RemoveAll(path)
+		require.Nil(t, err)
+	}()
+
+	// ping database should fail
+	err = drv.Ping(u)
+	require.EqualError(t, err, "unable to open database file")
+}