-
Notifications
You must be signed in to change notification settings - Fork 0
/
command.go
117 lines (97 loc) · 2.35 KB
/
command.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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package cmdutil
import (
graylog "github.com/gemnasium/logrus-graylog-hook/v3"
"github.com/google/uuid"
"github.com/sirupsen/logrus"
"github.com/spf13/cobra"
)
type Option func(*cobra.Command) error
func New(use, desc string, options ...Option) *cobra.Command {
cmd := &cobra.Command{
Use: use,
Short: desc,
}
var (
preRuns = make([]func(*cobra.Command, []string), 0)
persistentPreRuns = make([]func(*cobra.Command, []string), 0)
)
for _, o := range options {
err := o(cmd)
Must(err)
if cmd.PreRun != nil {
preRuns = append(preRuns, cmd.PreRun)
}
cmd.PreRun = nil
if cmd.PersistentPreRun != nil {
persistentPreRuns = append(persistentPreRuns, cmd.PersistentPreRun)
}
cmd.PersistentPreRun = nil
}
if len(persistentPreRuns) > 0 {
cmd.PersistentPreRun = func(cmd *cobra.Command, args []string) {
for _, run := range persistentPreRuns {
run(cmd, args)
}
}
}
cmd.PreRun = func(cmd *cobra.Command, args []string) {
for _, run := range preRuns {
run(cmd, args)
}
}
return cmd
}
func WithSubCommand(sub *cobra.Command) Option {
return func(parent *cobra.Command) error {
parent.AddCommand(sub)
return nil
}
}
func WithRun(run RunFuncWithContext) Option {
return func(cmd *cobra.Command) error {
cmd.Run = wrapRootConext(run)
return nil
}
}
type LoggerOption struct {
JSONFormatter bool
GELFLogger bool
}
func (o *LoggerOption) Bind(cmd *cobra.Command) error {
var (
flagJSON bool
flagGELFAddress string
)
// Bind json-logs flag, if enabled.
if o.JSONFormatter {
cmd.PersistentFlags().BoolVar(
&flagJSON, "json-logs", false, "Print the logs in JSON format")
}
// Bind gelf-address flag, if enabled.
if o.GELFLogger {
cmd.PersistentFlags().StringVar(
&flagGELFAddress, "gelf-address", "",
`Address to Graylog for logging (format: "ip:port")`)
}
cmd.PreRun = func(cmd *cobra.Command, args []string) {
if flagJSON {
logrus.SetFormatter(&logrus.JSONFormatter{
FieldMap: logrus.FieldMap{
logrus.FieldKeyTime: "time",
logrus.FieldKeyLevel: "level",
logrus.FieldKeyMsg: "message",
},
})
}
if flagGELFAddress != "" {
hook := graylog.NewGraylogHook(flagGELFAddress,
map[string]interface{}{
"uuid": uuid.New(),
"facility": Name,
})
hook.Level = logrus.DebugLevel
logrus.AddHook(hook)
}
}
return nil
}