iamolegga/enviper

package enviper

import (
    "reflect"
    "strings"

    "github.com/mitchellh/mapstructure"

    "github.com/spf13/viper"
)

// Enviper is a wrapper struct for viper,
// that makes it possible to unmarshal config to struct
// considering environment variables
type Enviper struct {
    *viper.Viper
    tagName string
}

// New returns an initialized Enviper instance
func New(v *viper.Viper) *Enviper {
    return &Enviper{
        Viper: v,
    }
}

const defaultTagName = "mapstructure"

// WithTagName sets custom tag name to be used instead of default `mapstructure`
func (e *Enviper) WithTagName(customTagName string) *Enviper {
    e.tagName = customTagName
    return e
}

// TagName returns currently used tag name (`mapstructure` by default)
func (e *Enviper) TagName() string {
    if e.tagName == "" {
        return defaultTagName
    }
    return e.tagName
}

// Unmarshal unmarshals the config into a Struct just like viper does.
// The difference between enviper and viper is in automatic overriding data from file by data from env variables
func (e *Enviper) Unmarshal(rawVal interface{}, opts ...viper.DecoderConfigOption) error {
    if e.TagName() != defaultTagName {
        opts = append(opts, func(c *mapstructure.DecoderConfig) {
            c.TagName = e.TagName()
        })
    }

    if err := e.Viper.ReadInConfig(); err != nil {
        switch err.(type) {
        case viper.ConfigFileNotFoundError:
            //     do nothing
        default:
            return err
        }
    }
    // We need to unmarshal before the env binding to make sure that keys of maps are bound just like the struct fields
    // We silence errors here because we'll unmarshal a second time
    _ = e.Viper.Unmarshal(rawVal, opts...)
    e.readEnvs(rawVal)
    return e.Viper.Unmarshal(rawVal, opts...)
}

func (e *Enviper) readEnvs(rawVal interface{}) {
    e.Viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
    e.bindEnvs(rawVal)
}

func (e *Enviper) bindEnvs(in interface{}, prev ...string) {
    ifv := reflect.ValueOf(in)
    if ifv.Kind() == reflect.Ptr {
        ifv = ifv.Elem()
    }

    switch ifv.Kind() {
    case reflect.Struct:
        for i := 0; i < ifv.NumField(); i++ {
            fv := ifv.Field(i)
            if fv.Kind() == reflect.Ptr {
                if fv.IsZero() {
                    fv = reflect.New(fv.Type().Elem()).Elem()
                } else {
                    fv = fv.Elem()
                }
            }
            t := ifv.Type().Field(i)
            tv, ok := t.Tag.Lookup(e.TagName())
            if ok {
                if index := strings.Index(tv, ","); index != -1 {
                    if tv[:index] == "-" {
                        continue
                    }

                    // If "squash" is specified in the tag, we squash the field down.
                    if strings.Contains(tv[index+1:], "squash") {
                        e.bindEnvs(fv.Interface(), prev...)
                        continue
                    }

                    tv = tv[:index]
                }

                if tv == "" {
                    tv = t.Name
                }
            } else {
                tv = t.Name
            }

            e.bindEnvs(fv.Interface(), append(prev, tv)...)
        }
    case reflect.Map:
        iter := ifv.MapRange()
        for iter.Next() {
            if key, ok := iter.Key().Interface().(string); ok {
                e.bindEnvs(iter.Value().Interface(), append(prev, key)...)
            }
        }
    default:
        env := strings.Join(prev, ".")
        // Viper.BindEnv will never return error
        // because env is always non empty string
        _ = e.Viper.BindEnv(env)
    }
}