flag.go 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166
  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. // Package flags implements command-line flag parsing.
  15. package flags
  16. import (
  17. "flag"
  18. "fmt"
  19. "net/url"
  20. "os"
  21. "strings"
  22. "github.com/coreos/pkg/capnslog"
  23. "github.com/spf13/pflag"
  24. )
  25. var (
  26. plog = capnslog.NewPackageLogger("github.com/coreos/etcd", "pkg/flags")
  27. )
  28. // DeprecatedFlag encapsulates a flag that may have been previously valid but
  29. // is now deprecated. If a DeprecatedFlag is set, an error occurs.
  30. type DeprecatedFlag struct {
  31. Name string
  32. }
  33. func (f *DeprecatedFlag) Set(_ string) error {
  34. return fmt.Errorf(`flag "-%s" is no longer supported.`, f.Name)
  35. }
  36. func (f *DeprecatedFlag) String() string {
  37. return ""
  38. }
  39. // IgnoredFlag encapsulates a flag that may have been previously valid but is
  40. // now ignored. If an IgnoredFlag is set, a warning is printed and
  41. // operation continues.
  42. type IgnoredFlag struct {
  43. Name string
  44. }
  45. // IsBoolFlag is defined to allow the flag to be defined without an argument
  46. func (f *IgnoredFlag) IsBoolFlag() bool {
  47. return true
  48. }
  49. func (f *IgnoredFlag) Set(s string) error {
  50. plog.Warningf(`flag "-%s" is no longer supported - ignoring.`, f.Name)
  51. return nil
  52. }
  53. func (f *IgnoredFlag) String() string {
  54. return ""
  55. }
  56. // SetFlagsFromEnv parses all registered flags in the given flagset,
  57. // and if they are not already set it attempts to set their values from
  58. // environment variables. Environment variables take the name of the flag but
  59. // are UPPERCASE, have the given prefix and any dashes are replaced by
  60. // underscores - for example: some-flag => ETCD_SOME_FLAG
  61. func SetFlagsFromEnv(prefix string, fs *flag.FlagSet) error {
  62. var err error
  63. alreadySet := make(map[string]bool)
  64. fs.Visit(func(f *flag.Flag) {
  65. alreadySet[FlagToEnv(prefix, f.Name)] = true
  66. })
  67. usedEnvKey := make(map[string]bool)
  68. fs.VisitAll(func(f *flag.Flag) {
  69. if serr := setFlagFromEnv(fs, prefix, f.Name, usedEnvKey, alreadySet, true); serr != nil {
  70. err = serr
  71. }
  72. })
  73. verifyEnv(prefix, usedEnvKey, alreadySet)
  74. return err
  75. }
  76. // SetPflagsFromEnv is similar to SetFlagsFromEnv. However, the accepted flagset type is pflag.FlagSet
  77. // and it does not do any logging.
  78. func SetPflagsFromEnv(prefix string, fs *pflag.FlagSet) error {
  79. var err error
  80. alreadySet := make(map[string]bool)
  81. usedEnvKey := make(map[string]bool)
  82. fs.VisitAll(func(f *pflag.Flag) {
  83. if f.Changed {
  84. alreadySet[FlagToEnv(prefix, f.Name)] = true
  85. }
  86. if serr := setFlagFromEnv(fs, prefix, f.Name, usedEnvKey, alreadySet, false); serr != nil {
  87. err = serr
  88. }
  89. })
  90. verifyEnv(prefix, usedEnvKey, alreadySet)
  91. return err
  92. }
  93. // FlagToEnv converts flag string to upper-case environment variable key string.
  94. func FlagToEnv(prefix, name string) string {
  95. return prefix + "_" + strings.ToUpper(strings.Replace(name, "-", "_", -1))
  96. }
  97. func verifyEnv(prefix string, usedEnvKey, alreadySet map[string]bool) {
  98. for _, env := range os.Environ() {
  99. kv := strings.SplitN(env, "=", 2)
  100. if len(kv) != 2 {
  101. plog.Warningf("found invalid env %s", env)
  102. }
  103. if usedEnvKey[kv[0]] {
  104. continue
  105. }
  106. if alreadySet[kv[0]] {
  107. // TODO: exit with error in v3.4
  108. plog.Warningf("recognized environment variable %s, but unused: shadowed by corresponding flag", kv[0])
  109. continue
  110. }
  111. if strings.HasPrefix(env, prefix+"_") {
  112. plog.Warningf("unrecognized environment variable %s", env)
  113. }
  114. }
  115. }
  116. type flagSetter interface {
  117. Set(fk string, fv string) error
  118. }
  119. func setFlagFromEnv(fs flagSetter, prefix, fname string, usedEnvKey, alreadySet map[string]bool, log bool) error {
  120. key := FlagToEnv(prefix, fname)
  121. if !alreadySet[key] {
  122. val := os.Getenv(key)
  123. if val != "" {
  124. usedEnvKey[key] = true
  125. if serr := fs.Set(fname, val); serr != nil {
  126. return fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
  127. }
  128. if log {
  129. plog.Infof("recognized and used environment variable %s=%s", key, val)
  130. }
  131. }
  132. }
  133. return nil
  134. }
  135. // URLsFromFlag returns a slices from url got from the flag.
  136. func URLsFromFlag(fs *flag.FlagSet, urlsFlagName string) []url.URL {
  137. return []url.URL(*fs.Lookup(urlsFlagName).Value.(*URLsValue))
  138. }
  139. func IsSet(fs *flag.FlagSet, name string) bool {
  140. set := false
  141. fs.Visit(func(f *flag.Flag) {
  142. if f.Name == name {
  143. set = true
  144. }
  145. })
  146. return set
  147. }