spider-scheduler/pkg/conf/conf.go

package conf

import (
	"encoding/json"
	"github.com/BurntSushi/toml"
	"github.com/eigeen/furryboard/spider-scheduler/pkg/log"
	"github.com/eigeen/furryboard/spider-scheduler/pkg/util"
	"github.com/mcuadros/go-defaults"
	"os"
)

type Config struct {
	Common     Common     `toml:"common"`
	Database   Database   `toml:"database"`
	SpiderCore SpiderCore `toml:"spider_core"`
}

type Common struct {
}

type Database struct {
	Host     string `toml:"host" default:"localhost"`
	Port     uint16 `toml:"port" default:"5432"`
	User     string `toml:"user" default:"user"`
	Password string `toml:"password" default:"password"`
	DB       string `toml:"db" default:"furryboard"`
}

type SpiderCore struct {
	Host       string `toml:"host" default:"localhost"`
	Port       uint16 `toml:"port" default:"9996"`
	BiliCookie string `toml:"bili_cookie"`
}

var (
	Conf *Config
)

// 释放默认配置文件
func releaseConfig(file string, config *Config) {
	f, err := util.CreateNestedFile(file)
	if err != nil {
		log.Logger().Fatalf("创建默认配置文件失败: %s", err.Error())
	}

	encoder := toml.NewEncoder(f)
	err = encoder.Encode(config)
	if err != nil {
		log.Logger().Fatalf("创建默认配置文件失败: %s", err.Error())
	}
}

func InitConfig(file string) {
	Conf = &Config{}
	defaults.SetDefaults(Conf)

	if util.NotExists(file) {
		// 默认配置文件不存在则创建
		if file == "config.toml" {
			releaseConfig(file, Conf)
			log.Logger().Infof("已创建默认配置文件: %s，请修改配置文件内容后重新启动", file)
			os.Exit(0)
		} else {
			log.Logger().Fatalf("配置文件不存在: %s", file)
		}
	}

	// 解析配置文件
	_, err := toml.DecodeFile(file, &Conf)
	if err != nil {
		log.Logger().Fatalf("配置文件解析出错: %s", err.Error())
	}

	// 配置预检查

	jsonConfig, err := json.Marshal(Conf)
	if err == nil {
		log.Logger().Debugf("配置文件已加载: %s", string(jsonConfig))
	}
}