Group bundle run flags by job and pipeline types

bundle/run/options.go

@@ -1,7 +1,8 @@
 package run
 
 import (
-	flag "github.com/spf13/pflag"
+	"github.com/databricks/cli/libs/cmdgroup"
+	"github.com/spf13/cobra"
 )
 
 type Options struct {
@@ -10,7 +11,11 @@ type Options struct {
 	NoWait   bool
 }
 
-func (o *Options) Define(fs *flag.FlagSet) {
-	o.Job.Define(fs)
-	o.Pipeline.Define(fs)
+func (o *Options) Define(cmd *cobra.Command) {
+	wrappedCmd := cmdgroup.NewCommandWithGroupFlag(cmd)
+	jobGroup := wrappedCmd.AddFlagGroup("Job")
+	o.Job.Define(jobGroup.FlagSet())
+
+	pipelineGroup := wrappedCmd.AddFlagGroup("Pipeline")
+	o.Pipeline.Define(pipelineGroup.FlagSet())
 }

cmd/bundle/run.go

@@ -24,7 +24,7 @@ func newRunCommand() *cobra.Command {
 	}
 
 	var runOptions run.Options
-	runOptions.Define(cmd.Flags())
+	runOptions.Define(cmd)
 
 	var noWait bool
 	cmd.Flags().BoolVar(&noWait, "no-wait", false, "Don't wait for the run to complete.")

libs/cmdgroup/command.go

@@ -0,0 +1,74 @@
+package cmdgroup
+
+import (
+	"io"
+	"strings"
+	"text/template"
+	"unicode"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+)
+
+type CommandWithGroupFlag struct {
+	cmd        *cobra.Command
+	flagGroups []*FlagGroup
+}
+
+func (c *CommandWithGroupFlag) Command() *cobra.Command {
+	return c.cmd
+}
+
+func (c *CommandWithGroupFlag) FlagGroups() []*FlagGroup {
+	return c.flagGroups
+}
+
+func NewCommandWithGroupFlag(cmd *cobra.Command) *CommandWithGroupFlag {
+	cmdWithFlagGroups := &CommandWithGroupFlag{cmd: cmd, flagGroups: make([]*FlagGroup, 0)}
+	cmd.SetUsageFunc(func(c *cobra.Command) error {
+		err := tmpl(c.OutOrStderr(), c.UsageTemplate(), cmdWithFlagGroups)
+		if err != nil {
+			c.PrintErrln(err)
+		}
+		return nil
+	})
+	cmd.SetUsageTemplate(usageTemplate)
+	return cmdWithFlagGroups
+}
+
+func (c *CommandWithGroupFlag) AddFlagGroup(name string) *FlagGroup {
+	fg := &FlagGroup{name: name, flagSet: pflag.NewFlagSet(name, pflag.ContinueOnError)}
+	c.flagGroups = append(c.flagGroups, fg)
+	return fg
+}
+
+type FlagGroup struct {
+	name    string
+	flagSet *pflag.FlagSet
+}
+
+func (c *FlagGroup) Name() string {
+	return c.name
+}
+
+func (c *FlagGroup) FlagSet() *pflag.FlagSet {
+	return c.flagSet
+}
+
+var templateFuncs = template.FuncMap{
+	"trim":                    strings.TrimSpace,
+	"trimRightSpace":          trimRightSpace,
+	"trimTrailingWhitespaces": trimRightSpace,
+}
+
+func trimRightSpace(s string) string {
+	return strings.TrimRightFunc(s, unicode.IsSpace)
+}
+
+// tmpl executes the given template text on data, writing the result to w.
+func tmpl(w io.Writer, text string, data interface{}) error {
+	t := template.New("top")
+	t.Funcs(templateFuncs)
+	template.Must(t.Parse(text))
+	return t.Execute(w, data)
+}

libs/cmdgroup/command_test.go

@@ -0,0 +1,51 @@
+package cmdgroup
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCommandFlagGrouping(t *testing.T) {
+	cmd := &cobra.Command{
+		Use:   "test [flags]",
+		Short: "test command",
+		Run: func(cmd *cobra.Command, args []string) {
+			// Do nothing
+		},
+	}
+
+	wrappedCmd := NewCommandWithGroupFlag(cmd)
+	jobGroup := wrappedCmd.AddFlagGroup("Job")
+	fs := jobGroup.FlagSet()
+	fs.String("job-name", "", "Name of the job")
+	fs.String("job-type", "", "Type of the job")
+
+	pipelineGroup := wrappedCmd.AddFlagGroup("Pipeline")
+	fs = pipelineGroup.FlagSet()
+	fs.String("pipeline-name", "", "Name of the pipeline")
+	fs.String("pipeline-type", "", "Type of the pipeline")
+
+	cmd.Flags().BoolP("bool", "b", false, "Bool flag")
+
+	buf := bytes.NewBuffer(nil)
+	cmd.SetOutput(buf)
+	cmd.Usage()
+
+	expected := `Usage:
+  test [flags]
+
+Job Flags:
+      --job-name string   Name of the job
+      --job-type string   Type of the job
+
+Pipeline Flags:
+      --pipeline-name string   Name of the pipeline
+      --pipeline-type string   Type of the pipeline
+
+Flags:
+  -b, --bool   Bool flag`
+	require.Equal(t, expected, buf.String())
+}

libs/cmdgroup/template.go

@@ -0,0 +1,13 @@
+package cmdgroup
+
+const usageTemplate = `Usage:{{if .Command.Runnable}}
+  {{.Command.UseLine}}{{end}}
+{{range .FlagGroups}}
+{{.Name}} Flags:
+{{.FlagSet.FlagUsages | trimTrailingWhitespaces}}
+{{end}}
+{{if .Command.HasAvailableLocalFlags}}Flags:
+{{.Command.LocalFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}{{if .Command.HasAvailableInheritedFlags}}
+
+Global Flags:
+{{.Command.InheritedFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}`