Add wait command (#35)

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")
+}