Repository: ariga/atlas
Pull request #1480 (2023-02-26 14:22:32)

cmd/atlas: add support for template_dir


Patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate.go b/cmd/atlas/internal/cmdapi/migrate.go
index bd07d808..0850773c 100644
--- a/cmd/atlas/internal/cmdapi/migrate.go
+++ b/cmd/atlas/internal/cmdapi/migrate.go
@@ -1280,12 +1280,16 @@ func dir(u string, create bool) (migrate.Dir, error) {
 
 // dirURL returns a migrate.Dir to use as migration directory. For now only local directories are supported.
 func dirURL(u *url.URL, create bool) (migrate.Dir, error) {
-	if u.Scheme != "file" {
-		return nil, fmt.Errorf("unsupported driver %q", u.Scheme)
-	}
 	path := filepath.Join(u.Host, u.Path)
-	if path == "" {
-		path = "migrations"
+	switch u.Scheme {
+	case "mem":
+		return migrate.OpenMemDir(path), nil
+	case "file":
+		if path == "" {
+			path = "migrations"
+		}
+	default:
+		return nil, fmt.Errorf("unsupported driver %q", u.Scheme)
 	}
 	fn := func() (migrate.Dir, error) { return migrate.NewLocalDir(path) }
 	switch f := u.Query().Get("format"); f {
diff --git a/cmd/atlas/internal/cmdext/cmdext.go b/cmd/atlas/internal/cmdext/cmdext.go
index daa49dc8..5bd83391 100644
--- a/cmd/atlas/internal/cmdext/cmdext.go
+++ b/cmd/atlas/internal/cmdext/cmdext.go
@@ -7,12 +7,16 @@
 package cmdext
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
+	"io/fs"
 	"net/url"
+	"os"
 	"path/filepath"
 	"strconv"
+	"text/template"
 	"time"
 
 	"ariga.io/atlas/schemahcl"
@@ -41,6 +45,7 @@ import (
 var DataSources = []schemahcl.Option{
 	schemahcl.WithDataSource("sql", QuerySrc),
 	schemahcl.WithDataSource("runtimevar", RuntimeVarSrc),
+	schemahcl.WithDataSource("template_dir", TemplateDir),
 }
 
 // RuntimeVarSrc exposes the gocloud.dev/runtimevar as a schemahcl datasource.
@@ -197,6 +202,109 @@ func QuerySrc(ctx *hcl.EvalContext, block *hclsyntax.Block) (cty.Value, error) {
 	return cty.ObjectVal(obj), nil
 }
 
+// TemplateDir implements migrate.Dir interface for template directories.
+//
+//	data "template_dir" "name" {
+//	  path = "path/to/directory"
+//	  vars = {
+//	    Env  = atlas.env
+//	    Seed = var.seed
+//	  }
+//	}
+//
+//	env "dev" {
+//	  url = "driver://path?query=param"
+//	  migration {
+//	    dir = data.template_dir.name.url
+//	  }
+//	}
+func TemplateDir(ctx *hcl.EvalContext, block *hclsyntax.Block) (cty.Value, error) {
+	var (
+		args struct {
+			Path   string   `hcl:"path"`
+			Remain hcl.Body `hcl:",remain"`
+		}
+		vars   = make(map[string]any)
+		errorf = blockError("data.template_dir", block)
+	)
+	if diags := gohcl.DecodeBody(block.Body, ctx, &args); diags.HasErrors() {
+		return cty.NilVal, errorf("decoding body: %v", diags)
+	}
+	attrs, diags := args.Remain.JustAttributes()
+	if diags.HasErrors() {
+		return cty.NilVal, errorf("getting attributes: %v", diags)
+	}
+	if vs, ok := attrs["vars"]; ok {
+		switch vs, diags := vs.Expr.Value(ctx); {
+		case diags.HasErrors():
+			return cty.NilVal, errorf(`evaluating "vars": %w`, diags)
+		case !vs.CanIterateElements():
+			return cty.NilVal, errorf(`attribute "vars" must be a map, got: %s`, vs.Type())
+		default:
+			for it := vs.ElementIterator(); it.Next(); {
+				k, v := it.Element()
+				switch v.Type() {
+				case cty.String:
+					vars[k.AsString()] = v.AsString()
+				case cty.Number:
+					f, _ := v.AsBigFloat().Float64()
+					vars[k.AsString()] = f
+				case cty.Bool:
+					vars[k.AsString()] = v.True()
+				default:
+					return cty.NilVal, errorf(`attribute "vars" must be a map of strings, numbers or booleans, got: %s`, v.Type())
+				}
+			}
+		}
+		delete(attrs, "vars")
+	}
+	if len(attrs) > 0 {
+		return cty.NilVal, errorf("unexpected attributes: %v", attrs)
+	}
+	t := template.New("template_dir")
+	err := filepath.Walk(args.Path, func(path string, d os.FileInfo, err error) error {
+		if err != nil {
+			return fmt.Errorf("walk path %s: %w", path, err)
+		}
+		if !d.IsDir() {
+			_, err = t.ParseFiles(path)
+		}
+		return err
+	})
+	if err != nil {
+		return cty.NilVal, errorf(err.Error())
+	}
+	dir := migrate.OpenMemDir(args.Path)
+	// Only top-level (template) files are treated as migrations.
+	matches, err := fs.Glob(os.DirFS(args.Path), "*.sql")
+	if err != nil {
+		return cty.NilVal, errorf("globbing templates: %w", err)
+	}
+	for _, m := range matches {
+		var b bytes.Buffer
+		if err := t.ExecuteTemplate(&b, m, vars); err != nil {
+			return cty.NilVal, errorf("executing template %q: %w", m, err)
+		}
+		if err := dir.WriteFile(m, b.Bytes()); err != nil {
+			return cty.NilVal, errorf("writing file %q: %w", m, err)
+		}
+	}
+	sum, err := dir.Checksum()
+	if err != nil {
+		return cty.NilVal, err
+	}
+	b, err := sum.MarshalText()
+	if err != nil {
+		return cty.NilVal, err
+	}
+	if err := dir.WriteFile(migrate.HashFileName, b); err != nil {
+		return cty.NilVal, err
+	}
+	return cty.ObjectVal(map[string]cty.Value{
+		"url": cty.StringVal(fmt.Sprintf("mem://%s", args.Path)),
+	}), nil
+}
+
 func blockError(name string, b *hclsyntax.Block) func(string, ...any) error {
 	return func(format string, args ...any) error {
 		return fmt.Errorf("%s.%s: %w", name, b.Labels[1], fmt.Errorf(format, args...))
diff --git a/sql/migrate/dir.go b/sql/migrate/dir.go
index dcf3f4b8..9579b2b9 100644
--- a/sql/migrate/dir.go
+++ b/sql/migrate/dir.go
@@ -18,6 +18,7 @@ import (
 	"regexp"
 	"sort"
 	"strings"
+	"sync"
 	"text/template"
 	"time"
 )
@@ -210,9 +211,39 @@ func (f LocalFile) Directive(name string) (ds []string) {
 	return ds
 }
 
-// MemDir provides an in-memory Dir implementation.
-type MemDir struct {
-	files map[string]File
+type (
+	// MemDir provides an in-memory Dir implementation.
+	MemDir struct {
+		files map[string]File
+	}
+	// An opened MemDir.
+	openedMem struct {
+		dir    *MemDir
+		numUse int
+	}
+)
+
+// A list of the opened memory-based directories.
+var memDirs struct {
+	sync.Mutex
+	opened map[string]*openedMem
+}
+
+// OpenMemDir opens an in-memory directory and registers it in the process namespace
+// with the given name. Hence, calling OpenMemDir with the same name will return the
+// same directory. The directory is deleted when the last reference of it is closed.
+func OpenMemDir(name string) *MemDir {
+	memDirs.Lock()
+	defer memDirs.Unlock()
+	if m, ok := memDirs.opened[name]; ok {
+		m.numUse++
+		return m.dir
+	}
+	if memDirs.opened == nil {
+		memDirs.opened = make(map[string]*openedMem)
+	}
+	memDirs.opened[name] = &openedMem{dir: &MemDir{}, numUse: 1}
+	return memDirs.opened[name].dir
 }
 
 // Open implements fs.FS.
@@ -226,6 +257,26 @@ func (d *MemDir) Open(name string) (fs.File, error) {
 	}, nil
 }
 
+// Close implements the io.Closer interface.
+func (d *MemDir) Close() error {
+	memDirs.Lock()
+	defer memDirs.Unlock()
+	var opened string
+	for name, m := range memDirs.opened {
+		switch {
+		case m.dir != d:
+		case opened != "":
+			return fmt.Errorf("dir was opened with different names: %q and %q", opened, name)
+		default:
+			opened = name
+			if m.numUse--; m.numUse == 0 {
+				delete(memDirs.opened, name)
+			}
+		}
+	}
+	return nil
+}
+
 // WriteFile adds a new file in-memory.
 func (d *MemDir) WriteFile(name string, data []byte) error {
 	if d.files == nil {


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate_test.go b/cmd/atlas/internal/cmdapi/migrate_test.go
index a3263bb7..161c56b3 100644
--- a/cmd/atlas/internal/cmdapi/migrate_test.go
+++ b/cmd/atlas/internal/cmdapi/migrate_test.go
@@ -537,6 +537,56 @@ env "local" {
 		// Empty list is expanded to zero blocks.
 		require.EqualError(t, err, `env "local" not defined in project file`)
 	})
+
+	t.Run("TemplateDir", func(t *testing.T) {
+		var (
+			h = `
+variable "path" {
+  type = string
+}
+
+data "template_dir" "migrations" {
+  path = var.path
+  vars = {
+    Env = atlas.env
+  }
+}
+
+env "dev" {
+  url = "sqlite://file://${atlas.env}?mode=memory&_fk=1"
+  migration {
+    dir = data.template_dir.migrations.url
+  }
+}
+
+env "prod" {
+  url = "sqlite://file://${atlas.env}?mode=memory&_fk=1"
+  migration {
+    dir = data.template_dir.migrations.url
+  }
+}
+`
+			p    = t.TempDir()
+			path = filepath.Join(p, "atlas.hcl")
+		)
+		err := os.WriteFile(path, []byte(h), 0600)
+		require.NoError(t, err)
+		for _, e := range []string{"dev", "prod"} {
+			cmd := migrateCmd()
+			cmd.AddCommand(migrateApplyCmd())
+			s, err := runCmd(
+				cmd, "apply",
+				"-c", "file://"+path,
+				"--env", e,
+				"--var", "path=testdata/templatedir",
+			)
+			require.NoError(t, err)
+			require.Contains(t, s, "Migrating to version 2 (2 migrations in total):")
+			require.Contains(t, s, fmt.Sprintf("create table %s1 (c text);", e))
+			require.Contains(t, s, fmt.Sprintf("create table %s2 (c text);", e))
+			require.Contains(t, s, fmt.Sprintf("create table users_%s2 (c text);", e))
+		}
+	})
 }
 
 func TestMigrate_ApplyTxMode(t *testing.T) {
diff --git a/cmd/atlas/internal/cmdapi/testdata/templatedir/1.sql b/cmd/atlas/internal/cmdapi/testdata/templatedir/1.sql
new file mode 100644
index 00000000..7956dc5f
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/templatedir/1.sql
@@ -0,0 +1,5 @@
+{{- if eq .Env "dev" }}
+    create table dev1 (c text);
+{{- else  }}
+    create table prod1 (c text);
+{{- end }}
\ No newline at end of file
diff --git a/cmd/atlas/internal/cmdapi/testdata/templatedir/2.sql b/cmd/atlas/internal/cmdapi/testdata/templatedir/2.sql
new file mode 100644
index 00000000..b8e8b01e
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/templatedir/2.sql
@@ -0,0 +1,7 @@
+{{- if eq .Env "dev" }}
+    create table dev2 (c text);
+    {{ template "shared/users" "dev2" }}
+{{- else  }}
+    create table prod2 (c text);
+    {{ template "shared/users" "prod2" }}
+{{- end }}
\ No newline at end of file
diff --git a/cmd/atlas/internal/cmdapi/testdata/templatedir/shared/users.sql b/cmd/atlas/internal/cmdapi/testdata/templatedir/shared/users.sql
new file mode 100644
index 00000000..0236e5f0
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/templatedir/shared/users.sql
@@ -0,0 +1,3 @@
+{{- define "shared/users" }}
+create table users_{{ $ }} (c text);
+{{- end}}
\ No newline at end of file
diff --git a/cmd/atlas/internal/cmdext/cmdext_test.go b/cmd/atlas/internal/cmdext/cmdext_test.go
index 385bed17..7ac63a77 100644
--- a/cmd/atlas/internal/cmdext/cmdext_test.go
+++ b/cmd/atlas/internal/cmdext/cmdext_test.go
@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 	"net/url"
+	"os"
 	"path/filepath"
 	"strings"
 	"testing"
@@ -21,6 +22,7 @@ import (
 
 	_ "github.com/mattn/go-sqlite3"
 	"github.com/stretchr/testify/require"
+	"github.com/zclconf/go-cty/cty"
 )
 
 func TestRuntimeVarSrc(t *testing.T) {
@@ -157,3 +159,40 @@ func TestEntLoader_MigrateDiff(t *testing.T) {
 		require.False(t, ok, "skipping schemas without globalid")
 	})
 }
+
+func TestTemplateDir(t *testing.T) {
+	var (
+		v struct {
+			Dir string `spec:"dir"`
+		}
+		dir   = t.TempDir()
+		state = schemahcl.New(cmdext.DataSources...)
+	)
+	err := os.WriteFile(filepath.Join(dir, "1.sql"), []byte("create table {{ .Schema }}.t(c int);"), 0644)
+	require.NoError(t, err)
+	err = state.EvalBytes([]byte(`
+variable "path" {
+  type = string
+}
+
+data "template_dir" "tenant" {
+  path = var.path
+  vars = {
+    Schema = "a8m"
+  }
+}
+
+dir = data.template_dir.tenant.url
+`), &v, map[string]cty.Value{
+		"path": cty.StringVal(dir),
+	})
+	require.NoError(t, err)
+	require.NotEmpty(t, v.Dir)
+	d := migrate.OpenMemDir(strings.TrimPrefix(v.Dir, "mem://"))
+	require.NoError(t, migrate.Validate(d))
+	files, err := d.Files()
+	require.NoError(t, err)
+	require.Len(t, files, 1)
+	require.Equal(t, "1.sql", files[0].Name())
+	require.Equal(t, "create table a8m.t(c int);", string(files[0].Bytes()))
+}
diff --git a/sql/migrate/dir_test.go b/sql/migrate/dir_test.go
index 5c69b014..70a7452a 100644
--- a/sql/migrate/dir_test.go
+++ b/sql/migrate/dir_test.go
@@ -226,6 +226,44 @@ func TestMemDir(t *testing.T) {
 	require.Len(t, files, 1) // 1.sql
 }
 
+func TestOpenMemDir(t *testing.T) {
+	dev1 := migrate.OpenMemDir("dev")
+	require.NoError(t, dev1.WriteFile("1.sql", []byte("create table t1(c int);")))
+	// Open the same dir.
+	dev2 := migrate.OpenMemDir("dev")
+	files2, err := dev2.Files()
+	require.NoError(t, err)
+	require.Len(t, files2, 1)
+	require.NoError(t, dev2.WriteFile("2.sql", []byte("create table t2(c int);")))
+	files1, err := dev1.Files()
+	require.NoError(t, err)
+	require.Len(t, files1, 2)
+	files2, err = dev2.Files()
+	require.NoError(t, err)
+	require.Len(t, files2, 2)
+	// Open a new dir.
+	etc := migrate.OpenMemDir("etc")
+	files, err := etc.Files()
+	require.NoError(t, err)
+	require.Empty(t, files)
+
+	// Closing dir and opening it should not
+	// clean it if there are active references.
+	require.NoError(t, dev1.Close())
+	dev1 = migrate.OpenMemDir("dev")
+	files1, err = dev1.Files()
+	require.NoError(t, err)
+	require.Len(t, files1, 2)
+
+	// Cleanup directory on close.
+	require.NoError(t, dev1.Close())
+	require.NoError(t, dev2.Close())
+	dev1 = migrate.OpenMemDir("dev")
+	files1, err = dev1.Files()
+	require.NoError(t, err)
+	require.Empty(t, files1)
+}
+
 func TestLocalFile_Directive(t *testing.T) {
 	f := migrate.NewLocalFile("1.sql", []byte(`-- atlas:lint ignore
 alter table users drop column id;