Repository: ariga/atlas
Pull request #844 (2022-06-01 19:31:52)

sql/migrate: refactor execution (in preparation for transactions)
This PR refactors the migration execution into two methods instead of one.

1. Compute the pending migration files
2. Execute a migration file onto a database

The above changes are needed to be able to wrap the execution of a migration file in a transaction. Splitting this into the two above methods also enables concurrent migration of multi-tenancy setups and keep them in sync before executing a new migration file.

Wrapping a migration execution in a transaction is should be controlled by the consumer instead of the executor itself. 


Patch:
diff --git a/cmd/atlas/doc.tmpl b/cmd/atlas/doc.tmpl
index 03ae307b..1e7ef325 100644
--- a/cmd/atlas/doc.tmpl
+++ b/cmd/atlas/doc.tmpl
@@ -17,10 +17,8 @@ For a more detailed introduction to the CLI capabilities, head over to the
 
 ## Distributed Binaries
 
-Starting [v0.3.0](https://github.com/ariga/atlas/releases/tag/v0.3.0),
-ֿthe distributed binaries include code for a [Management UI](ui/intro.md) wrapping the
-core atlas engine that is not currently released publicly. The binaries
-themselves are still released under the same [Apache License 2.0](https://github.com/ariga/atlas/blob/master/LICENSE).
+The binaries distributed in official releases are released under the  [Ariga End User License](https://ariga.io/legal/atlas/eula).
+If you would like to build Atlas from source follow the instructions [here](https://atlasgo.io/cli/reference#building-from-source).
 
 ### Building from Source
 
diff --git a/cmd/atlascmd/migrate/migrate.go b/cmd/atlascmd/migrate/migrate.go
index 9ee413d8..b7e32673 100644
--- a/cmd/atlascmd/migrate/migrate.go
+++ b/cmd/atlascmd/migrate/migrate.go
@@ -109,24 +109,18 @@ func (r *EntRevisions) ReadRevisions(ctx context.Context) (migrate.Revisions, er
 	return ret, nil
 }
 
-// WriteRevisions writes the revisions to the revisions table.
-func (r *EntRevisions) WriteRevisions(ctx context.Context, rs migrate.Revisions) error {
-	bulk := make([]*ent.RevisionCreate, len(rs))
-	for i, rev := range rs {
-		bulk[i] = r.ec.Revision.Create().
-			SetID(rev.Version).
-			SetDescription(rev.Description).
-			SetExecutionState(revision.ExecutionState(rev.ExecutionState)).
-			SetExecutedAt(rev.ExecutedAt).
-			SetExecutionTime(rev.ExecutionTime).
-			SetHash(rev.Hash).
-			SetOperatorVersion(rev.OperatorVersion).
-			SetMeta(rev.Meta)
-	}
-	return r.ec.Revision.CreateBulk(bulk...).
-		OnConflict(
-			sql.ConflictColumns(revision.FieldID),
-		).
+// WriteRevision writes a revision to the revisions table.
+func (r *EntRevisions) WriteRevision(ctx context.Context, rev *migrate.Revision) error {
+	return r.ec.Revision.Create().
+		SetID(rev.Version).
+		SetDescription(rev.Description).
+		SetExecutionState(revision.ExecutionState(rev.ExecutionState)).
+		SetExecutedAt(rev.ExecutedAt).
+		SetExecutionTime(rev.ExecutionTime).
+		SetHash(rev.Hash).
+		SetOperatorVersion(rev.OperatorVersion).
+		SetMeta(rev.Meta).
+		OnConflict(sql.ConflictColumns(revision.FieldID)).
 		UpdateNewValues().
 		Exec(ctx)
 }
diff --git a/doc/md/cli/reference.md b/doc/md/cli/reference.md
index 217e8c15..1b6952b6 100644
--- a/doc/md/cli/reference.md
+++ b/doc/md/cli/reference.md
@@ -14,10 +14,8 @@ For a more detailed introduction to the CLI capabilities, head over to the
 
 ## Distributed Binaries
 
-Starting [v0.3.0](https://github.com/ariga/atlas/releases/tag/v0.3.0),
-ֿthe distributed binaries include code for a [Management UI](ui/intro.md) wrapping the
-core atlas engine that is not currently released publicly. The binaries
-themselves are still released under the same [Apache License 2.0](https://github.com/ariga/atlas/blob/master/LICENSE).
+The binaries distributed in official releases are released under the  [Ariga End User License](https://ariga.io/legal/atlas/eula).
+If you would like to build Atlas from source follow the instructions [here](https://atlasgo.io/cli/reference#building-from-source).
 
 ### Building from Source
 
diff --git a/doc/md/getting-started/01-introduction.mdx b/doc/md/getting-started/01-introduction.mdx
index 4ff86567..716608d2 100644
--- a/doc/md/getting-started/01-introduction.mdx
+++ b/doc/md/getting-started/01-introduction.mdx
@@ -83,10 +83,9 @@ move the atlas binary to a file location on your system PATH.
 Latest release is updated twice a day and is based on the most recent tagged release.
 :::
 
-The distributed binaries include code for a [Management UI](../ui/intro.md) wrapping the
-core Atlas engine that is not currently released publicly. The binaries
-themselves are released under the [Ariga End User License](https://ariga.io/legal/atlas/eula).
-If you would like to build Atlas from source without the UI follow the instructions [here](../cli/reference.md#building-from-source).
+The binaries distributed in official releases are released under the  [Ariga End User License](https://ariga.io/legal/atlas/eula).
+If you would like to build Atlas from source follow the instructions [here](https://atlasgo.io/cli/reference#building-from-source).
+
 
 To verify that the CLI is installed correctly run:
 ```shell
diff --git a/sql/migrate/migrate.go b/sql/migrate/migrate.go
index cd860a63..bbfdd854 100644
--- a/sql/migrate/migrate.go
+++ b/sql/migrate/migrate.go
@@ -177,7 +177,7 @@ type (
 	// Atlas drivers provide a sql based implementation of this interface by default.
 	RevisionReadWriter interface {
 		ReadRevisions(context.Context) (Revisions, error)
-		WriteRevisions(context.Context, Revisions) error
+		WriteRevision(context.Context, *Revision) error
 	}
 
 	// Planner can plan the steps to take to migrate from one state to another. It uses the enclosed FS to write
@@ -340,6 +340,10 @@ func NewExecutor(drv Driver, dir Dir, rrw RevisionReadWriter, opts ...ExecutorOp
 	if _, ok := drv.(schema.Locker); !ok {
 		return nil, ErrLockUnsupported
 	}
+	// Check if the Dir implements Scanner interface.
+	if _, ok := dir.(Scanner); !ok {
+		return nil, errors.New("sql/migrate: execute: no scanner available")
+	}
 	p := &Executor{drv: drv, dir: dir, rrw: rrw}
 	for _, opt := range opts {
 		if err := opt(p); err != nil {
@@ -357,85 +361,138 @@ func WithLogger(log Logger) ExecutorOption {
 	}
 }
 
-// Execute executes n missing migration files on the database. Passing in n <= 0 means running all pending migrations.
-func (e *Executor) Execute(ctx context.Context, n int) (err error) {
+// Lock acquires a lock for the executor.
+// It is considered a user error to not call Lock before the Pending and Execute methods.
+func (e *Executor) Lock(ctx context.Context) (schema.UnlockFunc, error) {
 	unlock, err := e.drv.(schema.Locker).Lock(ctx, "atlas_migration_execute", 0)
 	if err != nil {
-		return fmt.Errorf("sql/migrate: acquiring database lock: %w", err)
+		return nil, fmt.Errorf("sql/migrate: acquiring database lock: %w", err)
 	}
-	defer unlock()
+	return unlock, nil
+}
+
+// Pending returns all pending (not applied) migration files in the migration directory.
+func (e *Executor) Pending(ctx context.Context) ([]File, error) {
 	// Don't operate with a broken migration directory.
-	// TODO(maseeelch): do not check here but let the caller check it before? Or let the caller decide if to skip this flag by using some flags.
 	if err := Validate(e.dir); err != nil {
-		return fmt.Errorf("sql/migrate: execute: validate migration directory: %w", err)
-	}
-	// Check if the Dir implements Scanner interface.
-	sc, ok := e.dir.(Scanner)
-	if !ok {
-		return errors.New("sql/migrate: execute: no scanner available")
+		return nil, fmt.Errorf("sql/migrate: execute: validate migration directory: %w", err)
 	}
 	// Read all applied database revisions.
-	revisions, err := e.rrw.ReadRevisions(ctx)
+	revs, err := e.rrw.ReadRevisions(ctx)
 	if err != nil {
-		return fmt.Errorf("sql/migrate: execute: read revisions: %w", err)
+		return nil, fmt.Errorf("sql/migrate: execute: read revisions: %w", err)
 	}
 	// Select the correct migration files.
+	sc := e.dir.(Scanner)
 	migrations, err := sc.Files()
 	if err != nil {
-		return fmt.Errorf("sql/migrate: execute: select migration files: %w", err)
+		return nil, fmt.Errorf("sql/migrate: execute: select migration files: %w", err)
 	}
 	// Check if the existing revisions did come from the migration directory and not from somewhere else.
 	hf, err := readHashFile(e.dir)
 	if err != nil && !errors.Is(err, os.ErrNotExist) {
-		return fmt.Errorf("sql/migrate: execute: read atlas.sum file: %w", err)
+		return nil, fmt.Errorf("sql/migrate: execute: read atlas.sum file: %w", err)
 	}
 	// If there is no atlas.sum file there are no migration files.
 	if errors.Is(err, os.ErrNotExist) {
-		return ErrNoPendingFiles
+		return nil, ErrNoPendingFiles
 	}
 	// For up to len(revisions) revisions the migration files must both match in order and content.
-	if len(revisions) > len(migrations) { // TODO(masseelch): keep compaction in mind.
-		return errors.New("sql/migrate: execute: revisions and migrations mismatch: more revisions than migrations")
+	if len(revs) > len(migrations) { // TODO(masseelch): keep compaction in mind.
+		return nil, errors.New("sql/migrate: execute: revisions and migrations mismatch: more revisions than migrations")
 	}
-	for i := range revisions {
+	for i := range revs {
 		m := migrations[i]
 		v, err := sc.Version(m)
 		if err != nil {
-			return fmt.Errorf("sql/migrate: execute: scan version from %q: %w", m.Name(), err)
+			return nil, fmt.Errorf("sql/migrate: execute: scan version from %q: %w", m.Name(), err)
 		}
 		d, err := sc.Desc(m)
 		if err != nil {
-			return fmt.Errorf("sql/migrate: execute: scan description from %q: %w", m.Name(), err)
+			return nil, fmt.Errorf("sql/migrate: execute: scan description from %q: %w", m.Name(), err)
 		}
-		r := revisions[i]
+		r := revs[i]
 		if v != r.Version || d != r.Description || hf[i].H != r.Hash {
-			return fmt.Errorf("sql/migrate: execute: revisions and migrations mismatch: rev %q <> file %q", v, r.Version)
+			return nil, fmt.Errorf("sql/migrate: execute: revisions and migrations mismatch: rev %q <> file %q", v, r.Version)
 		}
 	}
-	if len(migrations) == len(revisions) {
-		return ErrNoPendingFiles
+	if len(migrations) == len(revs) {
+		return nil, ErrNoPendingFiles
 	}
-	defer func(rrw RevisionReadWriter, ctx context.Context, revisions *Revisions) {
-		if dErr := rrw.WriteRevisions(ctx, *revisions); dErr != nil {
-			if err != nil {
-				err = fmt.Errorf("sql/migrate: execute: write revisions: %w: %v", dErr, err)
-				return
-			}
-			err = dErr
-		}
-	}(e.rrw, ctx, &revisions)
-	// Determine what migrations to run.
-	var pending []File
-	if n <= 0 {
-		// If n<=0 run every pending migrations.
-		pending = migrations[len(revisions):]
-	} else {
-		// Else run n pending migrations.
-		pc := len(migrations) - len(revisions)
-		if n > pc {
-			n = pc
+	return migrations[len(revs):], nil
+}
+
+// Execute executes the given migration file on the database.
+func (e *Executor) Execute(ctx context.Context, m File) (err error) {
+	r := &Revision{ExecutedAt: time.Now(), ExecutionState: stateOngoing}
+	if err := e.rrw.WriteRevision(ctx, r); err != nil {
+		return fmt.Errorf("sql/migrate: execute: storing revision: %w", err)
+	}
+	// Make sure to store the Revision information.
+	defer func(ctx context.Context, rrw RevisionReadWriter, r *Revision) {
+		if err2 := e.rrw.WriteRevision(ctx, r); err2 != nil {
+			err = wrap(fmt.Errorf("execute: write revision: %w", err2), err)
+		}
+	}(ctx, e.rrw, r)
+	sc := e.dir.(Scanner)
+	v, err := sc.Version(m)
+	if err != nil {
+		return r.setGoErr(fmt.Errorf("sql/migrate: execute: scan version from %q: %w", m.Name(), err))
+	}
+	r.Version = v
+	d, err := sc.Desc(m)
+	if err != nil {
+		return r.setGoErr(fmt.Errorf("sql/migrate: execute: scan description from %q: %w", m.Name(), err))
+
+	}
+	r.Description = d
+	if e.log != nil {
+		e.log.Log(LogFile{v, d})
+	}
+	hf, err := readHashFile(e.dir)
+	if err != nil {
+		return fmt.Errorf("sql/migrate: execute: read atlas.sum file: %w", err)
+	}
+	h, err := hf.sumByName(m.Name())
+	if err != nil {
+		return r.setGoErr(fmt.Errorf("sql/migrate: execute: scanning checksum for file %q: %w", m.Name(), err))
+	}
+	r.Hash = h
+	stmts, err := sc.Stmts(m)
+	if err != nil {
+		return r.setGoErr(fmt.Errorf("sql/migrate: execute: scanning statements from file %q: %w", m.Name(), err))
+	}
+	for _, stmt := range stmts {
+		if e.log != nil {
+			e.log.Log(LogStmt{stmt})
+		}
+		if _, err := e.drv.ExecContext(ctx, stmt); err != nil {
+			return r.setSQLErr(
+				fmt.Errorf("sql/migrate: execute: executing statement %q from version %q: %w", stmt, v, err),
+				stmt,
+			)
+		}
+		r.done(true)
+	}
+	return
+}
+
+// ExecuteN executes n pending migration files. If n<=0 all pending migration files are executed.
+func (e *Executor) ExecuteN(ctx context.Context, n int) error {
+	unlock, err := e.Lock(ctx)
+	if err != nil {
+		return err
+	}
+	defer unlock()
+	pending, err := e.Pending(ctx)
+	if err != nil {
+		return err
+	}
+	if n > 0 {
+		if n >= len(pending) {
+			n = len(pending)
 		}
-		pending = migrations[len(revisions) : len(revisions)+n]
+		pending = pending[:n]
 	}
 	if e.log != nil {
 		names := make([]string, len(pending))
@@ -444,46 +501,12 @@ func (e *Executor) Execute(ctx context.Context, n int) (err error) {
 		}
 		e.log.Log(LogExecution{names})
 	}
-	// TODO(masseelch): run in a transaction
 	for _, m := range pending {
-		r := &Revision{ExecutedAt: time.Now(), ExecutionState: stateOngoing}
-		revisions = append(revisions, r)
-		v, err := sc.Version(m)
-		if err != nil {
-			return r.setGoErr(fmt.Errorf("sql/migrate: execute: scan version from %q: %w", m.Name(), err))
-		}
-		r.Version = v
-		d, err := sc.Desc(m)
-		if err != nil {
-			return r.setGoErr(fmt.Errorf("sql/migrate: execute: scan description from %q: %w", m.Name(), err))
-		}
-		r.Description = d
-		if e.log != nil {
-			e.log.Log(LogFile{v, d})
-		}
-		h, err := hf.sumByName(m.Name())
-		if err != nil {
-			return r.setGoErr(fmt.Errorf("sql/migrate: execute: scanning checksum for file %q: %w", m.Name(), err))
-		}
-		r.Hash = h
-		stmts, err := sc.Stmts(m)
-		if err != nil {
-			return r.setGoErr(fmt.Errorf("sql/migrate: execute: scanning statements from file %q: %w", m.Name(), err))
-		}
-		for _, stmt := range stmts {
-			if e.log != nil {
-				e.log.Log(LogStmt{stmt})
-			}
-			if _, err := e.drv.ExecContext(ctx, stmt); err != nil {
-				return r.setSQLErr(
-					fmt.Errorf("sql/migrate: execute: executing statement %q from version %q: %w", stmt, v, err),
-					stmt,
-				)
-			}
-			r.done(true)
+		if err := e.Execute(ctx, m); err != nil {
+			return err
 		}
 	}
-	return
+	return nil
 }
 
 // ErrNotClean is returned when the connected dev-db is not in a clean state (aka it has schemas and tables).
@@ -542,12 +565,12 @@ func (e *Executor) ReadState(ctx context.Context) (realm *schema.Realm, err erro
 		for i, s := range realm.Schemas {
 			del[i] = &schema.DropSchema{S: s, Extra: []schema.Clause{&schema.IfExists{}}}
 		}
-		if derr := e.drv.ApplyChanges(ctx, del); derr != nil {
-			err = wrap(derr, err)
+		if err2 := e.drv.ApplyChanges(ctx, del); err2 != nil {
+			err = wrap(err2, err)
 		}
 	}()
 	// Replay the migration directory on the database.
-	if err := e.Execute(ctx, 0); err != nil && !errors.Is(err, ErrNoPendingFiles) {
+	if err := e.ExecuteN(ctx, 0); err != nil && !errors.Is(err, ErrNoPendingFiles) {
 		return nil, fmt.Errorf("sql/migrate: read migration directory state: %w", err)
 	}
 	// Inspect the database back and return the result.
@@ -564,11 +587,13 @@ func (NoopRevisionReadWriter) ReadRevisions(context.Context) (Revisions, error)
 	return nil, nil
 }
 
-// WriteRevisions implements RevisionsReadWriter.WriteRevisions,
-func (NoopRevisionReadWriter) WriteRevisions(context.Context, Revisions) error {
+// WriteRevision implements RevisionsReadWriter.WriteRevision,
+func (NoopRevisionReadWriter) WriteRevision(context.Context, *Revision) error {
 	return nil
 }
 
+var _ RevisionReadWriter = (*NoopRevisionReadWriter)(nil)
+
 // done computes and sets the ExecutionTime.
 func (r *Revision) done(ok bool) {
 	r.ExecutionTime = time.Now().Sub(r.ExecutedAt)
@@ -602,7 +627,7 @@ type (
 		logEntry()
 	}
 
-	// LogExecution is sent once when execution of the migration files has been started.
+	// LogExecution is sent once when execution of multiple migration files has been started.
 	// It holds the filenames of the pending migration files.
 	LogExecution struct {
 		Files []string
@@ -951,19 +976,6 @@ func readHashFile(dir Dir) (HashFile, error) {
 	return fh, nil
 }
 
-// reverse changes for the down migration.
-func reverse(changes []*Change) []*Change {
-	n := len(changes)
-	rev := make([]*Change, n)
-	if n%2 == 1 {
-		rev[n/2] = changes[n/2]
-	}
-	for i, j := 0, n-1; i < j; i, j = i+1, j-1 {
-		rev[i], rev[j] = changes[j], changes[i]
-	}
-	return rev
-}
-
 func wrap(err1, err2 error) error {
 	if err2 != nil {
 		return fmt.Errorf("sql/migrate: %w: %v", err2, err1)


Test patch:
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index 7b32489a..a57348a0 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -529,14 +529,14 @@ func testExecutor(t T) {
 
 	ex, err := migrate.NewExecutor(t.driver(), dir, t.revisionsStorage())
 	require.NoError(t, err)
-	require.NoError(t, ex.Execute(context.Background(), 2)) // usersT and postsT
+	require.NoError(t, ex.ExecuteN(context.Background(), 2)) // usersT and postsT
 	require.Len(t, *t.revisionsStorage().(*rrw), 2)
 	ensureNoChange(t, postsT, usersT)
-	require.NoError(t, ex.Execute(context.Background(), 1)) // petsT
+	require.NoError(t, ex.ExecuteN(context.Background(), 1)) // petsT
 	require.Len(t, *t.revisionsStorage().(*rrw), 3)
 	ensureNoChange(t, petsT, postsT, usersT)
 
-	require.ErrorIs(t, ex.Execute(context.Background(), 1), migrate.ErrNoPendingFiles)
+	require.ErrorIs(t, ex.ExecuteN(context.Background(), 1), migrate.ErrNoPendingFiles)
 }
 
 func plan(t T, name string, changes ...schema.Change) *migrate.Plan {
@@ -547,12 +547,18 @@ func plan(t T, name string, changes ...schema.Change) *migrate.Plan {
 
 type rrw migrate.Revisions
 
-func (r *rrw) WriteRevisions(ctx context.Context, revs migrate.Revisions) error {
-	*r = rrw(revs)
+func (r *rrw) WriteRevision(_ context.Context, rev *migrate.Revision) error {
+	for i, rev2 := range *r {
+		if rev2.Version == rev.Version {
+			(*r)[i] = rev
+			return nil
+		}
+	}
+	*r = append(*r, rev)
 	return nil
 }
 
-func (r *rrw) ReadRevisions(ctx context.Context) (migrate.Revisions, error) {
+func (r *rrw) ReadRevisions(context.Context) (migrate.Revisions, error) {
 	return migrate.Revisions(*r), nil
 }
 
diff --git a/sql/migrate/migrate_test.go b/sql/migrate/migrate_test.go
index e759970f..498cded3 100644
--- a/sql/migrate/migrate_test.go
+++ b/sql/migrate/migrate_test.go
@@ -202,6 +202,7 @@ func TestExecutor_ReadState(t *testing.T) {
 	require.Equal(t, []string{"DROP TABLE IF EXISTS t;", "CREATE TABLE t(c int);"}, drv.executed)
 	require.Equal(t, 2, drv.lockCounter)
 	require.Equal(t, 2, drv.unlockCounter)
+	require.True(t, drv.released())
 
 	// Does not work if locked.
 	drv.locks = map[string]struct{}{"atlas_migration_directory_state": {}}
@@ -209,6 +210,7 @@ func TestExecutor_ReadState(t *testing.T) {
 	require.EqualError(t, err, "sql/migrate: acquiring database lock: lockErr")
 	require.Equal(t, 2, drv.lockCounter)
 	require.Equal(t, 2, drv.unlockCounter)
+	require.False(t, drv.released())
 	drv.locks = make(map[string]struct{})
 
 	// Does not work if database is not clean.
@@ -217,6 +219,7 @@ func TestExecutor_ReadState(t *testing.T) {
 	require.EqualError(t, err, "sql/migrate: connected database is not clean")
 	require.Equal(t, 3, drv.lockCounter)
 	require.Equal(t, 3, drv.unlockCounter)
+	require.True(t, drv.released())
 
 	// Works.
 	edrv := &emptyMockDriver{drv}
@@ -232,6 +235,7 @@ func TestExecutor_ReadState(t *testing.T) {
 	}, edrv.applied)
 	require.Equal(t, 5, drv.lockCounter)
 	require.Equal(t, 5, drv.unlockCounter)
+	require.True(t, drv.released())
 }
 
 func TestLocalDir(t *testing.T) {
@@ -319,7 +323,7 @@ func TestExecutor(t *testing.T) {
 	ex, err = migrate.NewExecutor(&lockMockDriver{&mockDriver{}}, dir, &mockRevisionReadWriter{})
 	require.NoError(t, err)
 	require.NotNil(t, ex)
-	require.ErrorIs(t, ex.Execute(context.Background(), 0), migrate.ErrChecksumMismatch)
+	require.ErrorIs(t, ex.ExecuteN(context.Background(), 0), migrate.ErrChecksumMismatch)
 
 	// Prerequisites.
 	var (
@@ -345,7 +349,7 @@ func TestExecutor(t *testing.T) {
 	require.NoError(t, err)
 
 	// Applies all of them.
-	require.NoError(t, ex.Execute(context.Background(), 0))
+	require.NoError(t, ex.ExecuteN(context.Background(), 0))
 	require.Equal(t, drv.executed, []string{
 		"CREATE TABLE t_sub(c int);", "ALTER TABLE t_sub ADD c1 int;", "ALTER TABLE t_sub ADD c2 int;",
 	})
@@ -358,45 +362,66 @@ func TestExecutor(t *testing.T) {
 		migrate.LogFile{Version: "2.10.x-20", Desc: "description"},
 		migrate.LogStmt{SQL: "ALTER TABLE t_sub ADD c2 int;"},
 	}, []migrate.LogEntry(*log))
+	require.Equal(t, drv.lockCounter, 1)
+	require.Equal(t, drv.unlockCounter, 1)
+	require.True(t, drv.released())
 
 	// No pending files.
-	require.ErrorIs(t, ex.Execute(context.Background(), 0), migrate.ErrNoPendingFiles)
+	require.ErrorIs(t, ex.ExecuteN(context.Background(), 0), migrate.ErrNoPendingFiles)
 
 	// Apply one by one.
 	*rrw = mockRevisionReadWriter{}
 	*drv = lockMockDriver{&mockDriver{}}
 
-	require.NoError(t, ex.Execute(context.Background(), 1))
+	require.NoError(t, ex.ExecuteN(context.Background(), 1))
 	require.Equal(t, drv.executed, []string{"CREATE TABLE t_sub(c int);", "ALTER TABLE t_sub ADD c1 int;"})
 	requireEqualRevisions(t, migrate.Revisions{rev1}, migrate.Revisions(*rrw))
 
-	require.NoError(t, ex.Execute(context.Background(), 1))
+	require.NoError(t, ex.ExecuteN(context.Background(), 1))
 	require.Equal(t, drv.executed, []string{
 		"CREATE TABLE t_sub(c int);", "ALTER TABLE t_sub ADD c1 int;", "ALTER TABLE t_sub ADD c2 int;",
 	})
 	requireEqualRevisions(t, migrate.Revisions{rev1, rev2}, migrate.Revisions(*rrw))
-	require.ErrorIs(t, ex.Execute(context.Background(), 0), migrate.ErrNoPendingFiles)
+	require.ErrorIs(t, ex.ExecuteN(context.Background(), 0), migrate.ErrNoPendingFiles)
+
+	require.Equal(t, drv.lockCounter, 3)
+	require.Equal(t, drv.unlockCounter, 3)
+	require.True(t, drv.released())
 
 	// Suppose first revision is already executed, only execute second migration file.
 	*rrw = mockRevisionReadWriter(migrate.Revisions{rev1})
 	*drv = lockMockDriver{&mockDriver{}}
 
-	require.NoError(t, ex.Execute(context.Background(), 0))
+	require.NoError(t, ex.ExecuteN(context.Background(), 0))
 	require.Equal(t, drv.executed, []string{"ALTER TABLE t_sub ADD c2 int;"})
 	requireEqualRevisions(t, migrate.Revisions{rev1, rev2}, migrate.Revisions(*rrw))
-	require.ErrorIs(t, ex.Execute(context.Background(), 0), migrate.ErrNoPendingFiles)
+	require.ErrorIs(t, ex.ExecuteN(context.Background(), 0), migrate.ErrNoPendingFiles)
+
+	require.Equal(t, drv.lockCounter, 2)
+	require.Equal(t, drv.unlockCounter, 2)
+	require.True(t, drv.released())
 
 	// Unknown revision present.
 	*rrw = mockRevisionReadWriter(migrate.Revisions{&migrate.Revision{Version: "unknown"}})
-	require.EqualError(t, ex.Execute(context.Background(), 0),
+	*drv = lockMockDriver{&mockDriver{}}
+	require.EqualError(t, ex.ExecuteN(context.Background(), 0),
 		`sql/migrate: execute: revisions and migrations mismatch: rev "1.a" <> file "unknown"`,
 	)
 
+	require.Equal(t, drv.lockCounter, 1)
+	require.Equal(t, drv.unlockCounter, 1)
+	require.True(t, drv.released())
+
 	// More revisions than migration files.
 	*rrw = mockRevisionReadWriter(migrate.Revisions{&migrate.Revision{Version: "unknown"}, rev1, rev2})
-	require.EqualError(t, ex.Execute(context.Background(), 0),
+	*drv = lockMockDriver{&mockDriver{}}
+	require.EqualError(t, ex.ExecuteN(context.Background(), 0),
 		`sql/migrate: execute: revisions and migrations mismatch: more revisions than migrations`,
 	)
+
+	require.Equal(t, drv.lockCounter, 1)
+	require.Equal(t, drv.unlockCounter, 1)
+	require.True(t, drv.released())
 }
 
 type (
@@ -448,23 +473,32 @@ func (m *lockMockDriver) Lock(_ context.Context, name string, _ time.Duration) (
 		return nil
 	}, nil
 }
+func (m *lockMockDriver) released() bool {
+	return len(m.locks) == 0
+}
 func (m *emptyMockDriver) IsClean(context.Context) (bool, error) {
 	return true, nil
 }
 
 type mockRevisionReadWriter migrate.Revisions
 
-func (r *mockRevisionReadWriter) WriteRevisions(_ context.Context, revs migrate.Revisions) error {
-	*r = mockRevisionReadWriter(revs)
+func (rrw *mockRevisionReadWriter) WriteRevision(_ context.Context, r *migrate.Revision) error {
+	for i, rev := range *rrw {
+		if rev.Version == r.Version {
+			(*rrw)[i] = r
+			return nil
+		}
+	}
+	*rrw = append(*rrw, r)
 	return nil
 }
 
-func (r *mockRevisionReadWriter) ReadRevisions(_ context.Context) (migrate.Revisions, error) {
-	return migrate.Revisions(*r), nil
+func (rrw *mockRevisionReadWriter) ReadRevisions(_ context.Context) (migrate.Revisions, error) {
+	return migrate.Revisions(*rrw), nil
 }
 
-func (r *mockRevisionReadWriter) clean() {
-	*r = mockRevisionReadWriter(migrate.Revisions{})
+func (rrw *mockRevisionReadWriter) clean() {
+	*rrw = mockRevisionReadWriter(migrate.Revisions{})
 }
 
 type mockLogger []migrate.LogEntry