Merge pull request #23 from ShotaKitazawa/fix/22

fixed: required flag does not work well

Author: ShotaKitazawa

cmd/afterbench.go

@@ -10,6 +10,9 @@ import (
 var afterbenchCmd = &cobra.Command{
 	Use:   "afterbench",
 	Short: "Collect and parse profile data & Send to Slack",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkRequiredFlags(cmd.Flags())
+	},
 	RunE: func(c *cobra.Command, args []string) error {
 		executed = true
 		conf := cmd.ConfigAfterBench{

cmd/deploy.go

@@ -9,6 +9,9 @@ import (
 var deployCmd = &cobra.Command{
 	Use:   "deploy",
 	Short: "Deploy files from specified revision",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkRequiredFlags(cmd.Flags())
+	},
 	RunE: func(c *cobra.Command, args []string) error {
 		executed = true
 		conf := cmd.ConfigDeploy{
@@ -37,7 +40,7 @@ func init() {
 		"branch-name, tag-name, or commit-hash of deployed from Git remote-repo")
 	deployCmd.PersistentFlags().BoolVarP(&deployForce, "force", "f", false,
 		"force deploy")
-	deployCmd.PersistentFlags().StringVarP(&deploySlackToken, "slack-token", "t", getenvDefault("SLACK_TOKEN", ""),
+	deployCmd.PersistentFlags().StringVarP(&deploySlackToken, "slack-token", "t", getenvDefault("slack-token", ""),
 		"slack token of workspace where deployment notification will be sent")
-	_ = initCmd.MarkPersistentFlagRequired("slack-token")
+	setRequired(initCmd, "slack-token")
 }

cmd/import.go

@@ -10,6 +10,9 @@ import (
 var importCmd = &cobra.Command{
 	Use:   "import",
 	Short: "Import some files from hosts[].deploy.files[].target",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkRequiredFlags(cmd.Flags())
+	},
 	RunE: func(c *cobra.Command, args []string) error {
 		executed = true
 		conf := cmd.ConfigImport{

cmd/init.go

@@ -10,6 +10,9 @@ import (
 var initCmd = &cobra.Command{
 	Use:   "init",
 	Short: "Initialize local repository",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkRequiredFlags(cmd.Flags())
+	},
 	RunE: func(c *cobra.Command, args []string) error {
 		executed = true
 		conf := cmd.ConfigInit{
@@ -41,5 +44,5 @@ func init() {
 		"email of GitHub Account")
 	initCmd.PersistentFlags().StringVarP(&gitRemoteUrl, "remote-url", "r", getenvDefault("REMOTE_URL", ""),
 		"URL of remote repository (requirement)")
-	_ = initCmd.MarkPersistentFlagRequired("remote-url")
+	setRequired(initCmd, "remote-url")
 }

cmd/profiling.go

@@ -10,6 +10,9 @@ import (
 var profilingCmd = &cobra.Command{
 	Use:   "profiling",
 	Short: "Execute profiling command and wait to finish synchronously.",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkRequiredFlags(cmd.Flags())
+	},
 	RunE: func(c *cobra.Command, args []string) error {
 		executed = true
 		conf := cmd.ConfigProfiling{

cmd/push.go

@@ -9,6 +9,9 @@ import (
 var pushCmd = &cobra.Command{
 	Use:   "push",
 	Short: "Push local-repo to origin/${MAIN_BRANCH}",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkRequiredFlags(cmd.Flags())
+	},
 	RunE: func(c *cobra.Command, args []string) error {
 		executed = true
 		conf := cmd.ConfigPush{

cmd/root.go

@@ -1,12 +1,16 @@
 package cmd
 
 import (
+	"errors"
 	"flag"
 	"fmt"
+	"log"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 )
 
 // rootCmd represents the base command when called without any subcommands
@@ -46,10 +50,43 @@ func init() {
 		"local repository's path managed by isucontinuous")
 }
 
-func getenvDefault(key, defaultV string) string {
+// utils
+
+func getenvDefault(flagName, defaultV string) string {
+	key := strings.ToUpper(strings.ReplaceAll(flagName, "-", "_"))
 	result := os.Getenv(key)
 	if result == "" {
 		return defaultV
 	}
 	return result
 }
+
+const requiredFlagAnnotation = "isucontinuous/required"
+
+func setRequired(cmd *cobra.Command, flagNames ...string) {
+	for _, flagName := range flagNames {
+		if err := cmd.PersistentFlags().SetAnnotation(flagName, requiredFlagAnnotation, []string{"true"}); err != nil {
+			log.Fatal(err)
+		}
+	}
+}
+
+func checkRequiredFlags(flags *pflag.FlagSet) error {
+	requiredError := false
+	flagName := ""
+	flags.VisitAll(func(flag *pflag.Flag) {
+		requiredAnnotation := flag.Annotations[requiredFlagAnnotation]
+		if len(requiredAnnotation) == 0 {
+			return
+		}
+		flagRequired := requiredAnnotation[0] == "true"
+		if flagRequired && !flag.Changed && getenvDefault(flag.Name, "") == "" {
+			requiredError = true
+			flagName = flag.Name
+		}
+	})
+	if requiredError {
+		return errors.New("Required flag `" + flagName + "` has not been set")
+	}
+	return nil
+}

cmd/setup.go

@@ -10,6 +10,9 @@ import (
 var setupCmd = &cobra.Command{
 	Use:   "setup",
 	Short: "Install some softwares",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkRequiredFlags(cmd.Flags())
+	},
 	RunE: func(c *cobra.Command, args []string) error {
 		executed = true
 		conf := cmd.ConfigSetup{

cmd/sync.go

@@ -10,6 +10,9 @@ import (
 var syncCmd = &cobra.Command{
 	Use:   "sync",
 	Short: "synchronize local-repo with remote-repo",
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkRequiredFlags(cmd.Flags())
+	},
 	RunE: func(c *cobra.Command, args []string) error {
 		executed = true
 		conf := cmd.ConfigSync{