diff --git a/cmd/server.go b/cmd/server.go index 9b5821aa2..bdb27ab81 100644 --- a/cmd/server.go +++ b/cmd/server.go @@ -6,9 +6,11 @@ import ( "os" "os/signal" "runtime/pprof" + "strings" "time" "github.com/spf13/cobra" + flag "github.com/spf13/pflag" "github.com/spf13/viper" "github.com/pilosa/pilosa/server" @@ -24,22 +26,22 @@ var serveCmd = &cobra.Command{ It will load existing data from the configured directory, and start listening client connections on the configured port.`, - Run: func(cmd *cobra.Command, args []string) { + RunE: func(cmd *cobra.Command, args []string) error { + setupViper(cmd.Flags()) + fmt.Println(viper.AllSettings()) serve.Server.Handler.Version = Version fmt.Fprintf(serve.Stderr, "Pilosa %s, build time %s\n", Version, BuildTime) // Parse command line arguments. if err := serve.SetupConfig(args); err != nil { - fmt.Fprintln(serve.Stderr, err) - os.Exit(2) + return err } // Start CPU profiling. if serve.CPUProfile != "" { f, err := os.Create(serve.CPUProfile) if err != nil { - fmt.Fprintf(serve.Stderr, "create cpu profile: %v", err) - os.Exit(1) + return fmt.Errorf("create cpu profile: %v", err) } defer f.Close() @@ -54,8 +56,7 @@ on the configured port.`, // Execute the program. if err := serve.Run(); err != nil { - fmt.Fprintln(serve.Stderr, err) - os.Exit(1) + return err } // First SIGKILL causes server to shut down gracefully. @@ -68,22 +69,50 @@ on the configured port.`, go func() { <-c; os.Exit(1) }() if err := serve.Close(); err != nil { - fmt.Fprintln(serve.Stderr, err) - os.Exit(1) + return err } - + return nil }, } func init() { - serveCmd.Flags().StringVarP(&serve.ConfigPath, "config", "c", "", "Configuration file to read from") - serveCmd.Flags().StringVarP(&serve.CPUProfile, "cpuprofile", "", "", "Where to store CPU profile") - serveCmd.Flags().DurationVarP(&serve.CPUTime, "cputime", "", 30*time.Second, "CPU profile duration") + flags := serveCmd.Flags() - err := viper.BindPFlags(serveCmd.Flags()) - if err != nil { - log.Fatalf("Error binding server flags: %v", err) - } + flags.StringVarP(&serve.ConfigPath, "config", "c", "", "Configuration file to read from") + flags.StringVarP(&serve.CPUProfile, "cpu-profile", "", "", "Where to store CPU profile") + flags.DurationVarP(&serve.CPUTime, "cpu-time", "", 30*time.Second, "CPU profile duration") RootCmd.AddCommand(serveCmd) } + +func setupViper(flags *flag.FlagSet) { + // add cmd line flag def to viper + err := viper.BindPFlags(flags) + if err != nil { + log.Fatalf("Error binding server flags: %v", err) + } + // add env to viper + viper.SetEnvPrefix("PILOSA") + viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_")) + viper.AutomaticEnv() + + c := viper.GetString("config") + + // add config file to viper + if c != "" { + viper.AddConfigPath(c) + err := viper.ReadInConfig() + if err != nil { + log.Printf("Couldn't read config from '%s'", c) + } + } + flags.VisitAll(func(f *flag.Flag) { + log.Printf("Now visiting: %v with value '%s'", f.Name, f.Value) + value := viper.GetString(f.Name) + log.Printf("Setting to value: '%v'", value) + err := f.Value.Set(value) + if err != nil { + log.Printf("Error setting %s to '%s': %v", f.Name, value, err) + } + }) +}