kratos-utils/config/load.go

package config

import (
	"errors"
	"fmt"

	"dario.cat/mergo"
	"git.ikuban.com/server/yaml"
	"github.com/bytedance/sonic"
	"github.com/elliotchance/orderedmap/v3"
	"github.com/go-kratos/kratos/v2/config"
)

func Load(sources []Source, bc any) error {
	configSources := make([]config.Source, 0)
	sourceMap := orderedmap.NewOrderedMap[string, any]() // 有序map控制优先级

	for _, source := range sources {
		if !source.Validate() {
			continue
		}
		s, err := source.NewSource()
		if err != nil {
			return fmt.Errorf("%v, new source error: %v", source.String(), err)
		}

		kvs, err := s.Load()
		if err != nil {
			return fmt.Errorf("%v, source load error: %v", source.String(), err)
		}

		for _, v := range kvs {
			sourceMap.Set(v.Key, map[string]any{})
		}

		configSources = append(configSources, s)
	}

	if sourceMap.Len() == 0 {
		return errors.New("config source is empty")
	}

	c := config.New(
		config.WithSource(configSources...),
		config.WithDecoder(func(kv *config.KeyValue, v map[string]interface{}) error {
			ok := sourceMap.Has(kv.Key)
			if !ok {
				return fmt.Errorf("unknown config key,key:%s", kv.Key)
			}
			if kv.Format == "" {
				// 处理环境变量配置（Format 为空）
				// 环境变量的值已经是简单的字符串，不需要额外解析
				// 将环境变量值直接设置到 sourceMap 中对应键的位置
				// 这里需要将简单的键值对转换为嵌套结构
				envMap := buildNestedMap(kv.Key, string(kv.Value))
				sourceMap.Set(kv.Key, envMap)
			} else if kv.Format == Yaml.String() {
				next := map[string]any{}
				err := yaml.Unmarshal(kv.Value, next)
				if err != nil {
					return err
				}
				sourceMap.Set(kv.Key, convertMap(next))
			} else if kv.Format == Json.String() {
				next := map[string]any{}
				err := sonic.Unmarshal(kv.Value, &next)
				if err != nil {
					return err
				}
				sourceMap.Set(kv.Key, convertMap(next))
			}

			// 在每个配置源处理完成后，立即合并并解析到目标结构体
			// 这样确保配置热更新也能正常工作
			merged := make(map[string]any)
			for c := range sourceMap.Values() {
				err := mergo.Merge(&merged, c, mergo.WithOverride)
				if err != nil {
					return err
				}
			}

			// 将合并后的配置序列化为 YAML 再解析到目标结构体
			b, err := yaml.Marshal(merged)
			if err != nil {
				return err
			}

			err = yaml.Unmarshal(b, bc)
			if err != nil {
				return err
			}

			// 执行默认配置函数
			if defaultCfg != nil {
				defaultCfg(bc)
			}

			return nil
		}),
	)
	if err := c.Load(); err != nil {
		return err
	}

	return nil
}

type DefaultCfg func(bc any)

var defaultCfg DefaultCfg

func SetDefaultCfg(f DefaultCfg) {
	defaultCfg = f
}

func convertMap(src interface{}) interface{} {
	switch m := src.(type) {
	case map[string]interface{}:
		dst := make(map[string]interface{}, len(m))
		for k, v := range m {
			dst[k] = convertMap(v)
		}
		return dst
	case map[interface{}]interface{}:
		dst := make(map[interface{}]interface{}, len(m))
		for k, v := range m {
			dst[k] = convertMap(v)
		}
		return dst
	case []interface{}:
		dst := make([]interface{}, len(m))
		for k, v := range m {
			dst[k] = convertMap(v)
		}
		return dst
	case []byte:
		// there will be no binary data in the config data
		return string(m)
	default:
		return src
	}
}