sql/migrate: fixed delimiter on the same line with the line comment

ariga · Jan 21, 2025 · 9de8038 · 9de8038
1 parent 18b7894
commit 9de8038
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 1 deletion.
diff --git a/sql/migrate/dir.go b/sql/migrate/dir.go
@@ -552,7 +552,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 }}{{ $.Statement . }}{{ end }}`,
 			)),
 		},
 	}

diff --git a/sql/migrate/dir_test.go b/sql/migrate/dir_test.go
@@ -80,6 +80,32 @@ func TestHashSum(t *testing.T) {
 	require.NotContains(t, string(c), "exclude_2.sql")
 }
 
+func TestPlanDelimiter(t *testing.T) {
+	// Sum file gets created.
+	d := &migrate.MemDir{}
+	pl := migrate.NewPlanner(nil, d, migrate.PlanWithChecksum(false))
+	require.NotNil(t, pl)
+	plan := &migrate.Plan{
+		Version: "1",
+		Name:    "plan",
+		Changes: []*migrate.Change{{Cmd: "cmd\n-- comment"}},
+	}
+	require.NoError(t, pl.WritePlan(plan))
+
+	files, err := d.Files()
+	require.NoError(t, err)
+	require.Equal(t, 1, len(files))
+	require.Equal(t, "cmd\n-- comment\n;\n", string(files[0].Bytes()))
+
+	plan.Delimiter = "\nGO"
+	require.NoError(t, pl.WritePlan(plan))
+
+	files, err = d.Files()
+	require.NoError(t, err)
+	require.Equal(t, 1, len(files))
+	require.Equal(t, "-- atlas:delimiter \\nGO\n\ncmd\n-- comment\n\nGO\n", string(files[0].Bytes()))
+}
+
 var (
 	//go:embed testdata/migrate/atlas.sum
 	hash []byte

diff --git a/sql/migrate/migrate.go b/sql/migrate/migrate.go
@@ -59,6 +59,19 @@ type (
 	}
 )
 
+// Statement format the change's command
+func (p Plan) Statement(c *Change) string {
+	d := p.Delimiter
+	if d == "" {
+		d = ";"
+	}
+	if nl := strings.LastIndexByte(c.Cmd, '\n'); nl != -1 && strings.HasPrefix(c.Cmd[nl:], "\n-- ") {
+		// Put delimiter on the newline, to separate it from the line comment
+		return fmt.Sprintf("%s\n%s\n", c.Cmd, d)
+	}
+	return fmt.Sprintf("%s%s\n", c.Cmd, d)
+}
+
 // ReverseStmts returns the reverse statements of a Change, if any.
 func (c *Change) ReverseStmts() (cmd []string, err error) {
 	switch r := c.Reverse.(type) {