migrate: Print execution time of migrations

Debugging slow migrations, as in issue #3973[0] is difficult
because there is no indication which migration is being run. Admins
will need to check `pg_stat_activity`, then search the Gitaly source
for that query to find out which migration is misbehaving.

This commit updates the `sql-migrate` subcommand to execute migrations
one by one and displays their execution times.

This requires a bit of work as the `Exec` API exposed by `sql-migrate`
does not provide a simple way to execute only one migration. The obvious
option of passing a slice with a single migration does not work due to
how `ToApply()`[1] works. If `current` is not empty, which will be when one
or more migrations are up, then it will iterate over the migration slice
until `current` is found, returning the remainder of the slice. If
`current` is not present in the slice, then an empty slice is returned
and the target migration is not executed.

To work around this, we introduce the `MigrateSome()` function that
calculates all migrations prior to the target migration and passes them
along with the target to `sql-migrate.Exec()`.

Example output:

```
praefect sql-migrate: migrations to apply: 3

=  2022-02-01T17:01:37 20211025113445_remove_cancelled_replication_events: migrating
== 2022-02-01T17:01:37 20211025113445_remove_cancelled_replication_events: applied (9.010133ms)
=  2022-02-01T17:01:37 20211105100456_check_non_null_repository_id: migrating
== 2022-02-01T17:01:37 20211105100456_check_non_null_repository_id: applied (8.423652ms)
=  2022-02-01T17:01:37 20211105100458_non_nullable_repository_id: migrating
== 2022-02-01T17:01:37 20211105100458_non_nullable_repository_id: applied (12.438743ms)

praefect sql-migrate: OK (applied 3 migrations)
```

[0] https://gitlab.com/gitlab-org/gitaly/-/issues/3973
[1] https://github.com/rubenv/sql-migrate/blob/06338513c2375c000ad94f050dce5ee2a5bdcdc6/migrate.go#L628-L637

Changelog: changed

6 files changed, 240 insertions, 9 deletions

diff --git a/cmd/praefect/subcmd.go b/cmd/praefect/subcmd.go
index 8f2858e3e..3a4148244 100644
--- a/cmd/praefect/subcmd.go
+++ b/cmd/praefect/subcmd.go
@@ -33,7 +33,7 @@ const (
 
 var subcommands = map[string]subcmd{
 	sqlPingCmdName:                &sqlPingSubcommand{},
-	sqlMigrateCmdName:             &sqlMigrateSubcommand{},
+	sqlMigrateCmdName:             newSQLMigrateSubCommand(os.Stdout),
 	dialNodesCmdName:              newDialNodesSubcommand(os.Stdout),
 	sqlMigrateDownCmdName:         &sqlMigrateDownSubcommand{},
 	sqlMigrateStatusCmdName:       &sqlMigrateStatusSubcommand{},
diff --git a/cmd/praefect/subcmd_sql_migrate.go b/cmd/praefect/subcmd_sql_migrate.go
index 3bb4f9b11..27fd2b9ff 100644
--- a/cmd/praefect/subcmd_sql_migrate.go
+++ b/cmd/praefect/subcmd_sql_migrate.go
@@ -3,24 +3,36 @@ 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
 }
 
-func (s *sqlMigrateSubcommand) FlagSet() *flag.FlagSet {
+func newSQLMigrateSubCommand(writer io.Writer) *sqlMigrateSubcommand {
+	return &sqlMigrateSubcommand{w: writer}
+}
+
+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)")
 	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 +41,48 @@ 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()
+
+		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
 }
diff --git a/cmd/praefect/subcmd_sql_migrate_test.go b/cmd/praefect/subcmd_sql_migrate_test.go
new file mode 100644
index 000000000..c89965d6f
--- /dev/null
+++ b/cmd/praefect/subcmd_sql_migrate_test.go
@@ -0,0 +1,65 @@
+package main
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/migrations"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper/testdb"
+)
+
+func TestSubCmdSqlMigrate(t *testing.T) {
+	db := testdb.New(t)
+	dbCfg := testdb.GetConfig(t, db.Name)
+	cfg := config.Config{DB: dbCfg}
+
+	migrationCt := len(migrations.All())
+
+	for _, tc := range []struct {
+		desc           string
+		up             int
+		expectedOutput []string
+	}{
+		{
+			desc:           "All migrations up",
+			up:             migrationCt,
+			expectedOutput: []string{"praefect sql-migrate: all migrations are up"},
+		},
+		{
+			desc: "All migrations down",
+			up:   0,
+			expectedOutput: []string{
+				fmt.Sprintf("praefect sql-migrate: migrations to apply: %d", migrationCt),
+				"20200109161404_hello_world: migrating",
+				"20200109161404_hello_world: applied (",
+				fmt.Sprintf("praefect sql-migrate: OK (applied %d migrations)", migrationCt),
+			},
+		},
+		{
+			desc: "Some migrations down",
+			up:   10,
+			expectedOutput: []string{
+				fmt.Sprintf("praefect sql-migrate: migrations to apply: %d", migrationCt-10),
+				"20201126165633_repository_assignments_table: migrating",
+				"20201126165633_repository_assignments_table: applied (",
+				fmt.Sprintf("praefect sql-migrate: OK (applied %d migrations)", migrationCt-10),
+			},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			testdb.SetMigrations(t, db, cfg, tc.up)
+
+			var stdout bytes.Buffer
+			migrateCmd := sqlMigrateSubcommand{w: &stdout, ignoreUnknown: true}
+			assert.NoError(t, migrateCmd.Exec(flag.NewFlagSet("", flag.PanicOnError), cfg))
+
+			for _, out := range tc.expectedOutput {
+				assert.Contains(t, stdout.String(), out)
+			}
+		})
+	}
+}
diff --git a/internal/praefect/datastore/glsql/postgres.go b/internal/praefect/datastore/glsql/postgres.go
index 04983f059..24c4093d6 100644
--- a/internal/praefect/datastore/glsql/postgres.go
+++ b/internal/praefect/datastore/glsql/postgres.go
@@ -108,6 +108,41 @@ func Migrate(db *sql.DB, ignoreUnknown bool) (int, error) {
 	return migrationSet.Exec(db, "postgres", migrationSource, migrate.Up)
 }
 
+// MigrateSome will apply migration m and all unapplied migrations with earlier ids.
+// To ensure a single migration is executed, run sql-migrate.PlanMigration and call
+// MigrateSome for each migration returned.
+func MigrateSome(m *migrate.Migration, db *sql.DB, ignoreUnknown bool) (int, error) {
+	migrationSet := migrate.MigrationSet{
+		IgnoreUnknown: ignoreUnknown,
+		TableName:     migrations.MigrationTableName,
+	}
+
+	// sql-migrate.ToApply() expects all migrations prior to the final migration be present in the
+	// in the slice. If we pass in only the target migration it will not be executed.
+	migs := leadingMigrations(m)
+
+	migrationSource := &migrate.MemoryMigrationSource{
+		Migrations: migs,
+	}
+
+	return migrationSet.Exec(db, "postgres", migrationSource, migrate.Up)
+}
+
+// Create a slice of all migrations up to and including the one to be applied.
+func leadingMigrations(target *migrate.Migration) []*migrate.Migration {
+	allMigrations := migrations.All()
+
+	for i, m := range allMigrations {
+		if m.Id == target.Id {
+			return allMigrations[:i+1]
+		}
+	}
+
+	// Planned migration not found in migrations.All(), assume it is more recent
+	// and return all migrations.
+	return allMigrations
+}
+
 // Querier is an abstraction on *sql.DB and *sql.Tx that allows to use their methods without awareness about actual type.
 type Querier interface {
 	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
diff --git a/internal/praefect/datastore/glsql/postgres_test.go b/internal/praefect/datastore/glsql/postgres_test.go
index 8aac2694e..039145a5f 100644
--- a/internal/praefect/datastore/glsql/postgres_test.go
+++ b/internal/praefect/datastore/glsql/postgres_test.go
@@ -8,10 +8,12 @@ import (
 	"time"
 
 	"github.com/jackc/pgconn"
+	migrate "github.com/rubenv/sql-migrate"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"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"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper/testdb"
 )
@@ -321,3 +323,58 @@ func TestIsUniqueViolation(t *testing.T) {
 		})
 	}
 }
+
+func TestMigrateSome(t *testing.T) {
+	db := testdb.New(t)
+	dbCfg := testdb.GetConfig(t, db.Name)
+	cfg := config.Config{DB: dbCfg}
+
+	migs := migrations.All()
+	migrationCt := len(migs)
+
+	for _, tc := range []struct {
+		desc      string
+		up        int
+		executed  int
+		migration *migrate.Migration
+	}{
+		{
+			desc:      "All migrations up",
+			up:        migrationCt,
+			executed:  0,
+			migration: migs[migrationCt-1],
+		},
+		{
+			desc:      "Apply only first migration",
+			up:        0,
+			executed:  1,
+			migration: migs[0],
+		},
+		{
+			desc:      "Apply only last migration",
+			up:        migrationCt - 1,
+			executed:  1,
+			migration: migs[migrationCt-1],
+		},
+		{
+			desc:      "Apply only 10th migration",
+			up:        9,
+			executed:  1,
+			migration: migs[9],
+		},
+		{
+			desc:      "Apply 5th to 10th migrations",
+			up:        5,
+			executed:  5,
+			migration: migs[9],
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			testdb.SetMigrations(t, db, cfg, tc.up)
+
+			n, err := glsql.MigrateSome(tc.migration, db.DB, true)
+			assert.NoError(t, err)
+			assert.Equal(t, tc.executed, n)
+		})
+	}
+}
diff --git a/internal/testhelper/testdb/db.go b/internal/testhelper/testdb/db.go
index b9ea659ca..f085ebc58 100644
--- a/internal/testhelper/testdb/db.go
+++ b/internal/testhelper/testdb/db.go
@@ -19,6 +19,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"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"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 )
 
@@ -401,3 +402,27 @@ func WaitForBlockedQuery(ctx context.Context, t testing.TB, db glsql.Querier, qu
 		}
 	}
 }
+
+// SetMigrations ensures the requested number of migrations are up and the remainder are down.
+func SetMigrations(t testing.TB, db DB, cfg config.Config, up int) {
+	// Ensure all migrations are up first
+	_, err := glsql.Migrate(db.DB, true)
+	require.NoError(t, err)
+
+	migrationCt := len(migrations.All())
+
+	if up < migrationCt {
+		down := migrationCt - up
+
+		migrationSet := migrate.MigrationSet{
+			TableName: migrations.MigrationTableName,
+		}
+		ms := &migrate.MemoryMigrationSource{Migrations: migrations.All()}
+
+		// It would be preferable to use migrate.MigrateDown() here, but that introduces
+		// a circular dependency between testdb and datastore.
+		n, err := migrationSet.ExecMax(db.DB, "postgres", ms, migrate.Down, down)
+		require.NoError(t, err)
+		require.Equal(t, down, n)
+	}
+}