From cd9ef38d1377155b4147a91cba1b61f1097e9c4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8C=85=E5=AD=90?= <baozhecheng@foxmail.com>
Date: Mon, 4 Dec 2023 11:31:35 +0800
Subject: [PATCH] feat(config): config support custom merge config (#3106)

---
 config/config.go       | 12 ++++++------
 config/config_test.go  |  5 +++++
 config/options.go      | 14 +++++++++-----
 config/options_test.go | 11 +++++++++++
 config/reader.go       |  3 +--
 config/reader_test.go  | 11 +++++++++++
 6 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/config/config.go b/config/config.go
index 9efc90240..a7957d8d9 100644
--- a/config/config.go
+++ b/config/config.go
@@ -7,6 +7,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/imdario/mergo"
+
 	// init encoding
 	_ "github.com/go-kratos/kratos/v2/encoding/json"
 	_ "github.com/go-kratos/kratos/v2/encoding/proto"
@@ -17,12 +19,7 @@ import (
 
 var _ Config = (*config)(nil)
 
-var (
-	// ErrNotFound is key not found.
-	ErrNotFound = errors.New("key not found")
-	// ErrTypeAssert is type assert error.
-	ErrTypeAssert = errors.New("type assert error")
-)
+var ErrNotFound = errors.New("key not found") // ErrNotFound is key not found.
 
 // Observer is config observer.
 type Observer func(string, Value)
@@ -49,6 +46,9 @@ func New(opts ...Option) Config {
 	o := options{
 		decoder:  defaultDecoder,
 		resolver: defaultResolver,
+		merge: func(dst, src interface{}) error {
+			return mergo.Map(dst, src, mergo.WithOverride)
+		},
 	}
 	for _, opt := range opts {
 		opt(&o)
diff --git a/config/config_test.go b/config/config_test.go
index 9a166b00e..9af90a0e9 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -3,6 +3,8 @@ package config
 import (
 	"errors"
 	"testing"
+
+	"github.com/imdario/mergo"
 )
 
 const (
@@ -131,6 +133,9 @@ func TestConfig(t *testing.T) {
 		sources:  []Source{jSource},
 		decoder:  defaultDecoder,
 		resolver: defaultResolver,
+		merge: func(dst, src interface{}) error {
+			return mergo.Map(dst, src, mergo.WithOverride)
+		},
 	}
 	cf := &config{}
 	cf.opts = opts
diff --git a/config/options.go b/config/options.go
index 443ddaba1..fa833356b 100644
--- a/config/options.go
+++ b/config/options.go
@@ -6,7 +6,6 @@ import (
 	"strings"
 
 	"github.com/go-kratos/kratos/v2/encoding"
-	"github.com/go-kratos/kratos/v2/log"
 )
 
 // Decoder is config decoder.
@@ -15,6 +14,9 @@ type Decoder func(*KeyValue, map[string]interface{}) error
 // Resolver resolve placeholder in config.
 type Resolver func(map[string]interface{}) error
 
+// Merge is config merge func.
+type Merge func(dst, src interface{}) error
+
 // Option is config option.
 type Option func(*options)
 
@@ -22,6 +24,7 @@ type options struct {
 	sources  []Source
 	decoder  Decoder
 	resolver Resolver
+	merge    Merge
 }
 
 // WithSource with config source.
@@ -49,10 +52,11 @@ func WithResolver(r Resolver) Option {
 	}
 }
 
-// WithLogger with config logger.
-// Deprecated: use global logger instead.
-func WithLogger(_ log.Logger) Option {
-	return func(o *options) {}
+// WithMergeFunc with config merge func.
+func WithMergeFunc(m Merge) Option {
+	return func(o *options) {
+		o.merge = m
+	}
 }
 
 // defaultDecoder decode config from source KeyValue
diff --git a/config/options_test.go b/config/options_test.go
index c8cd2bf89..5b2c5acad 100644
--- a/config/options_test.go
+++ b/config/options_test.go
@@ -226,3 +226,14 @@ func TestExpand(t *testing.T) {
 		}
 	}
 }
+
+func TestWithMergeFunc(t *testing.T) {
+	c := &options{}
+	a := func(dst, src interface{}) error {
+		return nil
+	}
+	WithMergeFunc(a)(c)
+	if c.merge == nil {
+		t.Fatal("c.merge is nil")
+	}
+}
diff --git a/config/reader.go b/config/reader.go
index 2202de6b8..129705052 100644
--- a/config/reader.go
+++ b/config/reader.go
@@ -8,7 +8,6 @@ import (
 	"strings"
 	"sync"
 
-	"github.com/imdario/mergo"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
 
@@ -48,7 +47,7 @@ func (r *reader) Merge(kvs ...*KeyValue) error {
 			log.Errorf("Failed to config decode error: %v key: %s value: %s", err, kv.Key, string(kv.Value))
 			return err
 		}
-		if err := mergo.Map(&merged, convertMap(next), mergo.WithOverride); err != nil {
+		if err := r.opts.merge(&merged, convertMap(next)); err != nil {
 			log.Errorf("Failed to config merge error: %v key: %s value: %s", err, kv.Key, string(kv.Value))
 			return err
 		}
diff --git a/config/reader_test.go b/config/reader_test.go
index 11b2d990e..cfd49f778 100644
--- a/config/reader_test.go
+++ b/config/reader_test.go
@@ -6,6 +6,8 @@ import (
 	"testing"
 
 	"github.com/go-kratos/kratos/v2/encoding"
+
+	"github.com/imdario/mergo"
 )
 
 func TestReader_Merge(t *testing.T) {
@@ -21,6 +23,9 @@ func TestReader_Merge(t *testing.T) {
 			return fmt.Errorf("unsupported key: %s format: %s", kv.Key, kv.Format)
 		},
 		resolver: defaultResolver,
+		merge: func(dst, src interface{}) error {
+			return mergo.Map(dst, src, mergo.WithOverride)
+		},
 	}
 	r := newReader(opts)
 	err = r.Merge(&KeyValue{
@@ -82,6 +87,9 @@ func TestReader_Value(t *testing.T) {
 			return fmt.Errorf("unsupported key: %s format: %s", kv.Key, kv.Format)
 		},
 		resolver: defaultResolver,
+		merge: func(dst, src interface{}) error {
+			return mergo.Map(dst, src, mergo.WithOverride)
+		},
 	}
 
 	ymlval := `
@@ -184,6 +192,9 @@ func TestReader_Source(t *testing.T) {
 			return fmt.Errorf("unsupported key: %s format: %s", kv.Key, kv.Format)
 		},
 		resolver: defaultResolver,
+		merge: func(dst, src interface{}) error {
+			return mergo.Map(dst, src, mergo.WithOverride)
+		},
 	}
 	r := newReader(opts)
 	err = r.Merge(&KeyValue{