Praefect: add migration check

This commit adds a migration check that checks if all praefect
migrations have been run.

Changelog: added

2 files changed, 105 insertions, 0 deletions

diff --git a/internal/praefect/checks.go b/internal/praefect/checks.go
index 2da522ad9..f80108303 100644
--- a/internal/praefect/checks.go
+++ b/internal/praefect/checks.go
@@ -32,3 +32,40 @@ type Check struct {
 
 // CheckFunc is a function type that takes a praefect config and returns a Check
 type CheckFunc func(conf config.Config) *Check
+
+// NewPraefectMigrationCheck returns a Check that checks if all praefect migrations have run
+func NewPraefectMigrationCheck(conf config.Config) *Check {
+	return &Check{
+		Name:        "praefect migrations",
+		Description: "confirms whether or not all praefect migrations have run",
+		Run: func() error {
+			db, err := glsql.OpenDB(conf.DB)
+			if err != nil {
+				return err
+			}
+			defer db.Close()
+
+			migrationSet := migrate.MigrationSet{
+				TableName: migrations.MigrationTableName,
+			}
+
+			migrationSource := &migrate.MemoryMigrationSource{
+				Migrations: migrations.All(),
+			}
+
+			migrationsNeeded, _, err := migrationSet.PlanMigration(db, "postgres", migrationSource, migrate.Up, 0)
+			if err != nil {
+				return err
+			}
+
+			missingMigrations := len(migrationsNeeded)
+
+			if missingMigrations > 0 {
+				return fmt.Errorf("%d migrations have not been run", missingMigrations)
+			}
+
+			return nil
+		},
+		Severity: Fatal,
+	}
+}
diff --git a/internal/praefect/checks_test.go b/internal/praefect/checks_test.go
new file mode 100644
index 000000000..9e6a6f1b5
--- /dev/null
+++ b/internal/praefect/checks_test.go
@@ -0,0 +1,68 @@
+package praefect
+
+import (
+	"fmt"
+	"testing"
+
+	"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"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/glsql"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/migrations"
+)
+
+func TestPraefectMigrations_success(t *testing.T) {
+	testCases := []struct {
+		desc        string
+		prepare     func(cfg config.Config) error
+		expectedErr error
+	}{
+		{
+			desc: "no migrations have run",
+			prepare: func(cfg config.Config) error {
+				_, err := datastore.MigrateDown(cfg, len(migrations.All()))
+				if err != nil {
+					return err
+				}
+
+				return nil
+			},
+			expectedErr: fmt.Errorf("%d migrations have not been run", len(migrations.All())),
+		},
+		{
+			desc: "some migrations have run",
+			prepare: func(cfg config.Config) error {
+				_, err := datastore.MigrateDown(cfg, 3)
+				if err != nil {
+					return err
+				}
+
+				return nil
+			},
+			expectedErr: fmt.Errorf("3 migrations have not been run"),
+		},
+		{
+			desc: "all migrations have run",
+			prepare: func(cfg config.Config) error {
+				return nil
+			},
+			expectedErr: nil,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			var cfg config.Config
+			db := glsql.NewDB(t)
+			cfg.DB = glsql.GetDBConfig(t, db.Name)
+
+			require.NoError(t, tc.prepare(cfg))
+
+			migrationCheck := NewPraefectMigrationCheck(cfg)
+			assert.Equal(t, "praefect migrations", migrationCheck.Name)
+			assert.Equal(t, "confirms whether or not all praefect migrations have run", migrationCheck.Description)
+			assert.Equal(t, tc.expectedErr, migrationCheck.Run())
+		})
+	}
+}