dmctl: add command mode (#287) (#364)

cmd/dm-ctl/main.go

@@ -31,24 +31,72 @@ import (
 	"github.com/pingcap/dm/pkg/utils"
 )
 
+// output:
+// Usage: dmctl [global options] command [command options] [arguments...]
+//
+//Available Commands:
+//  break-ddl-lock        break-ddl-lock <-w worker ...> <task-name> [--remove-id] [--exec] [--skip]
+//  check-task            check-task <config-file>
+//  migrate-relay         migrate-relay <worker> <binlogName> <binlogPos>
+//  pause-relay           pause-relay <-w worker ...>
+//  pause-task            pause-task [-w worker ...] <task-name>
+//  purge-relay           purge-relay <-w worker> [--filename] [--sub-dir]
+//  query-error           query-error [-w worker ...] [task-name]
+//  query-status          query-status [-w worker ...] [task-name]
+//  refresh-worker-tasks  refresh-worker-tasks
+//  resume-relay          resume-relay <-w worker ...>
+//  resume-task           resume-task [-w worker ...] <task-name>
+//  show-ddl-locks        show-ddl-locks [-w worker ...] [task-name]
+//  sql-inject            sql-inject <-w worker> <task-name> <sql1;sql2;>
+//  sql-replace           sql-replace <-w worker> [-b binlog-pos] [-s sql-pattern] [--sharding] <task-name> <sql1;sql2;>
+//  sql-skip              sql-skip <-w worker> [-b binlog-pos] [-s sql-pattern] [--sharding] <task-name>
+//  start-task            start-task [-w worker ...] <config-file>
+//  stop-task             stop-task [-w worker ...] <task-name>
+//  switch-relay-master   switch-relay-master <-w worker ...>
+//  unlock-ddl-lock       unlock-ddl-lock [-w worker ...] <lock-ID>
+//  update-master-config  update-master-config <config-file>
+//  update-relay          update-relay [-w worker ...] <config-file>
+//  update-task           update-task [-w worker ...] <config-file>
+//
+//Special Commands:
+//  --encrypt encrypt plaintext to ciphertext
+//
+//Global Options:
+//  --V prints version and exit
+//  --config path to config file
+//  --master-addr master API server addr
+//  --rpc-timeout rpc timeout, default is 10m
+func helpUsage(cfg *common.Config) {
+	fmt.Println("Usage: dmctl [global options] command [command options] [arguments...]")
+	fmt.Println()
+	ctl.PrintUsage()
+	fmt.Println()
+	fmt.Println("Special Commands:")
+	f := cfg.FlagSet.Lookup(common.EncryptCmdName)
+	fmt.Println(fmt.Sprintf("  --%s %s", f.Name, f.Usage))
+	fmt.Println()
+	fmt.Println("Global Options:")
+	cfg.FlagSet.VisitAll(func(flag2 *flag.Flag) {
+		if flag2.Name == common.EncryptCmdName {
+			return
+		}
+		fmt.Println(fmt.Sprintf("  --%s %s", flag2.Name, flag2.Usage))
+	})
+}
+
 func main() {
 	cfg := common.NewConfig()
-	err := cfg.Parse(os.Args[1:])
-	switch errors.Cause(err) {
-	case nil:
-	case flag.ErrHelp:
+	args := os.Args[1:]
+
+	// no arguments: print help message about dmctl
+	if len(args) == 0 {
+		helpUsage(cfg)
 		os.Exit(0)
-	default:
-		fmt.Printf("parse cmd flags err: %s", err)
-		os.Exit(2)
 	}
 
-	utils.PrintInfo2("dmctl")
-	fmt.Println() // print a separater
-
 	// now, we use checker in dmctl while it using some pkg which log some thing when running
 	// to make dmctl output more clear, simply redirect log to file rather output to stdout
-	err = log.InitLogger(&log.Config{
+	err := log.InitLogger(&log.Config{
 		File:  "dmctl.log",
 		Level: "info",
 	})
@@ -57,11 +105,90 @@ func main() {
 		os.Exit(2)
 	}
 
+	// try to split one task operation from dmctl command
+	// because we allow user put task operation at last with two restrictions
+	// 1. one command one task operation
+	// 2. put task operation at last
+	cmdArgs := extractSubCommand(args)
+	lenArgs := len(args)
+	lenCmdArgs := len(cmdArgs)
+	if lenCmdArgs > 0 {
+		lenArgs = lenArgs - lenCmdArgs
+	}
+
+	finished, err := cfg.Parse(args[:lenArgs])
+	if finished {
+		os.Exit(0)
+	}
+
+	switch errors.Cause(err) {
+	case nil:
+	case flag.ErrHelp:
+		if lenCmdArgs > 0 {
+			// print help message about special subCommand
+			ctl.PrintHelp(cmdArgs)
+		} else {
+			// print help message about dmctl
+			helpUsage(cfg)
+		}
+		os.Exit(0)
+	default:
+		fmt.Printf("parse cmd flags err: %s", err)
+		os.Exit(2)
+	}
+
+	err = cfg.Validate()
+	if err != nil {
+		fmt.Printf("flags are not validate: %s", err)
+		os.Exit(2)
+	}
+
 	err = ctl.Init(cfg)
 	if err != nil {
 		fmt.Printf("init control error %v", errors.ErrorStack(err))
 		os.Exit(2)
 	}
+	if lenCmdArgs > 0 {
+		commandMode(cmdArgs)
+	} else {
+		interactionMode()
+	}
+}
+
+func extractSubCommand(args []string) []string {
+	collectedArgs := make([]string, 0, len(args))
+	subCommand := make([]string, 0, len(args))
+	for i := 0; i < len(args); i++ {
+		// check whether has multiple commands
+		if ctl.HasCommand(strings.ToLower(args[i])) {
+			subCommand = append(subCommand, args[i])
+		}
+	}
+	if len(subCommand) == 0 {
+		return collectedArgs
+	}
+	if len(subCommand) > 1 {
+		fmt.Printf("command mode only support one command at a time, find %d:", len(subCommand))
+		fmt.Println(subCommand)
+		os.Exit(1)
+	}
+
+	for i := 0; i < len(args); i++ {
+		if ctl.HasCommand(strings.ToLower(args[i])) {
+			collectedArgs = append(collectedArgs, args[i:]...)
+			break
+		}
+	}
+	return collectedArgs
+}
+
+func commandMode(args []string) {
+	ctl.Start(args)
+}
+
+func interactionMode() {
+	utils.PrintInfo2("dmctl")
+	fmt.Println() // print a separater
 
 	sc := make(chan os.Signal, 1)
 	signal.Notify(sc,

dm/ctl/common/config.go

@@ -28,19 +28,25 @@ import (
 
 const (
 	defaultRPCTimeout = "10m"
+
+	// EncryptCmdName is special command
+	EncryptCmdName = "encrypt"
 )
 
 // NewConfig creates a new base config for dmctl.
 func NewConfig() *Config {
 	cfg := &Config{}
 	cfg.FlagSet = flag.NewFlagSet("dmctl", flag.ContinueOnError)
+
+	// ignore default help usage
+	cfg.FlagSet.Usage = func() {}
 	fs := cfg.FlagSet
 
 	fs.BoolVar(&cfg.printVersion, "V", false, "prints version and exit")
 	fs.StringVar(&cfg.ConfigFile, "config", "", "path to config file")
 	fs.StringVar(&cfg.MasterAddr, "master-addr", "", "master API server addr")
 	fs.StringVar(&cfg.RPCTimeoutStr, "rpc-timeout", defaultRPCTimeout, fmt.Sprintf("rpc timeout, default is %s", defaultRPCTimeout))
-	fs.StringVar(&cfg.encrypt, "encrypt", "", "encrypt plaintext to ciphertext")
+	fs.StringVar(&cfg.encrypt, EncryptCmdName, "", "encrypt plaintext to ciphertext")
 
 	return cfg
 }
@@ -69,54 +75,60 @@ func (c *Config) String() string {
 }
 
 // Parse parses flag definitions from the argument list.
-func (c *Config) Parse(arguments []string) error {
-	// Parse first to get config file.
-	err := c.FlagSet.Parse(arguments)
+func (c *Config) Parse(arguments []string) (finish bool, err error) {
+	err = c.FlagSet.Parse(arguments)
 	if err != nil {
-		return errors.Trace(err)
+		return false, errors.Trace(err)
 	}
 
 	if c.printVersion {
 		fmt.Println(utils.GetRawInfo())
-		return flag.ErrHelp
+		return true, nil
 	}
 
 	if len(c.encrypt) > 0 {
 		ciphertext, err1 := utils.Encrypt(c.encrypt)
 		if err1 != nil {
-			fmt.Println(errors.ErrorStack(err1))
-		} else {
-			fmt.Println(ciphertext)
+			return true, err1
 		}
-		return flag.ErrHelp
+		fmt.Println(ciphertext)
+		return true, nil
 	}
 
 	// Load config file if specified.
 	if c.ConfigFile != "" {
 		err = c.configFromFile(c.ConfigFile)
 		if err != nil {
-			return errors.Trace(err)
+			return false, errors.Trace(err)
 		}
 	}
 
 	// Parse again to replace with command line options.
 	err = c.FlagSet.Parse(arguments)
 	if err != nil {
-		return errors.Trace(err)
+		return false, errors.Trace(err)
 	}
 
 	if len(c.FlagSet.Args()) != 0 {
-		return errors.Errorf("'%s' is an invalid flag", c.FlagSet.Arg(0))
+		return false, errors.Errorf("'%s' is an invalid flag", c.FlagSet.Arg(0))
+	}
+
+	if c.MasterAddr == "" {
+		return false, flag.ErrHelp
 	}
 
+	return false, errors.Trace(c.adjust())
+}
+
+// Validate check config is ready to execute commmand
+func (c *Config) Validate() error {
 	if c.MasterAddr == "" {
 		return errors.New("--master-addr not provided")
 	}
-	if err = validateAddr(c.MasterAddr); err != nil {
+	if err := validateAddr(c.MasterAddr); err != nil {
 		return errors.Annotatef(err, "specify master addr %s", c.MasterAddr)
 	}
-
-	return errors.Trace(c.adjust())
+	return nil
 }
 
 // configFromFile loads config from file.

dm/ctl/ctl.go

@@ -15,6 +15,7 @@ package ctl
 
 import (
 	"fmt"
+	"os"
 
 	"github.com/pingcap/dm/dm/ctl/common"
 	"github.com/pingcap/dm/dm/ctl/master"
@@ -27,27 +28,18 @@ import (
 
 var (
 	commandMasterFlags = CommandMasterFlags{}
+	rootCmd            = &cobra.Command{
+		Use:   "dmctl",
+		Short: "DM control",
+	}
 )
 
 // CommandMasterFlags are flags that used in all commands for dm-master
 type CommandMasterFlags struct {
 	workers []string // specify workers to control on these dm-workers
 }
 
-// Init initializes dm-control
-func Init(cfg *common.Config) error {
-	// set the log level temporarily
-	log.SetLevel(zapcore.InfoLevel)
-	return errors.Trace(common.InitUtils(cfg))
-}
-
-// Start starts running a command
-func Start(args []string) {
-	rootCmd := &cobra.Command{
-		Use:   "dmctl",
-		Short: "DM control",
-	}
-
+func init() {
 	// --worker worker1 -w worker2 --worker=worker3,worker4 -w=worker5,worker6
 	rootCmd.PersistentFlags().StringSliceVarP(&commandMasterFlags.workers, "worker", "w", []string{}, "DM-worker ID")
 	rootCmd.AddCommand(
@@ -74,9 +66,57 @@ func Start(args []string) {
 		master.NewPurgeRelayCmd(),
 		master.NewMigrateRelayCmd(),
 	)
+}
 
-	rootCmd.SetArgs(args)
+// Init initializes dm-control
+func Init(cfg *common.Config) error {
+	// set the log level temporarily
+	log.SetLevel(zapcore.InfoLevel)
+
+	return errors.Trace(common.InitUtils(cfg))
+}
+
+// PrintUsage prints usage
+func PrintUsage() {
+	maxCmdLen := 0
+	for _, cmd := range rootCmd.Commands() {
+		if maxCmdLen < len(cmd.Name()) {
+			maxCmdLen = len(cmd.Name())
+		}
+	}
+	fmt.Println("Available Commands:")
+	for _, cmd := range rootCmd.Commands() {
+		format := fmt.Sprintf("  %%-%ds\t%%s\n", maxCmdLen)
+		fmt.Printf(format, cmd.Name(), cmd.Use)
+	}
+}
 
+// HasCommand represent whether rootCmd has this command
+func HasCommand(name string) bool {
+	for _, cmd := range rootCmd.Commands() {
+		if name == cmd.Name() {
+			return true
+		}
+	}
+	return false
+}
+
+// PrintHelp print help message for special subCommand
+func PrintHelp(args []string) {
+	cmd, _, err := rootCmd.Find(args)
+	if err != nil {
+		fmt.Println(err)
+		rootCmd.SetOut(os.Stdout)
+		rootCmd.Usage()
+		return
+	}
+	cmd.SetOut(os.Stdout)
+	cmd.Usage()
+}
+
+// Start starts running a command
+func Start(args []string) {
+	rootCmd.SetArgs(args)
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(rootCmd.UsageString())
 	}

dm/ctl/master/break_ddl_lock.go

@@ -16,6 +16,7 @@ package master
 import (
 	"context"
 	"fmt"
+	"os"
 
 	"github.com/pingcap/dm/dm/ctl/common"
 	"github.com/pingcap/dm/dm/pb"
@@ -40,7 +41,8 @@ func NewBreakDDLLockCmd() *cobra.Command {
 // breakDDLLockFunc does break DDL lock
 func breakDDLLockFunc(cmd *cobra.Command, _ []string) {
 	if len(cmd.Flags().Args()) != 1 {
-		fmt.Println(cmd.Usage())
+		cmd.SetOut(os.Stdout)
+		cmd.Usage()
 		return
 	}
 	taskName := cmd.Flags().Arg(0)