使用Cobra构建命令行程序

February 22, 2020 (最后修改: August 21, 2021)

本文属于介绍 NWPC 消息平台 系列文章。

Cobra 是一款使用GOLANG实现的命令行库。官方给出的示例中使用全局变量设置参数，并大量使用init()函数添加命令逻辑。在各种编程规范中，都不建议使用全局变量，尤其是可变的全局变量。如果命令行程序比较复杂，使用init()初始化命令逻辑也不容易控制。

Cobra 作者开发的 Hugo 使用 Cobra 构建命令行程序，但 Hugo 对 Cobra 进行一定的封装，避免使用全局变量。

本文首先介绍 Cobra 默认的命令构建方式，再介绍一种参考 Hugo 模式简化而来的命令构建方式，在 nwpc-message-client 项目中使用。

Cobra 默认的命令构建方式

cobra 默认使用全局变量构建子程序和命令选项，并在每个子命令文件的init()函数中将子命令注册到主命令中。

下面的代码来自 nwpc-data-client 项目。

主命令在cmd包的root.go文件创建，并提供 Execute 函数供 main 文件调用命令。

package cmd

import (
	"fmt"
	"github.com/spf13/cobra"
	"os"
)

var rootCmd = &cobra.Command{
	Use:   "nwpc_data_client",
	Short: "Data client for NWPC.",
	Long:  "A data client for GRAPES models in NWPC.",
	Run: func(cmd *cobra.Command, args []string) {
		cmd.Help()
	},
}

func Execute() {
	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

每个子命令在单独的文件中创建，例如下面代码创建local 子命令的全局变量，使用全局变量定义命令选项，并在 init 函数中将该命令注册到主命令。

package cmd

import (
	"fmt"
	"github.com/nwpc-oper/nwpc-data-client/common"
	"github.com/nwpc-oper/nwpc-data-client/common/config"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"os"
	"path/filepath"
	"strings"
)

func init() {
	rootCmd.AddCommand(localCmd)

	localCmd.Flags().StringVar(&configDir, "config-dir", "",
		"Config dir.")

	localCmd.Flags().StringVar(&dataType, "data-type", "",
		"Data type used to locate config file path in config dir.")

	localCmd.Flags().StringVar(&locationLevels, "location-level", "",
		"Location levels, split by ',', such as 'runtime,archive'.")

	localCmd.Flags().StringVar(&startTimeSting, "start-time", "",
		"start time, YYYYMMDDHH, such as 2020021400")
	localCmd.Flags().StringVar(&forecastTimeString, "forecast-time", "",
		"forecast time, FFFh, such as 0h, 120h")
}

var localCmd = &cobra.Command{
	Use:   "local",
	Short: "Find local data path.",
	Run: func(cmd *cobra.Command, args []string) {
		// main code
	},
}

// ... skip ...

以同样的方式定义另一个子命令 hpc。

package cmd

import (
	"fmt"
	"github.com/nwpc-oper/nwpc-data-client/common"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"strings"
)

func init() {
	rootCmd.AddCommand(hpcCmd)

	hpcCmd.Flags().StringVar(&configDir, "config-dir", "",
		"Config dir")

	hpcCmd.Flags().StringVar(&dataType, "data-type", "",
		"Data type used to locate config file path in config dir.")

	hpcCmd.Flags().StringVar(&locationLevels, "location-level", "",
		"Location levels, split by ',', such as 'runtime,archive'.")

	hpcCmd.Flags().StringVar(&startTimeSting, "start-time", "",
		"start time, YYYYMMDDHH, such as 2020021400")
	hpcCmd.Flags().StringVar(&forecastTimeString, "forecast-time", "",
		"forecast time, FFFh, such as 0h, 120h")

	hpcCmd.Flags().StringVar(&storageUser, "storage-user", user, "user name for storage.")
	hpcCmd.Flags().StringVar(&storageHost, "storage-host", "10.40.140.44:22", "host for storage")
	hpcCmd.Flags().StringVar(&privateKeyFilePath, "private-key", fmt.Sprintf("%s/.ssh/id_rsa", home),
		"private key file path")
	hpcCmd.Flags().StringVar(&hostKeyFilePath, "host-key", fmt.Sprintf("%s/.ssh/known_hosts", home),
		"host key file path")
}

var hpcCmd = &cobra.Command{
	Use:   "hpc",
	Short: "Find data path on hpc.",
	Run: func(cmd *cobra.Command, args []string) {
		// main code
	},
}

// ...skip...

对比两个命令的命令参数，可以发现两者公用一些变量，也有单独使用的变量。这些变量都暴露在全局环境中，不利于不同命令间的隔离。即便两个命令使用相同的变量，将它们区分开也更利于后续的修改。在全局环境中为每个命令定义独立的参数变量显然太繁琐。而且这些变量仅在解析命令行参数的时候会使用到，使用范围仅限各自的子命令，没有必要暴露在全局环境中。

因此下面章节介绍仿照 Hugo 的命令行程序构建方法，将参数细节隐藏到子命令的struct中。

不使用全局变量的方式

上节的示例的全局环境中有三个部分：

主命令 rootCmd
子命令 localCmd 和 hpcCmd 等
子命令注册到主命令中，使用 init 函数实现

下面介绍如何将以上部分都从全局环境中隐藏。

主命令

将主命令封装到 commandsBuilder 中，并在该结构中保存所有子命令的列表。子命令符合 Commnad 接口，返回*cobra.Command对象，用于后续将子命令注册到主命令中。

type commandsBuilder struct {
	commands    []Command
	rootCommand *cobra.Command
}

type Command interface {
	getCommand() *cobra.Command
}

func (b *commandsBuilder) getCommand() *cobra.Command {
	return b.rootCommand
}

func (b *commandsBuilder) addCommands(commands ...Command) *commandsBuilder {
	b.commands = append(b.commands, commands...)
	return b
}

主程序使用 newCommandsBuilder 函数构建新的命令对象，避免创建全局变量。

func newCommandsBuilder() *commandsBuilder {
	return &commandsBuilder{
		rootCommand: &cobra.Command{
			Use:   appCommand,
			Short: "A client for NWPC message.",
			Long:  "A client for NWPC message.",
			Run: func(cmd *cobra.Command, args []string) {
			},
		},
	}
}

子命令

所有子命令都包含符合 Command 接口的 BaseCommand 结构。

type BaseCommand struct {
	cmd *cobra.Command
}

func (c *BaseCommand) getCommand() *cobra.Command {
	return c.cmd
}

下面的代码创建 ecflow-client 子命令，包含在 ecflowClientCommand 结构中。子命令需要的各项命令行参数也封装在该结构体中。

type ecflowClientCommand struct {
	BaseCommand

	commandOptions string

	rabbitmqServer string
	writeTimeout   time.Duration

	useBroker     bool
	brokerAddress string

	disableSend bool
}

func (ec *ecflowClientCommand) runCommand(cmd *cobra.Command, args []string) error {
	// ...skip...
}

使用 newEcflowClientCommand 函数创建 ecflowClientCommand 结构，并在该函数中绑定cobra需要的命令行参数等内容。

func newEcflowClientCommand() *ecflowClientCommand {
	ec := &ecflowClientCommand{
		writeTimeout: 2 * time.Second,
	}
	ecFlowClientCmd := &cobra.Command{
		Use:   "ecflow-client",
		Short: "send messages for ecflow_client command",
		RunE:  ec.runCommand,
	}

	ecFlowClientCmd.Flags().StringVar(&ec.commandOptions, "command-options", "",
		"ecflow_client command options")

	ecFlowClientCmd.Flags().StringVar(&ec.rabbitmqServer, "rabbitmq-server", "",
		"rabbitmq server, such as amqp://guest:guest@host:port")

	ecFlowClientCmd.Flags().BoolVar(&ec.useBroker, "with-broker", true,
		"deliver message using a broker, should set --broker-address when enabled.")
	ecFlowClientCmd.Flags().StringVar(&ec.brokerAddress, "broker-address", "",
		"broker address, work with --with-broker")

	ecFlowClientCmd.Flags().BoolVar(&ec.disableSend, "disable-send", false,
		"disable message deliver, just for debug.")

	ecFlowClientCmd.MarkFlagRequired("command-options")
	ecFlowClientCmd.MarkFlagRequired("rabbitmq-server")

	ec.cmd = ecFlowClientCmd
	return ec
}

其他子命令也使用类似的方式创建。

注册

commandsBuilder 的 addAll 函数将所有的子命令添加到主命令中。

func (b *commandsBuilder) addAll() *commandsBuilder {
	b.addCommands(
		newVersionCommand(),
		newProductionCommand(),
		newEcflowClientCommand(),
		newBrokerCommand(),
	)
	return b
}

commandsBuilder 的 build 函数将所有的子命令注册到主命令中。

func (b *commandsBuilder) build() *commandsBuilder {
	for _, command := range b.commands {
		b.rootCommand.AddCommand(command.getCommand())
	}
	return b
}

main程序只需要执行将上述步骤都封装好的 Execute 函数。

func Execute() {
	consumerCommand := newCommandsBuilder().addAll().build()
	rootCmd := consumerCommand.getCommand()

	if err := rootCmd.Execute(); err != nil {
		os.Exit(1)
	}
}

参考

nwpc-oper/nwpc-data-client

nwpc-oper/nwpc-message-client