1 files changed, 63 insertions, 8 deletions

diff --git a/cmd/praefect/subcmd_sql_migrate.go b/cmd/praefect/subcmd_sql_migrate.go
index 3bb4f9b11..1d4ee3420 100644
--- a/cmd/praefect/subcmd_sql_migrate.go
+++ b/cmd/praefect/subcmd_sql_migrate.go
@@ -3,24 +3,38 @@ package main
 import (
 	"flag"
 	"fmt"
+	"io"
+	"time"
 
+	migrate "github.com/rubenv/sql-migrate"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/glsql"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/migrations"
 )
 
-const sqlMigrateCmdName = "sql-migrate"
+const (
+	sqlMigrateCmdName = "sql-migrate"
+	timeFmt           = "2006-01-02T15:04:05"
+)
 
 type sqlMigrateSubcommand struct {
+	w             io.Writer
 	ignoreUnknown bool
+	verbose       bool
+}
+
+func newSQLMigrateSubCommand(writer io.Writer) *sqlMigrateSubcommand {
+	return &sqlMigrateSubcommand{w: writer}
 }
 
-func (s *sqlMigrateSubcommand) FlagSet() *flag.FlagSet {
+func (cmd *sqlMigrateSubcommand) FlagSet() *flag.FlagSet {
 	flags := flag.NewFlagSet(sqlMigrateCmdName, flag.ExitOnError)
-	flags.BoolVar(&s.ignoreUnknown, "ignore-unknown", true, "ignore unknown migrations (default is true)")
+	flags.BoolVar(&cmd.ignoreUnknown, "ignore-unknown", true, "ignore unknown migrations (default is true)")
+	flags.BoolVar(&cmd.verbose, "verbose", false, "show text of migration query (default is false)")
 	return flags
 }
 
-func (s *sqlMigrateSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
+func (cmd *sqlMigrateSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
 	const subCmd = progname + " " + sqlMigrateCmdName
 
 	db, clean, err := openDB(conf.DB)
@@ -29,11 +43,52 @@ func (s *sqlMigrateSubcommand) Exec(flags *flag.FlagSet, conf config.Config) err
 	}
 	defer clean()
 
-	n, err := glsql.Migrate(db, s.ignoreUnknown)
-	if err != nil {
-		return fmt.Errorf("%s: fail: %v", subCmd, err)
+	migrationSet := migrate.MigrationSet{
+		IgnoreUnknown: cmd.ignoreUnknown,
+		TableName:     migrations.MigrationTableName,
+	}
+
+	planSource := &migrate.MemoryMigrationSource{
+		Migrations: migrations.All(),
+	}
+
+	// Find all migrations that are currently down.
+	planMigrations, _, _ := migrationSet.PlanMigration(db, "postgres", planSource, migrate.Up, 0)
+
+	if len(planMigrations) == 0 {
+		fmt.Fprintf(cmd.w, "%s: all migrations are up\n", subCmd)
+		return nil
+	}
+	fmt.Fprintf(cmd.w, "%s: migrations to apply: %d\n\n", subCmd, len(planMigrations))
+
+	executed := 0
+	for _, mig := range planMigrations {
+		fmt.Fprintf(cmd.w, "=  %s %v: migrating\n", time.Now().Format(timeFmt), mig.Id)
+		start := time.Now()
+
+		if cmd.verbose {
+			fmt.Fprintf(cmd.w, "\t%v\n", mig.Up)
+		}
+
+		n, err := glsql.MigrateSome(mig.Migration, db, cmd.ignoreUnknown)
+		if err != nil {
+			return fmt.Errorf("%s: fail: %v", time.Now().Format(timeFmt), err)
+		}
+
+		if n > 0 {
+			fmt.Fprintf(cmd.w, "== %s %v: applied (%s)\n", time.Now().Format(timeFmt), mig.Id, time.Since(start))
+
+			// Additional migrations were run. No harm, but prevents us from tracking their execution duration.
+			if n > 1 {
+				fmt.Fprintf(cmd.w, "warning: %v additional migrations were applied successfully\n", n-1)
+			}
+		} else {
+			fmt.Fprintf(cmd.w, "== %s %v: skipped (%s)\n", time.Now().Format(timeFmt), mig.Id, time.Since(start))
+		}
+
+		executed += n
 	}
 
-	fmt.Printf("%s: OK (applied %d migrations)\n", subCmd, n)
+	fmt.Fprintf(cmd.w, "\n%s: OK (applied %d migrations)\n", subCmd, executed)
 	return nil
 }