Repository: ariga/atlas
Pull request #3027 (2024-08-06 06:32:15)

sql/migrate: add FormatFile for TemplateFormatter
Avoid `len` checks whenever `Format` is called on the default template. 


Patch:
diff --git a/cmd/atlas/internal/cmdext/reader.go b/cmd/atlas/internal/cmdext/reader.go
index 0957559f..99ac81f1 100644
--- a/cmd/atlas/internal/cmdext/reader.go
+++ b/cmd/atlas/internal/cmdext/reader.go
@@ -44,10 +44,11 @@ type (
 )
 
 // Close redirects calls to Close to the enclosed io.Closer.
-func (r *StateReadCloser) Close() {
+func (r *StateReadCloser) Close() error {
 	if r.Closer != nil {
-		r.Closer.Close()
+		return r.Closer.Close()
 	}
+	return nil
 }
 
 // StateReaderSQL returns a migrate.StateReader from an SQL file or a directory of migrations.
diff --git a/cmd/atlas/internal/cmdlog/cmdlog.go b/cmd/atlas/internal/cmdlog/cmdlog.go
index db62c404..eb99cbc8 100644
--- a/cmd/atlas/internal/cmdlog/cmdlog.go
+++ b/cmd/atlas/internal/cmdlog/cmdlog.go
@@ -31,6 +31,7 @@ import (
 
 var (
 	ColorCyan         = color.CyanString
+	ColorGray         = color.New(color.Attribute(90))
 	ColorGreen        = color.HiGreenString
 	ColorRed          = color.HiRedString
 	ColorRedBgWhiteFg = color.New(color.FgHiWhite, color.BgHiRed).SprintFunc()
diff --git a/cmd/atlas/internal/migratelint/run.go b/cmd/atlas/internal/migratelint/run.go
index 74568249..e2139988 100644
--- a/cmd/atlas/internal/migratelint/run.go
+++ b/cmd/atlas/internal/migratelint/run.go
@@ -16,7 +16,6 @@ import (
 	"text/template"
 	"time"
 
-	"ariga.io/atlas/cmd/atlas/internal/cmdlog"
 	"ariga.io/atlas/sql/migrate"
 	"ariga.io/atlas/sql/sqlcheck"
 	"ariga.io/atlas/sql/sqlclient"
@@ -187,7 +186,7 @@ func (r *Runner) analyze(ctx context.Context, files []*sqlcheck.File) error {
 
 var (
 	// TemplateFuncs are global functions available in templates.
-	TemplateFuncs = cmdlog.WithColorFuncs(template.FuncMap{
+	TemplateFuncs = template.FuncMap{
 		"json": func(v any, args ...string) (string, error) {
 			var (
 				b   []byte
@@ -226,7 +225,12 @@ var (
 			}
 			return append(lines, strings.Join(words[j:], " "))
 		},
-	})
+		"cyan":         color.CyanString,
+		"green":        color.HiGreenString,
+		"red":          color.HiRedString,
+		"redBgWhiteFg": color.New(color.FgHiWhite, color.BgHiRed).SprintFunc(),
+		"yellow":       color.YellowString,
+	}
 	// DefaultTemplate is the default template used by the CI job.
 	DefaultTemplate = template.Must(template.New("report").
 			Funcs(TemplateFuncs).
diff --git a/sql/migrate/dir.go b/sql/migrate/dir.go
index 6f0f89dc..6c160f3c 100644
--- a/sql/migrate/dir.go
+++ b/sql/migrate/dir.go
@@ -561,6 +561,18 @@ func (t TemplateFormatter) Format(plan *Plan) ([]File, error) {
 	return files, nil
 }
 
+// FormatFile is like Format, but expects and returns a single file.
+func (t TemplateFormatter) FormatFile(p *Plan) (File, error) {
+	files, err := t.Format(p)
+	if err != nil {
+		return nil, err
+	}
+	if len(files) != 1 {
+		return nil, fmt.Errorf("expected a single file, got %d", len(files))
+	}
+	return files[0], nil
+}
+
 // FormatTo calls Format and writes the files' content to the given writer.
 func (t TemplateFormatter) FormatTo(plan *Plan, w io.Writer) error {
 	files, err := t.Format(plan)


Test patch:
diff --git a/sql/migrate/dir_test.go b/sql/migrate/dir_test.go
index f2d06cfe..3170fc14 100644
--- a/sql/migrate/dir_test.go
+++ b/sql/migrate/dir_test.go
@@ -13,6 +13,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"text/template"
 	"time"
 
 	"ariga.io/atlas/sql/migrate"
@@ -546,6 +547,34 @@ create table t2(c int);
 `, b.String())
 }
 
+func TestDefaultFormatter_FormatFile(t *testing.T) {
+	f, err := migrate.DefaultFormatter.FormatFile(&migrate.Plan{
+		Changes: []*migrate.Change{
+			{Cmd: "create table t1(c int)"},
+			{Cmd: "create table t2(c int)", Comment: "create table"},
+		},
+	})
+	require.NoError(t, err)
+	require.Equal(t, `create table t1(c int);
+-- Create table
+create table t2(c int);
+`, string(f.Bytes()))
+
+	ft := migrate.TemplateFormatter{
+		{
+			N: template.Must(template.New("").Parse("name-1")),
+			C: template.Must(template.New("").Parse("cmd-1")),
+		},
+		{
+			N: template.Must(template.New("").Parse("name-2")),
+			C: template.Must(template.New("").Parse("cmd-2")),
+		},
+	}
+	f, err = ft.FormatFile(&migrate.Plan{})
+	require.Error(t, err, "expect only one file")
+	require.Nil(t, f)
+}
+
 func fileNames(r io.Reader) ([]string, error) {
 	var out []string
 	tr := tar.NewReader(r)