Viper and Cobra example

viper-and-cobra-example.go

package main

import (
	"fmt"
	"log"
	"net/http"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

type application struct {
	cmd *cobra.Command
	err error
}

func (a *application) AddCmdFlag(name string, value interface{}, usage string) {
	if a.err != nil {
		return
	}
	
	if value != nil {
		viper.SetDefault(name, value)
	}

	flags := a.cmd.Flags()

	switch v := value.(type) {
	case int:
		flags.Int(name, v, usage)
	case string:
		flags.String(name, v, usage)
	}

	a.err = viper.BindPFlag(name, flags.Lookup(name))
}

func (a *application) Run() error {
	if a.err != nil {
		return a.err
	}

	a.err = a.cmd.Execute()
	return a.err
}

func startServer(host string, port int) error {
	address := fmt.Sprintf("%s:%d", host, port)

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		_, err := w.Write([]byte("Hello!\n"))

		if err != nil {
			http.Error(w, "error sending response", 500)
		}
	})

	log.Printf("Listening on %s:%d\n", host, port)
	return http.ListenAndServe(address, nil)
}

func newServerCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "server",
		Short: "A simple HTTP server that reads configs from flags, env vars and a config file",
		RunE: func(cmd *cobra.Command, args []string) error {
			host := viper.GetString("host")
			port := viper.GetInt("port")
			return startServer(host, port)
		},
	}
}

func loadConfig() {
	if err := viper.ReadInConfig(); err != nil {
		log.Println("not using config file")
	}
}

func init() {
	viper.SetTypeByDefaultValue(true)

	viper.SetEnvPrefix("server")
	viper.AutomaticEnv()

	viper.SetConfigName("config")
	viper.SetConfigType("yaml")
	viper.AddConfigPath(".")
}

func main() {
	loadConfig()

	app := &application{cmd: newServerCommand()}
	app.AddCmdFlag("host", "localhost", "Set the application host")
	app.AddCmdFlag("port", 8000, "Set the application port")

	if err := app.Run(); err != nil {
		log.Fatalln(err)
	}
}

This example code shows how to use Viper and Cobra to create an application that gets its configurations from command line, env vars and/or a config file.

The app can be run as follows:

# read configs from command-line flags
$ go run . --host=0.0.0.0 --port=8000
2021/04/15 00:00:00 Listening on 0.0.0.0:8000

or

# read configs from env vars
$ SERVER_HOST=0.0.0.0 SERVER_PORT=8000 go run .
2021/04/15 00:00:00 Listening on 0.0.0.0:8000

or

# read configs from config.yaml
$ go run .
2021/04/15 00:00:00 Listening on 0.0.0.0:8000