Golang实践录：命令行cobra库实例

golang 1.5k 字 次

本文使用 cobra 库实现一个命令行工具，类似 git、docker、kubectl 这类的工具。
本文仅为一个初具模型的示例，但有实践参考意义。

起因

在编程中，很多时候，程序都会处理多个参数，特别是一些工具类的函数，需要整合较多功能，即使同一功能，也会有不同参数，利用配置文件或命令选项方式，可使程序具备通用性，也具扩展性。

简单介绍

cobra 功能较强大，在 golang 生态中有很多应用，如大名鼎鼎的 docker。其支持子命令执行，配置文件读写等，本文以实战为目的，不过多介绍。

整体结构

工程名为 cmdtool，见名知义。
工程目录及对应介绍如下： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
.
├── cmd ## 子命令总目录
│   ├── db    ## 子命令1实现目录
│   ├── misc  ## 子命令2实现目录
│   ├── rootCmd.go  ## 子命令入口
│   └── test  ## 子命令3实现目录
├── common ## 共用函数、变量
│   ├── conf
│   ├── constants
│   └── globalfunc.go
├── config.yaml ## 配置文件
├── go.mod
├── go.sum
├── main.go  ## 入口函数 
├── mybuild.sh ## 编译脚本
├── pkg  ## 库
│   ├── com
│   └── wait
├── README
└── vendor ## 依赖库
    ├── github.com
    ├── golang.org
    ├── gopkg.in
    ├── k8s.io
    └── xorm.io

其中 cmd 是所有子命令的入口目录，不同子命令，以不同子目录形式存在。common 目录存在共用的变量或初始化函数，等等。pkg 为个人总结积累的一些有用的库。
main.go 为主函数，调用了 cmd/rootCmd.go 的创建命令函数，由此进入 cobra 的处理框架中。
一般情况下，只需要扩展 cmd 目录下子命令，并补充 rootCmd.go 函数即可，其它即为业务程序的处理。
注：原本设计的思路是，在子命令包的 init 函数中自动注册到 rootCmd 中，但发现不一定符合逻辑，故舍弃，需手动在 rootCmd 添加。

工程分解

入口函数

主入口函数非常简单，实际调用了 rootCmd.go 中的执行函数。 

1
2
3
4
5
6
7
8
9
10
11
12
13
package main

import (
    _ "fmt"
    "os"
    rootCmd "github.com/latelee/cmdtool/cmd"
)

func main() {
trueif err := rootCmd.Execute(); err != nil {
truetrueos.Exit(1)
true}
}

命令行入口

rootCmd.go 源码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
package cmd

import (
true"os"
true"bytes"
    "path/filepath"
true"github.com/spf13/cobra"
true"github.com/spf13/viper"
true
true"github.com/fsnotify/fsnotify"

    "k8s.io/klog"
    test "github.com/latelee/cmdtool/cmd/test"
truemisc "github.com/latelee/cmdtool/cmd/misc"
truedb   "github.com/latelee/cmdtool/cmd/db"
trueconf "github.com/latelee/cmdtool/common/conf"
)

var (
    longDescription = `  database test tool.
  命令终端测试示例工具。
`
    example = `  comming soon...
`
)

var cfgFile string

var rootCmd = &cobra.Command{
trueUse:   filepath.Base(os.Args[0]),
trueShort: "database tool",
trueLong: longDescription,
trueExample: example,
trueVersion: "1.0",
}

func Execute() error {
    rootCmd.AddCommand(test.NewCmdTest())
truerootCmd.AddCommand(misc.NewCmdMisc())
truerootCmd.AddCommand(db.NewCmdDb())

truereturn rootCmd.Execute()
}

func init() {
truecobra.OnInitialize(initConfig)

truerootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (config.yaml)")

    rootCmd.PersistentFlags().BoolVar(&conf.FlagPrint, "print", false, "will print sth")

}

var yamlExample = []byte(
`dbserver:
  dbstr: helloooooo
  timeout:
    connect: 67s
    singleblock: 2s
  name:
    name: firstblood
`)

func initConfig() {
trueif cfgFile != "" {
truetrueviper.SetConfigFile(cfgFile)
true} else {
truetrueviper.AddConfigPath("./")
truetrueviper.SetConfigName("config")
truetrueviper.SetConfigType("yaml")
true}

trueviper.AutomaticEnv()

trueerr := viper.ReadInConfig();
trueif  err != nil {
truetrueklog.Println("not found config file. using default")
truetrueviper.ReadConfig(bytes.NewBuffer(yamlExample))
truetrueviper.SafeWriteConfig()
truetrue
true}
trueconf.FlagDBServer = viper.GetString("dbserver.dbstr")
trueconf.FlagTimeout = viper.GetString("dbserver.timeout.connect")
trueconf.FlagName = viper.GetString("dbserver.name.name")
trueklog.Println(conf.FlagDBServer, conf.FlagTimeout, conf.FlagName)

true//设置监听回调函数
trueviper.OnConfigChange(func(e fsnotify.Event) {
truetrueconf.FlagTimeout = viper.GetString("dbserver.timeout.connect")
true})

trueviper.WatchConfig()

}

其中 initConfig 函数作用是读取配置文件字段，如果没有文件则自动生成默认的配置。注意，该函数的 yamlExample 需要保持实际配置文件的格式（从 viper.GetString 函数参数可以看出 dbserver 为顶层字段）。
最后利用 viper 监听配置文件的变化。实际测试发现会触发2次，利用循环定时判断变量值可以解决。

子命令实现

子命令的实现形式大同小异，以 test 为例，源码如下： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
package cmd

import (
    "github.com/spf13/cobra"
true_ "github.com/spf13/pflag"
true
true"k8s.io/klog"
)

var (
    name = `test`
    shortDescription = `  test command`
    longDescription  = `  test...
`
    example = `  example comming up...
`
)

type UserCmdFunc struct {
truename string
truefn func(args []string)
}

func NewCmdTest() *cobra.Command{
true
    var cmd = &cobra.Command{
        Use:     name,
        Short:   shortDescription,
        Long:    longDescription,
        Example: example,
        RunE: func(cmd *cobra.Command, args []string) error {
truetruetrueif (len(args) == 0) {
truetruetruetrueklog.Warning("no args found")
truetruetruetruereturn nil
truetruetrue}

            if (args[0] == "foo"){
                foo(args)
            } else if (args[0] == "watch"){
                testWatch(args)
            } else {
truetruetruetrueklog.Printf("cmd '%v' not support", args[0])
truetruetruetruereturn nil
truetruetrue} 
            return nil
        },
    }

    return cmd
}

在 NewCmdTest 函数中创建 cobra.Command 并返回，在 RunE 中判断参数并真正执行业务函数。本例实现了参数监听功能，源码：

1
2
3
4
5
6
7
8
9
10
11
// 监听配置参数变化
func testWatch(args []string) {
truetimeout := conf.FlagTimeout
truefor {
truetrueif timeout != conf.FlagTimeout {
truetruetrueklog.Printf("param changed: %v\n", conf.FlagTimeout)
truetruetruetimeout = conf.FlagTimeout
truetrue}
truetruecom.Sleep(1000)
true}
}

当配置文件相应字段变化时，将其打印出来。

测试

默认输出帮助信息： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
$ ./cmdtool.exe
  database test tool.
  命令终端测试示例工具。

Usage:
  cmdtool.exe [command]

Examples:
  comming soon...


Available Commands:
  db            db command
  help        Help about any command
  misc          misc command
  test          test command

Flags:
  -h, --help      help for cmdtool.exe
      --print     will print sth
      --version   version for cmdtool.exe

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

执行子命令：

1
2
3
$ ./cmdtool.exe test foo
[2020-10-20 21:46:39.304 rootCmd.go:113] helloooooo 61s firstblood
[2020-10-20 21:46:39.305 busy.go:12] test foo.....

监听配置文件： 

1
2
3
$ ./cmdtool.exe test watch
[2020-10-20 21:47:14.408 rootCmd.go:113] helloooooo 61s firstblood
[2020-10-20 21:47:29.411 busy.go:20] param changed: 100s

源码

源码在此。

其它事项

利用viper.SafeWriteConfig()写配置文件时，发现 yamlExample 添加的注释会被删除，所以可以考虑直接将字符串通过ioutil.WriteFile写到文件。

viper 获取 yaml 参数的接口： 

1
2
获取数值、字符串、字符串数组、数值数组
GetInt GetInt32 GetInt64 GetUint GetUint32 GetUint64 GetString GetStringSlice GetIntSlice