From 9f229dde97c936c3cba07df4e00882cd5827374f Mon Sep 17 00:00:00 2001
From: Aneurin Barker Snook <aneurin@aneur.in>
Date: Mon, 6 Jul 2026 22:04:36 +0100
Subject: [PATCH] improve patch callback in migrations

---
 migration.go | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/migration.go b/migration.go
index 934c9cb..4f6e4df 100644
--- a/migration.go
+++ b/migration.go
@@ -4,11 +4,14 @@ import (
 	"sort"
 )
 
+// AfterPatchCallback is called after each time a migration is successfully applied (up) or reverted (down).
+type AfterPatchCallback func(applied, next *Version)
+
 // Migration maps provide a simple way to store and run versioned 'patch' functions.
 type Migration map[string]Patch
 
 // Down reverts all patches in reverse version order.
-func (m Migration) Down(after func(*Version)) error {
+func (m Migration) Down(after AfterPatchCallback) error {
 	versions, err := m.Versions()
 	if err != nil {
 		return err
@@ -22,7 +25,12 @@ func (m Migration) Down(after func(*Version)) error {
 		}
 
 		if after != nil {
-			after(v)
+			var next *Version
+			if i > 1 {
+				next = versions[i-1]
+			}
+
+			after(v, next)
 		}
 	}
 
@@ -47,20 +55,25 @@ func (m Migration) Match(constraints ...*Constraint) (Migration, error) {
 }
 
 // Up executes all patches in version order.
-func (m Migration) Up(after func(*Version)) error {
+func (m Migration) Up(after AfterPatchCallback) error {
 	versions, err := m.Versions()
 	if err != nil {
 		return err
 	}
 
-	for _, v := range versions {
+	for i, v := range versions {
 		patch := m[v.Text]
 		if err := patch.Up(); err != nil {
 			return err
 		}
 
 		if after != nil {
-			after(v)
+			var next *Version
+			if i < len(versions)-1 {
+				next = versions[i+1]
+			}
+
+			after(v, next)
 		}
 	}