Add 'UseMigrationLock' flag to acquire exclusive lock while performing migrations for Postgres

main.go

@@ -96,6 +96,10 @@ func NewApp() *cli.App {
 			Usage:   "timeout for --wait flag",
 			Value:   defaultDB.WaitTimeout,
 		},
+		&cli.BoolFlag{
+			Name:  "migration-lock",
+			Usage: "use a lock during the migration so other dbmate instances can not run migrations at the same time",
+		},
 	}
 
 	app.Commands = []*cli.Command{
@@ -127,6 +131,7 @@ func NewApp() *cli.App {
 			Action: action(func(db *dbmate.DB, c *cli.Context) error {
 				db.Strict = c.Bool("strict")
 				db.Verbose = c.Bool("verbose")
+				db.UseMigrationLock = c.Bool("migration-lock")
 				return db.CreateAndMigrate()
 			}),
 		},
@@ -163,6 +168,7 @@ func NewApp() *cli.App {
 			Action: action(func(db *dbmate.DB, c *cli.Context) error {
 				db.Strict = c.Bool("strict")
 				db.Verbose = c.Bool("verbose")
+				db.UseMigrationLock = c.Bool("migration-lock")
 				return db.Migrate()
 			}),
 		},

pkg/dbmate/db.go

@@ -60,6 +60,8 @@ type DB struct {
 	WaitInterval time.Duration
 	// WaitTimeout specifies maximum time for connection attempts
 	WaitTimeout time.Duration
+	// UseMigrationLock uses an exclusive lock while performing migrations
+	UseMigrationLock bool
 }
 
 // StatusResult represents an available migration status
@@ -83,6 +85,7 @@ func New(databaseURL *url.URL) *DB {
 		WaitBefore:          false,
 		WaitInterval:        time.Second,
 		WaitTimeout:         60 * time.Second,
+		UseMigrationLock:    false,
 	}
 }
 
@@ -153,12 +156,36 @@ func (db *DB) Wait() error {
 }
 
 // CreateAndMigrate creates the database (if necessary) and runs migrations
-func (db *DB) CreateAndMigrate() error {
+func (db *DB) CreateAndMigrate() (returnErr error) {
 	drv, err := db.Driver()
 	if err != nil {
 		return err
 	}
 
+	// try and acquire a lock for the duration of the migration.
+	// this is to prevent multiple instances performing the same migration in parallel.
+	if db.UseMigrationLock {
+		drvLock, ok := drv.(DriverMigrationLock)
+		if !ok {
+			return fmt.Errorf("driver does not support the use of a migration lock")
+		}
+
+		if err := drvLock.Lock(); err != nil {
+			return err
+		}
+
+		defer func() {
+			err := drvLock.Unlock()
+			if err != nil {
+				if returnErr != nil {
+					returnErr = fmt.Errorf("failed to unlock: %v: %w", err, returnErr)
+					return
+				}
+				returnErr = fmt.Errorf("failed to unlock: %w", err)
+			}
+		}()
+	}
+
 	// create database if it does not already exist
 	// skip this step if we cannot determine status
 	// (e.g. user does not have list database permission)
@@ -333,12 +360,37 @@ func (db *DB) openDatabaseForMigration(drv Driver) (*sql.DB, error) {
 }
 
 // Migrate migrates database to the latest version
-func (db *DB) Migrate() error {
+func (db *DB) Migrate() (returnErr error) {
 	drv, err := db.Driver()
 	if err != nil {
 		return err
 	}
 
+	if db.UseMigrationLock {
+		drvLock, ok := drv.(DriverMigrationLock)
+		if !ok {
+			return fmt.Errorf("driver does not support the use of a migration lock")
+		}
+
+		// only try and lock if we haven't already done so, e.g called CreateAndMigrate.
+		if !drvLock.IsLocked() {
+			if err := drvLock.Lock(); err != nil {
+				return err
+			}
+
+			defer func() {
+				err := drvLock.Unlock()
+				if err != nil {
+					if returnErr != nil {
+						returnErr = fmt.Errorf("failed to unlock: %v: %w", err, returnErr)
+						return
+					}
+					returnErr = fmt.Errorf("failed to unlock: %w", err)
+				}
+			}()
+		}
+	}
+
 	migrations, err := db.FindMigrations()
 	if err != nil {
 		return err

pkg/dbmate/db_test.go

@@ -59,6 +59,7 @@ func TestNew(t *testing.T) {
 	require.False(t, db.WaitBefore)
 	require.Equal(t, time.Second, db.WaitInterval)
 	require.Equal(t, 60*time.Second, db.WaitTimeout)
+	require.Equal(t, false, db.UseMigrationLock)
 }
 
 func TestGetDriver(t *testing.T) {

pkg/dbmate/driver.go

@@ -25,6 +25,12 @@ type Driver interface {
 	QueryError(string, error) error
 }
 
+type DriverMigrationLock interface {
+	Lock() error
+	Unlock() error
+	IsLocked() bool
+}
+
 // DriverConfig holds configuration passed to driver constructors
 type DriverConfig struct {
 	DatabaseURL         *url.URL

pkg/driver/postgres/postgres.go

@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"bytes"
+	"context"
 	"database/sql"
 	"fmt"
 	"io"
@@ -28,6 +29,8 @@ type Driver struct {
 	migrationsTableName string
 	databaseURL         *url.URL
 	log                 io.Writer
+
+	migrationLockTx *sql.Tx
 }
 
 // NewDriver initializes the driver
@@ -456,3 +459,48 @@ func (drv *Driver) quotedMigrationsTableNameParts(db dbutil.Transaction) (string
 	// if more than one part, we already have a schema
 	return quotedNameParts[0], strings.Join(quotedNameParts[1:], "."), nil
 }
+
+const lockKey = 48372615
+
+func (drv *Driver) Lock() error {
+	if drv.migrationLockTx != nil {
+		return fmt.Errorf("already locked")
+	}
+
+	db, err := drv.Open()
+	if err != nil {
+		return err
+	}
+
+	tx, err := db.BeginTx(context.Background(), nil)
+	if err != nil {
+		return err
+	}
+
+	drv.migrationLockTx = tx
+
+	_, err = tx.Exec("SELECT pg_advisory_xact_lock($1)", lockKey)
+	if err != nil {
+		return fmt.Errorf("failed to acquire lock: %w", err)
+	}
+
+	return nil
+}
+
+func (drv *Driver) Unlock() error {
+	if drv.migrationLockTx == nil {
+		return fmt.Errorf("not locked")
+	}
+
+	if err := drv.migrationLockTx.Rollback(); err != nil {
+		return err
+	}
+
+	drv.migrationLockTx = nil
+
+	return nil
+}
+
+func (drv *Driver) IsLocked() bool {
+	return drv.migrationLockTx != nil
+}

pkg/driver/postgres/postgres_test.go

@@ -6,6 +6,7 @@ import (
 	"net/url"
 	"runtime"
 	"testing"
+	"time"
 
 	"github.com/amacneil/dbmate/v2/pkg/dbmate"
 	"github.com/amacneil/dbmate/v2/pkg/dbtest"
@@ -790,3 +791,38 @@ func TestPostgresMigrationsTableExists(t *testing.T) {
 		require.Equal(t, true, exists)
 	})
 }
+
+func TestPostgresMigrationLock(t *testing.T) {
+	t.Run("lock and unlock", func(t *testing.T) {
+		drv := testPostgresDriver(t)
+		err := drv.Lock()
+		require.NoError(t, err)
+
+		err = drv.Lock()
+		require.Error(t, err, "Should not be able to lock again without unlock")
+
+		err = drv.Unlock()
+		require.NoError(t, err, "Should be able to unlock")
+	})
+
+	t.Run("lock on one instance should block lock attempt on another", func(t *testing.T) {
+		drv1 := testPostgresDriver(t)
+		err1 := drv1.Lock()
+		require.NoError(t, err1)
+
+		var isUnlocked bool
+		go func() {
+			time.Sleep(10 * time.Millisecond)
+			err := drv1.Unlock()
+			require.NoError(t, err, "Should be able to unlock")
+			isUnlocked = true
+		}()
+
+		drv2 := testPostgresDriver(t)
+		err2 := drv2.Lock()
+		require.NoError(t, err2)
+		require.Equal(t, true, isUnlocked)
+		err2 = drv2.Unlock()
+		require.NoError(t, err2, "Should be able to unlock")
+	})
+}