Skip to content

Commit

Permalink
update yaml v2 as yaml parser
Browse files Browse the repository at this point in the history
Signed-off-by: xianliang.li <[email protected]>
  • Loading branch information
foxspy committed Nov 5, 2024
1 parent 0645d46 commit 84fbb77
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 61 deletions.
72 changes: 33 additions & 39 deletions pkg/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)
Expand Down Expand Up @@ -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
}
}
}
29 changes: 7 additions & 22 deletions pkg/config/file_source.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,15 @@
package config

import (
"bytes"
"fmt"
"io"
"os"
"path/filepath"
"sync"

"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"
)
Expand Down Expand Up @@ -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)
}

Expand All @@ -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)
Expand Down

0 comments on commit 84fbb77

Please sign in to comment.