diff --git a/main.go b/main.go index 4f909d8b7..9ef258413 100644 --- a/main.go +++ b/main.go @@ -48,6 +48,8 @@ func init() { func main() { flag.Parse() + setFlagsFromEnv() + if *proxyMode { startProxy() } else { @@ -198,3 +200,25 @@ func (as *Addrs) Set(s string) error { func (as *Addrs) String() string { return strings.Join(*as, ",") } + +// setFlagsFromEnv parses all registered flags in the global flagset, +// and if they are not already set it attempts to set their values from +// environment variables. Environment variables take the name of the flag but +// are UPPERCASE, have the prefix "ETCD_", and any dashes are replaced by +// underscores - for example: some-flag => ETCD_SOME_FLAG +func setFlagsFromEnv() { + alreadySet := make(map[string]bool) + flag.Visit(func(f *flag.Flag) { + alreadySet[f.Name] = true + }) + flag.VisitAll(func(f *flag.Flag) { + if !alreadySet[f.Name] { + key := "ETCD_" + strings.ToUpper(strings.Replace(f.Name, "-", "_", -1)) + val := os.Getenv(key) + if val != "" { + flag.Set(f.Name, val) + } + } + + }) +} diff --git a/main_test.go b/main_test.go new file mode 100644 index 000000000..f3c976cd1 --- /dev/null +++ b/main_test.go @@ -0,0 +1,43 @@ +package main + +import "os" +import "flag" +import "testing" + +func TestSetFlagsFromEnv(t *testing.T) { + os.Clearenv() + // flags should be settable using env vars + os.Setenv("ETCD_DATA_DIR", "/foo/bar") + // and command-line flags + if err := flag.Set("peer-bind-addr", "1.2.3.4"); err != nil { + t.Fatal(err) + } + // command-line flags take precedence over env vars + os.Setenv("ETCD_ID", "woof") + if err := flag.Set("id", "quack"); err != nil { + t.Fatal(err) + } + + // first verify that flags are as expected before reading the env + for f, want := range map[string]string{ + "data-dir": "", + "peer-bind-addr": "1.2.3.4", + "id": "quack", + } { + if got := flag.Lookup(f).Value.String(); got != want { + t.Fatalf("flag %q=%q, want %q", f, got, want) + } + } + + // now read the env and verify flags were updated as expected + setFlagsFromEnv() + for f, want := range map[string]string{ + "data-dir": "/foo/bar", + "peer-bind-addr": "1.2.3.4", + "id": "quack", + } { + if got := flag.Lookup(f).Value.String(); got != want { + t.Errorf("flag %q=%q, want %q", f, got, want) + } + } +}