Repository: ariga/atlas
Pull request #566 (2022-02-11 08:53:12)

sql/migrate: add template formatter
This PR adds `TemplateFormatter` which implements `Formatter` and uses `text/template.Template`s to define how migrations files are written.

A default formatter is given, that is compatible with [golang-migrate/migrate](https://github.com/golang-migrate/migrate).


Patch:
diff --git a/doc/md/ddl/sql.md b/doc/md/ddl/sql.md
index b679ca15..628423f2 100644
--- a/doc/md/ddl/sql.md
+++ b/doc/md/ddl/sql.md
@@ -14,18 +14,23 @@ state of a SQL database.
 ### Schema
 
 ```hcl
-schema "default" {}
+# Empty schema.
+schema "public" {}
+
+# Schema with attributes.
+schema "main" {
+  charset   = "utf8mb4"
+  collation = "utf8mb4_0900_ai_ci"
+}
 ```
 
 ### Table
 
-A `table` describes a table in a SQL database. 
-
-#### Example
+A `table` describes a table in a SQL database.
 
 ```hcl
 table "users" {
-  schema = schema.default
+  schema = schema.public
   column "id" {
     type = int
   }
@@ -175,12 +180,18 @@ index "idx_name" {
     }
     unique = true
 }
+
+index "idx_name" {
+  type = GIN
+  columns = [column.data]
+}
 ```
 
 #### Properties
 
-| Name      | Kind      | Type                    | Description                                                  |
-|-----------|-----------|-------------------------|--------------------------------------------------------------|
-| unique    | attribute | boolean                 | Defines whether a uniqueness constraint is set on the index. |
-| columns   | attribute | reference (list)        | The columns that comprise the index.                         |
-| on        | attribute | schema.IndexPart (list) | The index parts that comprise the index.                     |
+| Name      | Kind      | Type                    | Description                                                    |
+|-----------|-----------|-------------------------|----------------------------------------------------------------|
+| unique    | attribute | boolean                 | Defines whether a uniqueness constraint is set on the index.   |
+| type      | attribute | IndexType (enum)        | Defines the index type. e.g. `HASH`, `GIN`, `FULLTEXT`.        |
+| columns   | attribute | reference (list)        | The columns that comprise the index.                           |
+| on        | resource  | schema.IndexPart (list) | The index parts that comprise the index.                       |
diff --git a/sql/migrate/migrate.go b/sql/migrate/migrate.go
index 3c3e4db4..366254bb 100644
--- a/sql/migrate/migrate.go
+++ b/sql/migrate/migrate.go
@@ -5,12 +5,16 @@
 package migrate
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
+	"strings"
+	"text/template"
+	"time"
 
 	"ariga.io/atlas/sql/schema"
 )
@@ -159,28 +163,6 @@ func New(drv Driver, dir Dir, fmt Formatter) *Planner {
 	}
 }
 
-// LocalDir implements Dir for a local path.
-type LocalDir struct {
-	dir string
-}
-
-// Open implements Dir.Open.
-func (dir *LocalDir) Open(name string) (File, error) {
-	return os.Create(filepath.Join(dir.dir, name))
-}
-
-// NewLocalDir returns a new the Dir used by a Planner to work on the given local path.
-func NewLocalDir(path string) (*LocalDir, error) {
-	fi, err := os.Stat(path)
-	if err != nil {
-		return nil, err
-	}
-	if !fi.IsDir() {
-		return nil, fmt.Errorf("sql/migrate: %q is not a dir", path)
-	}
-	return &LocalDir{dir: path}, nil
-}
-
 // Plan calculates the migration Plan required for moving the current state (from) state to
 // the next state (to). A StateReader can be a directory, static schema elements or a Driver connection.
 func (p *Planner) Plan(ctx context.Context, name string, from StateReader, to StateReader) (*Plan, error) {
@@ -202,8 +184,7 @@ func (p *Planner) Plan(ctx context.Context, name string, from StateReader, to St
 	return p.drv.PlanChanges(ctx, name, changes)
 }
 
-// WritePlan writes the given plan to the directory
-// based on the given Write configuration.
+// WritePlan writes the given plan to the Dir based on the given Formatter.
 func (p *Planner) WritePlan(plan *Plan) error {
 	fs, err := p.fmt.Format(plan)
 	if err != nil {
@@ -222,3 +203,122 @@ func (p *Planner) WritePlan(plan *Plan) error {
 	}
 	return nil
 }
+
+// LocalDir implements Dir for a local path.
+type LocalDir struct {
+	dir string
+}
+
+// Open implements Dir.Open.
+func (d *LocalDir) Open(name string) (File, error) {
+	return os.Create(filepath.Join(d.dir, name))
+}
+
+// NewLocalDir returns a new the Dir used by a Planner to work on the given local path.
+func NewLocalDir(path string) (*LocalDir, error) {
+	fi, err := os.Stat(path)
+	if err != nil {
+		return nil, err
+	}
+	if !fi.IsDir() {
+		return nil, fmt.Errorf("sql/migrate: %q is not a dir", path)
+	}
+	return &LocalDir{dir: path}, nil
+}
+
+// TemplateFormatter implements Formatter by using templates.
+type TemplateFormatter struct {
+	templates []struct{ N, C *template.Template }
+}
+
+var (
+	// templateFuncs contains the template.FuncMap for the DefaultFormatter.
+	templateFuncs = template.FuncMap{
+		"now": time.Now,
+		"sem": ensureSemicolonSuffix,
+		"rev": reverse,
+	}
+	// DefaultFormatter is a default implementation for Formatter. Compatible with golang-migrate/migrate.
+	DefaultFormatter = &TemplateFormatter{
+		templates: []struct{ N, C *template.Template }{
+			{
+				N: template.Must(template.New("").Funcs(templateFuncs).Parse(
+					"{{ now.Unix }}{{ with .Name }}_{{ . }}{{ end }}.up.sql",
+				)),
+				C: template.Must(template.New("").Funcs(templateFuncs).Parse(
+					"{{ range .Changes }}{{ println (sem .Cmd) }}{{ end }}",
+				)),
+			},
+			{
+				N: template.Must(template.New("").Funcs(templateFuncs).Parse(
+					"{{ now.Unix }}{{ with .Name }}_{{ . }}{{ end }}.down.sql",
+				)),
+				C: template.Must(template.New("").Funcs(templateFuncs).Parse(
+					"{{ range rev .Changes }}{{ println (sem .Reverse) }}{{ end }}",
+				)),
+			},
+		},
+	}
+)
+
+// NewTemplateFormatter creates a new Formatter working with the given templates.
+//
+//	migrate.NewTemplateFormatter(
+//		template.Must(template.New("").Parse("{{now.Unix}}{{.Name}}.sql")),                 // name template
+//		template.Must(template.New("").Parse("{{range .Changes}}{{println .Cmd}}{{end}}")), // content template
+//	)
+//
+func NewTemplateFormatter(templates ...*template.Template) (*TemplateFormatter, error) {
+	if n := len(templates); n == 0 || n%2 == 1 {
+		return nil, fmt.Errorf("zero or odd number of templates given")
+	}
+	t := new(TemplateFormatter)
+	for i := 0; i < len(templates); i += 2 {
+		t.templates = append(t.templates, struct{ N, C *template.Template }{templates[i], templates[i+1]})
+	}
+	return t, nil
+}
+
+// Format implements the Formatter interface.
+func (t *TemplateFormatter) Format(plan *Plan) ([]File, error) {
+	fs := make([]File, 0, len(t.templates))
+	for _, tpl := range t.templates {
+		var n, c bytes.Buffer
+		if err := tpl.N.Execute(&n, plan); err != nil {
+			return nil, err
+		}
+		if err := tpl.C.Execute(&c, plan); err != nil {
+			return nil, err
+		}
+		fs = append(fs, &templateFile{
+			Buffer: &c,
+			n:      n.String(),
+		})
+	}
+	return fs, nil
+}
+
+type templateFile struct {
+	*bytes.Buffer
+	n string
+}
+
+// Name implements the File interface.
+func (f *templateFile) Name() string { return f.n }
+
+// reverse changes for the down migration.
+func reverse(changes []*Change) []*Change {
+	rev := make([]*Change, len(changes))
+	for i, j := 0, len(changes)-1; i < j; i, j = i+1, j-1 {
+		rev[i], rev[j] = changes[j], changes[i]
+	}
+	return rev
+}
+
+// ensure an SQL statement has a trailing ";".
+func ensureSemicolonSuffix(s string) string {
+	if !strings.HasSuffix(s, ";") {
+		return s + ";"
+	}
+	return s
+}


Test patch:
diff --git a/sql/migrate/migrate_test.go b/sql/migrate/migrate_test.go
new file mode 100644
index 00000000..6873c36b
--- /dev/null
+++ b/sql/migrate/migrate_test.go
@@ -0,0 +1,93 @@
+// 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 migrate_test
+
+import (
+	"bytes"
+	"io"
+	"strconv"
+	"testing"
+	"text/template"
+	"time"
+
+	"ariga.io/atlas/sql/migrate"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPlanner_WritePlan(t *testing.T) {
+	var mfs = &mockFS{}
+	plan := &migrate.Plan{
+		Name: "add_t1_and_t2",
+		Changes: []*migrate.Change{
+			{Cmd: "CREATE TABLE t1(c int);", Reverse: "DROP TABLE t1 IF EXISTS"},
+			{Cmd: "CREATE TABLE t2(c int)", Reverse: "DROP TABLE t2"},
+		},
+	}
+
+	// DefaultFormatter
+	pl := migrate.New(nil, mfs, migrate.DefaultFormatter)
+	require.NotNil(t, pl)
+	require.NoError(t, pl.WritePlan(plan))
+	v := strconv.FormatInt(time.Now().Unix(), 10)
+	require.Len(t, mfs.files, 2)
+	require.Equal(t, &file{
+		n: v + "_add_t1_and_t2.up.sql",
+		b: bytes.NewBufferString("CREATE TABLE t1(c int);\nCREATE TABLE t2(c int);\n"),
+	}, mfs.files[0])
+	require.Equal(t, &file{
+		n: v + "_add_t1_and_t2.down.sql",
+		b: bytes.NewBufferString("DROP TABLE t2;\nDROP TABLE t1 IF EXISTS;\n"),
+	}, mfs.files[1])
+
+	// Custom formatter (creates only "up" migration files).
+	fmt, err := migrate.NewTemplateFormatter(
+		template.Must(template.New("").Parse("{{.Name}}.sql")),
+		template.Must(template.New("").Parse("{{range .Changes}}{{println .Cmd}}{{end}}")),
+	)
+	require.NoError(t, err)
+	pl = migrate.New(nil, mfs, fmt)
+	require.NotNil(t, pl)
+	require.NoError(t, pl.WritePlan(plan))
+	require.Len(t, mfs.files, 3)
+	require.Equal(t, &file{
+		n: "add_t1_and_t2.sql",
+		b: bytes.NewBufferString("CREATE TABLE t1(c int);\nCREATE TABLE t2(c int)\n"),
+	}, mfs.files[2])
+}
+
+type (
+	mockFS struct {
+		files []*file
+	}
+	file struct {
+		n string
+		b *bytes.Buffer
+	}
+)
+
+func (f *file) Read(b []byte) (int, error) {
+	return f.b.Read(b)
+}
+
+func (f *file) Write(b []byte) (int, error) {
+	return f.b.Write(b)
+}
+
+func (f *file) Close() error { return nil }
+
+func (f *file) Name() string { return f.n }
+
+func (fs *mockFS) Open(name string) (io.ReadWriteCloser, error) {
+	for _, f := range fs.files {
+		if f.n == name {
+			return f, nil
+		}
+	}
+	f := &file{n: name, b: new(bytes.Buffer)}
+	fs.files = append(fs.files, f)
+	return f, nil
+}
+
+var _ io.ReadWriteCloser = (*file)(nil)