From 1022d75c736dd00cb382b1c93ef7a04d13959f05 Mon Sep 17 00:00:00 2001 From: spf13 Date: Mon, 22 Dec 2014 18:31:11 -0500 Subject: [PATCH] Adding Support for Environment variable prefixes --- README.md | 2 +- viper.go | 20 ++++++++++- viper_test.go | 93 ++++++++++++++++++++++++++++++++++++++++----------- 3 files changed, 93 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index e9b489c..99c60c6 100644 --- a/README.md +++ b/README.md @@ -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. -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 value) and call BindEnv on that key. It does not simply import all ENV variables. Because of this behavior it’s diff --git a/viper.go b/viper.go index f942d92..d075a95 100644 --- a/viper.go +++ b/viper.go @@ -74,6 +74,7 @@ type viper struct { configName string configFile string configType string + envPrefix string config 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 func ConfigFileUsed() string { return v.ConfigFileUsed() } 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 // ENV variables are case sensitive // 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 (v *viper) BindEnv(input ...string) (err error) { var key, envkey string @@ -352,7 +370,7 @@ func (v *viper) BindEnv(input ...string) (err error) { key = strings.ToLower(input[0]) if len(input) == 1 { - envkey = strings.ToUpper(key) + envkey = strings.ToUpper(v.mergeWithEnvPrefix(key)) } else { envkey = input[1] } diff --git a/viper_test.go b/viper_test.go index ea45a31..ce1aa3c 100644 --- a/viper_test.go +++ b/viper_test.go @@ -67,6 +67,49 @@ var remoteExample = []byte(`{ "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 type stringValue string @@ -139,34 +182,23 @@ func TestAliasInConfigFile(t *testing.T) { } func TestYML(t *testing.T) { - Reset() - SetConfigType("yml") - r := bytes.NewReader(yamlExample) - - marshalReader(r, v.config) + initYAML() assert.Equal(t, "steve", Get("name")) } func TestJSON(t *testing.T) { - SetConfigType("json") - r := bytes.NewReader(jsonExample) - - marshalReader(r, v.config) + initJSON() assert.Equal(t, "0001", Get("id")) } func TestTOML(t *testing.T) { - SetConfigType("toml") - r := bytes.NewReader(tomlExample) - - marshalReader(r, v.config) + initTOML() assert.Equal(t, "TOML Example", Get("title")) } func TestRemotePrecedence(t *testing.T) { - SetConfigType("json") - r := bytes.NewReader(jsonExample) - marshalReader(r, v.config) + initJSON() + remote := bytes.NewReader(remoteExample) assert.Equal(t, "0001", Get("id")) marshalReader(remote, v.kvstore) @@ -180,9 +212,8 @@ func TestRemotePrecedence(t *testing.T) { } func TestEnv(t *testing.T) { - SetConfigType("json") - r := bytes.NewReader(jsonExample) - marshalReader(r, v.config) + initJSON() + BindEnv("id") BindEnv("f", "FOOD") @@ -199,10 +230,32 @@ func TestEnv(t *testing.T) { 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) { + initConfigs() + 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") - 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 allkeys = AllKeys()