Add configurable string based migration splitting

database/postgres/README.md

@@ -37,3 +37,18 @@ In PostgreSQL running multiple SQL statements in one `Exec` executes them inside
 behavior is not desirable because some statements can be only run outside of transaction (e.g.
 `CREATE INDEX CONCURRENTLY`). If you want to use `CREATE INDEX CONCURRENTLY` without activating multi-statement mode
 you have to put such statements in a separate migration files.
+
+### Custom migration splitter
+
+If your migrations contain SQL bodies with embedded semicolons (e.g. `DO $$ BEGIN ... END $$;`), the
+semicolon-based splitting used by `x-multi-statement` will produce incorrect statement boundaries. Instead, you
+can use the CLI `-migration-splitter` flag to split a migration into steps using a delimiter line you control:
+
+```
+migrate -database "$DATABASE_URL" -source "file://migrations" \
+  -migration-splitter '---' up
+```
+
+Place a line containing only `---` between each step in your migration file. The splitter matches when
+a whole line exactly equals the token (supports both LF and CRLF line endings), and each step is
+executed sequentially in a separate database call.

internal/cli/main.go

@@ -68,21 +68,23 @@ func Main(version string) {
 	pathPtr := flag.String("path", "", "")
 	databasePtr := flag.String("database", "", "")
 	sourcePtr := flag.String("source", "", "")
+	migrationSplitterPtr := flag.String("migration-splitter", "", "")
 
 	flag.Usage = func() {
 		fmt.Fprintf(os.Stderr,
 			`Usage: migrate OPTIONS COMMAND [arg...]
        migrate [ -version | -help ]
 
 Options:
-  -source          Location of the migrations (driver://url)
-  -path            Shorthand for -source=file://path
-  -database        Run migrations against this database (driver://url)
-  -prefetch N      Number of migrations to load in advance before executing (default 10)
-  -lock-timeout N  Allow N seconds to acquire database lock (default 15)
-  -verbose         Print verbose logging
-  -version         Print version
-  -help            Print usage
+  -source               Location of the migrations (driver://url)
+  -path                 Shorthand for -source=file://path
+  -database             Run migrations against this database (driver://url)
+  -prefetch N           Number of migrations to load in advance before executing (default 10)
+  -lock-timeout N       Allow N seconds to acquire database lock (default 15)
+  -migration-splitter   Split migrations into sequential steps when a whole line equals this string
+  -verbose              Print verbose logging
+  -version              Print version
+  -help                 Print usage
 
 Commands:
   %s
@@ -134,6 +136,9 @@ Database drivers: `+strings.Join(database.List(), ", ")+"\n", createUsage, gotoU
 		migrater.Log = log
 		migrater.PrefetchMigrations = *prefetchPtr
 		migrater.LockTimeout = time.Duration(int64(*lockTimeoutPtr)) * time.Second
+		if *migrationSplitterPtr != "" {
+			migrater.MigrationSplitter = []byte(*migrationSplitterPtr)
+		}
 
 		// handle Ctrl+c
 		signals := make(chan os.Signal, 1)

migrate.go

@@ -5,8 +5,10 @@
 package migrate
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"sync"
 	"time"
@@ -78,6 +80,12 @@ type Migrate struct {
 	// LockTimeout defaults to DefaultLockTimeout,
 	// but can be set per Migrate instance.
 	LockTimeout time.Duration
+
+	// MigrationSplitter, if non-empty, splits each migration into sequential steps.
+	// A splitter matches only when a full logical line equals this token.
+	// Both LF and CRLF input line endings are recognized.
+	// If nil or empty, each migration is passed to the driver as-is.
+	MigrationSplitter []byte
 }
 
 // New returns a new Migrate instance from a source URL and a database URL.
@@ -741,8 +749,20 @@ func (m *Migrate) runMigrations(ret <-chan interface{}) error {
 
 			if migr.Body != nil {
 				m.logVerbosePrintf("Read and execute %v\n", migr.LogString())
-				if err := m.databaseDrv.Run(migr.BufferedBody); err != nil {
-					return err
+				if len(m.MigrationSplitter) == 0 {
+					if err := m.databaseDrv.Run(migr.BufferedBody); err != nil {
+						return err
+					}
+				} else {
+					content, err := io.ReadAll(migr.BufferedBody)
+					if err != nil {
+						return err
+					}
+					for _, step := range splitMigrationSteps(content, m.MigrationSplitter) {
+						if err := m.databaseDrv.Run(bytes.NewReader(step)); err != nil {
+							return err
+						}
+					}
 				}
 			}
 
@@ -771,6 +791,33 @@ func (m *Migrate) runMigrations(ret <-chan interface{}) error {
 	return nil
 }
 
+func splitMigrationSteps(content, splitter []byte) [][]byte {
+	if len(splitter) == 0 {
+		return [][]byte{content}
+	}
+
+	steps := make([][]byte, 0, 1)
+	stepStart := 0
+	offset := 0
+
+	for offset < len(content) {
+		lineStart := offset
+		lineEnd := bytes.IndexByte(content[offset:], '\n')
+		if lineEnd < 0 {
+			offset = len(content)
+		} else {
+			offset += lineEnd + 1
+		}
+
+		if bytes.Equal(bytes.TrimRight(content[lineStart:offset], "\r\n"), splitter) {
+			steps = append(steps, content[stepStart:lineStart])
+			stepStart = offset
+		}
+	}
+
+	return append(steps, content[stepStart:])
+}
+
 // versionExists checks the source if either the up or down migration for
 // the specified migration version exists.
 func (m *Migrate) versionExists(version uint) (result error) {

migrate_test.go

@@ -1409,3 +1409,193 @@ func equalDbSeq(t *testing.T, i int, expected migrationSequence, got *dStub.Stub
 		t.Fatalf("\nexpected sequence %v,\ngot               %v, in %v", bs, got.MigrationSequence, i)
 	}
 }
+
+// newMigrateWithContent creates a Migrate instance backed by stub drivers.
+// The single up migration at version 1 has the given content as its body.
+func newMigrateWithContent(t *testing.T, content string) (*Migrate, *dStub.Stub) {
+	t.Helper()
+
+	migrations := source.NewMigrations()
+	migrations.Append(&source.Migration{Version: 1, Direction: source.Up, Identifier: content})
+
+	m, err := New("stub://", "stub://")
+	if err != nil {
+		t.Fatal(err)
+	}
+	m.sourceDrv.(*sStub.Stub).Migrations = migrations
+
+	dbDrv := m.databaseDrv.(*dStub.Stub)
+	return m, dbDrv
+}
+
+var errEmptyMigrationBody = errors.New("empty migration body")
+
+type emptyRejectingStub struct {
+	*dStub.Stub
+}
+
+func (s *emptyRejectingStub) Run(migration io.Reader) error {
+	m, err := io.ReadAll(migration)
+	if err != nil {
+		return err
+	}
+	if len(bytes.TrimSpace(m)) == 0 {
+		return errEmptyMigrationBody
+	}
+	return s.Stub.Run(bytes.NewReader(m))
+}
+
+// TestMigrationSplitterSplitsIntoSteps verifies that when MigrationSplitter is
+// set, each step between delimiter lines is passed to the driver as a
+// separate Run call.
+func TestMigrationSplitterSplitsIntoSteps(t *testing.T) {
+	const content = "stmt1;\n---\nstmt2;"
+	m, dbDrv := newMigrateWithContent(t, content)
+	m.MigrationSplitter = []byte("---")
+
+	if err := m.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	want := []string{"stmt1;\n", "stmt2;"}
+	if !dbDrv.EqualSequence(want) {
+		t.Errorf("MigrationSequence = %q, want %q", dbDrv.MigrationSequence, want)
+	}
+}
+
+// TestMigrationSplitterDollarQuotedBody verifies that a migration containing a
+// DO $$ ... $$; block (which contains semicolons) is split correctly when a
+// newline-anchored delimiter is used, so the inner semicolons are never split on.
+func TestMigrationSplitterDollarQuotedBody(t *testing.T) {
+	const content = "CREATE TABLE foo (id INT);\n---\nDO $$ BEGIN RAISE NOTICE 'hi'; END $$;"
+	m, dbDrv := newMigrateWithContent(t, content)
+	m.MigrationSplitter = []byte("---")
+
+	if err := m.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	want := []string{
+		"CREATE TABLE foo (id INT);\n",
+		"DO $$ BEGIN RAISE NOTICE 'hi'; END $$;",
+	}
+	if !dbDrv.EqualSequence(want) {
+		t.Errorf("MigrationSequence = %q, want %q", dbDrv.MigrationSequence, want)
+	}
+}
+
+// TestMigrationSplitterNilPreservesExistingBehavior verifies that when
+// MigrationSplitter is nil (the default), the migration body is passed to the
+// driver in a single Run call with its content unchanged.
+func TestMigrationSplitterNilPreservesExistingBehavior(t *testing.T) {
+	const content = "stmt1; stmt2;"
+	m, dbDrv := newMigrateWithContent(t, content)
+	// MigrationSplitter is nil by default — no split should occur.
+
+	if err := m.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	want := []string{content}
+	if !dbDrv.EqualSequence(want) {
+		t.Errorf("MigrationSequence = %q, want %q", dbDrv.MigrationSequence, want)
+	}
+}
+
+// TestMigrationSplitterNilDoesNotSkipEmptyMigration verifies that when
+// MigrationSplitter is nil, empty migration bodies are still passed through.
+func TestMigrationSplitterNilDoesNotSkipEmptyMigration(t *testing.T) {
+	m, _ := newMigrateWithContent(t, "")
+	m.databaseDrv = &emptyRejectingStub{Stub: m.databaseDrv.(*dStub.Stub)}
+
+	err := m.Up()
+	if !errors.Is(err, errEmptyMigrationBody) {
+		t.Fatalf("expected %v, got %v", errEmptyMigrationBody, err)
+	}
+}
+
+// TestMigrationSplitterAbsentInContent verifies that when the splitter is set
+// but does not appear in the migration body, the entire body is passed as a
+// single statement.
+func TestMigrationSplitterAbsentInContent(t *testing.T) {
+	const content = "CREATE TABLE bar (id INT);"
+	m, dbDrv := newMigrateWithContent(t, content)
+	m.MigrationSplitter = []byte("---")
+
+	if err := m.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	want := []string{content}
+	if !dbDrv.EqualSequence(want) {
+		t.Errorf("MigrationSequence = %q, want %q", dbDrv.MigrationSequence, want)
+	}
+}
+
+// TestMigrationSplitterEmptyPreservesExistingBehavior verifies that an empty
+// splitter is treated as disabled and the full migration body is executed once.
+func TestMigrationSplitterEmptyPreservesExistingBehavior(t *testing.T) {
+	const content = "stmt1; stmt2;"
+	m, dbDrv := newMigrateWithContent(t, content)
+	m.MigrationSplitter = []byte{}
+
+	if err := m.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	want := []string{content}
+	if !dbDrv.EqualSequence(want) {
+		t.Errorf("MigrationSequence = %q, want %q", dbDrv.MigrationSequence, want)
+	}
+}
+
+// TestMigrationSplitterIncludesWhitespaceOnlySteps verifies that fragments
+// containing only whitespace between splitter lines are executed as-is.
+func TestMigrationSplitterIncludesWhitespaceOnlySteps(t *testing.T) {
+	const content = "stmt1;\n---\n \t \r\n---\nstmt2;"
+	m, dbDrv := newMigrateWithContent(t, content)
+	m.MigrationSplitter = []byte("---")
+
+	if err := m.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	want := []string{"stmt1;\n", " \t \r\n", "stmt2;"}
+	if !dbDrv.EqualSequence(want) {
+		t.Errorf("MigrationSequence = %q, want %q", dbDrv.MigrationSequence, want)
+	}
+}
+
+// TestMigrationSplitterSupportsCRLF verifies that splitter lines are detected
+// in CRLF-formatted files.
+func TestMigrationSplitterSupportsCRLF(t *testing.T) {
+	const content = "stmt1;\r\n---\r\nstmt2;"
+	m, dbDrv := newMigrateWithContent(t, content)
+	m.MigrationSplitter = []byte("---")
+
+	if err := m.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	want := []string{"stmt1;\r\n", "stmt2;"}
+	if !dbDrv.EqualSequence(want) {
+		t.Errorf("MigrationSequence = %q, want %q", dbDrv.MigrationSequence, want)
+	}
+}
+
+// TestMigrationSplitterAtFileBoundaries verifies splitter lines are recognized
+// at both the start and end of a file, including empty boundary fragments.
+func TestMigrationSplitterAtFileBoundaries(t *testing.T) {
+	const content = "---\nstmt1;\n---\nstmt2;\n---"
+	m, dbDrv := newMigrateWithContent(t, content)
+	m.MigrationSplitter = []byte("---")
+
+	if err := m.Up(); err != nil {
+		t.Fatal(err)
+	}
+
+	want := []string{"", "stmt1;\n", "stmt2;\n", ""}
+	if !dbDrv.EqualSequence(want) {
+		t.Errorf("MigrationSequence = %q, want %q", dbDrv.MigrationSequence, want)
+	}
+}