od-database-crawler/config.go

package main

import (
	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"time"
)

var config struct {
	ServerUrl string
	Token     string
	Retries   int
	Workers   int
	Tasks     int32
	StatsInterval time.Duration
	Verbose   bool
}

const (
	ConfServerUrl = "server.url"
	ConfToken     = "server.token"
	ConfTasks     = "crawl.tasks"
	ConfRetries   = "crawl.retries"
	ConfWorkers   = "crawl.connections"
	ConfStatsInterval = "output.stats_interval"
	ConfVerbose   = "output.verbose"
)

func prepareConfig() {
	viper.SetDefault(ConfRetries, 5)
	viper.SetDefault(ConfWorkers, 2)
	viper.SetDefault(ConfTasks, 3)
	viper.SetDefault(ConfStatsInterval, 3 * time.Second)
	viper.SetDefault(ConfVerbose, false)
}

func readConfig() {
	viper.AddConfigPath(".")
	viper.SetConfigName("config")
	err := viper.ReadInConfig()
	if err != nil {
		logrus.Fatal(err)
	}

	config.ServerUrl = viper.GetString(ConfServerUrl)
	if config.ServerUrl == "" {
		configMissing(ConfServerUrl)
	}

	config.Token = viper.GetString(ConfToken)
	if config.Token == "" {
		configMissing(ConfToken)
	}

	config.Retries = viper.GetInt(ConfRetries)
	if config.Retries < 0 {
		config.Retries = 1 << 31
	}

	config.Workers = viper.GetInt(ConfWorkers)
	if config.Workers <= 0 {
		configOOB(ConfWorkers, config.Workers)
	}

	config.Tasks = viper.GetInt32(ConfTasks)
	if config.Tasks <= 0 {
		configOOB(ConfTasks, int(config.Tasks))
	}

	config.StatsInterval = viper.GetDuration(ConfStatsInterval)

	config.Verbose = viper.GetBool(ConfVerbose)
	if config.Verbose {
		logrus.SetLevel(logrus.DebugLevel)
	}
}

func configMissing(key string) {
	logrus.Fatalf("config: %s not set!", key)
}

func configOOB(key string, v int) {
	logrus.Fatal("config: illegal value %d for %key!", v, key)
}