github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/conf/env_var.go (about)

     1  package conf
     2  
     3  import (
     4  	"go/ast"
     5  	"reflect"
     6  	"sort"
     7  	"strings"
     8  
     9  	"github.com/sirupsen/logrus"
    10  
    11  	"github.com/artisanhe/tools/strutil"
    12  )
    13  
    14  var TagConf = "conf"
    15  
    16  type TagConfOption struct {
    17  	CanConfig     bool `opt:"env"`
    18  	IsUpstream    bool `opt:"upstream"`
    19  	FallbackValue interface{}
    20  }
    21  
    22  func GetTagConfOption(opt string) (tagConfOption TagConfOption) {
    23  	if opt == "" {
    24  		return
    25  	}
    26  
    27  	optionList := strings.Split(opt, ",")
    28  	options := make(map[string]bool)
    29  	for _, opt := range optionList {
    30  		options[opt] = true
    31  	}
    32  
    33  	rv := reflect.Indirect(reflect.ValueOf(&tagConfOption))
    34  	tpe := rv.Type()
    35  
    36  	for i := 0; i < tpe.NumField(); i++ {
    37  		field := tpe.Field(i)
    38  		if options[field.Tag.Get("opt")] {
    39  			rv.Field(i).SetBool(true)
    40  		}
    41  	}
    42  
    43  	return
    44  }
    45  
    46  type IHasDockerDefaults interface {
    47  	DockerDefaults() DockerDefaults
    48  }
    49  
    50  type DockerDefaults map[string]interface{}
    51  
    52  func (dockerDefaults DockerDefaults) Merge(nextDockerDefaults DockerDefaults) DockerDefaults {
    53  	finalDockerDefaults := DockerDefaults{}
    54  	for key, value := range dockerDefaults {
    55  		finalDockerDefaults[key] = value
    56  	}
    57  	for key, value := range nextDockerDefaults {
    58  		finalDockerDefaults[key] = value
    59  	}
    60  	return finalDockerDefaults
    61  }
    62  
    63  func collectEnvVars(rv reflect.Value, envVarKey string, tagConfOption TagConfOption, dockerDefaults DockerDefaults) (envVars EnvVars, err error) {
    64  	envVars = EnvVars{}
    65  	if rv.Kind() == reflect.Ptr && rv.IsNil() {
    66  		return
    67  	}
    68  	rv = reflect.Indirect(rv)
    69  	if hasDockerDefaults, ok := rv.Interface().(IHasDockerDefaults); ok {
    70  		// parent should overwrite child as anonymous
    71  		dockerDefaults = hasDockerDefaults.DockerDefaults().Merge(dockerDefaults)
    72  	}
    73  	rv = reflect.Indirect(rv)
    74  	switch rv.Kind() {
    75  	case reflect.Func:
    76  		// skip func
    77  	case reflect.Struct:
    78  		walkStructField(rv, false, func(fieldValue reflect.Value, field reflect.StructField) bool {
    79  			nextEnvKey := resolveEnvVarKeyByField(envVarKey, field)
    80  			tagConfOption = GetTagConfOption(field.Tag.Get(TagConf))
    81  			if v, exists := dockerDefaults[field.Name]; exists {
    82  				tagConfOption.FallbackValue = v
    83  			}
    84  			subEnvVars, errForCollect := collectEnvVars(fieldValue, nextEnvKey, tagConfOption, dockerDefaults)
    85  			envVars.Merge(subEnvVars)
    86  			return errForCollect == nil
    87  		})
    88  	default:
    89  		envVars.Set(envVarKey, EnvVar{
    90  			TagConfOption: tagConfOption,
    91  			Value:         rv,
    92  		})
    93  	}
    94  	return
    95  }
    96  
    97  func walkStructField(rv reflect.Value, forSet bool, fn func(fieldValue reflect.Value, field reflect.StructField) bool) {
    98  	reflectType := rv.Type()
    99  	for i := 0; i < reflectType.NumField(); i++ {
   100  		field := reflectType.Field(i)
   101  		fieldValue := rv.Field(i)
   102  		if !ast.IsExported(field.Name) {
   103  			continue
   104  		}
   105  
   106  		if forSet && (!fieldValue.IsValid() || !fieldValue.CanSet()) {
   107  			continue
   108  		}
   109  
   110  		next := fn(fieldValue, field)
   111  		if !next {
   112  			break
   113  		}
   114  	}
   115  }
   116  
   117  func resolveEnvVarKeyByField(pre string, field reflect.StructField) string {
   118  	if field.Anonymous {
   119  		return pre
   120  	}
   121  	if pre == "" {
   122  		return strings.ToUpper(field.Name)
   123  	}
   124  	return strings.ToUpper(pre + "_" + field.Name)
   125  }
   126  
   127  func CollectEnvVars(rv reflect.Value, envVarKey string) (envVars EnvVars, err error) {
   128  	return collectEnvVars(rv, envVarKey, TagConfOption{}, DockerDefaults{})
   129  }
   130  
   131  type ISecurityStringer interface {
   132  	SecurityString() string
   133  }
   134  
   135  type EnvVars map[string]EnvVar
   136  
   137  type EnvVar struct {
   138  	Value reflect.Value
   139  	TagConfOption
   140  }
   141  
   142  func stringValueOf(rv reflect.Value, security bool) string {
   143  	v := rv.Interface()
   144  	switch rv.Kind() {
   145  	case reflect.Array, reflect.Slice:
   146  		values := make([]string, 0)
   147  		for i := 0; i < rv.Len(); i++ {
   148  			values = append(values, stringValueOf(rv.Index(i), security))
   149  		}
   150  		return strings.Join(values, ",")
   151  	default:
   152  		if security {
   153  			if securityStringer, ok := v.(ISecurityStringer); ok {
   154  				return securityStringer.SecurityString()
   155  			}
   156  		}
   157  		s, err := strutil.ConvertToStr(v)
   158  		if err != nil {
   159  			panic(err)
   160  		}
   161  		return s
   162  	}
   163  }
   164  
   165  func (envVar EnvVar) GetFallbackValue(security bool) string {
   166  	// zero value will return if the fallbackv alue is nil
   167  	if envVar.FallbackValue == nil {
   168  		return ""
   169  	}
   170  	return stringValueOf(reflect.ValueOf(envVar.FallbackValue), security)
   171  }
   172  
   173  func (envVar EnvVar) GetValue(security bool) string {
   174  	rv := envVar.Value
   175  	switch rv.Kind() {
   176  	case reflect.Slice, reflect.Array:
   177  		values := make([]string, 0)
   178  		for i := 0; i < rv.Len(); i++ {
   179  			values = append(values, stringValueOf(rv.Index(i), security))
   180  		}
   181  		return strings.Join(values, ",")
   182  	default:
   183  		return stringValueOf(rv, security)
   184  	}
   185  }
   186  
   187  func (envVars EnvVars) Set(key string, envVar EnvVar) {
   188  	envVars[key] = envVar
   189  }
   190  
   191  func (envVars EnvVars) Merge(nextEnvVars EnvVars) {
   192  	for key, envVar := range nextEnvVars {
   193  		envVars.Set(key, envVar)
   194  	}
   195  }
   196  
   197  func (envVars EnvVars) Print() {
   198  	keysWarning := make([]string, 0)
   199  	keysNormal := make([]string, 0)
   200  
   201  	for key, envVar := range envVars {
   202  		if envVar.CanConfig {
   203  			keysWarning = append(keysWarning, key)
   204  		} else {
   205  			keysNormal = append(keysNormal, key)
   206  		}
   207  	}
   208  
   209  	sort.Strings(keysWarning)
   210  	sort.Strings(keysNormal)
   211  
   212  	fields := logrus.Fields{}
   213  
   214  	logger := logrus.New()
   215  	logger.Formatter = &logrus.TextFormatter{}
   216  
   217  	getLogger := func(key string) *logrus.Entry {
   218  		envVar := envVars[key]
   219  		fields[key] = envVar.GetValue(true)
   220  		l := logger.WithField(key, fields[key])
   221  		if envVar.FallbackValue != nil {
   222  			l = l.WithField("fallback", envVars[key].GetFallbackValue(true))
   223  		}
   224  		return l
   225  	}
   226  
   227  	for _, key := range keysWarning {
   228  		getLogger(key).Warning()
   229  	}
   230  	for _, key := range keysNormal {
   231  		getLogger(key).Info()
   232  	}
   233  
   234  	logger.Formatter = &logrus.JSONFormatter{}
   235  	logger.WithField("config", fields).Warning()
   236  }