Make pgroll pull pull only migrations that don't exist in target directory (#811)

Make the `pgroll pull` command pull only those migrations from the
target database that are missing in the local directory.

This ensures that existing migrations in the local directory are not
overwritten or modified when pulling from the target database, for
example with formatting or order-of-field changes.

Author: andrew-farries

cmd/pull.go

@@ -4,11 +4,11 @@ package cmd
 
 import (
 	"fmt"
-
-	"github.com/xataio/pgroll/cmd/flags"
-	"github.com/xataio/pgroll/pkg/state"
+	"os"
+	"path/filepath"
 
 	"github.com/spf13/cobra"
+	"github.com/xataio/pgroll/pkg/migrations"
 )
 
 func pullCmd() *cobra.Command {
@@ -27,34 +27,51 @@ func pullCmd() *cobra.Command {
 			ctx := cmd.Context()
 			targetDir := args[0]
 
-			state, err := state.New(ctx, flags.PostgresURL(), flags.StateSchema())
+			m, err := NewRoll(ctx)
 			if err != nil {
 				return err
 			}
-			defer state.Close()
+			defer m.Close()
 
 			// Ensure that pgroll is initialized
-			ok, err := state.IsInitialized(cmd.Context())
+			ok, err := m.State().IsInitialized(cmd.Context())
 			if err != nil {
 				return err
 			}
 			if !ok {
 				return errPGRollNotInitialized
 			}
 
-			migs, err := state.SchemaHistory(ctx, flags.Schema())
+			// Ensure that the target directory is valid, creating it if it doesn't
+			// exist
+			_, err = os.Stat(targetDir)
+			if err != nil {
+				if os.IsNotExist(err) {
+					err := os.MkdirAll(targetDir, 0o755)
+					if err != nil {
+						return fmt.Errorf("failed to create target directory: %w", err)
+					}
+				} else {
+					return fmt.Errorf("failed to stat directory: %w", err)
+				}
+			}
+
+			// Get the list of missing migrations (those that have been applied to
+			// the target database but are missing in the local directory).
+			migs, err := m.MissingMigrations(ctx, os.DirFS(targetDir))
 			if err != nil {
-				return fmt.Errorf("failed to read schema history: %w", err)
+				return fmt.Errorf("failed to read migrations from target directory: %w", err)
 			}
 
+			// Write the missing migrations to the target directory
 			for i, mig := range migs {
 				prefix := ""
 				if withPrefixes {
 					prefix = fmt.Sprintf("%04d", i+1) + "_"
 				}
-				err := mig.WriteToFile(targetDir, prefix, useJSON)
+				err := writeMigrationToFile(mig, targetDir, prefix, useJSON)
 				if err != nil {
-					return fmt.Errorf("failed to write migration %q: %w", mig.Migration.Name, err)
+					return fmt.Errorf("failed to write migration %q: %w", mig.Name, err)
 				}
 			}
 			return nil
@@ -66,3 +83,33 @@ func pullCmd() *cobra.Command {
 
 	return pullCmd
 }
+
+// WriteToFile writes the migration to a file in `targetDir`, prefixing the
+// filename with `prefix`. The output format defaults to YAML, but can
+// be changed to JSON by setting `useJSON` to true.
+func writeMigrationToFile(m *migrations.Migration, targetDir, prefix string, useJSON bool) error {
+	err := os.MkdirAll(targetDir, 0o755)
+	if err != nil {
+		return err
+	}
+
+	suffix := "yaml"
+	if useJSON {
+		suffix = "json"
+	}
+
+	fileName := fmt.Sprintf("%s%s.%s", prefix, m.Name, suffix)
+	filePath := filepath.Join(targetDir, fileName)
+
+	file, err := os.Create(filePath)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	if useJSON {
+		return m.WriteAsJSON(file)
+	} else {
+		return m.WriteAsYAML(file)
+	}
+}

pkg/migrations/migrations.go

@@ -4,9 +4,12 @@ package migrations
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"io"
 
 	_ "github.com/lib/pq"
+	"sigs.k8s.io/yaml"
 
 	"github.com/xataio/pgroll/pkg/db"
 	"github.com/xataio/pgroll/pkg/schema"
@@ -102,3 +105,22 @@ func (m *Migration) ContainsRawSQLOperation() bool {
 	}
 	return false
 }
+
+// WriteAsJSON writes the migration to the given writer in JSON format
+func (m *Migration) WriteAsJSON(w io.Writer) error {
+	encoder := json.NewEncoder(w)
+	encoder.SetIndent("", "  ")
+
+	return encoder.Encode(m)
+}
+
+// WriteAsYAML writes the migration to the given writer in YAML format
+func (m *Migration) WriteAsYAML(w io.Writer) error {
+	yml, err := yaml.Marshal(m)
+	if err != nil {
+		return err
+	}
+
+	_, err = w.Write(yml)
+	return err
+}

pkg/roll/missing.go

@@ -0,0 +1,61 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package roll
+
+import (
+	"context"
+	"fmt"
+	"io/fs"
+
+	"github.com/xataio/pgroll/pkg/migrations"
+)
+
+// MissingMigrations returns the slice of migrations that have been applied to
+// the target database but are missing from the local migrations directory
+// `dir`.
+func (m *Roll) MissingMigrations(ctx context.Context, dir fs.FS) ([]*migrations.Migration, error) {
+	// Determine the latest version of the database
+	latestVersion, err := m.State().LatestVersion(ctx, m.Schema())
+	if err != nil {
+		return nil, fmt.Errorf("determining latest version: %w", err)
+	}
+
+	// If no migrations are applied, return a nil slice
+	if latestVersion == nil {
+		return nil, nil
+	}
+
+	// Collect all migration files from the directory
+	files, err := migrations.CollectFilesFromDir(dir)
+	if err != nil {
+		return nil, fmt.Errorf("reading migration files: %w", err)
+	}
+
+	// Create a set of local migration names for fast lookup
+	localMigNames := make(map[string]struct{}, len(files))
+	for _, file := range files {
+		mig, err := migrations.ReadMigration(dir, file)
+		if err != nil {
+			return nil, fmt.Errorf("reading migration file %s: %w", file, err)
+		}
+		localMigNames[mig.Name] = struct{}{}
+	}
+
+	// Get the full schema history from the database
+	history, err := m.State().SchemaHistory(ctx, m.Schema())
+	if err != nil {
+		return nil, fmt.Errorf("reading schema history: %w", err)
+	}
+
+	// Find all migrations that have been applied to the database but are missing
+	// from the local directory
+	migs := make([]*migrations.Migration, 0, len(history))
+	for _, h := range history {
+		if _, ok := localMigNames[h.Migration.Name]; ok {
+			continue
+		}
+		migs = append(migs, &h.Migration)
+	}
+
+	return migs, nil
+}