-
Notifications
You must be signed in to change notification settings - Fork 2
/
cli.go
169 lines (152 loc) · 4.5 KB
/
cli.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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
package cli
import (
"encoding/json"
"fmt"
"os"
"os/signal"
"strings"
"syscall"
"time"
"github.com/spf13/pflag"
"github.com/spf13/viper"
logger "github.com/textileio/go-log/v2"
)
var (
log = logger.Logger("cli")
)
// Flag describes a configuration flag.
type Flag struct {
Name string
DefValue interface{}
Description string
}
// ConfigureCLI configures a Viper environment with flags and envs.
func ConfigureCLI(v *viper.Viper, envPrefix string, flags []Flag, flagSet *pflag.FlagSet) {
v.SetEnvPrefix(envPrefix)
v.AutomaticEnv()
v.SetEnvKeyReplacer(strings.NewReplacer(".", "_", "-", "_"))
for _, flag := range flags {
switch defval := flag.DefValue.(type) {
case string:
flagSet.String(flag.Name, defval, flag.Description)
v.SetDefault(flag.Name, defval)
case []string:
flagSet.StringSlice(flag.Name, defval, flag.Description+"; repeatable")
v.SetDefault(flag.Name, defval)
case bool:
flagSet.Bool(flag.Name, defval, flag.Description)
v.SetDefault(flag.Name, defval)
case int:
flagSet.Int(flag.Name, defval, flag.Description)
v.SetDefault(flag.Name, defval)
case int64:
flagSet.Int64(flag.Name, defval, flag.Description)
v.SetDefault(flag.Name, defval)
case uint64:
flagSet.Uint64(flag.Name, defval, flag.Description)
v.SetDefault(flag.Name, defval)
case time.Duration:
flagSet.Duration(flag.Name, defval, flag.Description)
v.SetDefault(flag.Name, defval)
default:
log.Fatalf("unknown flag type: %T", flag)
}
if err := v.BindPFlag(flag.Name, flagSet.Lookup(flag.Name)); err != nil {
log.Fatalf("binding flag %s: %s", flag.Name, err)
}
}
}
// ConfigureLogging configures the default logger with the right setup depending flag/envs.
// If logLevels is not nil, only logLevels values will be configured to Info/Debug depending
// on viper flags. if logLevels is nil, all sub-logs will be configured.
func ConfigureLogging(v *viper.Viper, logLevels []string) error {
var format logger.LogFormat
if v.GetBool("log-json") {
format = logger.JSONOutput
} else if v.GetBool("log-plaintext") {
format = logger.PlaintextOutput
} else {
format = logger.ColorizedOutput
}
logger.SetupLogging(logger.Config{
Format: format,
Level: logger.LevelError,
Stderr: false,
Stdout: true,
})
logLevel := logger.LevelInfo
if v.GetBool("log-debug") {
logLevel = logger.LevelDebug
}
if len(logLevels) == 0 {
logger.SetAllLoggers(logLevel)
return nil
}
mapLevel := make(map[string]logger.LogLevel, len(logLevels))
for i := range logLevels {
mapLevel[logLevels[i]] = logLevel
}
if err := logger.SetLogLevels(mapLevel); err != nil {
return fmt.Errorf("set log levels: %s", err)
}
return nil
}
// ParseStringSlice returns a single slice of values that may have been set by either repeating
// a flag or using comma separation in a single flag.
// This is used to enable repeated flags as well as env vars that can't be repeated.
// In either case, Viper understands how to write the config entry as a list.
func ParseStringSlice(v *viper.Viper, key string) []string {
vals := make([]string, 0)
for _, val := range v.GetStringSlice(key) {
parts := strings.Split(val, ",")
for _, p := range parts {
if p != "" {
vals = append(vals, p)
}
}
}
return vals
}
// ExpandEnvVars expands env vars present in the config.
func ExpandEnvVars(v *viper.Viper, settings map[string]interface{}) {
for name, val := range settings {
if str, ok := val.(string); ok {
v.Set(name, os.ExpandEnv(str))
}
}
}
// CheckErr ends in a fatal log if err is not nil.
func CheckErr(err error) {
if err != nil {
log.Fatal(err)
}
}
// CheckErrf ends in a fatal log if err is not nil.
func CheckErrf(format string, err error) {
if err != nil {
log.Fatalf(format, err)
}
}
// MarshalConfig marshals a *viper.Viper config to JSON. pretty controls if the
// result is indented or not. It replaces the masked fields with three
// asterisks, if they are present.
func MarshalConfig(v *viper.Viper, pretty bool, maskedFields ...string) ([]byte, error) {
all := v.AllSettings()
for _, f := range maskedFields {
if _, exists := all[f]; exists {
all[f] = "***"
}
}
if pretty {
return json.MarshalIndent(all, "", " ")
}
return json.Marshal(all)
}
// HandleInterrupt attempts to cleanup while allowing the user to force stop the process.
func HandleInterrupt(cleanup func()) {
quit := make(chan os.Signal, 1)
signal.Notify(quit, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
<-quit
log.Info("Gracefully stopping... (press Ctrl+C again to force)")
cleanup()
}