golang使用Cobra快速创建命令行

1、flag 获取参数

先看下我们将要创建的目录结构:

├── cmd
│   └── main.go
├── config
│   ├── config.go
│   ├── config.yaml
│   └── parse.go
├── go.mod
└── go.sum

我们需要从 config.yaml 中读取应用的配置到结构体中。

首先我们创建了 config.yaml 配置文件,里面包含了 server 和 data 的相关配置:

server:
  http:
    addr: 0.0.0.0:8012
data:
  database:
    driver: "mysql"
    connection: root:liufutian@tcp(localhost:3306)/answer
  cache:
    file_path: "./data/tmp/cache/cache.db"

接着我们需要定义下转出的结构体 config.go:

package config

// AllConfig all config
type AllConfig struct {
	Debug  bool    `json:"debug" mapstructure:"debug"`
	Data   *Data   `json:"data" mapstructure:"data"`
	Server *Server `json:"server" mapstructure:"server"`
}

// Server server config
type Server struct {
	HTTP *HTTP `json:"http" mapstructure:"http"`
}

// HTTP http config
type HTTP struct {
	Addr string `json:"addr" mapstructure:"addr"`
}

// Data data config
type Data struct {
	Database *Database  `json:"database" mapstructure:"database"`
	Cache    *CacheConf `json:"cache" mapstructure:"cache"`
}

// Database database config
type Database struct {
	Driver          string `json:"driver" mapstructure:"driver"`
	Connection      string `json:"connection" mapstructure:"connection"`
	ConnMaxLifeTime int    `json:"conn_max_life_time" mapstructure:"conn_max_life_time"`
	MaxOpenConn     int    `json:"max_open_conn" mapstructure:"max_open_conn"`
	MaxIdleConn     int    `json:"max_idle_conn" mapstructure:"max_idle_conn"`
}

// CacheConf cache
type CacheConf struct {
	FilePath string `json:"file_path" mapstructure:"file_path"`
}

接着我们定义下解析配置的代码 parse.go:

package config

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

// Config interface determines the common methods for parsing configuration from specified resources
type Config interface { // 1
	Parse(any) error
}

type FileConfig struct { // 2
	Path   string `json:"path"`
	parser *viper.Viper
}

// NewWithPath to create config from specified path
func NewWithPath(filePath string) (c *FileConfig, err error) { // 3
	var stat os.FileInfo

	stat, err = os.Stat(filePath)
	if err != nil {
		return
	}

	if !stat.Mode().IsRegular() {
		return nil, fmt.Errorf("%s is not a regular file", filePath)
	}

	p := viper.New()
	p.SetConfigFile(filePath)

	err = p.ReadInConfig()
	if err != nil {
		return
	}

	return &FileConfig{Path: filePath, parser: p}, nil
}

// Parse parses the configuration by object pointer
func (c *FileConfig) Parse(obj any) error { // 4
	return c.parser.Unmarshal(obj)
}
  1. 首先我们定义了一个 Config 接口,有一个 Parse 方法,接受一个 any 类型也就是 interface
  2. 接着我们定义一个具体的结构 FileConfig,实现 Config 接口。有两个字段 Path 文件路径和 viper 的引用。我们需要通过 viper 将文件读入一个 map
  3. 我们定义了一个工厂方法,通过传入的文件路径,将文件中的配置读入 viper 的 map 中,并返回实例化后的 FileConfig
  4. 最后我们实现 Config 的 Parse 接口,通过 viper 的 Unmarshal 将 map 中的配置解析到传入结构 obj 中

最后我们看下 main.go 中的逻辑:

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"go-demo/cobra/config"
)

var DefaultConfigFile = "./config/config.yaml"
var configFile string

func main() {
	c, err := readConfig() // 1
	if err != nil {
		panic(err)
	}
	marshal, err := json.Marshal(c)
	if err != nil {
		panic(err)
	}
	fmt.Println(string(marshal))
}

func readConfig() (c *config.AllConfig, err error) {
	flag.StringVar(&configFile, "c", "", "go run cmd/main.go -c xxx") // 2
	flag.Parse()

	if len(configFile) == 0 {
		configFile = DefaultConfigFile
	}
	c = &config.AllConfig{}

	conf, err := config.NewWithPath(configFile) // 3
	if err != nil {
		return nil, err
	}
	if err = conf.Parse(&c); err != nil { // 4
		return nil, err
	}
	return c, nil
}
  1. 首先我们定义了 readConfig 方法用于读取配置
  2. readConfig 中我们使用 go 中的 flag 包解析出命令行中的自定义配置文件路径 c
  3. 通过获取的文件路径初始化 FileConfig
  4. 解析文件中的配置到 AllConfig 结构中

我们执行下上面的代码:

go build cmd/main.go 
./main       
{"debug":false,"data":{"database":{"driver":"mysql","connection":"root:liufutian@tcp(localhost:3306)/answer","conn_max_life_time":0,"max_open_conn":0,"max_idle_conn":0},"cache":{"file_path":"./data/tmp/cache/cache.db"}},"server":{"http":{"addr":"0.0.0.0:8012"}}}

./main -c
flag needs an argument: -c
Usage of ./main:
  -c string
        go run cmd/main.go -c xxx

这个应该是我们大部分开发者一般的做法。但是这么做有一些缺点就是,界面太丑不清晰,无法灵活扩展嵌套命令,没有帮助文档等。

1、为什么使用 Cobra

Cobra 是一个强大的 Golang 库和工具,用于创建 CLI(命令行界面)应用程序。 Cobra 通过提供自动化流程的工具并提供可提高开发人员生产力的关键抽象来做到这一点。

使用 Cobra 的优点:

2、Cobra 的安装和使用

cobra 中有对应的客户端命令,我们可以像下面这样安装:

go install github.com/spf13/cobra-cli@latest

然后我们可以通过初始化命令创建一个 cobra 模板:

cobra-cli init 

然后可以用 add 命令添加你的命令:

cobra-cli add 你的命令

我们修改下我们上面的例子,使用 cobra 构造我们的客户端命令,首先在 cmd 文件夹下面增加一个 command.go 文件:

package main

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

// rootCmd 表示在没有任何子命令的情况下调用时的基本命令
var rootCmd = &cobra.Command{
	Use:     "main",
	Version: "0.0.1",
	Short:   "这里是短简介.",
	Long: `为了运行她, 使用:
	- 'main run' 启动这个程序.`,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	// Run: func(cmd *cobra.Command, args []string) { },
}

// Execute 将所有子命令添加到根命令并适当地设置标志。
// 这由 main.main() 调用。 它只需要在 rootCmd 上发生一次。
func Execute() {
	err := rootCmd.Execute()
	if err != nil {
		os.Exit(1)
	}
}

func init() {
	// 在这里您将定义您的标志和配置设置。
	// Cobra 支持持久化标志,如果在这里定义,
	// 对您的应用程序来说是全局的。

	// rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.cobra.yaml)")

	// Cobra 也支持本地标志,当这个动作被直接调用时它才会运行。
	rootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}

cobra.Command 命令中有几个参数介绍下:

然后 main.go 中只需要调用上面的 Excecute 方法:

package main

func main() {
	Execute()
}

接着我们构建下程序:

go build -o main ./cmd

运行下 main 看下效果:

./main              
为了运行她, 使用:
        - 'main run' 启动这个程序.

我们还可以添加一个 run 的子命令,用来运行应用程序。我们在 command.go 中添加子命令的代码:

package main

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

// rootCmd 表示在没有任何子命令的情况下调用时的基本命令
var rootCmd = &cobra.Command{
	Use:   "main",
	Short: "这里是短简介.",
	Long: `为了运行她, 使用:
	- 'main run' 启动这个程序.`,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	// Run: func(cmd *cobra.Command, args []string) { },
}

// runCmd represents the run command
var runCmd = &cobra.Command{ // 1
	Use:   "run",
	Short: "运行程序",
	Long:  `运行程序`,
	Run: func(cmd *cobra.Command, args []string) {
		runApp()
	},
}

// Execute 将所有子命令添加到根命令并适当地设置标志。
// 这由 main.main() 调用。 它只需要在 rootCmd 上发生一次。
func Execute() {
	err := rootCmd.Execute()
	if err != nil {
		os.Exit(1)
	}
}

func init() {
	// 在这里您将定义您的标志和配置设置。
	// Cobra 支持持久化标志,如果在这里定义,
	// 对您的应用程序来说是全局的。

	// rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.cobra.yaml)")

	// Cobra 也支持本地标志,当这个动作被直接调用时它才会运行。
	rootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")

	runCmd.Flags().StringVarP(&configFile, "config", "c", "", "可以指定自定义的配置文件") // 2

	rootCmd.AddCommand(runCmd) // 3
}
  1. 我们增加了一个 runCmd 子命令,runApp 里处理具体的运行逻辑
  2. 我们给 runCmd 增加了一个字符串类型的参数,并赋值给了我们定义的变量 configFile
  3. 根命令添加 runCmd 子命令

接着我们修改下 main.go 中的逻辑,去掉之前 flag 包获取配置的方法:

package main

import (
	"encoding/json"
	"fmt"
	"go-demo/cobra/config"
)

var DefaultConfigFile = "./config/config.yaml"
var configFile string // 2

func main() {
	Execute()
}

func runApp() { // 1
	c, err := readConfig()
	if err != nil {
		panic(err)
	}
	marshal, err := json.Marshal(c)
	if err != nil {
		panic(err)
	}
	fmt.Println(string(marshal))
}

func readConfig() (c *config.AllConfig, err error) {

	if len(configFile) == 0 {
		configFile = DefaultConfigFile
	}
	c = &config.AllConfig{}

	conf, err := config.NewWithPath(configFile)
	if err != nil {
		return nil, err
	}
	if err = conf.Parse(&c); err != nil {
		return nil, err
	}
	return c, nil
}
  1. runApp 运行程序的具体逻辑
  2. configFile 已经通过 cobra 的 flags 参数赋值

我们编译下 main 函数,并执行看下命令行的输出:

go build -o main ./cmd

./main   # 父命令            
为了运行她, 使用:
        - 'main run' 启动这个程序.

Usage:
  main [command]

Available Commands:
  completion  Generate the autocompletion script for the specified shell
  help        Help about any command
  run         运行程序

Flags:
  -h, --help     help for main
  -t, --toggle   Help message for toggle

Use "main [command] --help" for more information about a command.

./main -v   # 查看版本  
main version 0.1.1


 ./main run  # 执行run子命令       
{"debug":false,"data":{"database":{"driver":"mysql","connection":"root:liufutian@tcp(localhost:3306)/answer","conn_max_life_time":0,"max_open_conn":0,"max_idle_conn":0},"cache":{"file_path":"./data/tmp/cache/cache.db"}},"server":{"http":{"addr":"0.0.0.0:8012"}}}


./main run -h  #执行子命令help程序
运行程序

Usage:
  main run [flags]

Flags:
  -c, --config string   可以指定自定义的配置文件
  -h, --help            help for run