Merge branch 'ps-praefect-cmd-fix' into 'master'

praefect: Fix inconsistencies in praefect sub-commands

See merge request gitlab-org/gitaly!4026

15 files changed, 229 insertions, 179 deletions

diff --git a/cmd/praefect/main_test.go b/cmd/praefect/main_test.go
index 9a4e8949f..cdbea8f2a 100644
--- a/cmd/praefect/main_test.go
+++ b/cmd/praefect/main_test.go
@@ -21,61 +21,6 @@ func TestNoConfigFlag(t *testing.T) {
 	assert.Equal(t, err, errNoConfigFile)
 }
 
-func TestFlattenNodes(t *testing.T) {
-	for _, tt := range []struct {
-		desc   string
-		conf   config.Config
-		expect map[string]*nodePing
-	}{
-		{
-			desc: "Flatten common address between storages",
-			conf: config.Config{
-				VirtualStorages: []*config.VirtualStorage{
-					{
-						Name: "meow",
-						Nodes: []*config.Node{
-							{
-								Storage: "foo",
-								Address: "tcp://example.com",
-								Token:   "abc",
-							},
-						},
-					},
-					{
-						Name: "woof",
-						Nodes: []*config.Node{
-							{
-								Storage: "bar",
-								Address: "tcp://example.com",
-								Token:   "abc",
-							},
-						},
-					},
-				},
-			},
-			expect: map[string]*nodePing{
-				"tcp://example.com": {
-					address: "tcp://example.com",
-					storages: map[gitalyStorage][]virtualStorage{
-						"foo": {"meow"},
-						"bar": {"woof"},
-					},
-					vStorages: map[virtualStorage]struct{}{
-						"meow": {},
-						"woof": {},
-					},
-					token: "abc",
-				},
-			},
-		},
-	} {
-		t.Run(tt.desc, func(t *testing.T) {
-			actual := flattenNodes(tt.conf)
-			require.Equal(t, tt.expect, actual)
-		})
-	}
-}
-
 func TestGetStarterConfigs(t *testing.T) {
 	for _, tc := range []struct {
 		desc   string
diff --git a/cmd/praefect/subcmd.go b/cmd/praefect/subcmd.go
index a85f0b772..2f09c7587 100644
--- a/cmd/praefect/subcmd.go
+++ b/cmd/praefect/subcmd.go
@@ -14,7 +14,6 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/client"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect"
 	"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"
 	"google.golang.org/grpc"
 )
@@ -24,15 +23,17 @@ type subcmd interface {
 	Exec(flags *flag.FlagSet, config config.Config) error
 }
 
+const defaultDialTimeout = 30 * time.Second
+
 var subcommands = map[string]subcmd{
-	"sql-ping":                    &sqlPingSubcommand{},
-	"sql-migrate":                 &sqlMigrateSubcommand{},
-	"dial-nodes":                  &dialNodesSubcommand{},
-	"sql-migrate-down":            &sqlMigrateDownSubcommand{},
-	"sql-migrate-status":          &sqlMigrateStatusSubcommand{},
-	"dataloss":                    newDatalossSubcommand(),
-	"accept-dataloss":             &acceptDatalossSubcommand{},
-	"set-replication-factor":      newSetReplicatioFactorSubcommand(os.Stdout),
+	sqlPingCmdName:                &sqlPingSubcommand{},
+	sqlMigrateCmdName:             &sqlMigrateSubcommand{},
+	dialNodesCmdName:              &dialNodesSubcommand{},
+	sqlMigrateDownCmdName:         &sqlMigrateDownSubcommand{},
+	sqlMigrateStatusCmdName:       &sqlMigrateStatusSubcommand{},
+	datalossCmdName:               newDatalossSubcommand(),
+	acceptDatalossCmdName:         &acceptDatalossSubcommand{},
+	setReplicationFactorCmdName:   newSetReplicatioFactorSubcommand(os.Stdout),
 	removeRepositoryCmdName:       newRemoveRepository(logger),
 	trackRepositoryCmdName:        newTrackRepository(logger),
 	listUntrackedRepositoriesName: newListUntrackedRepositories(logger, os.Stdout),
@@ -83,57 +84,6 @@ func getNodeAddress(cfg config.Config) (string, error) {
 	}
 }
 
-type sqlPingSubcommand struct{}
-
-func (s *sqlPingSubcommand) FlagSet() *flag.FlagSet {
-	return flag.NewFlagSet("sql-ping", flag.ExitOnError)
-}
-
-func (s *sqlPingSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
-	const subCmd = progname + " sql-ping"
-
-	db, clean, err := openDB(conf.DB)
-	if err != nil {
-		return err
-	}
-	defer clean()
-
-	if err := datastore.CheckPostgresVersion(db); err != nil {
-		return fmt.Errorf("%s: fail: %v", subCmd, err)
-	}
-
-	fmt.Printf("%s: OK\n", subCmd)
-	return nil
-}
-
-type sqlMigrateSubcommand struct {
-	ignoreUnknown bool
-}
-
-func (s *sqlMigrateSubcommand) FlagSet() *flag.FlagSet {
-	flags := flag.NewFlagSet("sql-migrate", flag.ExitOnError)
-	flags.BoolVar(&s.ignoreUnknown, "ignore-unknown", true, "ignore unknown migrations (default is true)")
-	return flags
-}
-
-func (s *sqlMigrateSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
-	const subCmd = progname + " sql-migrate"
-
-	db, clean, err := openDB(conf.DB)
-	if err != nil {
-		return err
-	}
-	defer clean()
-
-	n, err := glsql.Migrate(db, s.ignoreUnknown)
-	if err != nil {
-		return fmt.Errorf("%s: fail: %v", subCmd, err)
-	}
-
-	fmt.Printf("%s: OK (applied %d migrations)\n", subCmd, n)
-	return nil
-}
-
 func openDB(conf config.DB) (*sql.DB, func(), error) {
 	db, err := glsql.OpenDB(conf)
 	if err != nil {
@@ -153,8 +103,8 @@ func printfErr(format string, a ...interface{}) (int, error) {
 	return fmt.Fprintf(os.Stderr, format, a...)
 }
 
-func subCmdDial(addr, token string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
-	ctx, cancel := context.WithTimeout(context.TODO(), 30*time.Second)
+func subCmdDial(ctx context.Context, addr, token string, timeout time.Duration, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
+	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
 	opts = append(opts,
@@ -171,3 +121,15 @@ func subCmdDial(addr, token string, opts ...grpc.DialOption) (*grpc.ClientConn,
 
 	return client.DialContext(ctx, addr, opts)
 }
+
+type requiredParameterError string
+
+func (p requiredParameterError) Error() string {
+	return fmt.Sprintf("%q is a required parameter", string(p))
+}
+
+type unexpectedPositionalArgsError struct{ Command string }
+
+func (err unexpectedPositionalArgsError) Error() string {
+	return fmt.Sprintf("%s doesn't accept positional arguments", err.Command)
+}
diff --git a/cmd/praefect/subcmd_accept_dataloss.go b/cmd/praefect/subcmd_accept_dataloss.go
index 0c4ea015a..04b076743 100644
--- a/cmd/praefect/subcmd_accept_dataloss.go
+++ b/cmd/praefect/subcmd_accept_dataloss.go
@@ -9,13 +9,8 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
 )
 
-type requiredParameterError string
-
-func (p requiredParameterError) Error() string {
-	return fmt.Sprintf("%q is a required parameter", string(p))
-}
-
 const (
+	acceptDatalossCmdName     = "accept-dataloss"
 	paramVirtualStorage       = "virtual-storage"
 	paramRelativePath         = "repository"
 	paramAuthoritativeStorage = "authoritative-storage"
@@ -28,7 +23,7 @@ type acceptDatalossSubcommand struct {
 }
 
 func (cmd *acceptDatalossSubcommand) FlagSet() *flag.FlagSet {
-	fs := flag.NewFlagSet("accept-dataloss", flag.ContinueOnError)
+	fs := flag.NewFlagSet(acceptDatalossCmdName, flag.ContinueOnError)
 	fs.StringVar(&cmd.virtualStorage, paramVirtualStorage, "", "name of the repository's virtual storage")
 	fs.StringVar(&cmd.relativePath, paramRelativePath, "", "repository to accept data loss for")
 	fs.StringVar(&cmd.authoritativeStorage, paramAuthoritativeStorage, "", "storage with the repository to consider as authoritative")
@@ -50,15 +45,16 @@ func (cmd *acceptDatalossSubcommand) Exec(flags *flag.FlagSet, cfg config.Config
 	if err != nil {
 		return err
 	}
+	ctx := context.TODO()
 
-	conn, err := subCmdDial(nodeAddr, cfg.Auth.Token)
+	conn, err := subCmdDial(ctx, nodeAddr, cfg.Auth.Token, defaultDialTimeout)
 	if err != nil {
 		return fmt.Errorf("error dialing: %w", err)
 	}
 	defer conn.Close()
 
 	client := gitalypb.NewPraefectInfoServiceClient(conn)
-	if _, err := client.SetAuthoritativeStorage(context.TODO(), &gitalypb.SetAuthoritativeStorageRequest{
+	if _, err := client.SetAuthoritativeStorage(ctx, &gitalypb.SetAuthoritativeStorageRequest{
 		VirtualStorage:       cmd.virtualStorage,
 		RelativePath:         cmd.relativePath,
 		AuthoritativeStorage: cmd.authoritativeStorage,
diff --git a/cmd/praefect/subcmd_dataloss.go b/cmd/praefect/subcmd_dataloss.go
index 286e696b6..5eb88f4e6 100644
--- a/cmd/praefect/subcmd_dataloss.go
+++ b/cmd/praefect/subcmd_dataloss.go
@@ -14,11 +14,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
 )
 
-type unexpectedPositionalArgsError struct{ Command string }
-
-func (err unexpectedPositionalArgsError) Error() string {
-	return fmt.Sprintf("%s doesn't accept positional arguments", err.Command)
-}
+const datalossCmdName = "dataloss"
 
 type datalossSubcommand struct {
 	output                    io.Writer
@@ -31,7 +27,7 @@ func newDatalossSubcommand() *datalossSubcommand {
 }
 
 func (cmd *datalossSubcommand) FlagSet() *flag.FlagSet {
-	fs := flag.NewFlagSet("dataloss", flag.ContinueOnError)
+	fs := flag.NewFlagSet(datalossCmdName, flag.ContinueOnError)
 	fs.StringVar(&cmd.virtualStorage, "virtual-storage", "", "virtual storage to check for data loss")
 	fs.BoolVar(&cmd.includePartiallyAvailable, "partially-unavailable", false, strings.TrimSpace(`
 Additionally include repositories which are available but some assigned replicas
@@ -65,7 +61,8 @@ func (cmd *datalossSubcommand) Exec(flags *flag.FlagSet, cfg config.Config) erro
 		return err
 	}
 
-	conn, err := subCmdDial(nodeAddr, cfg.Auth.Token)
+	ctx := context.Background()
+	conn, err := subCmdDial(ctx, nodeAddr, cfg.Auth.Token, defaultDialTimeout)
 	if err != nil {
 		return fmt.Errorf("error dialing: %v", err)
 	}
@@ -78,7 +75,7 @@ func (cmd *datalossSubcommand) Exec(flags *flag.FlagSet, cfg config.Config) erro
 	client := gitalypb.NewPraefectInfoServiceClient(conn)
 
 	for _, vs := range virtualStorages {
-		resp, err := client.DatalossCheck(context.Background(), &gitalypb.DatalossCheckRequest{
+		resp, err := client.DatalossCheck(ctx, &gitalypb.DatalossCheckRequest{
 			VirtualStorage:             vs,
 			IncludePartiallyReplicated: cmd.includePartiallyAvailable,
 		})
diff --git a/cmd/praefect/subcmd_pingnodes.go b/cmd/praefect/subcmd_dial_nodes.go
index 0b0873484..0b22a0555 100644
--- a/cmd/praefect/subcmd_pingnodes.go
+++ b/cmd/praefect/subcmd_dial_nodes.go
@@ -15,6 +15,37 @@ import (
 	"google.golang.org/grpc/health/grpc_health_v1"
 )
 
+const dialNodesCmdName = "dial-nodes"
+
+type dialNodesSubcommand struct{}
+
+func (s *dialNodesSubcommand) FlagSet() *flag.FlagSet {
+	return flag.NewFlagSet(dialNodesCmdName, flag.ExitOnError)
+}
+
+func (s *dialNodesSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
+	nodes := flattenNodes(conf)
+
+	var wg sync.WaitGroup
+	for _, n := range nodes {
+		wg.Add(1)
+		go func(n *nodePing) {
+			defer wg.Done()
+			n.checkNode()
+		}(n)
+	}
+	wg.Wait()
+
+	var err error
+	for _, n := range nodes {
+		if n.err != nil {
+			err = n.err
+		}
+	}
+
+	return err
+}
+
 type (
 	virtualStorage string
 	gitalyStorage  string
@@ -58,37 +89,8 @@ func flattenNodes(conf config.Config) map[string]*nodePing {
 	return nodeByAddress
 }
 
-type dialNodesSubcommand struct{}
-
-func (s *dialNodesSubcommand) FlagSet() *flag.FlagSet {
-	return flag.NewFlagSet("dial-nodes", flag.ExitOnError)
-}
-
-func (s *dialNodesSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
-	nodes := flattenNodes(conf)
-
-	var wg sync.WaitGroup
-	for _, n := range nodes {
-		wg.Add(1)
-		go func(n *nodePing) {
-			defer wg.Done()
-			n.checkNode()
-		}(n)
-	}
-	wg.Wait()
-
-	var err error
-	for _, n := range nodes {
-		if n.err != nil {
-			err = n.err
-		}
-	}
-
-	return err
-}
-
 func (npr *nodePing) dial() (*grpc.ClientConn, error) {
-	return subCmdDial(npr.address, npr.token)
+	return subCmdDial(context.TODO(), npr.address, npr.token, defaultDialTimeout)
 }
 
 func (npr *nodePing) healthCheck(cc *grpc.ClientConn) (grpc_health_v1.HealthCheckResponse_ServingStatus, error) {
diff --git a/cmd/praefect/subcmd_pingnodes_test.go b/cmd/praefect/subcmd_dial_nodes_test.go
index a207a4a80..df013916f 100644
--- a/cmd/praefect/subcmd_pingnodes_test.go
+++ b/cmd/praefect/subcmd_dial_nodes_test.go
@@ -120,3 +120,58 @@ func TestSubCmdDialNodes(t *testing.T) {
 		})
 	}
 }
+
+func TestFlattenNodes(t *testing.T) {
+	for _, tt := range []struct {
+		desc   string
+		conf   config.Config
+		expect map[string]*nodePing
+	}{
+		{
+			desc: "Flatten common address between storages",
+			conf: config.Config{
+				VirtualStorages: []*config.VirtualStorage{
+					{
+						Name: "meow",
+						Nodes: []*config.Node{
+							{
+								Storage: "foo",
+								Address: "tcp://example.com",
+								Token:   "abc",
+							},
+						},
+					},
+					{
+						Name: "woof",
+						Nodes: []*config.Node{
+							{
+								Storage: "bar",
+								Address: "tcp://example.com",
+								Token:   "abc",
+							},
+						},
+					},
+				},
+			},
+			expect: map[string]*nodePing{
+				"tcp://example.com": {
+					address: "tcp://example.com",
+					storages: map[gitalyStorage][]virtualStorage{
+						"foo": {"meow"},
+						"bar": {"woof"},
+					},
+					vStorages: map[virtualStorage]struct{}{
+						"meow": {},
+						"woof": {},
+					},
+					token: "abc",
+				},
+			},
+		},
+	} {
+		t.Run(tt.desc, func(t *testing.T) {
+			actual := flattenNodes(tt.conf)
+			require.Equal(t, tt.expect, actual)
+		})
+	}
+}
diff --git a/cmd/praefect/subcmd_list_untracked_repositories.go b/cmd/praefect/subcmd_list_untracked_repositories.go
index b1e37db1b..e74f02315 100644
--- a/cmd/praefect/subcmd_list_untracked_repositories.go
+++ b/cmd/praefect/subcmd_list_untracked_repositories.go
@@ -7,6 +7,7 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"time"
 
 	"github.com/sirupsen/logrus"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect"
@@ -63,7 +64,7 @@ func (cmd listUntrackedRepositories) Exec(flags *flag.FlagSet, cfg config.Config
 	logger.Debugf("starting %s command", cmd.FlagSet().Name())
 
 	logger.Debug("dialing to gitaly nodes...")
-	nodeSet, err := dialGitalyStorages(cfg)
+	nodeSet, err := dialGitalyStorages(ctx, cfg, defaultDialTimeout)
 	if err != nil {
 		return fmt.Errorf("dial nodes: %w", err)
 	}
@@ -97,11 +98,11 @@ func (cmd listUntrackedRepositories) Exec(flags *flag.FlagSet, cfg config.Config
 	return nil
 }
 
-func dialGitalyStorages(cfg config.Config) (praefect.NodeSet, error) {
+func dialGitalyStorages(ctx context.Context, cfg config.Config, timeout time.Duration) (praefect.NodeSet, error) {
 	nodeSet := praefect.NodeSet{}
 	for _, vs := range cfg.VirtualStorages {
 		for _, node := range vs.Nodes {
-			conn, err := subCmdDial(node.Address, node.Token)
+			conn, err := subCmdDial(ctx, node.Address, node.Token, timeout)
 			if err != nil {
 				return nil, fmt.Errorf("dial with %q gitaly at %q", node.Storage, node.Address)
 			}
diff --git a/cmd/praefect/subcmd_remove_repository.go b/cmd/praefect/subcmd_remove_repository.go
index 016a10daa..b9e0bfb07 100644
--- a/cmd/praefect/subcmd_remove_repository.go
+++ b/cmd/praefect/subcmd_remove_repository.go
@@ -28,10 +28,11 @@ type removeRepository struct {
 	logger         logrus.FieldLogger
 	virtualStorage string
 	relativePath   string
+	dialTimeout    time.Duration
 }
 
 func newRemoveRepository(logger logrus.FieldLogger) *removeRepository {
-	return &removeRepository{logger: logger}
+	return &removeRepository{logger: logger, dialTimeout: defaultDialTimeout}
 }
 
 func (cmd *removeRepository) FlagSet() *flag.FlagSet {
@@ -136,7 +137,7 @@ func (cmd *removeRepository) removeRepositoryFromDatabase(ctx context.Context, d
 }
 
 func (cmd *removeRepository) removeRepository(ctx context.Context, repo *gitalypb.Repository, addr, token string) (bool, error) {
-	conn, err := subCmdDial(addr, token)
+	conn, err := subCmdDial(ctx, addr, token, cmd.dialTimeout)
 	if err != nil {
 		return false, fmt.Errorf("error dialing: %w", err)
 	}
@@ -162,9 +163,15 @@ func (cmd *removeRepository) removeReplicationEvents(ctx context.Context, logger
 	// As some of them could be a repository creation jobs we need to remove those newly created
 	// repositories after replication finished.
 	start := time.Now()
+	var tick helper.Ticker
 	for found := true; found; {
-		ticker.Reset()
-		<-ticker.C()
+		if tick != nil {
+			tick.Reset()
+			<-tick.C()
+		} else {
+			tick = ticker
+		}
+
 		if int(time.Since(start).Seconds())%5 == 0 {
 			logger.Debug("awaiting for the repository in_progress replication jobs to complete...")
 		}
diff --git a/cmd/praefect/subcmd_remove_repository_test.go b/cmd/praefect/subcmd_remove_repository_test.go
index 519ab4764..f1ffbd3eb 100644
--- a/cmd/praefect/subcmd_remove_repository_test.go
+++ b/cmd/praefect/subcmd_remove_repository_test.go
@@ -127,6 +127,7 @@ func TestRemoveRepository_Exec(t *testing.T) {
 			logger:         testhelper.NewTestLogger(t),
 			virtualStorage: repo.StorageName,
 			relativePath:   repo.RelativePath,
+			dialTimeout:    time.Second,
 		}
 		require.NoError(t, cmd.Exec(flag.NewFlagSet("", flag.PanicOnError), conf))
 
@@ -153,6 +154,7 @@ func TestRemoveRepository_Exec(t *testing.T) {
 			logger:         logrus.NewEntry(logger),
 			virtualStorage: praefectStorage,
 			relativePath:   repo.RelativePath,
+			dialTimeout:    time.Second,
 		}
 		require.NoError(t, cmd.Exec(flag.NewFlagSet("", flag.PanicOnError), conf))
 		var found bool
@@ -179,6 +181,7 @@ func TestRemoveRepository_Exec(t *testing.T) {
 			logger:         logrus.NewEntry(logger),
 			virtualStorage: praefectStorage,
 			relativePath:   repo.RelativePath,
+			dialTimeout:    100 * time.Millisecond,
 		}
 
 		for {
diff --git a/cmd/praefect/subcmd_set_replication_factor.go b/cmd/praefect/subcmd_set_replication_factor.go
index ff2cc4562..0da3ecdc6 100644
--- a/cmd/praefect/subcmd_set_replication_factor.go
+++ b/cmd/praefect/subcmd_set_replication_factor.go
@@ -11,7 +11,10 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
 )
 
-const paramReplicationFactor = "replication-factor"
+const (
+	setReplicationFactorCmdName = "set-replication-factor"
+	paramReplicationFactor      = "replication-factor"
+)
 
 type setReplicationFactorSubcommand struct {
 	stdout            io.Writer
@@ -25,7 +28,7 @@ func newSetReplicatioFactorSubcommand(stdout io.Writer) *setReplicationFactorSub
 }
 
 func (cmd *setReplicationFactorSubcommand) FlagSet() *flag.FlagSet {
-	fs := flag.NewFlagSet("set-replication-factor", flag.ContinueOnError)
+	fs := flag.NewFlagSet(setReplicationFactorCmdName, flag.ContinueOnError)
 	fs.StringVar(&cmd.virtualStorage, paramVirtualStorage, "", "name of the repository's virtual storage")
 	fs.StringVar(&cmd.relativePath, paramRelativePath, "", "repository to set the replication factor for")
 	fs.IntVar(&cmd.replicationFactor, paramReplicationFactor, -1, "desired replication factor")
@@ -48,7 +51,7 @@ func (cmd *setReplicationFactorSubcommand) Exec(flags *flag.FlagSet, cfg config.
 		return err
 	}
 
-	conn, err := subCmdDial(nodeAddr, cfg.Auth.Token)
+	conn, err := subCmdDial(context.TODO(), nodeAddr, cfg.Auth.Token, defaultDialTimeout)
 	if err != nil {
 		return fmt.Errorf("error dialing: %w", err)
 	}
diff --git a/cmd/praefect/subcmd_sql_migrate.go b/cmd/praefect/subcmd_sql_migrate.go
new file mode 100644
index 000000000..3bb4f9b11
--- /dev/null
+++ b/cmd/praefect/subcmd_sql_migrate.go
@@ -0,0 +1,39 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/glsql"
+)
+
+const sqlMigrateCmdName = "sql-migrate"
+
+type sqlMigrateSubcommand struct {
+	ignoreUnknown bool
+}
+
+func (s *sqlMigrateSubcommand) FlagSet() *flag.FlagSet {
+	flags := flag.NewFlagSet(sqlMigrateCmdName, flag.ExitOnError)
+	flags.BoolVar(&s.ignoreUnknown, "ignore-unknown", true, "ignore unknown migrations (default is true)")
+	return flags
+}
+
+func (s *sqlMigrateSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
+	const subCmd = progname + " " + sqlMigrateCmdName
+
+	db, clean, err := openDB(conf.DB)
+	if err != nil {
+		return err
+	}
+	defer clean()
+
+	n, err := glsql.Migrate(db, s.ignoreUnknown)
+	if err != nil {
+		return fmt.Errorf("%s: fail: %v", subCmd, err)
+	}
+
+	fmt.Printf("%s: OK (applied %d migrations)\n", subCmd, n)
+	return nil
+}
diff --git a/cmd/praefect/subcmd_sqldown.go b/cmd/praefect/subcmd_sql_migrate_down.go
index 78977fe57..14e6c510b 100644
--- a/cmd/praefect/subcmd_sqldown.go
+++ b/cmd/praefect/subcmd_sql_migrate_down.go
@@ -10,16 +10,18 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore"
 )
 
+const sqlMigrateDownCmdName = "sql-migrate-down"
+
 type sqlMigrateDownSubcommand struct {
 	force bool
 }
 
 func (s *sqlMigrateDownSubcommand) FlagSet() *flag.FlagSet {
-	flags := flag.NewFlagSet("sql-migrate-down", flag.ExitOnError)
+	flags := flag.NewFlagSet(sqlMigrateDownCmdName, flag.ExitOnError)
 	flags.Usage = func() {
 		flag.PrintDefaults()
-		printfErr("  MAX_MIGRATIONS\n")
-		printfErr("\tNumber of migrations to roll back\n")
+		_, _ = printfErr("  MAX_MIGRATIONS\n")
+		_, _ = printfErr("\tNumber of migrations to roll back\n")
 	}
 	flags.BoolVar(&s.force, "f", false, "apply down-migrations (default is dry run)")
 	return flags
diff --git a/cmd/praefect/subcmd_sqlstatus.go b/cmd/praefect/subcmd_sql_migrate_status.go
index debdbd2b0..760beb22c 100644
--- a/cmd/praefect/subcmd_sqlstatus.go
+++ b/cmd/praefect/subcmd_sql_migrate_status.go
@@ -10,10 +10,12 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore"
 )
 
+const sqlMigrateStatusCmdName = "sql-migrate-status"
+
 type sqlMigrateStatusSubcommand struct{}
 
 func (s *sqlMigrateStatusSubcommand) FlagSet() *flag.FlagSet {
-	return flag.NewFlagSet("sql-migrate-status", flag.ExitOnError)
+	return flag.NewFlagSet(sqlMigrateStatusCmdName, flag.ExitOnError)
 }
 
 func (s *sqlMigrateStatusSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
diff --git a/cmd/praefect/subcmd_sql_ping.go b/cmd/praefect/subcmd_sql_ping.go
new file mode 100644
index 000000000..5fcf6c772
--- /dev/null
+++ b/cmd/praefect/subcmd_sql_ping.go
@@ -0,0 +1,36 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore"
+)
+
+const (
+	sqlPingCmdName = "sql-ping"
+)
+
+type sqlPingSubcommand struct{}
+
+func (s *sqlPingSubcommand) FlagSet() *flag.FlagSet {
+	return flag.NewFlagSet(sqlPingCmdName, flag.ExitOnError)
+}
+
+func (s *sqlPingSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
+	const subCmd = progname + " " + sqlPingCmdName
+
+	db, clean, err := openDB(conf.DB)
+	if err != nil {
+		return err
+	}
+	defer clean()
+
+	if err := datastore.CheckPostgresVersion(db); err != nil {
+		return fmt.Errorf("%s: fail: %v", subCmd, err)
+	}
+
+	fmt.Printf("%s: OK\n", subCmd)
+	return nil
+}
diff --git a/cmd/praefect/subcmd_track_repository.go b/cmd/praefect/subcmd_track_repository.go
index 0be14d778..17b0d4781 100644
--- a/cmd/praefect/subcmd_track_repository.go
+++ b/cmd/praefect/subcmd_track_repository.go
@@ -187,7 +187,7 @@ func (cmd *trackRepository) trackRepository(
 }
 
 func (cmd *trackRepository) repositoryExists(ctx context.Context, repo *gitalypb.Repository, addr, token string) (bool, error) {
-	conn, err := subCmdDial(addr, token)
+	conn, err := subCmdDial(ctx, addr, token, defaultDialTimeout)
 	if err != nil {
 		return false, fmt.Errorf("error dialing: %w", err)
 	}