Adding Support for Environment variable prefixes

This commit is contained in:
spf13 2014-12-22 18:31:11 -05:00
parent 2909239689
commit 1022d75c73
3 changed files with 93 additions and 22 deletions

View file

@ -170,7 +170,7 @@ the value will be read each time it is accessed. It does not fix the
value when the BindEnv is called. value when the BindEnv is called.
AutomaticEnv is intended to be a convience helper. It will look for all AutomaticEnv is intended to be a convenience helper. It will look for all
keys that have been set (via defaults, config file, flag, or remote key keys that have been set (via defaults, config file, flag, or remote key
value) and call BindEnv on that key. It does value) and call BindEnv on that key. It does
not simply import all ENV variables. Because of this behavior its not simply import all ENV variables. Because of this behavior its

View file

@ -74,6 +74,7 @@ type viper struct {
configName string configName string
configFile string configFile string
configType string configType string
envPrefix string
config map[string]interface{} config map[string]interface{}
override map[string]interface{} override map[string]interface{}
@ -125,6 +126,22 @@ func (v *viper) SetConfigFile(in string) {
} }
} }
// Define a prefix that ENVIRONMENT variables will use.
func SetEnvPrefix(in string) { v.SetEnvPrefix(in) }
func (v *viper) SetEnvPrefix(in string) {
if in != "" {
v.envPrefix = in
}
}
func (v *viper) mergeWithEnvPrefix(in string) string {
if v.envPrefix != "" {
return v.envPrefix + "_" + in
}
return in
}
// Return the config file used // Return the config file used
func ConfigFileUsed() string { return v.ConfigFileUsed() } func ConfigFileUsed() string { return v.ConfigFileUsed() }
func (v *viper) ConfigFileUsed() string { return v.configFile } func (v *viper) ConfigFileUsed() string { return v.configFile }
@ -342,6 +359,7 @@ func (v *viper) BindPFlag(key string, flag *pflag.Flag) (err error) {
// Binds a viper key to a ENV variable // Binds a viper key to a ENV variable
// ENV variables are case sensitive // ENV variables are case sensitive
// If only a key is provided, it will use the env key matching the key, uppercased. // If only a key is provided, it will use the env key matching the key, uppercased.
// EnvPrefix will be used when set when env name is not provided.
func BindEnv(input ...string) (err error) { return v.BindEnv(input...) } func BindEnv(input ...string) (err error) { return v.BindEnv(input...) }
func (v *viper) BindEnv(input ...string) (err error) { func (v *viper) BindEnv(input ...string) (err error) {
var key, envkey string var key, envkey string
@ -352,7 +370,7 @@ func (v *viper) BindEnv(input ...string) (err error) {
key = strings.ToLower(input[0]) key = strings.ToLower(input[0])
if len(input) == 1 { if len(input) == 1 {
envkey = strings.ToUpper(key) envkey = strings.ToUpper(v.mergeWithEnvPrefix(key))
} else { } else {
envkey = input[1] envkey = input[1]
} }

View file

@ -67,6 +67,49 @@ var remoteExample = []byte(`{
"newkey":"remote" "newkey":"remote"
}`) }`)
func initConfigs() {
Reset()
SetConfigType("yaml")
r := bytes.NewReader(yamlExample)
marshalReader(r, v.config)
SetConfigType("json")
r = bytes.NewReader(jsonExample)
marshalReader(r, v.config)
SetConfigType("toml")
r = bytes.NewReader(tomlExample)
marshalReader(r, v.config)
SetConfigType("json")
remote := bytes.NewReader(remoteExample)
marshalReader(remote, v.kvstore)
}
func initYAML() {
Reset()
SetConfigType("yaml")
r := bytes.NewReader(yamlExample)
marshalReader(r, v.config)
}
func initJSON() {
Reset()
SetConfigType("json")
r := bytes.NewReader(jsonExample)
marshalReader(r, v.config)
}
func initTOML() {
Reset()
SetConfigType("toml")
r := bytes.NewReader(tomlExample)
marshalReader(r, v.config)
}
//stubs for PFlag Values //stubs for PFlag Values
type stringValue string type stringValue string
@ -139,34 +182,23 @@ func TestAliasInConfigFile(t *testing.T) {
} }
func TestYML(t *testing.T) { func TestYML(t *testing.T) {
Reset() initYAML()
SetConfigType("yml")
r := bytes.NewReader(yamlExample)
marshalReader(r, v.config)
assert.Equal(t, "steve", Get("name")) assert.Equal(t, "steve", Get("name"))
} }
func TestJSON(t *testing.T) { func TestJSON(t *testing.T) {
SetConfigType("json") initJSON()
r := bytes.NewReader(jsonExample)
marshalReader(r, v.config)
assert.Equal(t, "0001", Get("id")) assert.Equal(t, "0001", Get("id"))
} }
func TestTOML(t *testing.T) { func TestTOML(t *testing.T) {
SetConfigType("toml") initTOML()
r := bytes.NewReader(tomlExample)
marshalReader(r, v.config)
assert.Equal(t, "TOML Example", Get("title")) assert.Equal(t, "TOML Example", Get("title"))
} }
func TestRemotePrecedence(t *testing.T) { func TestRemotePrecedence(t *testing.T) {
SetConfigType("json") initJSON()
r := bytes.NewReader(jsonExample)
marshalReader(r, v.config)
remote := bytes.NewReader(remoteExample) remote := bytes.NewReader(remoteExample)
assert.Equal(t, "0001", Get("id")) assert.Equal(t, "0001", Get("id"))
marshalReader(remote, v.kvstore) marshalReader(remote, v.kvstore)
@ -180,9 +212,8 @@ func TestRemotePrecedence(t *testing.T) {
} }
func TestEnv(t *testing.T) { func TestEnv(t *testing.T) {
SetConfigType("json") initJSON()
r := bytes.NewReader(jsonExample)
marshalReader(r, v.config)
BindEnv("id") BindEnv("id")
BindEnv("f", "FOOD") BindEnv("f", "FOOD")
@ -199,10 +230,32 @@ func TestEnv(t *testing.T) {
assert.Equal(t, "crunk", Get("name")) assert.Equal(t, "crunk", Get("name"))
} }
func TestEnvPrefix(t *testing.T) {
initJSON()
SetEnvPrefix("foo") // will be uppercased automatically
BindEnv("id")
BindEnv("f", "FOOD") // not using prefix
os.Setenv("FOO_ID", "13")
os.Setenv("FOOD", "apple")
os.Setenv("FOO_NAME", "crunk")
assert.Equal(t, "13", Get("id"))
assert.Equal(t, "apple", Get("f"))
assert.Equal(t, "Cake", Get("name"))
AutomaticEnv()
assert.Equal(t, "crunk", Get("name"))
}
func TestAllKeys(t *testing.T) { func TestAllKeys(t *testing.T) {
initConfigs()
ks := sort.StringSlice{"title", "newkey", "owner", "name", "beard", "ppu", "batters", "hobbies", "clothing", "age", "hacker", "id", "type", "eyes"} ks := sort.StringSlice{"title", "newkey", "owner", "name", "beard", "ppu", "batters", "hobbies", "clothing", "age", "hacker", "id", "type", "eyes"}
dob, _ := time.Parse(time.RFC3339, "1979-05-27T07:32:00Z") dob, _ := time.Parse(time.RFC3339, "1979-05-27T07:32:00Z")
all := map[string]interface{}{"hacker": true, "beard": true, "newkey": "remote", "batters": map[string]interface{}{"batter": []interface{}{map[string]interface{}{"type": "Regular"}, map[string]interface{}{"type": "Chocolate"}, map[string]interface{}{"type": "Blueberry"}, map[string]interface{}{"type": "Devil's Food"}}}, "hobbies": []interface{}{"skateboarding", "snowboarding", "go"}, "ppu": 0.55, "clothing": map[interface{}]interface{}{"jacket": "leather", "trousers": "denim"}, "name": "crunk", "owner": map[string]interface{}{"organization": "MongoDB", "Bio": "MongoDB Chief Developer Advocate & Hacker at Large", "dob": dob}, "id": "13", "title": "TOML Example", "age": 35, "type": "donut", "eyes": "brown"} all := map[string]interface{}{"owner": map[string]interface{}{"organization": "MongoDB", "Bio": "MongoDB Chief Developer Advocate & Hacker at Large", "dob": dob}, "title": "TOML Example", "ppu": 0.55, "eyes": "brown", "clothing": map[interface{}]interface{}{"trousers": "denim", "jacket": "leather"}, "id": "0001", "batters": map[string]interface{}{"batter": []interface{}{map[string]interface{}{"type": "Regular"}, map[string]interface{}{"type": "Chocolate"}, map[string]interface{}{"type": "Blueberry"}, map[string]interface{}{"type": "Devil's Food"}}}, "hacker": true, "beard": true, "hobbies": []interface{}{"skateboarding", "snowboarding", "go"}, "age": 35, "type": "donut", "newkey": "remote", "name": "Cake"}
var allkeys sort.StringSlice var allkeys sort.StringSlice
allkeys = AllKeys() allkeys = AllKeys()