sonet/pkg/config/loader.go

package config

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"github.com/spf13/viper"
	"os"
	"sonet/pkg/utils/logger"
	"strings"
)

func parseConfPathFlag(confPath string) (filePath, fileName, confName, confType string) {
	idx := strings.LastIndex(confPath, "/")
	filePath = confPath[:idx+1]
	fileName = confPath[idx+1:]
	idx2 := strings.LastIndex(fileName, ".")
	confName = fileName[:idx2]
	confType = fileName[idx2+1:]
	return
}

// LoadConfig load config.toml
// appConf service custom config
// return common service configuration
func LoadConfig(appConf any, confPathArg ...string) *Configuration {
	confPath := ""
	confName := "config"
	confType := "toml"
	envPrefix := "SO"

	if len(confPathArg) > 0 {
		confPath = confPathArg[0]
	}

	// go run xx -conf=config/xx.toml
	confPathFlag := flag.String("conf", "", "config file path.")
	envPrefixFlag := flag.String("envPrefix", "", "env config key prefix.")
	flag.Parse()
	if *confPathFlag != "" {
		confPath, _, confName, confType = parseConfPathFlag(*confPathFlag)
	}
	if *envPrefixFlag != "" {
		envPrefix = *envPrefixFlag
	}
	if confPath == "" {
		confPath = "./"
	}
	logger.Infof("use config file: %s/%s.%s, env prefix=%s\n", confPath, confName, confType, envPrefix)

	viper.AddConfigPath(confPath)
	viper.SetConfigName(confName)
	viper.SetConfigType(confType)
	viper.SetEnvPrefix(envPrefix)
	viper.AutomaticEnv()
	viper.AllowEmptyEnv(true)

	conf := &Configuration{}

	if err := viper.ReadInConfig(); err != nil {
		panic(errors.New("viper read config fail: " + err.Error()))
	}
	if err := viper.Unmarshal(conf); err != nil {
		panic(errors.New("viper unmarshal config failed: " + err.Error()))
	}

	// 读取appConf
	if appConf != nil && len(conf.App) > 0 {
		app, err := json.Marshal(conf.App)
		if err != nil {
			panic(err)
		}
		// 不能 json.Unmarshal 类型转换问题, viper.UnmarshalKey 无法读取环境变量
		viper.Reset()
		viper.SetConfigType("json")
		if err = viper.ReadConfig(bytes.NewBuffer(app)); err != nil {
			panic(err)
		}
		if err = viper.Unmarshal(appConf); err != nil {
			panic(err)
		}
	}

	InitLogger(conf.Grpc.Log)
	return conf
}

func LoadIdlPath(idlPath string) (string, error) {
	idl, err := os.ReadFile(idlPath)
	if err != nil {
		return "", err
	}
	return string(idl), nil
}