From 84fbb775a7ca2a4eec5aa50d38f36850cb19e727 Mon Sep 17 00:00:00 2001
From: "xianliang.li" <xianliang.li@zilliz.com>
Date: Mon, 4 Nov 2024 20:35:36 +0800
Subject: [PATCH] update yaml v2 as yaml parser

Signed-off-by: xianliang.li <xianliang.li@zilliz.com>
---
 pkg/config/config.go      | 72 ++++++++++++++++++---------------------
 pkg/config/file_source.go | 29 ++++------------
 2 files changed, 40 insertions(+), 61 deletions(-)

diff --git a/pkg/config/config.go b/pkg/config/config.go
index 2f2a5ef6f6b69..387692378483f 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -17,10 +17,11 @@
 package config
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/cockroachdb/errors"
-	"gopkg.in/yaml.v3"
+	"github.com/spf13/cast"
 
 	"github.com/milvus-io/milvus/pkg/util/typeutil"
 )
@@ -80,48 +81,41 @@ func formatKey(key string) string {
 	return result
 }
 
-func flattenNode(node *yaml.Node, parentKey string, result map[string]string) {
-	// The content of the node should contain key-value pairs in a MappingNode
-	if node.Kind == yaml.MappingNode {
-		for i := 0; i < len(node.Content); i += 2 {
-			keyNode := node.Content[i]
-			valueNode := node.Content[i+1]
-
-			key := keyNode.Value
-			// Construct the full key with parent hierarchy
-			fullKey := key
-			if parentKey != "" {
-				fullKey = parentKey + "." + key
-			}
+func flattenAndMergeMap(prefix string, m map[string]interface{}, result map[string]string) {
+	for k, v := range m {
+		fullKey := k
+		if prefix != "" {
+			fullKey = prefix + "." + k
+		}
 
-			switch valueNode.Kind {
-			case yaml.ScalarNode:
-				// handle null value
-				if valueNode.Tag == "!!null" {
-					result[lowerKey(fullKey)] = ""
-					result[formatKey(fullKey)] = ""
-				} else {
-					// Scalar value, store it as a string
-					result[lowerKey(fullKey)] = valueNode.Value
-					result[formatKey(fullKey)] = valueNode.Value
+		switch val := v.(type) {
+		case map[string]interface{}:
+			flattenAndMergeMap(fullKey, val, result)
+		case map[interface{}]interface{}:
+			flattenAndMergeMap(fullKey, cast.ToStringMap(val), result)
+		case []interface{}:
+			str := ""
+			for i, item := range val {
+				itemStr, err := cast.ToStringE(item)
+				if err != nil {
+					continue
 				}
-			case yaml.MappingNode:
-				// Nested map, process recursively
-				flattenNode(valueNode, fullKey, result)
-			case yaml.SequenceNode:
-				// List (sequence), process elements
-				var listStr string
-				for j, item := range valueNode.Content {
-					if j > 0 {
-						listStr += ","
-					}
-					if item.Kind == yaml.ScalarNode {
-						listStr += item.Value
-					}
+				if i == 0 {
+					str = itemStr
+				} else {
+					str = str + "," + itemStr
 				}
-				result[lowerKey(fullKey)] = listStr
-				result[formatKey(fullKey)] = listStr
 			}
+			result[lowerKey(fullKey)] = str
+			result[formatKey(fullKey)] = str
+		default:
+			str, err := cast.ToStringE(val)
+			if err != nil {
+				fmt.Printf("cast to string failed %s, error = %s\n", fullKey, err.Error())
+				continue
+			}
+			result[lowerKey(fullKey)] = str
+			result[formatKey(fullKey)] = str
 		}
 	}
 }
diff --git a/pkg/config/file_source.go b/pkg/config/file_source.go
index 535e87d5c2963..6b158fcd2fdfa 100644
--- a/pkg/config/file_source.go
+++ b/pkg/config/file_source.go
@@ -17,9 +17,7 @@
 package config
 
 import (
-	"bytes"
 	"fmt"
-	"io"
 	"os"
 	"path/filepath"
 	"sync"
@@ -27,7 +25,7 @@ import (
 	"github.com/cockroachdb/errors"
 	"github.com/samber/lo"
 	"go.uber.org/zap"
-	"gopkg.in/yaml.v3"
+	"gopkg.in/yaml.v2"
 
 	"github.com/milvus-io/milvus/pkg/log"
 )
@@ -131,7 +129,7 @@ func (fs *FileSource) loadFromFile() error {
 		}
 
 		ext := filepath.Ext(configFile)
-		if len(ext) == 0 || ext[1:] != "yaml" {
+		if len(ext) == 0 || (ext[1:] != "yaml" && ext[1:] != "yml") {
 			return fmt.Errorf("Unsupported Config Type: " + ext)
 		}
 
@@ -140,26 +138,13 @@ func (fs *FileSource) loadFromFile() error {
 			return errors.Wrap(err, "Read config failed: "+configFile)
 		}
 
-		// handle empty file
-		if len(data) == 0 {
-			continue
-		}
-
-		var node yaml.Node
-		decoder := yaml.NewDecoder(bytes.NewReader(data))
-		if err := decoder.Decode(&node); err != nil && !errors.Is(err, io.EOF) {
-			return errors.Wrap(err, "YAML unmarshal failed: "+configFile)
+		var config map[string]interface{}
+		err = yaml.Unmarshal(data, &config)
+		if err != nil {
+			return errors.Wrap(err, "unmarshal yaml file "+configFile+" failed")
 		}
 
-		if node.Kind == yaml.DocumentNode && len(node.Content) > 0 {
-			// Get the content of the Document Node
-			contentNode := node.Content[0]
-
-			// Recursively process the content of the Document Node
-			flattenNode(contentNode, "", newConfig)
-		} else if node.Kind == yaml.MappingNode {
-			flattenNode(&node, "", newConfig)
-		}
+		flattenAndMergeMap("", config, newConfig)
 	}
 
 	return fs.update(newConfig)