Skip to content

Commit cbefdf6

Browse files
flarcoflarco
committed
refactor: enhance YAML marshaling in EnvFile to preserve comments and key order; add comprehensive tests for comment preservation
1 parent 8ad4ecb commit cbefdf6

2 files changed

Lines changed: 376 additions & 80 deletions

File tree

‎core/env/envfile.go‎

Lines changed: 194 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,14 @@
11
package env
22

33
import (
4+
"bytes"
45
"os"
56
"path"
6-
"sort"
77
"strings"
88

99
"github.com/flarco/g"
1010
cmap "github.com/orcaman/concurrent-map/v2"
11-
"github.com/samber/lo"
12-
"github.com/spf13/cast"
13-
"gopkg.in/yaml.v2"
11+
"gopkg.in/yaml.v3"
1412
)
1513

1614
type EnvFile struct {
@@ -23,49 +21,6 @@ type EnvFile struct {
2321
Body string `json:"-" yaml:"-"`
2422
}
2523

26-
// marshalEnvFileBytes marshals the EnvFile into formatted YAML bytes
27-
func (ef *EnvFile) marshalEnvFileBytes() ([]byte, error) {
28-
connsMap := yaml.MapSlice{}
29-
30-
// order connections names
31-
names := lo.Keys(ef.Connections)
32-
sort.Strings(names)
33-
for _, name := range names {
34-
keyMap := ef.Connections[name]
35-
// order connection keys (type first)
36-
cMap := yaml.MapSlice{}
37-
keys := lo.Keys(keyMap)
38-
sort.Strings(keys)
39-
if v, ok := keyMap["type"]; ok {
40-
cMap = append(cMap, yaml.MapItem{Key: "type", Value: v})
41-
}
42-
43-
for _, k := range keys {
44-
if k == "type" {
45-
continue // already put first
46-
}
47-
k = cast.ToString(k)
48-
cMap = append(cMap, yaml.MapItem{Key: k, Value: keyMap[k]})
49-
}
50-
51-
// add to connection map
52-
connsMap = append(connsMap, yaml.MapItem{Key: name, Value: cMap})
53-
}
54-
55-
efMap := yaml.MapSlice{
56-
{Key: "connections", Value: connsMap},
57-
{Key: "variables", Value: ef.Env},
58-
}
59-
60-
envBytes, err := yaml.Marshal(efMap)
61-
if err != nil {
62-
return nil, g.Error(err, "could not marshal into YAML")
63-
}
64-
65-
output := formatYAML([]byte(ef.TopComment + string(envBytes)))
66-
return output, nil
67-
}
68-
6924
func (ef *EnvFile) WriteEnvFile() (err error) {
7025
output, err := ef.marshalEnvFileBytes()
7126
if err != nil {
@@ -91,43 +46,93 @@ func (ef *EnvFile) MarshalBody() (string, error) {
9146
return string(output), nil
9247
}
9348

94-
func formatYAML(input []byte) []byte {
95-
newOutput := []byte{}
96-
pIndent := 0
97-
indent := 0
98-
inIndent := true
99-
prevC := byte('-')
100-
for _, c := range input {
101-
add := false
102-
if c == ' ' && inIndent {
103-
indent++
104-
add = true
105-
} else if c == '\n' {
106-
pIndent = indent
107-
indent = 0
108-
add = true
109-
inIndent = true
110-
} else if prevC == '\n' {
111-
newOutput = append(newOutput, '\n') // add extra space
112-
add = true
113-
} else if prevC == ' ' && pIndent > indent && inIndent {
114-
newOutput = append(newOutput, '\n') // add extra space
115-
for i := 0; i < indent; i++ {
116-
newOutput = append(newOutput, ' ')
117-
}
118-
add = true
119-
inIndent = false
120-
} else {
121-
add = true
122-
inIndent = false
49+
func (ef *EnvFile) freshRoot() *yaml.Node {
50+
root := &yaml.Node{
51+
Kind: yaml.DocumentNode,
52+
Content: []*yaml.Node{{
53+
Kind: yaml.MappingNode,
54+
}},
55+
}
56+
if ef.TopComment != "" {
57+
root.Content[0].HeadComment = strings.TrimRight(ef.TopComment, "\n")
58+
}
59+
return root
60+
}
61+
62+
// marshalEnvFileBytes renders the EnvFile as YAML, preserving comments, key
63+
// order, and unmanaged top-level keys from the file at ef.Path.
64+
func (ef *EnvFile) marshalEnvFileBytes() ([]byte, error) {
65+
original, err := ef.loadRootNode()
66+
if err != nil {
67+
return nil, err
68+
}
69+
70+
newRoot, err := ef.structToRootNode(original)
71+
if err != nil {
72+
return nil, err
73+
}
74+
75+
merged := mergeNode(original, newRoot)
76+
77+
var buf bytes.Buffer
78+
enc := yaml.NewEncoder(&buf)
79+
enc.SetIndent(2)
80+
if err := enc.Encode(merged); err != nil {
81+
_ = enc.Close()
82+
return nil, g.Error(err, "could not marshal into YAML")
83+
}
84+
if err := enc.Close(); err != nil {
85+
return nil, g.Error(err, "could not finalize YAML encoder")
86+
}
87+
return buf.Bytes(), nil
88+
}
89+
90+
// structToRootNode marshals the EnvFile struct into a DocumentNode, mirroring
91+
// unmanaged top-level keys from original so the merge doesn't drop them.
92+
func (ef *EnvFile) structToRootNode(original *yaml.Node) (*yaml.Node, error) {
93+
b, err := yaml.Marshal(ef)
94+
if err != nil {
95+
return nil, g.Error(err, "could not marshal env file")
96+
}
97+
var doc yaml.Node
98+
if uerr := yaml.Unmarshal(b, &doc); uerr != nil {
99+
return nil, g.Error(uerr, "could not re-parse env file node")
100+
}
101+
if doc.Kind == 0 {
102+
doc = yaml.Node{
103+
Kind: yaml.DocumentNode,
104+
Content: []*yaml.Node{{Kind: yaml.MappingNode}},
123105
}
106+
}
107+
if len(doc.Content) == 0 || doc.Content[0].Kind != yaml.MappingNode {
108+
doc.Content = []*yaml.Node{{Kind: yaml.MappingNode}}
109+
}
124110

125-
if add {
126-
newOutput = append(newOutput, c)
111+
managed := map[string]struct{}{
112+
"connections": {}, "variables": {}, "env": {},
113+
}
114+
if original != nil && len(original.Content) > 0 && original.Content[0].Kind == yaml.MappingNode {
115+
newMap := doc.Content[0]
116+
newKeys := map[string]struct{}{}
117+
for i := 0; i < len(newMap.Content); i += 2 {
118+
newKeys[newMap.Content[i].Value] = struct{}{}
119+
}
120+
origMap := original.Content[0]
121+
for i := 0; i < len(origMap.Content); i += 2 {
122+
k := origMap.Content[i].Value
123+
if _, isManaged := managed[k]; isManaged {
124+
continue
125+
}
126+
if _, alreadyInNew := newKeys[k]; alreadyInNew {
127+
continue
128+
}
129+
keyCopy := *origMap.Content[i]
130+
valCopy := *origMap.Content[i+1]
131+
newMap.Content = append(newMap.Content, &keyCopy, &valCopy)
127132
}
128-
prevC = c
129133
}
130-
return newOutput
134+
135+
return &doc, nil
131136
}
132137

133138
var dotEnvMap = cmap.New[string]()
@@ -256,6 +261,11 @@ func LoadEnvFile(path string) (ef EnvFile) {
256261

257262
for k, v := range ef.Env {
258263
if _, found := envMap[k]; !found {
264+
// non-scalar values (e.g. SLING_ASSIST) are read from ef.Env, not os.Getenv
265+
switch v.(type) {
266+
case map[string]any, map[any]any, []any:
267+
continue
268+
}
259269
os.Setenv(k, g.CastToString(v))
260270
}
261271
}
@@ -265,3 +275,107 @@ func LoadEnvFile(path string) (ef EnvFile) {
265275
func GetEnvFilePath(dir string) string {
266276
return CleanWindowsPath(path.Join(dir, "env.yaml"))
267277
}
278+
279+
// mergeNode deep-merges newNode into original, keeping original's comments and
280+
// key order. Adapted from pulumi/pulumi's yamlutil.editNodes (Apache 2.0).
281+
func mergeNode(original, newNode *yaml.Node) *yaml.Node {
282+
if original == nil {
283+
out := *newNode
284+
return &out
285+
}
286+
if newNode == nil {
287+
out := *original
288+
return &out
289+
}
290+
if original.Kind != newNode.Kind {
291+
out := *newNode
292+
return &out
293+
}
294+
295+
ret := *original
296+
ret.Tag = newNode.Tag
297+
ret.Value = newNode.Value
298+
299+
switch original.Kind {
300+
case yaml.DocumentNode, yaml.SequenceNode:
301+
minLen := len(newNode.Content)
302+
if len(original.Content) < minLen {
303+
minLen = len(original.Content)
304+
}
305+
content := make([]*yaml.Node, 0, len(newNode.Content))
306+
for i := 0; i < minLen; i++ {
307+
content = append(content, mergeNode(original.Content[i], newNode.Content[i]))
308+
}
309+
content = append(content, newNode.Content[minLen:]...)
310+
ret.Content = content
311+
case yaml.MappingNode:
312+
ret.Content = mergeMappingContent(original, newNode)
313+
case yaml.ScalarNode, yaml.AliasNode:
314+
ret.Content = newNode.Content
315+
}
316+
return &ret
317+
}
318+
319+
// mergeMappingContent merges two mapping nodes: original keys keep their
320+
// position and comments; new-only keys append at the end; dropped keys are
321+
// removed.
322+
func mergeMappingContent(original, newNode *yaml.Node) []*yaml.Node {
323+
origIdx := map[string]int{}
324+
newIdx := map[string]int{}
325+
var origOrder, newOnly []string
326+
327+
for i := 0; i < len(original.Content); i += 2 {
328+
k := original.Content[i].Value
329+
origIdx[k] = i
330+
origOrder = append(origOrder, k)
331+
}
332+
for i := 0; i < len(newNode.Content); i += 2 {
333+
k := newNode.Content[i].Value
334+
newIdx[k] = i
335+
if _, ok := origIdx[k]; !ok {
336+
newOnly = append(newOnly, k)
337+
}
338+
}
339+
340+
content := make([]*yaml.Node, 0, len(newNode.Content))
341+
for _, k := range origOrder {
342+
ni, present := newIdx[k]
343+
if !present {
344+
continue
345+
}
346+
oi := origIdx[k]
347+
key := mergeNode(original.Content[oi], newNode.Content[ni])
348+
val := mergeNode(original.Content[oi+1], newNode.Content[ni+1])
349+
content = append(content, key, val)
350+
}
351+
for _, k := range newOnly {
352+
ni := newIdx[k]
353+
key := *newNode.Content[ni]
354+
val := *newNode.Content[ni+1]
355+
content = append(content, &key, &val)
356+
}
357+
return content
358+
}
359+
360+
// loadRootNode parses ef.Path into a yaml.Node tree, returning a fresh root
361+
// if the file is missing, empty, or not a mapping at the top.
362+
func (ef *EnvFile) loadRootNode() (*yaml.Node, error) {
363+
root := &yaml.Node{}
364+
if ef.Path == "" {
365+
return ef.freshRoot(), nil
366+
}
367+
data, rerr := os.ReadFile(ef.Path)
368+
if rerr != nil || len(bytes.TrimSpace(data)) == 0 {
369+
return ef.freshRoot(), nil
370+
}
371+
if uerr := yaml.Unmarshal(data, root); uerr != nil {
372+
return nil, g.Error(uerr, "could not parse %s", ef.Path)
373+
}
374+
if root.Kind == 0 {
375+
return ef.freshRoot(), nil
376+
}
377+
if len(root.Content) == 0 || root.Content[0].Kind != yaml.MappingNode {
378+
root.Content = []*yaml.Node{{Kind: yaml.MappingNode}}
379+
}
380+
return root, nil
381+
}

0 commit comments

Comments
 (0)