Allow setting multiple environment variables for the same key. (#14)

* Allow setting more than one env. variable into BindEnv.
* Allow setting value transformers which permit the modification of
an environment variable before being assigned to a key.
This commit is contained in:
Gabriel Aszalos 2020-09-14 11:27:04 +03:00 committed by GitHub
parent fa4e03eaf5
commit 4ee0750c91
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 67 additions and 14 deletions

View file

@ -194,7 +194,8 @@ type Viper struct {
defaults map[string]interface{} defaults map[string]interface{}
kvstore map[string]interface{} kvstore map[string]interface{}
pflags map[string]FlagValue pflags map[string]FlagValue
env map[string]string env map[string][]string
envTransform map[string]func(string) interface{}
aliases map[string]string aliases map[string]string
knownKeys map[string]interface{} knownKeys map[string]interface{}
typeByDefValue bool typeByDefValue bool
@ -217,7 +218,8 @@ func New() *Viper {
v.defaults = make(map[string]interface{}) v.defaults = make(map[string]interface{})
v.kvstore = make(map[string]interface{}) v.kvstore = make(map[string]interface{})
v.pflags = make(map[string]FlagValue) v.pflags = make(map[string]FlagValue)
v.env = make(map[string]string) v.env = make(map[string][]string)
v.envTransform = make(map[string]func(string) interface{})
v.aliases = make(map[string]string) v.aliases = make(map[string]string)
v.knownKeys = make(map[string]interface{}) v.knownKeys = make(map[string]interface{})
v.typeByDefValue = false v.typeByDefValue = false
@ -368,6 +370,13 @@ func (v *Viper) SetEnvPrefix(in string) {
} }
} }
// SetEnvKeyTransformer allows defining a transformer function which decides
// how an environment variables value gets assigned to key.
func SetEnvKeyTransformer(key string, fn func(string) interface{}) { v.SetEnvKeyTransformer(key, fn) }
func (v *Viper) SetEnvKeyTransformer(key string, fn func(string) interface{}) {
v.envTransform[strings.ToLower(key)] = fn
}
func (v *Viper) mergeWithEnvPrefix(in string) string { func (v *Viper) mergeWithEnvPrefix(in string) string {
if v.envPrefix != "" { if v.envPrefix != "" {
return strings.ToUpper(v.envPrefix + "_" + in) return strings.ToUpper(v.envPrefix + "_" + in)
@ -1023,21 +1032,20 @@ func (v *Viper) BindFlagValue(key string, flag FlagValue) error {
// EnvPrefix will be used when set when env name is not provided. // EnvPrefix will be used when set when env name is not provided.
func BindEnv(input ...string) error { return v.BindEnv(input...) } func BindEnv(input ...string) error { return v.BindEnv(input...) }
func (v *Viper) BindEnv(input ...string) error { func (v *Viper) BindEnv(input ...string) error {
var key, envkey string
if len(input) == 0 { if len(input) == 0 {
return fmt.Errorf("BindEnv missing key to bind to") return fmt.Errorf("BindEnv missing key to bind to")
} }
key = strings.ToLower(input[0]) key := strings.ToLower(input[0])
var envkeys []string
if len(input) == 1 { if len(input) == 1 {
envkey = v.mergeWithEnvPrefix(key) envkeys = []string{v.mergeWithEnvPrefix(key)}
} else { } else {
envkey = input[1] envkeys = input[1:]
} }
v.env[key] = envkey v.env[key] = append(v.env[key], envkeys...)
v.SetKnown(key) v.SetKnown(key)
return nil return nil
@ -1108,10 +1116,15 @@ func (v *Viper) find(lcaseKey string) interface{} {
return nil return nil
} }
} }
envkey, exists := v.env[lcaseKey] envkeys, exists := v.env[lcaseKey]
if exists { if exists {
if val, ok := v.getEnv(envkey); ok { for _, key := range envkeys {
return val if val, ok := v.getEnv(key); ok {
if fn, ok := v.envTransform[lcaseKey]; ok {
return fn(val)
}
return val
}
} }
} }
if nested && v.isPathShadowedInFlatMap(path, v.env) != "" { if nested && v.isPathShadowedInFlatMap(path, v.env) != "" {
@ -1629,6 +1642,14 @@ func castToMapStringInterface(
return tgt return tgt
} }
func castMapStringSliceToMapInterface(src map[string][]string) map[string]interface{} {
tgt := map[string]interface{}{}
for k, v := range src {
tgt[k] = v
}
return tgt
}
func castMapStringToMapInterface(src map[string]string) map[string]interface{} { func castMapStringToMapInterface(src map[string]string) map[string]interface{} {
tgt := map[string]interface{}{} tgt := map[string]interface{}{}
for k, v := range src { for k, v := range src {
@ -1795,7 +1816,7 @@ func (v *Viper) AllKeys() []string {
m = v.flattenAndMergeMap(m, castMapStringToMapInterface(v.aliases), "") m = v.flattenAndMergeMap(m, castMapStringToMapInterface(v.aliases), "")
m = v.flattenAndMergeMap(m, v.override, "") m = v.flattenAndMergeMap(m, v.override, "")
m = v.mergeFlatMap(m, castMapFlagToMapInterface(v.pflags)) m = v.mergeFlatMap(m, castMapFlagToMapInterface(v.pflags))
m = v.mergeFlatMap(m, castMapStringToMapInterface(v.env)) m = v.mergeFlatMap(m, castMapStringSliceToMapInterface(v.env))
m = v.flattenAndMergeMap(m, v.config, "") m = v.flattenAndMergeMap(m, v.config, "")
m = v.flattenAndMergeMap(m, v.kvstore, "") m = v.flattenAndMergeMap(m, v.kvstore, "")
m = v.flattenAndMergeMap(m, v.defaults, "") m = v.flattenAndMergeMap(m, v.defaults, "")
@ -1850,7 +1871,7 @@ func (v *Viper) flattenAndMergeMap(shadow map[string]bool, m map[string]interfac
func (v *Viper) mergeFlatMap(shadow map[string]bool, m map[string]interface{}) map[string]bool { func (v *Viper) mergeFlatMap(shadow map[string]bool, m map[string]interface{}) map[string]bool {
// scan keys // scan keys
outer: outer:
for k, _ := range m { for k := range m {
path := strings.Split(k, v.keyDelim) path := strings.Split(k, v.keyDelim)
// scan intermediate paths // scan intermediate paths
var parentKey string var parentKey string

View file

@ -383,10 +383,11 @@ func TestEnv(t *testing.T) {
initJSON(v) initJSON(v)
v.BindEnv("id") v.BindEnv("id")
v.BindEnv("f", "FOOD") v.BindEnv("f", "FOOD", "DEPRECATED_FOOD")
os.Setenv("ID", "13") os.Setenv("ID", "13")
os.Setenv("FOOD", "apple") os.Setenv("FOOD", "apple")
os.Setenv("DEPRECATED_FOOD", "banana")
os.Setenv("NAME", "crunk") os.Setenv("NAME", "crunk")
assert.Equal(t, "13", v.Get("id")) assert.Equal(t, "13", v.Get("id"))
@ -396,7 +397,38 @@ func TestEnv(t *testing.T) {
v.AutomaticEnv() v.AutomaticEnv()
assert.Equal(t, "crunk", v.Get("name")) assert.Equal(t, "crunk", v.Get("name"))
}
func TestEnvTransformer(t *testing.T) {
v := New()
initJSON(v)
v.BindEnv("id", "MY_ID")
v.SetEnvKeyTransformer("id", func(in string) interface{} { return "transformed" })
os.Setenv("MY_ID", "14")
assert.Equal(t, "transformed", v.Get("id"))
}
func TestMultipleEnv(t *testing.T) {
v := New()
initJSON(v)
v.BindEnv("id")
v.BindEnv("f", "FOOD", "DEPRECATED_FOOD")
os.Setenv("ID", "13")
os.Unsetenv("FOOD")
os.Setenv("DEPRECATED_FOOD", "banana")
os.Setenv("NAME", "crunk")
assert.Equal(t, "13", v.Get("id"))
assert.Equal(t, "banana", v.Get("f"))
assert.Equal(t, "Cake", v.Get("name"))
v.AutomaticEnv()
assert.Equal(t, "crunk", v.Get("name"))
} }
func TestEmptyEnv(t *testing.T) { func TestEmptyEnv(t *testing.T) {