Repository: ariga/atlas
Pull request #1358 (2022-12-21 12:08:29)

cmd/atlas: extend sql state-reader to accept single files


Patch:
diff --git a/cmd/atlas/internal/cmdapi/cmdapi.go b/cmd/atlas/internal/cmdapi/cmdapi.go
index 520482d3..656808ab 100644
--- a/cmd/atlas/internal/cmdapi/cmdapi.go
+++ b/cmd/atlas/internal/cmdapi/cmdapi.go
@@ -427,76 +427,15 @@ func stateReader(ctx context.Context, config *stateReaderConfig) (*stateReadClos
 	switch scheme {
 	// "file" scheme is valid for both migration directory and HCL paths.
 	case "file":
-		if len(config.urls) > 1 {
-			// Consider urls being HCL paths.
-			return hclStateReader(ctx, config, parsed)
-		}
-		path := filepath.Join(parsed[0].Host, parsed[0].Path)
-		fi, err := os.Stat(path)
-		if err != nil {
-			return nil, err
-		}
-		if !fi.IsDir() {
-			// If there is only one url given, and it is a file, consider it an HCL path.
-			return hclStateReader(ctx, config, parsed)
-		}
-		files, err := os.ReadDir(path)
-		if err != nil {
+		switch ext, err := filesExt(parsed); {
+		case err != nil:
 			return nil, err
-		}
-		// Check all files, if we find both HCL and SQL files, abort. Otherwise, proceed accordingly.
-		var hcl, sql bool
-		for _, f := range files {
-			ext := filepath.Ext(f.Name())
-			switch {
-			case hcl && ext == ".sql", sql && ext == ".hcl":
-				return nil, fmt.Errorf("ambiguos files: %q contains both SQL and HCL files", path)
-			case ext == ".hcl":
-				hcl = true
-			case ext == ".sql":
-				sql = true
-			default:
-				// unknown extension, we don't care
-			}
-		}
-		switch {
-		case hcl:
+		case ext == extHCL:
 			return hclStateReader(ctx, config, parsed)
-		case sql:
-			// Replaying a migration directory requires a dev connection.
-			if config.dev == nil {
-				return nil, errors.New("--dev-url cannot be empty")
-			}
-			dir, err := dirURL(parsed[0], false)
-			if err != nil {
-				return nil, err
-			}
-			ex, err := migrate.NewExecutor(config.dev.Driver, dir, migrate.NopRevisionReadWriter{})
-			if err != nil {
-				return nil, err
-			}
-			var opts []migrate.ReplayOption
-			if v := parsed[0].Query().Get("version"); v != "" {
-				opts = append(opts, migrate.ReplayToVersion(v))
-			}
-			sr, err := ex.Replay(ctx, func() migrate.StateReader {
-				if config.dev.URL.Schema != "" {
-					return migrate.SchemaConn(config.dev, "", nil)
-				}
-				return migrate.RealmConn(config.dev, &schema.InspectRealmOption{
-					Schemas: config.schemas,
-					Exclude: config.exclude,
-				})
-			}(), opts...)
-			if err != nil && !errors.Is(err, migrate.ErrNoPendingFiles) {
-				return nil, fmt.Errorf("replaying the migration directory: %w", err)
-			}
-			return &stateReadCloser{
-				StateReader: migrate.Realm(sr),
-				schema:      config.dev.URL.Schema,
-			}, nil
+		case ext == extSQL:
+			return sqlStateReader(ctx, config, parsed)
 		default:
-			return nil, fmt.Errorf("%q contains neither SQL nor HCL files", path)
+			panic("unreachable") // checked by filesExt.
 		}
 	default:
 		// All other schemes are database (or docker) connections.
@@ -522,7 +461,7 @@ func stateReader(ctx context.Context, config *stateReaderConfig) (*stateReadClos
 	}
 }
 
-// hclStateReadr returns a migrate.StateReader that reads the state from the given HCL paths urls.
+// hclStateReadr returns a StateReader that reads the state from the given HCL paths urls.
 func hclStateReader(ctx context.Context, config *stateReaderConfig, urls []*url.URL) (*stateReadCloser, error) {
 	var client *sqlclient.Client
 	switch {
@@ -576,9 +515,122 @@ func hclStateReader(ctx context.Context, config *stateReaderConfig, urls []*url.
 	return t, nil
 }
 
+func sqlStateReader(ctx context.Context, config *stateReaderConfig, urls []*url.URL) (*stateReadCloser, error) {
+	if len(urls) != 1 {
+		return nil, fmt.Errorf("the provided SQL state must be either a single schema file or a migration directory, but %d paths were found", len(urls))
+	}
+	// Replaying a migration directory requires a dev connection.
+	if config.dev == nil {
+		return nil, errors.New("--dev-url cannot be empty")
+	}
+	var (
+		dir  migrate.Dir
+		opts []migrate.ReplayOption
+		path = filepath.Join(urls[0].Host, urls[0].Path)
+	)
+	switch fi, err := os.Stat(path); {
+	case err != nil:
+		return nil, err
+	// A single schema file.
+	case !fi.IsDir():
+		b, err := os.ReadFile(path)
+		if err != nil {
+			return nil, err
+		}
+		dir = &migrate.MemDir{}
+		if err := dir.WriteFile(fi.Name(), b); err != nil {
+			return nil, err
+		}
+	// A migration directory.
+	default:
+		if dir, err = dirURL(urls[0], false); err != nil {
+			return nil, err
+		}
+		if v := urls[0].Query().Get("version"); v != "" {
+			opts = append(opts, migrate.ReplayToVersion(v))
+		}
+	}
+	ex, err := migrate.NewExecutor(config.dev.Driver, dir, migrate.NopRevisionReadWriter{})
+	if err != nil {
+		return nil, err
+	}
+	sr, err := ex.Replay(ctx, func() migrate.StateReader {
+		if config.dev.URL.Schema != "" {
+			return migrate.SchemaConn(config.dev, "", nil)
+		}
+		return migrate.RealmConn(config.dev, &schema.InspectRealmOption{
+			Schemas: config.schemas,
+			Exclude: config.exclude,
+		})
+	}(), opts...)
+	if err != nil && !errors.Is(err, migrate.ErrNoPendingFiles) {
+		return nil, err
+	}
+	return &stateReadCloser{
+		StateReader: migrate.Realm(sr),
+		schema:      config.dev.URL.Schema,
+	}, nil
+}
+
 // Close redirects calls to Close to the enclosed io.Closer.
 func (sr *stateReadCloser) Close() {
 	if sr.Closer != nil {
 		sr.Closer.Close()
 	}
 }
+
+const (
+	extHCL = ".hcl"
+	extSQL = ".sql"
+)
+
+func filesExt(urls []*url.URL) (string, error) {
+	var path, ext string
+	set := func(curr string) error {
+		switch e := filepath.Ext(curr); {
+		case e != extHCL && e != extSQL:
+			return fmt.Errorf("unknown schema file: %q", curr)
+		case ext != "" && ext != e:
+			return fmt.Errorf("ambiguous schema: both SQL and HCL files found: %q, %q", path, curr)
+		default:
+			path, ext = curr, e
+			return nil
+		}
+	}
+	for _, u := range urls {
+		path := filepath.Join(u.Host, u.Path)
+		switch fi, err := os.Stat(path); {
+		case err != nil:
+			return "", err
+		case fi.IsDir():
+			files, err := os.ReadDir(path)
+			if err != nil {
+				return "", err
+			}
+			for _, f := range files {
+				switch filepath.Ext(f.Name()) {
+				// Ignore unknown extensions in case we read directories.
+				case extHCL, extSQL:
+					if err := set(f.Name()); err != nil {
+						return "", err
+					}
+				}
+			}
+		default:
+			if err := set(fi.Name()); err != nil {
+				return "", err
+			}
+		}
+	}
+	switch {
+	case ext != "":
+	case len(urls) == 1 && (urls[0].Host != "" || urls[0].Path != ""):
+		return "", fmt.Errorf(
+			"%q contains neither SQL nor HCL files",
+			filepath.Base(filepath.Join(urls[0].Host, urls[0].Path)),
+		)
+	default:
+		return "", errors.New("schema contains neither SQL nor HCL files")
+	}
+	return ext, nil
+}
diff --git a/sql/migrate/dir.go b/sql/migrate/dir.go
index 641a982a..ab4e1634 100644
--- a/sql/migrate/dir.go
+++ b/sql/migrate/dir.go
@@ -117,28 +117,11 @@ func (d *LocalDir) Files() ([]File, error) {
 
 // Checksum implements Dir.Checksum. By default, it calls Files() and creates a checksum from them.
 func (d *LocalDir) Checksum() (HashFile, error) {
-	var (
-		hs HashFile
-		h  = sha256.New()
-	)
 	files, err := d.Files()
 	if err != nil {
 		return nil, err
 	}
-	for _, f := range files {
-		if _, err = h.Write([]byte(f.Name())); err != nil {
-			return nil, err
-		}
-		// Check if this file contains an "atlas:sum" directive and if so, act to it.
-		if mode, ok := directive(string(f.Bytes()), directiveSum); ok && mode == sumModeIgnore {
-			continue
-		}
-		if _, err = h.Write(f.Bytes()); err != nil {
-			return nil, err
-		}
-		hs = append(hs, struct{ N, H string }{f.Name(), base64.StdEncoding.EncodeToString(h.Sum(nil))})
-	}
-	return hs, nil
+	return NewHashFile(files)
 }
 
 // LocalFile is used by LocalDir to implement the Scanner interface.
@@ -196,6 +179,65 @@ func (f LocalFile) Bytes() []byte {
 	return f.b
 }
 
+// MemDir provides an in-memory Dir implementation.
+type MemDir struct {
+	files map[string]File
+}
+
+// Open implements fs.FS.
+func (d *MemDir) Open(name string) (fs.File, error) {
+	var b []byte
+	switch name {
+	case HashFileName:
+		h, err := d.Checksum()
+		if err != nil {
+			return nil, err
+		}
+		if b, err = h.MarshalText(); err != nil {
+			return nil, err
+		}
+	default:
+		f, ok := d.files[name]
+		if !ok {
+			return nil, fs.ErrNotExist
+		}
+		b = f.Bytes()
+	}
+	return &memFile{
+		ReadCloser: io.NopCloser(bytes.NewReader(b)),
+	}, nil
+}
+
+// WriteFile adds a new file in-memory.
+func (d *MemDir) WriteFile(name string, data []byte) error {
+	if d.files == nil {
+		d.files = make(map[string]File)
+	}
+	d.files[name] = NewLocalFile(name, data)
+	return nil
+}
+
+// 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))
+	for _, f := range d.files {
+		files = append(files, f)
+	}
+	sort.Slice(files, func(i, j int) bool {
+		return files[i].Name() < files[j].Name()
+	})
+	return files, nil
+}
+
+// Checksum implements Dir.Checksum.
+func (d *MemDir) Checksum() (HashFile, error) {
+	files, err := d.Files()
+	if err != nil {
+		return nil, err
+	}
+	return NewHashFile(files)
+}
+
 var (
 	// templateFuncs contains the template.FuncMap for the DefaultFormatter.
 	templateFuncs = template.FuncMap{"now": func() string { return time.Now().UTC().Format("20060102150405") }}
@@ -261,11 +303,26 @@ const HashFileName = "atlas.sum"
 // HashFile represents the integrity sum file of the migration dir.
 type HashFile []struct{ N, H string }
 
-// HashSum reads the whole dir, sorts the files by name and creates a HashSum from its contents.
-//
-// Deprecated: Use Dir.Checksum instead.
-func HashSum(dir Dir) (HashFile, error) {
-	return dir.Checksum()
+// NewHashFile computes and returns a HashFile from the given directory's files.
+func NewHashFile(files []File) (HashFile, error) {
+	var (
+		hs HashFile
+		h  = sha256.New()
+	)
+	for _, f := range files {
+		if _, err := h.Write([]byte(f.Name())); err != nil {
+			return nil, err
+		}
+		// Check if this file contains an "atlas:sum" directive and if so, act to it.
+		if mode, ok := directive(string(f.Bytes()), directiveSum); ok && mode == sumModeIgnore {
+			continue
+		}
+		if _, err := h.Write(f.Bytes()); err != nil {
+			return nil, err
+		}
+		hs = append(hs, struct{ N, H string }{f.Name(), base64.StdEncoding.EncodeToString(h.Sum(nil))})
+	}
+	return hs, nil
 }
 
 // WriteSumFile writes the given HashFile to the Dir. If the file does not exist, it is created.
@@ -413,3 +470,15 @@ func readHashFile(dir Dir) (HashFile, error) {
 	}
 	return fh, nil
 }
+
+// memFile implements the File interface for a file in memory.
+type memFile struct{ io.ReadCloser }
+
+// Stat returns a zero FileInfo.
+func (m *memFile) Stat() (fs.FileInfo, error) { return m, nil }
+func (m *memFile) Name() string               { return "" }
+func (m *memFile) Size() int64                { return 0 }
+func (m *memFile) Mode() fs.FileMode          { return 0 }
+func (m *memFile) ModTime() time.Time         { return time.Time{} }
+func (m *memFile) IsDir() bool                { return false }
+func (m *memFile) Sys() interface{}           { return nil }


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/schema_test.go b/cmd/atlas/internal/cmdapi/schema_test.go
index 522bfdab..33f9aaf3 100644
--- a/cmd/atlas/internal/cmdapi/schema_test.go
+++ b/cmd/atlas/internal/cmdapi/schema_test.go
@@ -229,7 +229,7 @@ table "users" {
 	})
 }
 
-func TestCmdSchemaApply(t *testing.T) {
+func TestSchema_Apply(t *testing.T) {
 	const drvName = "checknormalizer"
 	// If no dev-database is given, there must not be a call to Driver.Normalize.
 	sqlclient.Register(
@@ -464,6 +464,87 @@ env "local" {
 	}, lines)
 }
 
+func TestSchema_ApplySQL(t *testing.T) {
+	t.Run("File", func(t *testing.T) {
+		db := openSQLite(t, "")
+		p := filepath.Join(t.TempDir(), "schema.sql")
+		require.NoError(t, os.WriteFile(p, []byte(`create table t1 (id int NOT NULL);`), 0600))
+		s, err := runCmd(
+			schemaApplyCmd(),
+			"-u", db,
+			"--dev-url", openSQLite(t, ""),
+			"--to", "file://"+p,
+			"--auto-approve",
+		)
+		require.NoError(t, err)
+		require.Equal(t, "-- Planned Changes:\n-- Create \"t1\" table\nCREATE TABLE `t1` (`id` int NOT NULL)\n", s)
+
+		s, err = runCmd(
+			schemaApplyCmd(),
+			"-u", db,
+			"--dev-url", openSQLite(t, ""),
+			"--to", "file://"+p,
+			"--auto-approve",
+		)
+		require.NoError(t, err)
+		require.Equal(t, "Schema is synced, no changes to be made.\n", s)
+	})
+	t.Run("Dir", func(t *testing.T) {
+		db := openSQLite(t, "")
+		s, err := runCmd(
+			schemaApplyCmd(),
+			"-u", db,
+			"--dev-url", openSQLite(t, ""),
+			"--to", "file://testdata/sqlite",
+			"--auto-approve",
+		)
+		require.NoError(t, err)
+		require.Equal(t, "-- Planned Changes:\n-- Create \"tbl\" table\nCREATE TABLE `tbl` (`col` int NOT NULL, `col_2` bigint NULL)\n", s)
+
+		s, err = runCmd(
+			schemaApplyCmd(),
+			"-u", db,
+			"--dev-url", openSQLite(t, ""),
+			"--to", "file://testdata/sqlite",
+			"--auto-approve",
+		)
+		require.NoError(t, err)
+		require.Equal(t, "Schema is synced, no changes to be made.\n", s)
+	})
+	t.Run("Error", func(t *testing.T) {
+		_, err := runCmd(
+			schemaApplyCmd(),
+			"-u", openSQLite(t, ""),
+			"--dev-url", openSQLite(t, ""),
+			"--to", "file://testdata/sqlite",
+			"--to", "file://testdata/sqlite2",
+			"--auto-approve",
+		)
+		require.EqualError(t, err, `the provided SQL state must be either a single schema file or a migration directory, but 2 paths were found`)
+
+		_, err = runCmd(
+			schemaApplyCmd(),
+			"-u", openSQLite(t, ""),
+			"--dev-url", openSQLite(t, ""),
+			"--to", "file://"+t.TempDir(),
+			"--auto-approve",
+		)
+		require.ErrorContains(t, err, `contains neither SQL nor HCL files`)
+
+		p := t.TempDir()
+		require.NoError(t, os.WriteFile(filepath.Join(p, "schema.sql"), []byte(`create table t1 (id int NOT NULL);`), 0600))
+		require.NoError(t, os.WriteFile(filepath.Join(p, "schema.hcl"), []byte(`schema "main" {}`), 0600))
+		_, err = runCmd(
+			schemaApplyCmd(),
+			"-u", openSQLite(t, ""),
+			"--dev-url", openSQLite(t, ""),
+			"--to", "file://"+p,
+			"--auto-approve",
+		)
+		require.EqualError(t, err, `ambiguous schema: both SQL and HCL files found: "schema.hcl", "schema.sql"`)
+	})
+}
+
 func TestSchema_InspectLog(t *testing.T) {
 	db := openSQLite(t, "create table t1 (id integer primary key);create table t2 (name text);")
 	cmd := schemaCmd()
diff --git a/sql/migrate/dir_test.go b/sql/migrate/dir_test.go
index d7c45868..d46b1a41 100644
--- a/sql/migrate/dir_test.go
+++ b/sql/migrate/dir_test.go
@@ -195,3 +195,23 @@ func TestLocalDir(t *testing.T) {
 	require.Equal(t, "2.10.x-20", files[1].Version())
 	require.Equal(t, "description", files[1].Desc())
 }
+
+func TestMemDir(t *testing.T) {
+	var d migrate.MemDir
+	files, err := d.Files()
+	require.NoError(t, err)
+	require.Empty(t, files)
+	require.NoError(t, d.WriteFile("1.sql", []byte("create table t(c int);")))
+	files, err = d.Files()
+	require.NoError(t, err)
+	require.Len(t, files, 1)
+	require.Equal(t, "1.sql", files[0].Name())
+	require.Equal(t, "1", files[0].Version())
+	require.Equal(t, "", files[0].Desc())
+	require.EqualValues(t, "create table t(c int);", files[0].Bytes())
+	hs1, err := d.Checksum()
+	require.NoError(t, err)
+	hs2, err := migrate.NewHashFile(files)
+	require.NoError(t, err)
+	require.Equal(t, hs1, hs2)
+}