Automatic cobra command registration with fx

Cobra is a popular Go package for creating CLIs. It provides a lot of functionality for creating commands, subcommands, and flags. However, it can be tedious to manually register all of your commands.

fx is a Go package that provides a dependency injection framework. It can be used to automatically register your application components, including but not limited to Cobra commands.

In this blog post, I will show you how you can use fx to automatically register your Cobra commands.

This code was written 5 minutes ago, but works, and I imagine it could help folks bootstrap more complex CLIs rapidly.

The problem

When you create a Cobra command, you need to manually register it with the root command. This can be tedious, especially if you have a lot of commands.

For example, the following code registers a command called hello:

func NewHelloCmd() *cobra.Command {
  cmd := &cobra.Command{
    Use:   "hello",
    Short: "Say hello",
    Run: func(cmd *cobra.Command, args []string) {
      cmd.Println("Hello, world!")
    },
  }
  return cmd
}

func main() {
  rootCmd := &cobra.Command{
    Use:   "myapp",
    Short: "My application",
  }
  rootCmd.AddCommand(NewHelloCmd())
  rootCmd.Execute()
}Code language: Go (go)

As you can see, the NewHelloCmd() function creates a cobra.Command object. The main() function then registers this command with the rootCmd object.

As CLIs get more complex the number of registrations that need to happen in the right place can amass.

A Solution

The fx package provides a way to automatically register your Cobra commands. This can be done by using the fx.Provide() function.

Taking some inspiration from the docs site and registering “many handlers” for a web app, a similar principle can be applied to cobra.

An interface and a common type for the interface will be used.

type CobraCommand interface {
	Command() *cobra.Command
	FullName() string
}

type GenericCommand struct {
	cmd      *cobra.Command
	fullName string
}

func (g GenericCommand) Command() *cobra.Command {
	return g.cmd
}

func (g GenericCommand) FullName() string {
	return g.fullName
}Code language: Go (go)

These enable a command to be defined that contains both the command definition but also where it should site within the fully registered cobra application.

For example, fullName might be root hello for a hello world command. Deeper commands would have more components.

An example hello world task definition would perhaps look as follows:

func NewHelloCmd() GenericCommand {
	cmd := cobra.Command{
		Use:   "hello",
		Short: "hello - just says hello",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Println("Hello, World!")
		},
	}
	return GenericCommand{cmd: &cmd, fullName: "root hello"}
}
Code language: Go (go)

The real magic happens alongside the definition of the root command.

EDIT: Check the comments for a suggested better loop.

func NewRootCmd(commands []CobraCommand, lc fx.Lifecycle, shutdowner fx.Shutdowner) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "root",
		Short: "root - a simple CLI for this example",
		Run: func(cmd *cobra.Command, args []string) {
			cmd.Help()
		},
	}

	// Loop through commands adding them to other commands
	done := make(map[string]bool)
	for _, c1 := range commands {
		for _, c2 := range commands {
			if c1.FullName() == c2.FullName() || done[c1.FullName()] {
				continue
			}
			fullParent := strings.Join(strings.Split(c1.FullName(), " ")[:len(strings.Split(c1.FullName(), " "))-1], " ")
			if fullParent == c2.FullName() {
				// Add the child to the parent
				c2.Command().AddCommand(c1.Command())
				done[c1.FullName()] = true
				continue
			}
			if fullParent == "root" {
				cmd.AddCommand(c1.Command())
				done[c1.FullName()] = true
				continue
			}
		}
	}

	lc.Append(fx.Hook{
		OnStart: func(ctx context.Context) error {
			if err := cmd.Execute(); err != nil {
				fmt.Fprintf(os.Stderr, "Whoops. There was an error while executing the CLI '%s'", err)
				shutdowner.Shutdown(fx.ExitCode(1))
			}
			shutdowner.Shutdown()
			return nil
		},
	})
	return cmd
}
Code language: Go (go)

Firstly you see a regular root command definition

Next, you find a rather large set of loops that takes the full names defined by each command and compares them against each other to find the parent commands for each command.

When found this is registered and that command is marked as done.

The root command is a special case that needs to be looked for separately, as it is not in the main list of subcommands provided by fx.

The latter part of the method when you see lc.Append defines how the command should be run within the lifecycle of the fx application.

Finally, we just need to define our application

func New() *fx.App {
	return fx.New(
		fx.Provide(
			AsCommand(NewHelloCmd),
			AsCommand(NewPingerPing),
			AsCommand(NewPingerPong),
			fx.Annotate(
				commands.NewRootCmd,
				fx.ParamTags(`group:"commands"`),
			),
		),
		fx.NopLogger, // Note, this currently removes all fx logs, even on error
		fx.Invoke(func(*cobra.Command) {}),
	)
}

func AsCommand(f any) any {
	return fx.Annotate(
		f,
		fx.As(new(commands.CobraCommand)),
		fx.ResultTags(`group:"commands"`),
	)
}Code language: Go (go)

And run it, currently I use a 30 minute timeout for my case, but you might want to change this.

package main

import "addshore.com/fx-cobra-registration-demo/app"

func main() {
	app := app.New()

	startCtx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
	defer cancel()
	if err := app.Start(startCtx); err != nil {
		log.Fatal(err)
	}
}
Code language: Go (go)

Conclusion

Hopefully, this can provide you with an easy copy-and-paste starting point to work on your own automatic registration system.

I’m sure the logic in the loop can be cleaned up somewhat and also that there might be a better way to manage all of this within fx, but this first attempt works for me right now.