Let Viper handle the config file

This commit is contained in:
Maddie Zhan 2020-11-10 16:33:19 +08:00
parent c815103e20
commit fc6e99b45c
2 changed files with 8 additions and 36 deletions

View File

@ -1,8 +1,6 @@
package config package config
import ( import (
"os"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
"github.com/spf13/viper" "github.com/spf13/viper"
) )
@ -30,6 +28,7 @@ type Config struct {
} }
var ( var (
configFile string = ""
loadedConfig *Config = nil loadedConfig *Config = nil
) )
@ -51,8 +50,12 @@ func init() {
viper.AddConfigPath(".") viper.AddConfigPath(".")
} }
func Load() Config { func Load(configPath string) Config {
var conf Config var conf Config
configFile = configPath
viper.SetConfigFile(configPath)
if err := viper.ReadInConfig(); err != nil { if err := viper.ReadInConfig(); err != nil {
if _, ok := err.(viper.ConfigFileNotFoundError); ok { if _, ok := err.(viper.ConfigFileNotFoundError); ok {
log.Warnf("No config file found in search paths, using default values") log.Warnf("No config file found in search paths, using default values")
@ -70,33 +73,9 @@ func Load() Config {
return conf return conf
} }
func LoadFile(configFile string) Config {
var conf Config
f, err := os.OpenFile(configFile, os.O_RDONLY, 0444)
if err != nil {
log.Fatalf("Failed to open config file: %s", err)
}
defer f.Close()
if err := viper.ReadConfig(f); err != nil {
log.Fatalf("Error reading config: %s", err)
}
if err := viper.Unmarshal(&conf); err != nil {
log.Fatalf("Error parsing config: %s", err)
}
loadedConfig = &conf
return conf
}
func LoadedConfig() *Config { func LoadedConfig() *Config {
if loadedConfig == nil { if loadedConfig == nil {
Load() Load(configFile)
} }
return loadedConfig return loadedConfig
} }

View File

@ -17,14 +17,7 @@ var (
func main() { func main() {
flag.Parse() flag.Parse()
conf := config.Load(*optConfig)
var conf config.Config
if *optConfig != "" {
conf = config.LoadFile(*optConfig)
} else {
conf = config.Load()
}
web.SetServerLocation(&conf) web.SetServerLocation(&conf)
results.Initialize(&conf) results.Initialize(&conf)
database.SetDBInfo(&conf) database.SetDBInfo(&conf)