flag.go

Documentation: go.etcd.io/etcd/pkg/v3/flags

     1  // Copyright 2015 The etcd Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package flags implements command-line flag parsing.
    16  package flags
    17  
    18  import (
    19  	"flag"
    20  	"fmt"
    21  	"os"
    22  	"strings"
    23  
    24  	"github.com/spf13/pflag"
    25  	"go.uber.org/zap"
    26  )
    27  
    28  // SetFlagsFromEnv parses all registered flags in the given flagset,
    29  // and if they are not already set it attempts to set their values from
    30  // environment variables. Environment variables take the name of the flag but
    31  // are UPPERCASE, have the given prefix  and any dashes are replaced by
    32  // underscores - for example: some-flag => ETCD_SOME_FLAG
    33  func SetFlagsFromEnv(lg *zap.Logger, prefix string, fs *flag.FlagSet) error {
    34  	var err error
    35  	alreadySet := make(map[string]bool)
    36  	fs.Visit(func(f *flag.Flag) {
    37  		alreadySet[FlagToEnv(prefix, f.Name)] = true
    38  	})
    39  	usedEnvKey := make(map[string]bool)
    40  	fs.VisitAll(func(f *flag.Flag) {
    41  		if serr := setFlagFromEnv(lg, fs, prefix, f.Name, usedEnvKey, alreadySet, true); serr != nil {
    42  			err = serr
    43  		}
    44  	})
    45  	verifyEnv(lg, prefix, usedEnvKey, alreadySet)
    46  	return err
    47  }
    48  
    49  // SetPflagsFromEnv is similar to SetFlagsFromEnv. However, the accepted flagset type is pflag.FlagSet
    50  // and it does not do any logging.
    51  func SetPflagsFromEnv(lg *zap.Logger, prefix string, fs *pflag.FlagSet) error {
    52  	var err error
    53  	alreadySet := make(map[string]bool)
    54  	usedEnvKey := make(map[string]bool)
    55  	fs.VisitAll(func(f *pflag.Flag) {
    56  		if f.Changed {
    57  			alreadySet[FlagToEnv(prefix, f.Name)] = true
    58  		}
    59  		if serr := setFlagFromEnv(lg, fs, prefix, f.Name, usedEnvKey, alreadySet, false); serr != nil {
    60  			err = serr
    61  		}
    62  	})
    63  	verifyEnv(lg, prefix, usedEnvKey, alreadySet)
    64  	return err
    65  }
    66  
    67  // FlagToEnv converts flag string to upper-case environment variable key string.
    68  func FlagToEnv(prefix, name string) string {
    69  	return prefix + "_" + strings.ToUpper(strings.Replace(name, "-", "_", -1))
    70  }
    71  
    72  func verifyEnv(lg *zap.Logger, prefix string, usedEnvKey, alreadySet map[string]bool) {
    73  	for _, env := range os.Environ() {
    74  		kv := strings.SplitN(env, "=", 2)
    75  		if len(kv) != 2 {
    76  			if lg != nil {
    77  				lg.Warn("found invalid environment variable", zap.String("environment-variable", env))
    78  			}
    79  		}
    80  		if usedEnvKey[kv[0]] {
    81  			continue
    82  		}
    83  		if alreadySet[kv[0]] {
    84  			if lg != nil {
    85  				lg.Fatal(
    86  					"conflicting environment variable is shadowed by corresponding command-line flag (either unset environment variable or disable flag))",
    87  					zap.String("environment-variable", kv[0]),
    88  				)
    89  			}
    90  		}
    91  		if strings.HasPrefix(env, prefix+"_") {
    92  			if lg != nil {
    93  				lg.Warn("unrecognized environment variable", zap.String("environment-variable", env))
    94  			}
    95  		}
    96  	}
    97  }
    98  
    99  type flagSetter interface {
   100  	Set(fk string, fv string) error
   101  }
   102  
   103  func setFlagFromEnv(lg *zap.Logger, fs flagSetter, prefix, fname string, usedEnvKey, alreadySet map[string]bool, log bool) error {
   104  	key := FlagToEnv(prefix, fname)
   105  	if !alreadySet[key] {
   106  		val := os.Getenv(key)
   107  		if val != "" {
   108  			usedEnvKey[key] = true
   109  			if serr := fs.Set(fname, val); serr != nil {
   110  				return fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
   111  			}
   112  			if log && lg != nil {
   113  				lg.Info(
   114  					"recognized and used environment variable",
   115  					zap.String("variable-name", key),
   116  					zap.String("variable-value", val),
   117  				)
   118  			}
   119  		}
   120  	}
   121  	return nil
   122  }
   123  
   124  func IsSet(fs *flag.FlagSet, name string) bool {
   125  	set := false
   126  	fs.Visit(func(f *flag.Flag) {
   127  		if f.Name == name {
   128  			set = true
   129  		}
   130  	})
   131  	return set
   132  }
   133
View as plain text