Merge branch 'pks-praefect-subcmd' into 'master'

Streamline handling of subcommands for praefect executable

See merge request gitlab-org/gitaly!2020

5 files changed, 123 insertions, 89 deletions

diff --git a/cmd/praefect/main.go b/cmd/praefect/main.go
index 9814d795b..dff299354 100644
--- a/cmd/praefect/main.go
+++ b/cmd/praefect/main.go
@@ -83,6 +83,17 @@ var (
 const progname = "praefect"
 
 func main() {
+	flag.Usage = func() {
+		cmds := []string{}
+		for k, _ := range subcommands {
+			cmds = append(cmds, k)
+		}
+
+		printfErr("Usage of %s:\n", progname)
+		flag.PrintDefaults()
+		printfErr("  subcommand (optional)\n")
+		printfErr("\tOne of %s\n", strings.Join(cmds, ", "))
+	}
 	flag.Parse()
 
 	// If invoked with -version
diff --git a/cmd/praefect/subcmd.go b/cmd/praefect/subcmd.go
index a29882a90..4ebba9ad7 100644
--- a/cmd/praefect/subcmd.go
+++ b/cmd/praefect/subcmd.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"database/sql"
+	"flag"
 	"fmt"
 	"os"
 	"os/signal"
@@ -16,6 +17,21 @@ import (
 	"google.golang.org/grpc"
 )
 
+type subcmd interface {
+	FlagSet() *flag.FlagSet
+	Exec(flags *flag.FlagSet, config config.Config) error
+}
+
+var (
+	subcommands = map[string]subcmd{
+		"sql-ping":         &sqlPingSubcommand{},
+		"sql-migrate":      &sqlMigrateSubcommand{},
+		"dial-nodes":       &dialNodesSubcommand{},
+		"reconcile":        &reconcileSubcommand{},
+		"sql-migrate-down": &sqlMigrateDownSubcommand{},
+	}
+)
+
 const invocationPrefix = progname + " -config CONFIG_TOML"
 
 // subCommand returns an exit code, to be fed into os.Exit.
@@ -28,65 +44,78 @@ func subCommand(conf config.Config, arg0 string, argRest []string) int {
 		os.Exit(130) // indicates program was interrupted
 	}()
 
-	switch arg0 {
-	case "sql-ping":
-		return sqlPing(conf)
-	case "sql-migrate":
-		return sqlMigrate(conf)
-	case subCmdSQLMigrateDown:
-		return sqlMigrateDown(conf, argRest)
-	case "dial-nodes":
-		return dialNodes(conf)
-	case "reconcile":
-		return reconcile(conf, argRest)
-	default:
+	subcmd, ok := subcommands[arg0]
+	if !ok {
 		printfErr("%s: unknown subcommand: %q\n", progname, arg0)
 		return 1
 	}
+
+	flags := subcmd.FlagSet()
+
+	if err := flags.Parse(argRest); err != nil {
+		printfErr("%s\n", err)
+		return 1
+	}
+
+	if err := subcmd.Exec(flags, conf); err != nil {
+		printfErr("%s\n", err)
+		return 1
+	}
+
+	return 0
+}
+
+type sqlPingSubcommand struct{}
+
+func (s *sqlPingSubcommand) FlagSet() *flag.FlagSet {
+	return flag.NewFlagSet("sql-ping", flag.ExitOnError)
 }
 
-func sqlPing(conf config.Config) int {
+func (s *sqlPingSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
 	const subCmd = progname + " sql-ping"
 
-	db, clean, code := openDB(conf.DB)
-	if code != 0 {
-		return code
+	db, clean, err := openDB(conf.DB)
+	if err != nil {
+		return err
 	}
 	defer clean()
 
 	if err := datastore.CheckPostgresVersion(db); err != nil {
-		printfErr("%s: fail: %v\n", subCmd, err)
-		return 1
+		return fmt.Errorf("%s: fail: %v\n", subCmd, err)
 	}
 
 	fmt.Printf("%s: OK\n", subCmd)
-	return 0
+	return nil
 }
 
-func sqlMigrate(conf config.Config) int {
+type sqlMigrateSubcommand struct{}
+
+func (s *sqlMigrateSubcommand) FlagSet() *flag.FlagSet {
+	return flag.NewFlagSet("sql-migrate", flag.ExitOnError)
+}
+
+func (s *sqlMigrateSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
 	const subCmd = progname + " sql-migrate"
 
-	db, clean, code := openDB(conf.DB)
-	if code != 0 {
-		return code
+	db, clean, err := openDB(conf.DB)
+	if err != nil {
+		return err
 	}
 	defer clean()
 
 	n, err := glsql.Migrate(db)
 	if err != nil {
-		printfErr("%s: fail: %v\n", subCmd, err)
-		return 1
+		return fmt.Errorf("%s: fail: %v\n", subCmd, err)
 	}
 
 	fmt.Printf("%s: OK (applied %d migrations)\n", subCmd, n)
-	return 0
+	return nil
 }
 
-func openDB(conf config.DB) (*sql.DB, func(), int) {
+func openDB(conf config.DB) (*sql.DB, func(), error) {
 	db, err := glsql.OpenDB(conf)
 	if err != nil {
-		printfErr("sql open: %v\n", err)
-		return nil, nil, 1
+		return nil, nil, fmt.Errorf("sql open: %v\n", err)
 	}
 
 	clean := func() {
@@ -95,7 +124,7 @@ func openDB(conf config.DB) (*sql.DB, func(), int) {
 		}
 	}
 
-	return db, clean, 0
+	return db, clean, nil
 }
 
 func printfErr(format string, a ...interface{}) (int, error) {
diff --git a/cmd/praefect/subcmd_pingnodes.go b/cmd/praefect/subcmd_pingnodes.go
index b3bf24fec..742f8eb12 100644
--- a/cmd/praefect/subcmd_pingnodes.go
+++ b/cmd/praefect/subcmd_pingnodes.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"errors"
+	"flag"
 	"fmt"
 	"log"
 	"sync"
@@ -45,7 +46,13 @@ func flattenNodes(conf config.Config) map[string]*nodePing {
 	return nodeByAddress
 }
 
-func dialNodes(conf config.Config) int {
+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
@@ -58,14 +65,14 @@ func dialNodes(conf config.Config) int {
 	}
 	wg.Wait()
 
-	exitCode := 0
+	var err error
 	for _, n := range nodes {
 		if n.err != nil {
-			exitCode = 1
+			err = n.err
 		}
 	}
 
-	return exitCode
+	return err
 }
 
 func (npr *nodePing) dial() (*grpc.ClientConn, error) {
diff --git a/cmd/praefect/subcmd_reconcile.go b/cmd/praefect/subcmd_reconcile.go
index 21f3ed7fc..bb2b12b6c 100644
--- a/cmd/praefect/subcmd_reconcile.go
+++ b/cmd/praefect/subcmd_reconcile.go
@@ -19,32 +19,33 @@ type nodeReconciler struct {
 	referenceStorage string
 }
 
-func reconcile(conf config.Config, subCmdArgs []string) int {
-	var (
-		fs = flag.NewFlagSet("reconcile", flag.ExitOnError)
-		vs = fs.String("virtual", "", "virtual storage for target storage")
-		t  = fs.String("target", "", "target storage to reconcile")
-		r  = fs.String("reference", "", "reference storage to reconcile (optional)")
-	)
-
-	if err := fs.Parse(subCmdArgs); err != nil {
-		log.Printf("unable to parse args %v: %s", subCmdArgs, err)
-		return 1
-	}
+type reconcileSubcommand struct {
+	virtual   string
+	target    string
+	reference string
+}
 
+func (s *reconcileSubcommand) FlagSet() *flag.FlagSet {
+	fs := flag.NewFlagSet("reconcile", flag.ExitOnError)
+	fs.StringVar(&s.virtual, "virtual", "", "virtual storage for target storage")
+	fs.StringVar(&s.target, "target", "", "target storage to reconcile")
+	fs.StringVar(&s.reference, "reference", "", "reference storage to reconcile (optional)")
+	return fs
+}
+
+func (s *reconcileSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
 	nr := nodeReconciler{
 		conf:             conf,
-		virtualStorage:   *vs,
-		targetStorage:    *t,
-		referenceStorage: *r,
+		virtualStorage:   s.virtual,
+		targetStorage:    s.target,
+		referenceStorage: s.reference,
 	}
 
 	if err := nr.reconcile(); err != nil {
-		log.Print("unable to reconcile: ", err)
-		return 1
+		return fmt.Errorf("unable to reconcile: %s", err)
 	}
 
-	return 0
+	return nil
 }
 
 func (nr nodeReconciler) reconcile() error {
diff --git a/cmd/praefect/subcmd_sqldown.go b/cmd/praefect/subcmd_sqldown.go
index f34395d6c..2748de0d0 100644
--- a/cmd/praefect/subcmd_sqldown.go
+++ b/cmd/praefect/subcmd_sqldown.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"flag"
 	"fmt"
 	"strconv"
@@ -9,43 +10,28 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/praefect/datastore"
 )
 
-const subCmdSQLMigrateDown = "sql-migrate-down"
-
-func sqlMigrateDown(conf config.Config, args []string) int {
-	cmd := &sqlMigrateDownCmd{Config: conf}
-	return cmd.Run(args)
+type sqlMigrateDownSubcommand struct {
+	force bool
 }
 
-type sqlMigrateDownCmd struct{ config.Config }
-
-func (*sqlMigrateDownCmd) prefix() string { return progname + " " + subCmdSQLMigrateDown }
-
-func (*sqlMigrateDownCmd) invocation() string { return invocationPrefix + " " + subCmdSQLMigrateDown }
-
-func (smd *sqlMigrateDownCmd) Run(args []string) int {
-	flagset := flag.NewFlagSet(smd.prefix(), flag.ExitOnError)
-	flagset.Usage = func() {
-		printfErr("usage:  %s [-f] MAX_MIGRATIONS\n", smd.invocation())
-	}
-	force := flagset.Bool("f", false, "apply down-migrations (default is dry run)")
-
-	_ = flagset.Parse(args) // No error check because flagset is set to ExitOnError
-
-	if flagset.NArg() != 1 {
-		flagset.Usage()
-		return 1
+func (s *sqlMigrateDownSubcommand) FlagSet() *flag.FlagSet {
+	flags := flag.NewFlagSet("sql-migrate-down", flag.ExitOnError)
+	flags.Usage = func() {
+		flag.PrintDefaults()
+		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
+}
 
-	if err := smd.run(*force, flagset.Arg(0)); err != nil {
-		printfErr("%s: fail: %v\n", smd.prefix(), err)
-		return 1
+func (s *sqlMigrateDownSubcommand) Exec(flags *flag.FlagSet, conf config.Config) error {
+	if flags.NArg() != 1 {
+		flags.Usage()
+		return errors.New("invalid usage")
 	}
 
-	return 0
-}
-
-func (smd *sqlMigrateDownCmd) run(force bool, maxString string) error {
-	maxMigrations, err := strconv.Atoi(maxString)
+	maxMigrations, err := strconv.Atoi(flags.Arg(0))
 	if err != nil {
 		return err
 	}
@@ -54,26 +40,26 @@ func (smd *sqlMigrateDownCmd) run(force bool, maxString string) error {
 		return fmt.Errorf("number of migrations to roll back must be 1 or more")
 	}
 
-	if force {
-		n, err := datastore.MigrateDown(smd.Config, maxMigrations)
+	if s.force {
+		n, err := datastore.MigrateDown(conf, maxMigrations)
 		if err != nil {
 			return err
 		}
 
-		fmt.Printf("%s: OK (applied %d \"down\" migrations)\n", smd.prefix(), n)
+		fmt.Printf("OK (applied %d \"down\" migrations)\n", n)
 		return nil
 	}
 
-	planned, err := datastore.MigrateDownPlan(smd.Config, maxMigrations)
+	planned, err := datastore.MigrateDownPlan(conf, maxMigrations)
 	if err != nil {
 		return err
 	}
 
-	fmt.Printf("%s: DRY RUN -- would roll back:\n\n", smd.prefix())
+	fmt.Printf("DRY RUN -- would roll back:\n\n")
 	for _, id := range planned {
 		fmt.Printf("- %s\n", id)
 	}
-	fmt.Printf("\nTo apply these migrations run: %s -f %d\n", smd.invocation(), maxMigrations)
+	fmt.Printf("\nTo apply these migrations run with -f\n")
 
 	return nil
 }