sg: Add sg migration revert <commit> (sourcegraph#30772)

Author: efritz

dev/sg/internal/db/db.go

@@ -65,6 +65,12 @@ var (
 	DefaultDatabase = databases[0]
 )
 
+func Databases() []Database {
+	c := make([]Database, len(databases))
+	copy(c, databases)
+	return c
+}
+
 func DatabaseNames() []string {
 	databaseNames := make([]string, 0, len(databases))
 	for _, database := range databases {

dev/sg/internal/migration/add.go

@@ -7,13 +7,15 @@ import (
 	"time"
 
 	"github.com/sourcegraph/sourcegraph/dev/sg/internal/db"
+	"github.com/sourcegraph/sourcegraph/dev/sg/internal/stdout"
+	"github.com/sourcegraph/sourcegraph/lib/output"
 )
 
-const metadataFileTemplate = `name: %s
+const newMetadataFileTemplate = `name: %s
 parents: [%s]
 `
 
-const upMigrationFileTemplate = `BEGIN;
+const newUpMigrationFileTemplate = `BEGIN;
 
 -- Perform migration here.
 --
@@ -29,7 +31,7 @@ const upMigrationFileTemplate = `BEGIN;
 COMMIT;
 `
 
-const downMigrationFileTemplate = `BEGIN;
+const newDownMigrationFileTemplate = `BEGIN;
 
 -- Undo the changes made in the up migration
 
@@ -38,10 +40,14 @@ COMMIT;
 
 // Add creates a new directory with stub migration files in the given schema and returns the
 // names of the newly created files. If there was an error, the filesystem is rolled-back.
-func Add(database db.Database, migrationName string) (up, down, metadata string, _ error) {
+func Add(database db.Database, migrationName string) error {
+	return add(database, migrationName, newUpMigrationFileTemplate, newDownMigrationFileTemplate)
+}
+
+func add(database db.Database, migrationName, upMigrationFileTemplate, downMigrationFileTemplate string) error {
 	definitions, err := readDefinitions(database)
 	if err != nil {
-		return "", "", "", err
+		return err
 	}
 
 	leaves := definitions.Leaves()
@@ -50,23 +56,27 @@ func Add(database db.Database, migrationName string) (up, down, metadata string,
 		parents = append(parents, leaf.ID)
 	}
 
-	id := int(time.Now().UTC().Unix())
-
-	upPath, downPath, metadataPath, err := makeMigrationFilenames(database, id)
+	files, err := makeMigrationFilenames(database, int(time.Now().UTC().Unix()))
 	if err != nil {
-		return "", "", "", err
+		return err
 	}
 
 	contents := map[string]string{
-		upPath:       upMigrationFileTemplate,
-		downPath:     downMigrationFileTemplate,
-		metadataPath: fmt.Sprintf(metadataFileTemplate, migrationName, strings.Join(intsToStrings(parents), ", ")),
+		files.UpFile:       upMigrationFileTemplate,
+		files.DownFile:     downMigrationFileTemplate,
+		files.MetadataFile: fmt.Sprintf(newMetadataFileTemplate, migrationName, strings.Join(intsToStrings(parents), ", ")),
 	}
 	if err := writeMigrationFiles(contents); err != nil {
-		return "", "", "", err
+		return err
 	}
 
-	return upPath, downPath, metadataPath, nil
+	block := stdout.Out.Block(output.Linef("", output.StyleBold, "Migration files created"))
+	block.Writef("Up query file: %s", files.UpFile)
+	block.Writef("Down query file: %s", files.DownFile)
+	block.Writef("Metadata file: %s", files.MetadataFile)
+	block.Close()
+
+	return nil
 }
 
 func intsToStrings(ints []int) []string {

dev/sg/internal/migration/revert.go

@@ -0,0 +1,97 @@
+package migration
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/sourcegraph/sourcegraph/dev/sg/internal/db"
+	"github.com/sourcegraph/sourcegraph/dev/sg/internal/run"
+	"github.com/sourcegraph/sourcegraph/dev/sg/internal/stdout"
+	"github.com/sourcegraph/sourcegraph/internal/database/migration/definition"
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+	"github.com/sourcegraph/sourcegraph/lib/output"
+)
+
+// Revert creates a new migration that reverts the set of migrations from a target commit.
+func Revert(databases []db.Database, commit string) error {
+	versionsByDatabase := make(map[string][]int, len(databases))
+	for _, database := range databases {
+		definitions, err := readDefinitions(database)
+		if err != nil {
+			return err
+		}
+
+		versions, err := selectMigrationsDefinedInCommit(database, definitions, commit)
+		if err != nil {
+			return err
+		}
+
+		versionsByDatabase[database.Name] = versions
+	}
+
+	redacted := false
+	for name, versions := range versionsByDatabase {
+		if len(versions) == 0 {
+			continue
+		}
+		redacted = true
+
+		var (
+			database, _ = db.DatabaseByName(name)
+			upPaths     = make([]string, 0, len(versions))
+			downQueries = make([]string, 0, len(versions))
+		)
+
+		for _, version := range versions {
+			files, err := makeMigrationFilenames(database, version)
+			if err != nil {
+				return err
+			}
+
+			downQuery, err := os.ReadFile(files.DownFile)
+			if err != nil {
+				return err
+			}
+			upPaths = append(upPaths, files.UpFile)
+			downQueries = append(downQueries, string(downQuery))
+
+			contents := map[string]string{
+				files.UpFile: "-- REDACTED\n",
+			}
+			if err := writeMigrationFiles(contents); err != nil {
+				return err
+			}
+		}
+
+		block := stdout.Out.Block(output.Linef("", output.StyleBold, "Migration files redacted"))
+		for _, path := range upPaths {
+			block.Writef("Up query file: %s", path)
+		}
+		block.Close()
+
+		if err := add(database, fmt.Sprintf("revert %s", commit), strings.Join(downQueries, "\n\n"), "-- No-op\n"); err != nil {
+			return err
+		}
+	}
+	if !redacted {
+		return errors.Newf("No migrations defined on commit %q", commit)
+	}
+
+	return nil
+}
+
+// selectMigrationsDefinedInCommit returns the identifiers of migrations defined in the given
+// commit for the given schema.a
+func selectMigrationsDefinedInCommit(database db.Database, ds *definition.Definitions, commit string) ([]int, error) {
+	migrationsDir := filepath.Join("migrations", database.Name)
+
+	output, err := run.GitCmd("diff", "--name-only", commit+".."+commit+"~1", migrationsDir)
+	if err != nil {
+		return nil, err
+	}
+
+	versions := parseVersions(strings.Split(output, "\n"), migrationsDir)
+	return versions, nil
+}

dev/sg/internal/migration/squash.go

@@ -54,24 +54,24 @@ func Squash(database db.Database, commit string) error {
 	block := stdout.Out.Block(output.Linef("", output.StyleBold, "Updated filesystem"))
 	defer block.Close()
 
-	upPath, downPath, metadataPath, err := makeMigrationFilenames(database, newRoot)
+	files, err := makeMigrationFilenames(database, newRoot)
 	if err != nil {
 		return err
 	}
 
 	contents := map[string]string{
-		upPath:       squashedUpMigration,
-		downPath:     squashedDownMigration,
-		metadataPath: "name: 'squashed migrations'\n",
+		files.UpFile:       squashedUpMigration,
+		files.DownFile:     squashedDownMigration,
+		files.MetadataFile: "name: 'squashed migrations'\n",
 	}
 
 	if err := writeMigrationFiles(contents); err != nil {
 		return err
 	}
 
-	block.Writef("Created: %s", upPath)
-	block.Writef("Created: %s", downPath)
-	block.Writef("Created: %s", metadataPath)
+	block.Writef("Created: %s", files.UpFile)
+	block.Writef("Created: %s", files.DownFile)
+	block.Writef("Created: %s", files.MetadataFile)
 
 	// Remove the migration file pairs that were just squashed
 	filenames, err := removeAncestorsOf(database, definitions, newRoot)
@@ -97,16 +97,8 @@ func selectNewRootMigration(database db.Database, ds *definition.Definitions, co
 	if err != nil {
 		return 0, false, err
 	}
-	lines := strings.Split(output, "\n")
 
-	versions := make([]int, 0, len(lines))
-	for _, filename := range lines {
-		if version, err := strconv.Atoi(strings.Split(filename, "_")[0]); err == nil {
-			versions = append(versions, version)
-		}
-	}
-
-	ds, err = ds.Filter(versions)
+	ds, err = ds.Filter(parseVersions(strings.Split(output, "\n"), migrationsDir))
 	if err != nil {
 		return 0, false, err
 	}

dev/sg/internal/migration/util.go

@@ -4,6 +4,9 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"sort"
+	"strconv"
+	"strings"
 
 	"github.com/sourcegraph/sourcegraph/dev/sg/internal/db"
 	"github.com/sourcegraph/sourcegraph/dev/sg/root"
@@ -20,17 +23,23 @@ func readDefinitions(database db.Database) (*definition.Definitions, error) {
 	return definition.ReadDefinitions(fs)
 }
 
+type MigrationFiles struct {
+	UpFile       string
+	DownFile     string
+	MetadataFile string
+}
+
 // makeMigrationFilenames makes a pair of (absolute) paths to migration files with the given migration index.
-func makeMigrationFilenames(database db.Database, migrationIndex int) (up, down, metadata string, _ error) {
+func makeMigrationFilenames(database db.Database, migrationIndex int) (MigrationFiles, error) {
 	baseDir, err := migrationDirectoryForDatabase(database)
 	if err != nil {
-		return "", "", "", err
+		return MigrationFiles{}, err
 	}
 
 	upPath := filepath.Join(baseDir, fmt.Sprintf("%d/up.sql", migrationIndex))
 	downPath := filepath.Join(baseDir, fmt.Sprintf("%d/down.sql", migrationIndex))
 	metadataPath := filepath.Join(baseDir, fmt.Sprintf("%d/metadata.yaml", migrationIndex))
-	return upPath, downPath, metadataPath, nil
+	return MigrationFiles{upPath, downPath, metadataPath}, nil
 }
 
 // migrationDirectoryForDatabase returns the directory where migration files are stored for the
@@ -67,3 +76,39 @@ func writeMigrationFiles(contents map[string]string) (err error) {
 
 	return nil
 }
+
+// parseVersions takes a list of filepaths (the output of some git command) and a base
+// migrations directory and returns the versions of migrations present in the list.
+func parseVersions(lines []string, migrationsDir string) []int {
+	var (
+		pathSeparator       = string(os.PathSeparator)
+		prefixesToTrim      = []string{migrationsDir, pathSeparator}
+		separatorsToSplitBy = []string{pathSeparator, "_"}
+	)
+
+	versionMap := make(map[int]struct{}, len(lines))
+	for _, rawVersion := range lines {
+		// Remove leading migration directory if it exists
+		for _, prefix := range prefixesToTrim {
+			rawVersion = strings.TrimPrefix(rawVersion, prefix)
+		}
+
+		// Remove trailing filepath (if dir) or name prefix (if old migration)
+		for _, separator := range separatorsToSplitBy {
+			rawVersion = strings.Split(rawVersion, separator)[0]
+		}
+
+		// Should be left with only a version number
+		if version, err := strconv.Atoi(rawVersion); err == nil {
+			versionMap[version] = struct{}{}
+		}
+	}
+
+	versions := make([]int, 0, len(versionMap))
+	for version := range versionMap {
+		versions = append(versions, version)
+	}
+	sort.Ints(versions)
+
+	return versions
+}