-
Notifications
You must be signed in to change notification settings - Fork 23
/
main.go
128 lines (106 loc) · 4.54 KB
/
main.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
package main
import (
"fmt"
"os"
"strings"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"github.com/spf13/viper"
)
// I'm declaring as vars so I can test easier, I recommend declaring these as constants
var (
// The name of our config file, without the file extension because viper supports many different config file languages.
defaultConfigFilename = "stingoftheviper"
// The environment variable prefix of all environment variables bound to our command line flags.
// For example, --number is bound to STING_NUMBER.
envPrefix = "STING"
// Replace hyphenated flag names with camelCase in the config file
replaceHyphenWithCamelCase = false
)
func main() {
cmd := NewRootCommand()
if err := cmd.Execute(); err != nil {
os.Exit(1)
}
}
// Build the cobra command that handles our command line tool.
func NewRootCommand() *cobra.Command {
// Store the result of binding cobra flags and viper config. In a
// real application these would be data structures, most likely
// custom structs per command. This is simplified for the demo app and is
// not recommended that you use one-off variables. The point is that we
// aren't retrieving the values directly from viper or flags, we read the values
// from standard Go data structures.
color := ""
number := 0
// Define our command
rootCmd := &cobra.Command{
Use: "stingoftheviper",
Short: "Cober and Viper together at last",
Long: `Demonstrate how to get cobra flags to bind to viper properly`,
PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
// You can bind cobra and viper in a few locations, but PersistencePreRunE on the root command works well
return initializeConfig(cmd)
},
Run: func(cmd *cobra.Command, args []string) {
// Working with OutOrStdout/OutOrStderr allows us to unit test our command easier
out := cmd.OutOrStdout()
// Print the final resolved value from binding cobra flags and viper config
fmt.Fprintln(out, "Your favorite color is:", color)
fmt.Fprintln(out, "The magic number is:", number)
},
}
// Define cobra flags, the default value has the lowest (least significant) precedence
rootCmd.Flags().IntVarP(&number, "number", "n", 7, "What is the magic number?")
rootCmd.Flags().StringVarP(&color, "favorite-color", "c", "red", "Should come from flag first, then env var STING_FAVORITE_COLOR then the config file, then the default last")
return rootCmd
}
func initializeConfig(cmd *cobra.Command) error {
v := viper.New()
// Set the base name of the config file, without the file extension.
v.SetConfigName(defaultConfigFilename)
// Set as many paths as you like where viper should look for the
// config file. We are only looking in the current working directory.
v.AddConfigPath(".")
// Attempt to read the config file, gracefully ignoring errors
// caused by a config file not being found. Return an error
// if we cannot parse the config file.
if err := v.ReadInConfig(); err != nil {
// It's okay if there isn't a config file
if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
return err
}
}
// When we bind flags to environment variables expect that the
// environment variables are prefixed, e.g. a flag like --number
// binds to an environment variable STING_NUMBER. This helps
// avoid conflicts.
v.SetEnvPrefix(envPrefix)
// Environment variables can't have dashes in them, so bind them to their equivalent
// keys with underscores, e.g. --favorite-color to STING_FAVORITE_COLOR
v.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
// Bind to environment variables
// Works great for simple config names, but needs help for names
// like --favorite-color which we fix in the bindFlags function
v.AutomaticEnv()
// Bind the current command's flags to viper
bindFlags(cmd, v)
return nil
}
// Bind each cobra flag to its associated viper configuration (config file and environment variable)
func bindFlags(cmd *cobra.Command, v *viper.Viper) {
cmd.Flags().VisitAll(func(f *pflag.Flag) {
// Determine the naming convention of the flags when represented in the config file
configName := f.Name
// If using camelCase in the config file, replace hyphens with a camelCased string.
// Since viper does case-insensitive comparisons, we don't need to bother fixing the case, and only need to remove the hyphens.
if replaceHyphenWithCamelCase {
configName = strings.ReplaceAll(f.Name, "-", "")
}
// Apply the viper config value to the flag when the flag is not set and viper has a value
if !f.Changed && v.IsSet(configName) {
val := v.Get(configName)
cmd.Flags().Set(f.Name, fmt.Sprintf("%v", val))
}
})
}