From 2cac8fb75c7d39ba6ca1749eee078e2664883b8f Mon Sep 17 00:00:00 2001
From: "Giau. Tran Minh" <hello@giautm.dev>
Date: Wed, 22 Jan 2025 01:11:16 +0700
Subject: [PATCH] sql/migrate: fixed delimiter on the same line with the line
 comment

---
 sql/migrate/dir.go      | 18 +++++++++++++++++-
 sql/migrate/dir_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/sql/migrate/dir.go b/sql/migrate/dir.go
index 9b3cba4b601..f48b34197b9 100644
--- a/sql/migrate/dir.go
+++ b/sql/migrate/dir.go
@@ -538,12 +538,28 @@ func delim(s string) string {
 	return fmt.Sprintf("-- atlas:%s %s", directiveDelimiter, s)
 }
 
+// stmt returns a statement with the given delimiter.
+func stmt(cmd, delim string) string {
+	if cmd == "" {
+		return ""
+	}
+	if delim == "" {
+		delim = delimiter
+	}
+	if nl := strings.LastIndexByte(cmd, '\n'); nl != -1 && strings.HasPrefix(cmd[nl+1:], directivePrefixSQL) {
+		// Put delimiter on the newline, to separate it from the line comment
+		return fmt.Sprintf("%s\n%s\n", cmd, delim)
+	}
+	return fmt.Sprintf("%s%s\n", cmd, delim)
+}
+
 var (
 	// templateFunc contains the template.FuncMap for the DefaultFormatter.
 	templateFuncs = template.FuncMap{
 		"upper": strings.ToUpper,
 		"now":   NewVersion,
 		"delim": delim,
+		"stmt":  stmt,
 	}
 	// DefaultFormatter is a default implementation for Formatter.
 	DefaultFormatter = TemplateFormatter{
@@ -552,7 +568,7 @@ var (
 				"{{ with .Version }}{{ . }}{{ else }}{{ now }}{{ end }}{{ with .Name }}_{{ . }}{{ end }}.sql",
 			)),
 			C: template.Must(template.New("").Funcs(templateFuncs).Parse(
-				`{{ with .Delimiter }}{{ delim . | printf "%s\n\n" }}{{ end }}{{ range .Changes }}{{ with .Comment }}{{ printf "-- %s%s\n" (slice . 0 1 | upper ) (slice . 1) }}{{ end }}{{ printf "%s%s\n" .Cmd (or $.Delimiter ";") }}{{ end }}`,
+				`{{ with .Delimiter }}{{ delim . | printf "%s\n\n" }}{{ end }}{{ range .Changes }}{{ with .Comment }}{{ printf "-- %s%s\n" (slice . 0 1 | upper ) (slice . 1) }}{{ end }}{{ stmt .Cmd $.Delimiter }}{{ end }}`,
 			)),
 		},
 	}
diff --git a/sql/migrate/dir_test.go b/sql/migrate/dir_test.go
index d330dff27e0..3cd8ca35f08 100644
--- a/sql/migrate/dir_test.go
+++ b/sql/migrate/dir_test.go
@@ -80,6 +80,37 @@ func TestHashSum(t *testing.T) {
 	require.NotContains(t, string(c), "exclude_2.sql")
 }
 
+func TestPlanDelimiter(t *testing.T) {
+	var cases = []struct {
+		Cmd       string
+		Delimiter string
+		Expected  string
+	}{
+		{"", "", ""},
+		{"\n", "", "\n;\n"},
+		{"cmd\n-- comment", "", "cmd\n-- comment\n;\n"},
+		{"cmd\n-- comment", "\nGO", "-- atlas:delimiter \\nGO\n\ncmd\n-- comment\n\nGO\n"},
+	}
+	d := &migrate.MemDir{}
+	pl := migrate.NewPlanner(nil, d, migrate.PlanWithChecksum(false))
+	require.NotNil(t, pl)
+	for _, c := range cases {
+		plan := &migrate.Plan{
+			Version:   "1",
+			Name:      "plan",
+			Delimiter: c.Delimiter,
+			Changes: []*migrate.Change{{
+				Cmd: c.Cmd,
+			}},
+		}
+		require.NoError(t, pl.WritePlan(plan))
+		files, err := d.Files()
+		require.NoError(t, err)
+		require.Equal(t, 1, len(files))
+		require.Equal(t, c.Expected, string(files[0].Bytes()))
+	}
+}
+
 var (
 	//go:embed testdata/migrate/atlas.sum
 	hash []byte