Repository: ariga/atlas
Pull request #3024 (2024-08-05 14:33:13)

cmd/atlas/internal/cmdapi: run upgrade suggestion in logic


Patch:
diff --git a/cmd/atlas/internal/cmdapi/cmdapi_oss.go b/cmd/atlas/internal/cmdapi/cmdapi_oss.go
index dfef514a..1f2a2ffc 100644
--- a/cmd/atlas/internal/cmdapi/cmdapi_oss.go
+++ b/cmd/atlas/internal/cmdapi/cmdapi_oss.go
@@ -40,9 +40,6 @@ func init() {
 		schemaInspectCmd(),
 		unsupportedCommand("schema", "test"),
 	)
-	schemaCmd.PersistentPreRun = func(cmd *cobra.Command, _ []string) {
-		maySuggestUpgrade(cmd)
-	}
 	Root.AddCommand(schemaCmd)
 	migrateCmd := migrateCmd()
 	migrateCmd.AddCommand(
@@ -63,9 +60,6 @@ func init() {
 		unsupportedCommand("migrate", "push"),
 		unsupportedCommand("migrate", "test"),
 	)
-	migrateCmd.PersistentPreRun = func(cmd *cobra.Command, _ []string) {
-		maySuggestUpgrade(cmd)
-	}
 	Root.AddCommand(migrateCmd)
 }
 
diff --git a/cmd/atlas/internal/cmdapi/schema.go b/cmd/atlas/internal/cmdapi/schema.go
index 51c1ba71..34485f78 100644
--- a/cmd/atlas/internal/cmdapi/schema.go
+++ b/cmd/atlas/internal/cmdapi/schema.go
@@ -186,6 +186,7 @@ func schemaApplyRun(cmd *cobra.Command, flags schemaApplyFlags, env *Env) error
 	if err != nil {
 		return err
 	}
+	maySuggestUpgrade(cmd)
 	// Returning at this stage should
 	// not trigger the help message.
 	cmd.SilenceUsage = true
@@ -438,6 +439,7 @@ func schemaDiffRun(cmd *cobra.Command, _ []string, flags schemaDiffFlags, env *E
 	if err != nil {
 		return err
 	}
+	maySuggestUpgrade(cmd)
 	return format.Execute(cmd.OutOrStdout(),
 		cmdlog.NewSchemaDiff(ctx, c, diff.from, diff.to, diff.changes),
 	)
@@ -544,6 +546,7 @@ func schemaInspectRun(cmd *cobra.Command, _ []string, flags schemaInspectFlags,
 	if err != nil {
 		return err
 	}
+	maySuggestUpgrade(cmd)
 	return format.Execute(cmd.OutOrStdout(), cmdlog.NewSchemaInspect(ctx, client, s))
 }
 
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 430dcb83..eb99cbc8 100644
--- a/cmd/atlas/internal/cmdlog/cmdlog.go
+++ b/cmd/atlas/internal/cmdlog/cmdlog.go
@@ -10,10 +10,13 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/url"
+	"reflect"
 	"slices"
 	"sort"
 	"strings"
+	"sync"
 	"text/template"
 	"time"
 
@@ -28,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()
@@ -1119,3 +1123,20 @@ func indentLn(input string, indent int) string {
 	pad := strings.Repeat(" ", indent)
 	return strings.ReplaceAll(input, "\n", "\n"+pad)
 }
+
+// WarnOnce allow writing warning messages to the given writer,
+// but ensures only one message will be written in process run.
+func WarnOnce(w io.Writer, text string) error {
+	if !reflect.TypeOf(w).Comparable() {
+		_, err := w.Write([]byte(text))
+		return err
+	}
+	if _, loaded := oneWrites.LoadOrStore(w, true); !loaded {
+		_, err := w.Write([]byte(text))
+		return err
+	}
+	return nil
+}
+
+// oneWrites per writer.
+var oneWrites sync.Map
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/cmd/atlas/main.go b/cmd/atlas/main.go
index 98f32ece..160c002a 100644
--- a/cmd/atlas/main.go
+++ b/cmd/atlas/main.go
@@ -16,6 +16,7 @@ import (
 
 	"ariga.io/atlas/cmd/atlas/internal/cmdapi"
 	"ariga.io/atlas/cmd/atlas/internal/cmdapi/vercheck"
+	"ariga.io/atlas/cmd/atlas/internal/cmdlog"
 	_ "ariga.io/atlas/cmd/atlas/internal/docker"
 	_ "ariga.io/atlas/sql/mysql"
 	_ "ariga.io/atlas/sql/mysql/mysqlcheck"
@@ -58,7 +59,7 @@ func main() {
 	update := checkForUpdate(ctx)
 	err = cmdapi.Root.ExecuteContext(ctx)
 	if u := update(); u != "" {
-		fmt.Fprintln(os.Stderr, u)
+		_ = cmdlog.WarnOnce(os.Stderr, u)
 	}
 	done(err)
 	if err != nil {
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index ada57f79..904ef49c 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -5,7 +5,6 @@
 package specutil
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
 	"slices"
@@ -683,7 +682,7 @@ func FromView(v *schema.View, colFn ViewColumnSpecFunc, idxFn IndexSpecFunc) (*s
 	}
 	// In case the view definition is multi-line,
 	// format it as indented heredoc with two spaces.
-	as := MightHeredoc(v.Def)
+	as := sqlspec.MightHeredoc(v.Def)
 	embed := &schemahcl.Resource{
 		Attrs: []*schemahcl.Attr{
 			schemahcl.StringAttr("as", as),
@@ -706,39 +705,6 @@ func FromView(v *schema.View, colFn ViewColumnSpecFunc, idxFn IndexSpecFunc) (*s
 	return spec, nil
 }
 
-// normalizeCRLF for heredoc strings that inspected and printed in the HCL as-is to
-// avoid having mixed endings in the printed file - Unix-like (default) and Windows-like.
-func normalizeCRLF(s string) string {
-	if strings.Contains(s, "\r\n") {
-		return strings.ReplaceAll(s, "\r\n", "\n")
-	}
-	return s
-}
-
-// MightHeredoc returns the string as an indented heredoc if it has multiple lines.
-func MightHeredoc(s string) string {
-	s = normalizeCRLF(strings.TrimSpace(s))
-	// In case the given definition is multi-line,
-	// format it as indented heredoc with two spaces.
-	if lines := strings.Split(s, "\n"); len(lines) > 1 {
-		var b bytes.Buffer
-		b.Grow(len(s))
-		b.WriteString("<<-SQL\n")
-		for _, l := range lines {
-			// Skip spaces-only lines, as editors stripped these spaces off,
-			// and HCL parser results a different string for them.
-			if strings.TrimSpace(l) != "" {
-				b.WriteString("  ")
-				b.WriteString(l)
-			}
-			b.WriteByte('\n')
-		}
-		b.WriteString("  SQL")
-		s = b.String()
-	}
-	return s
-}
-
 // dependsOn returns the depends_on attribute for the given objects.
 func dependsOn(realm *schema.Realm, objects []schema.Object) (*schemahcl.Attr, bool) {
 	var (
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)
diff --git a/sql/sqlspec/sqlspec.go b/sql/sqlspec/sqlspec.go
index 4060da35..12bca2f3 100644
--- a/sql/sqlspec/sqlspec.go
+++ b/sql/sqlspec/sqlspec.go
@@ -5,6 +5,9 @@
 package sqlspec
 
 import (
+	"bytes"
+	"strings"
+
 	"ariga.io/atlas/schemahcl"
 
 	"github.com/zclconf/go-cty/cty"
@@ -192,3 +195,36 @@ func init() {
 	schemahcl.Register("sequence", &Sequence{})
 	schemahcl.Register("schema", &Schema{})
 }
+
+// normalizeCRLF for heredoc strings that inspected and printed in the HCL as-is to
+// avoid having mixed endings in the printed file - Unix-like (default) and Windows-like.
+func normalizeCRLF(s string) string {
+	if strings.Contains(s, "\r\n") {
+		return strings.ReplaceAll(s, "\r\n", "\n")
+	}
+	return s
+}
+
+// MightHeredoc returns the string as an indented heredoc if it has multiple lines.
+func MightHeredoc(s string) string {
+	s = normalizeCRLF(strings.TrimSpace(s))
+	// In case the given definition is multi-line,
+	// format it as indented heredoc with two spaces.
+	if lines := strings.Split(s, "\n"); len(lines) > 1 {
+		var b bytes.Buffer
+		b.Grow(len(s))
+		b.WriteString("<<-SQL\n")
+		for _, l := range lines {
+			// Skip spaces-only lines, as editors stripped these spaces off,
+			// and HCL parser results a different string for them.
+			if strings.TrimSpace(l) != "" {
+				b.WriteString("  ")
+				b.WriteString(l)
+			}
+			b.WriteByte('\n')
+		}
+		b.WriteString("  SQL")
+		s = b.String()
+	}
+	return s
+}


Test patch:
diff --git a/cmd/atlas/internal/cmdlog/cmdlog_test.go b/cmd/atlas/internal/cmdlog/cmdlog_test.go
index ccc6bb6f..14493aac 100644
--- a/cmd/atlas/internal/cmdlog/cmdlog_test.go
+++ b/cmd/atlas/internal/cmdlog/cmdlog_test.go
@@ -9,8 +9,10 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"io"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"text/template"
 	"time"
@@ -909,3 +911,38 @@ func TestReporter_OutOfOrder(t *testing.T) {
   ERROR: migration files 1.5.sql, 1.6.sql were added out of order. See: https://atlasgo.io/versioned/apply#non-linear-error
 `, buf.String())
 }
+
+func TestWarnOnce(t *testing.T) {
+	b := &strings.Builder{}
+	require.NoError(t, cmdlog.WarnOnce(b, "one"))
+	require.NoError(t, cmdlog.WarnOnce(b, "two"))
+	require.Equal(t, "one", b.String())
+
+	var wg sync.WaitGroup
+	b = &strings.Builder{}
+	wg.Add(5)
+	for i := 0; i < 5; i++ {
+		go func() {
+			defer wg.Done()
+			require.NoError(t, cmdlog.WarnOnce(b, "one"))
+		}()
+	}
+	wg.Wait()
+	require.Equal(t, "one", b.String())
+
+	// Ensure the type is not assignable.
+	require.Panics(t, func() {
+		var m sync.Map
+		m.LoadOrStore(unassignable{}, "text")
+	})
+	b.Reset()
+	require.NoError(t, cmdlog.WarnOnce(unassignable{Writer: b}, "done"))
+	require.Equal(t, "done", b.String())
+}
+
+type unassignable struct {
+	_ func()
+	io.Writer
+}
+
+func (a unassignable) Write(p []byte) (n int, err error) { return a.Writer.Write(p) }
diff --git a/sql/internal/specutil/convert_test.go b/sql/internal/specutil/convert_test.go
index dabf339c..70dee7b3 100644
--- a/sql/internal/specutil/convert_test.go
+++ b/sql/internal/specutil/convert_test.go
@@ -180,49 +180,3 @@ func TestFromView(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, "<<-SQL\n  SELECT * FROM users\n   WHERE c NOT LIKE \"\\r\\n\"\n  SQL", s)
 }
-
-func TestMightHeredoc(t *testing.T) {
-	for _, tt := range []struct {
-		input    string
-		expected string
-	}{
-		{
-			input:    "SELECT * FROM users",
-			expected: "SELECT * FROM users",
-		},
-		{
-			input: `
-SELECT
-  *
-  FROM users
-  WHERE active`,
-			expected: `<<-SQL
-  SELECT
-    *
-    FROM users
-    WHERE active
-  SQL`,
-		},
-		{
-			input: `
--- The line below includes spaces.
-  
-	
-SELECT
-  *
-  FROM users
-  WHERE active`,
-			expected: `<<-SQL
-  -- The line below includes spaces.
-
-
-  SELECT
-    *
-    FROM users
-    WHERE active
-  SQL`,
-		},
-	} {
-		require.Equal(t, tt.expected, MightHeredoc(tt.input))
-	}
-}
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)
diff --git a/sql/sqlspec/sqlspec_test.go b/sql/sqlspec/sqlspec_test.go
new file mode 100644
index 00000000..a8425236
--- /dev/null
+++ b/sql/sqlspec/sqlspec_test.go
@@ -0,0 +1,57 @@
+// Copyright 2021-present The Atlas Authors. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package sqlspec
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMightHeredoc(t *testing.T) {
+	for _, tt := range []struct {
+		input    string
+		expected string
+	}{
+		{
+			input:    "SELECT * FROM users",
+			expected: "SELECT * FROM users",
+		},
+		{
+			input: `
+SELECT
+  *
+  FROM users
+  WHERE active`,
+			expected: `<<-SQL
+  SELECT
+    *
+    FROM users
+    WHERE active
+  SQL`,
+		},
+		{
+			input: `
+-- The line below includes spaces.
+  
+	
+SELECT
+  *
+  FROM users
+  WHERE active`,
+			expected: `<<-SQL
+  -- The line below includes spaces.
+
+
+  SELECT
+    *
+    FROM users
+    WHERE active
+  SQL`,
+		},
+	} {
+		require.Equal(t, tt.expected, MightHeredoc(tt.input))
+	}
+}