atlasexec: handle migrate.Dir when copy

atlasexec/working_dir.go

@@ -7,6 +7,8 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+
+	"ariga.io/atlas/sql/migrate"
 )
 
 type (
@@ -137,26 +139,54 @@ func (ce *WorkingDir) CreateFile(name string, fn func(w io.Writer) error) error
 // in the temporary directory.
 // If source is nil, an error is returned.
 func (cs *WorkingDir) CopyFS(name string, src fs.FS) error {
-	if src == nil {
-		return errors.New("atlasexec: source is nil")
-	}
 	dst := cs.Path(name)
 	// Ensure destination directory exists.
 	if err := os.MkdirAll(dst, 0700); err != nil {
 		return err
 	}
-	return fs.WalkDir(src, ".", func(path string, d fs.DirEntry, err error) error {
-		if err != nil || path == "." {
+	switch dir := src.(type) {
+	case nil:
+		return errors.New("atlasexec: source is nil")
+	case migrate.Dir:
+		// The migrate.MemDir doesn't 100% compatible with fs.FS.
+		// It returns fs.ErrNotExist error when open "." directory.
+		// So, we need to handle it separately using the Files method.
+		files, err := dir.Files()
+		if err != nil {
 			return err
 		}
-		name := filepath.Join(dst, path)
-		if d.IsDir() {
-			return os.Mkdir(name, 0700)
+		for _, f := range files {
+			name := filepath.Join(dst, f.Name())
+			if err := os.WriteFile(name, f.Bytes(), 0644); err != nil {
+				return err
+			}
 		}
-		data, err := fs.ReadFile(src, path)
-		if err != nil {
-			return err
+		// If the atlas.sum file exists, copy it to the destination directory.
+		if hf, err := dir.Open(migrate.HashFileName); err == nil {
+			data, err := io.ReadAll(hf)
+			if err != nil {
+				return err
+			}
+			name := filepath.Join(dst, migrate.HashFileName)
+			if err := os.WriteFile(name, data, 0644); err != nil {
+				return err
+			}
 		}
-		return os.WriteFile(name, data, 0644)
-	})
+		return nil
+	default:
+		return fs.WalkDir(dir, ".", func(path string, d fs.DirEntry, err error) error {
+			if err != nil || path == "." {
+				return err
+			}
+			name := filepath.Join(dst, path)
+			if d.IsDir() {
+				return os.Mkdir(name, 0700)
+			}
+			data, err := fs.ReadFile(dir, path)
+			if err != nil {
+				return err
+			}
+			return os.WriteFile(name, data, 0644)
+		})
+	}
 }

atlasexec/working_dir_test.go

@@ -10,6 +10,7 @@ import (
 	"testing/fstest"
 	"text/template"
 
+	"ariga.io/atlas/sql/migrate"
 	"github.com/stretchr/testify/require"
 )
 
@@ -37,6 +38,16 @@ func TestContextExecer(t *testing.T) {
 	checkFileContent(t, filepath.Join("migrations", "bar"), "bar-content")
 	require.NoError(t, ce.Close())
 
+	// Test WithMigrations - MemDir.
+	dir := &migrate.MemDir{}
+	require.NoError(t, dir.WriteFile("1.sql", []byte("-- only .sql files are copied\nmem-content")))
+	require.NoError(t, dir.WriteFile(migrate.HashFileName, []byte("-- And the atlas.sum")))
+	ce, err = NewWorkingDir(WithMigrations(dir))
+	require.NoError(t, err)
+	checkFileContent(t, filepath.Join("migrations", "1.sql"), "-- only .sql files are copied\nmem-content")
+	checkFileContent(t, filepath.Join("migrations", migrate.HashFileName), "-- And the atlas.sum")
+	require.NoError(t, ce.Close())
+
 	// Test WithAtlasHCL.
 	ce, err = NewWorkingDir(
 		WithAtlasHCL(func(w io.Writer) error {

go.mod

@@ -5,7 +5,7 @@ go 1.22
 toolchain go1.22.0
 
 require (
-	ariga.io/atlas v0.20.1-0.20240321075817-75fd3b1accbf
+	ariga.io/atlas v0.21.2-0.20240418081819-02b3f6239b04
 	github.com/mattn/go-sqlite3 v1.14.17
 	github.com/stretchr/testify v1.8.4
 )

go.sum

@@ -1,5 +1,5 @@
-ariga.io/atlas v0.20.1-0.20240321075817-75fd3b1accbf h1:jpDQmb4YyO6SADrfS2+13p+dKOBfNJB80K8sTLQMr9g=
-ariga.io/atlas v0.20.1-0.20240321075817-75fd3b1accbf/go.mod h1:VPlcXdd4w2KqKnH54yEZcry79UAhpaWaxEsmn5JRNoE=
+ariga.io/atlas v0.21.2-0.20240418081819-02b3f6239b04 h1:YF3qiqtnhn+y4tfhZKTfZKfizpjqHYt7rWPUb+eA4ZA=
+ariga.io/atlas v0.21.2-0.20240418081819-02b3f6239b04/go.mod h1:VPlcXdd4w2KqKnH54yEZcry79UAhpaWaxEsmn5JRNoE=
 github.com/DATA-DOG/go-sqlmock v1.5.0 h1:Shsta01QNfFxHCfpW6YH2STWB0MudeXXEWMr20OEh60=
 github.com/DATA-DOG/go-sqlmock v1.5.0/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM=
 github.com/agext/levenshtein v1.2.3 h1:YB2fHEn0UJagG8T1rrWknE3ZQzWM06O8AMAatNn7lmo=