Repository: ariga/atlas
Pull request #1764 (2023-06-20 15:18:13)

cmd/atlas/internal: support writing to template files


Patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate.go b/cmd/atlas/internal/cmdapi/migrate.go
index a465707e..cb61590c 100644
--- a/cmd/atlas/internal/cmdapi/migrate.go
+++ b/cmd/atlas/internal/cmdapi/migrate.go
@@ -166,11 +166,11 @@ func migrateApplyRun(cmd *cobra.Command, args []string, flags migrateApplyFlags,
 		}
 	}
 	// Open and validate the migration directory.
-	migrationDir, err := cmdmigrate.Dir(flags.dirURL, false)
+	dir, err := cmdmigrate.Dir(flags.dirURL, false)
 	if err != nil {
 		return err
 	}
-	if err := migrate.Validate(migrationDir); err != nil {
+	if err := migrate.Validate(dir); err != nil {
 		printChecksumError(cmd)
 		return err
 	}
@@ -205,7 +205,7 @@ func migrateApplyRun(cmd *cobra.Command, args []string, flags migrateApplyFlags,
 		return err
 	}
 	// Setup reporting info.
-	report := cmdlog.NewMigrateApply(client, migrationDir)
+	report := cmdlog.NewMigrateApply(client, dir)
 	mr.Init(client, report, rrw.(*cmdmigrate.EntRevisions))
 	// If cloud reporting is enabled, and we cannot obtain the current
 	// target identifier, abort and report it to the user.
@@ -214,7 +214,7 @@ func migrateApplyRun(cmd *cobra.Command, args []string, flags migrateApplyFlags,
 	}
 	// Determine pending files.
 	opts := append(flags.migrateOptions(), migrate.WithOperatorVersion(operatorVersion()), migrate.WithLogger(report))
-	ex, err := migrate.NewExecutor(client.Driver, migrationDir, rrw, opts...)
+	ex, err := migrate.NewExecutor(client.Driver, dir, rrw, opts...)
 	if err != nil {
 		return err
 	}
@@ -253,7 +253,7 @@ func migrateApplyRun(cmd *cobra.Command, args []string, flags migrateApplyFlags,
 		if drv, rrw, err = mux.driverFor(ctx, f); err != nil {
 			break
 		}
-		if ex, err = migrate.NewExecutor(drv, migrationDir, rrw, opts...); err != nil {
+		if ex, err = migrate.NewExecutor(drv, dir, rrw, opts...); err != nil {
 			return fmt.Errorf("unexpected executor creation error: %w", err)
 		}
 		if err = mux.mayRollback(ex.Execute(ctx, f)); err != nil {
diff --git a/cmd/atlas/internal/cmdext/cmdext.go b/cmd/atlas/internal/cmdext/cmdext.go
index e3b87583..2ceb3341 100644
--- a/cmd/atlas/internal/cmdext/cmdext.go
+++ b/cmd/atlas/internal/cmdext/cmdext.go
@@ -16,6 +16,7 @@ import (
 	"net/url"
 	"os"
 	"os/exec"
+	"path"
 	"path/filepath"
 	"reflect"
 	"strconv"
@@ -268,6 +269,13 @@ func TemplateDir(ctx *hcl.EvalContext, block *hclsyntax.Block) (cty.Value, error
 	if len(attrs) > 0 {
 		return cty.NilVal, errorf("unexpected attributes: %v", attrs)
 	}
+	dirname := path.Join(args.Path, block.Labels[1])
+	dir := migrate.OpenMemDir(dirname)
+	// Clear existing directories in case the config was called
+	// multiple times with different variables.
+	if files, err := dir.Files(); err != nil || len(files) > 0 {
+		dir.Reset()
+	}
 	t := template.New("template_dir")
 	err := filepath.Walk(args.Path, func(path string, d os.FileInfo, err error) error {
 		if err != nil {
@@ -281,7 +289,6 @@ func TemplateDir(ctx *hcl.EvalContext, block *hclsyntax.Block) (cty.Value, error
 	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 {
@@ -300,15 +307,29 @@ func TemplateDir(ctx *hcl.EvalContext, block *hclsyntax.Block) (cty.Value, error
 	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 {
+	if err := migrate.WriteSumFile(dir, sum); err != nil {
 		return cty.NilVal, err
 	}
+	// Sync template dir to local filesystem.
+	dir.SyncWrites(func(name string, data []byte) error {
+		if name == migrate.HashFileName {
+			return nil
+		}
+		l, err := migrate.NewLocalDir(args.Path)
+		if err != nil {
+			return err
+		}
+		if err := l.WriteFile(name, data); err != nil {
+			return err
+		}
+		sum, err := l.Checksum()
+		if err != nil {
+			return err
+		}
+		return migrate.WriteSumFile(l, sum)
+	})
 	return cty.ObjectVal(map[string]cty.Value{
-		"url": cty.StringVal(fmt.Sprintf("mem://%s", args.Path)),
+		"url": cty.StringVal(fmt.Sprintf("mem://%s", dirname)),
 	}), nil
 }
 
diff --git a/doc/md/atlas-schema/projects.mdx b/doc/md/atlas-schema/projects.mdx
index d14e7422..8cf36222 100644
--- a/doc/md/atlas-schema/projects.mdx
+++ b/doc/md/atlas-schema/projects.mdx
@@ -662,7 +662,10 @@ have the `.sql` file extension, and generating an in-memory migration directory
 
 - `url` - A [URL](../concepts/url.mdx) to the generated migration directory.
 
-```hcl title="atlas.hcl" {6-14}
+<Tabs>
+<TabItem label="Read only templates" value="read-only">
+
+```hcl title="atlas.hcl" {6-14,19}
 variable "path" {
   type        = string
   description = "A path to the template directory"
@@ -686,6 +689,41 @@ env "dev" {
 }
 ```
 
+</TabItem>
+<TabItem label="Variables shared between HCL and directory" value="diff">
+
+```hcl title="atlas.hcl" {13-18,24}
+variable "schema_name" {
+  type    = string
+  default = "Database schema name injected to both migrations directory and HCL schema"
+}
+
+data "hcl_schema" "app" {
+  path = "path/to/schema.hcl"
+  vars = {
+    schema_name = var.schema_name
+  }
+}
+
+data "template_dir" "migrations" {
+  path = "path/to/directory"
+  vars = {
+    schema_name = var.schema_name
+  }
+}
+
+env "local" {
+  src = data.hcl_schema.app.url
+  dev = "sqlite://file?mode=memory&_fk=1"
+  migration {
+    dir = data.template_dir.migrations.url
+  }
+}
+```
+
+</TabItem>
+</Tabs>
+
 ### Environments
 
 The `env` block defines an environment block that can be selected by using the `--env` flag.
diff --git a/sql/migrate/dir.go b/sql/migrate/dir.go
index 09e555fb..06bef373 100644
--- a/sql/migrate/dir.go
+++ b/sql/migrate/dir.go
@@ -219,7 +219,8 @@ func (f LocalFile) Directive(name string) (ds []string) {
 type (
 	// MemDir provides an in-memory Dir implementation.
 	MemDir struct {
-		files map[string]File
+		files  map[string]File
+		syncTo []func(string, []byte) error
 	}
 	// An opened MemDir.
 	openedMem struct {
@@ -262,6 +263,12 @@ func (d *MemDir) Open(name string) (fs.File, error) {
 	}, nil
 }
 
+// Reset the in-memory directory to its initial state.
+func (d *MemDir) Reset() {
+	d.files = nil
+	d.syncTo = nil
+}
+
 // Close implements the io.Closer interface.
 func (d *MemDir) Close() error {
 	memDirs.Lock()
@@ -288,9 +295,20 @@ func (d *MemDir) WriteFile(name string, data []byte) error {
 		d.files = make(map[string]File)
 	}
 	d.files[name] = NewLocalFile(name, data)
+	for _, f := range d.syncTo {
+		if err := f(name, data); err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
+// SyncWrites allows syncing writes from in-memory directory
+// the underlying storage.
+func (d *MemDir) SyncWrites(fs ...func(string, []byte) error) {
+	d.syncTo = append(d.syncTo, fs...)
+}
+
 // Files returns a set of files stored in-memory to be executed on a database.
 func (d *MemDir) Files() ([]File, error) {
 	files := make([]File, 0, len(d.files))


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate_test.go b/cmd/atlas/internal/cmdapi/migrate_test.go
index 818974c7..d35a49ac 100644
--- a/cmd/atlas/internal/cmdapi/migrate_test.go
+++ b/cmd/atlas/internal/cmdapi/migrate_test.go
@@ -1353,6 +1353,147 @@ env "local" {
 		require.NoError(t, err)
 		require.Len(t, files, 2)
 	})
+
+	t.Run("TemplateDir", func(t *testing.T) {
+		var (
+			h = `
+variable "default" {
+  type    = string
+  default = "Default"
+}
+
+variable "schema_path" {
+  type = string
+}
+
+variable "migrations_path" {
+  type = string
+}
+
+data "hcl_schema" "app" {
+  path = var.schema_path
+  vars = {
+    default = var.default
+  }
+}
+
+data "template_dir" "app" {
+  path = var.migrations_path
+  vars = {
+    default = var.default
+  }
+}
+
+env "local" {
+  src = data.hcl_schema.app.url
+  dev = "sqlite://file?mode=memory&_fk=1"
+  migration {
+    dir = data.template_dir.app.url
+  }
+}
+`
+			p, md      = t.TempDir(), t.TempDir()
+			cfg, state = filepath.Join(p, "atlas.hcl"), filepath.Join(p, "schema.hcl")
+		)
+		err := os.WriteFile(cfg, []byte(h), 0600)
+		require.NoError(t, err)
+		err = os.WriteFile(state, []byte(`variable "default" { type = string  }
+schema "main" {}
+table "users" {
+  schema = schema.main
+  column "c" {
+    type = text
+    default = var.default
+  }
+}
+`), 0600)
+		require.NoError(t, err)
+		dir, err := migrate.NewLocalDir(md)
+		require.NoError(t, err)
+		err = dir.WriteFile("1.sql", []byte(`create table users (c text default "{{ .default }}" NOT NULL);`))
+		require.NoError(t, err)
+
+		cmd := migrateCmd()
+		cmd.AddCommand(migrateHashCmd())
+		_, err = runCmd(
+			cmd, "hash",
+			"--dir", "file://"+md,
+		)
+		require.NoError(t, err)
+		f, err := dir.Open(migrate.HashFileName)
+		require.NoError(t, err)
+		b, err := io.ReadAll(f)
+		require.NoError(t, err)
+		require.NotEmpty(t, b)
+
+		cmd = migrateCmd()
+		cmd.AddCommand(migrateDiffCmd())
+		s, err := runCmd(
+			cmd, "diff",
+			"-c", "file://"+cfg,
+			"--env", "local",
+			"--var", "migrations_path="+md,
+			"--var", "schema_path="+state,
+		)
+		// Desired state and migration directory are in sync.
+		require.NoError(t, err)
+		require.Equal(t, "The migration directory is synced with the desired state, no changes to be made\n", s)
+		f, err = dir.Open(migrate.HashFileName)
+		require.NoError(t, err)
+		nb, err := io.ReadAll(f)
+		require.NoError(t, err)
+		require.Equal(t, b, nb, "hash file should not be updated")
+
+		// Update the desired state and run diff again.
+		err = os.WriteFile(state, []byte(`variable "default" { type = string  }
+schema "main" {}
+table "users" {
+  schema = schema.main
+  column "c" {
+    type = text
+    default = var.default
+  }
+  column "d" {
+    type = text
+  }
+}
+`), 0600)
+		require.NoError(t, err)
+		_, err = runCmd(
+			cmd, "diff",
+			"-c", "file://"+cfg,
+			"--env", "local",
+			"--var", "migrations_path="+md,
+			"--var", "schema_path="+state,
+		)
+		require.NoError(t, err)
+
+		// Check files.
+		files, err := dir.Files()
+		require.NoError(t, err)
+		require.Len(t, files, 2)
+		require.Equal(t, "1.sql", files[0].Name())
+		require.Equal(t, `create table users (c text default "{{ .default }}" NOT NULL);`, string(files[0].Bytes()), "should not update template files")
+		require.Equal(t, "-- Add column \"d\" to table: \"users\"\nALTER TABLE `users` ADD COLUMN `d` text NOT NULL;\n", string(files[1].Bytes()))
+
+		// Ensure the sum file is consistent.
+		f, err = dir.Open(migrate.HashFileName)
+		require.NoError(t, err)
+		before, err := io.ReadAll(f)
+		require.NoError(t, err)
+		cmd = migrateCmd()
+		cmd.AddCommand(migrateHashCmd())
+		_, err = runCmd(
+			cmd, "hash",
+			"--dir", "file://"+md,
+		)
+		require.NoError(t, err)
+		f, err = dir.Open(migrate.HashFileName)
+		require.NoError(t, err)
+		after, err := io.ReadAll(f)
+		require.NoError(t, err)
+		require.Equal(t, before, after)
+	})
 }
 
 func TestMigrate_StatusJSON(t *testing.T) {
diff --git a/cmd/atlas/internal/cmdapi/testdata/templatedir/atlas.sum b/cmd/atlas/internal/cmdapi/testdata/templatedir/atlas.sum
new file mode 100644
index 00000000..d015a018
--- /dev/null
+++ b/cmd/atlas/internal/cmdapi/testdata/templatedir/atlas.sum
@@ -0,0 +1,3 @@
+h1:qhwlEKNEXRVSUiyBXmbiSiS26un12DoxHCR3WDzrDdA=
+1.sql h1:b/5P45x6+CVoyVbBJ/BVc5Z2cI46XfoW92QH3T1kqxY=
+2.sql h1:El2EOcXWK2PfPV6yJg4hBVtc4aHheV/MFEPjho9SEmQ=
diff --git a/sql/migrate/dir_test.go b/sql/migrate/dir_test.go
index f786c054..443c7a13 100644
--- a/sql/migrate/dir_test.go
+++ b/sql/migrate/dir_test.go
@@ -227,6 +227,18 @@ func TestMemDir(t *testing.T) {
 	files, err = d.Files()
 	require.NoError(t, err)
 	require.Len(t, files, 1) // 1.sql
+
+	// Sync with additional directory.
+	var d2, d3 migrate.MemDir
+	d.SyncWrites(d2.WriteFile, d3.WriteFile)
+	require.NoError(t, d.WriteFile("2.sql", []byte("create table t2(c int);")))
+	for _, d1 := range []migrate.MemDir{d, d2, d3} {
+		f, err := d1.Open("2.sql")
+		require.NoError(t, err)
+		c, err := io.ReadAll(f)
+		require.NoError(t, err)
+		require.Equal(t, "create table t2(c int);", string(c))
+	}
 }
 
 func TestOpenMemDir(t *testing.T) {