sql/migrate: support custom Delimiter in Plan (#3238)

* sql/migrate: support custom Delimiter in Plan

* chore: fixed review comment

Author: giautm

sql/migrate/dir.go

@@ -528,11 +528,22 @@ func CheckVersion(v string) error {
 	return nil
 }
 
+// delim returns a directive with the given delimiter.
+func delim(s string) string {
+	if s == "" {
+		return ""
+	}
+	// Escape delimiters. e.g. "\n" => "\\n".
+	s = strings.NewReplacer("\n", `\n`, "\r", `\r`, "\t", `\t`).Replace(s)
+	return fmt.Sprintf("-- atlas:%s %s", directiveDelimiter, s)
+}
+
 var (
 	// templateFunc contains the template.FuncMap for the DefaultFormatter.
 	templateFuncs = template.FuncMap{
 		"upper": strings.ToUpper,
 		"now":   NewVersion,
+		"delim": delim,
 	}
 	// DefaultFormatter is a default implementation for Formatter.
 	DefaultFormatter = TemplateFormatter{
@@ -541,7 +552,7 @@ var (
 				"{{ with .Version }}{{ . }}{{ else }}{{ now }}{{ end }}{{ with .Name }}_{{ . }}{{ end }}.sql",
 			)),
 			C: template.Must(template.New("").Funcs(templateFuncs).Parse(
-				`{{ range .Changes }}{{ with .Comment }}{{ printf "-- %s%s\n" (slice . 0 1 | upper ) (slice . 1) }}{{ end }}{{ printf "%s;\n" .Cmd }}{{ end }}`,
+				`{{ 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 }}`,
 			)),
 		},
 	}

sql/migrate/migrate.go

@@ -34,6 +34,9 @@ type (
 
 		// Changes defines the list of changeset in the plan.
 		Changes []*Change
+
+		// Delimiter to use for separating statements.
+		Delimiter string
 	}
 
 	// A Change of migration.

sql/migrate/migrate_test.go

@@ -74,6 +74,15 @@ func TestPlanner_WritePlan(t *testing.T) {
 	require.Equal(t, countFiles(t, d), 3)
 	requireFileEqual(t, d, "add_t1_and_t2.up.sql", "CREATE TABLE t1(c int)\nCREATE TABLE t2(c int)\n")
 	requireFileEqual(t, d, "add_t1_and_t2.down.sql", "DROP TABLE t1 IF EXISTS\nDROP TABLE t2\n")
+
+	// With custom delimiter.
+	plan.Delimiter = "\nGO"
+	pl = migrate.NewPlanner(nil, d, migrate.PlanWithChecksum(false))
+	require.NotNil(t, pl)
+	require.NoError(t, pl.WritePlan(plan))
+	v = time.Now().UTC().Format("20060102150405")
+	require.Equal(t, countFiles(t, d), 3)
+	requireFileEqual(t, d, v+"_add_t1_and_t2.sql", "-- atlas:delimiter \\nGO\n\nCREATE TABLE t1(c int)\nGO\nCREATE TABLE t2(c int)\nGO\n")
 }
 
 func TestPlanner_WriteCheckpoint(t *testing.T) {