From f4f200cf0407bb2ce4ec742263450c8f4849fe24 Mon Sep 17 00:00:00 2001
From: Ming Deng <mingflycash@gmail.com>
Date: Mon, 31 Aug 2020 13:02:22 +0000
Subject: [PATCH] enhance yaml

---
 pkg/infrastructure/config/ini.go            |  4 ++
 pkg/infrastructure/config/yaml/yaml.go      | 60 ++++++++++++++++++++-
 pkg/infrastructure/config/yaml/yaml_test.go | 35 ++++++++++++
 3 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/pkg/infrastructure/config/ini.go b/pkg/infrastructure/config/ini.go
index 92ed8df8..cc67e4cd 100644
--- a/pkg/infrastructure/config/ini.go
+++ b/pkg/infrastructure/config/ini.go
@@ -66,6 +66,10 @@ func (ini *IniConfig) parseData(dir string, data []byte) (*IniConfigContainer, e
 		keyComment:     make(map[string]string),
 		RWMutex:        sync.RWMutex{},
 	}
+
+	cfg.BaseConfiger = NewBaseConfiger(func(ctx context.Context, key string) (string, error) {
+		return cfg.getdata(key), nil
+	})
 	cfg.Lock()
 	defer cfg.Unlock()
 
diff --git a/pkg/infrastructure/config/yaml/yaml.go b/pkg/infrastructure/config/yaml/yaml.go
index 1f4f1d23..61ea45b9 100644
--- a/pkg/infrastructure/config/yaml/yaml.go
+++ b/pkg/infrastructure/config/yaml/yaml.go
@@ -42,8 +42,10 @@ import (
 	"sync"
 
 	"github.com/beego/goyaml2"
+	"gopkg.in/yaml.v2"
 
 	"github.com/astaxie/beego/pkg/infrastructure/config"
+	"github.com/astaxie/beego/pkg/infrastructure/logs"
 )
 
 // Config is a yaml config parser and implements Config interface.
@@ -120,11 +122,61 @@ func parseYML(buf []byte) (cnf map[string]interface{}, err error) {
 
 // ConfigContainer is a config which represents the yaml configuration.
 type ConfigContainer struct {
-	config.BaseConfiger
 	data map[string]interface{}
 	sync.RWMutex
 }
 
+// Unmarshaler is similar to Sub
+func (c *ConfigContainer) Unmarshaler(ctx context.Context, prefix string, obj interface{}, opt ...config.DecodeOption) error {
+	sub, err := c.sub(ctx, prefix)
+	if err != nil {
+		return err
+	}
+
+	bytes, err := yaml.Marshal(sub)
+	if err != nil {
+		return err
+	}
+	return yaml.Unmarshal(bytes, obj)
+}
+
+func (c *ConfigContainer) Sub(ctx context.Context, key string) (config.Configer, error) {
+	sub, err := c.sub(ctx, key)
+	if err != nil {
+		return nil, err
+	}
+	return &ConfigContainer{
+		data: sub,
+	}, nil
+}
+
+func (c *ConfigContainer) sub(ctx context.Context, key string) (map[string]interface{}, error) {
+	tmpData := c.data
+	keys := strings.Split(key, ".")
+	for idx, k := range keys {
+		if v, ok := tmpData[k]; ok {
+			switch v.(type) {
+			case map[string]interface{}:
+				{
+					tmpData = v.(map[string]interface{})
+					if idx == len(keys)-1 {
+						return tmpData, nil
+					}
+				}
+			default:
+				return nil, errors.New(fmt.Sprintf("the key is invalid: %s", key))
+			}
+		}
+	}
+
+	return tmpData, nil
+}
+
+func (c *ConfigContainer) OnChange(ctx context.Context, key string, fn func(value string)) {
+	// do nothing
+	logs.Warn("Unsupported operation: OnChange")
+}
+
 // Bool returns the boolean value for a given key.
 func (c *ConfigContainer) Bool(ctx context.Context, key string) (bool, error) {
 	v, err := c.getData(key)
@@ -291,7 +343,7 @@ func (c *ConfigContainer) getData(key string) (interface{}, error) {
 	c.RLock()
 	defer c.RUnlock()
 
-	keys := strings.Split(key, ".")
+	keys := strings.Split(c.key(key), ".")
 	tmpData := c.data
 	for idx, k := range keys {
 		if v, ok := tmpData[k]; ok {
@@ -314,6 +366,10 @@ func (c *ConfigContainer) getData(key string) (interface{}, error) {
 	return nil, fmt.Errorf("not exist key %q", key)
 }
 
+func (c *ConfigContainer) key(key string) string {
+	return key
+}
+
 func init() {
 	config.Register("yaml", &Config{})
 }
diff --git a/pkg/infrastructure/config/yaml/yaml_test.go b/pkg/infrastructure/config/yaml/yaml_test.go
index 197b68e4..1fd4e894 100644
--- a/pkg/infrastructure/config/yaml/yaml_test.go
+++ b/pkg/infrastructure/config/yaml/yaml_test.go
@@ -15,10 +15,13 @@
 package yaml
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/astaxie/beego/pkg/infrastructure/config"
 )
 
@@ -37,6 +40,9 @@ func TestYaml(t *testing.T) {
 "path1": ${GOPATH}
 "path2": ${GOPATH||/home/go}
 "empty": "" 
+"user":
+  "name": "tom"
+  "age": 13
 `
 
 		keyValue = map[string]interface{}{
@@ -114,4 +120,33 @@ func TestYaml(t *testing.T) {
 		t.Fatal("get name error")
 	}
 
+	sub, err := yamlconf.Sub(context.Background(), "user")
+	assert.Nil(t, err)
+	assert.NotNil(t, sub)
+	name, err := sub.String(context.Background(), "name")
+	assert.Nil(t, err)
+	assert.Equal(t, "tom", name)
+
+	age, err := sub.Int(context.Background(), "age")
+	assert.Nil(t, err)
+	assert.Equal(t, 13, age)
+
+	user := &User{}
+
+	err = sub.Unmarshaler(context.Background(), "", user)
+	assert.Nil(t, err)
+	assert.Equal(t, "tom", user.Name)
+	assert.Equal(t, 13, user.Age)
+
+	user = &User{}
+
+	err = yamlconf.Unmarshaler(context.Background(), "user", user)
+	assert.Nil(t, err)
+	assert.Equal(t, "tom", user.Name)
+	assert.Equal(t, 13, user.Age)
+}
+
+type User struct {
+	Name string `yaml:"name"`
+	Age  int    `yaml:"age"`
 }