44 "context"
55 "fmt"
66 "path/filepath"
7+ "regexp"
78 "strings"
89 "sync"
910 "text/template"
@@ -21,7 +22,11 @@ import (
2122
2223 "github.com/argoproj/argo-cd/v2/pkg/apiclient/application"
2324 "github.com/argoproj/argo-cd/v2/pkg/apis/application/v1alpha1"
24- "gopkg.in/yaml.v2"
25+
26+ "github.com/goccy/go-yaml"
27+ "github.com/goccy/go-yaml/ast"
28+ "github.com/goccy/go-yaml/parser"
29+ "github.com/goccy/go-yaml/token"
2530)
2631
2732// Stores some statistics about the results of a run
@@ -459,8 +464,7 @@ func marshalParamsOverride(app *v1alpha1.Application, originalData []byte) ([]by
459464 if strings .HasPrefix (app .Annotations [common .WriteBackTargetAnnotation ], common .HelmPrefix ) {
460465 images := GetImagesAndAliasesFromApplication (app )
461466
462- helmNewValues := yaml.MapSlice {}
463- err = yaml .Unmarshal (originalData , & helmNewValues )
467+ helmNewValues , err := parser .ParseBytes (originalData , parser .ParseComments )
464468 if err != nil {
465469 return nil , err
466470 }
@@ -488,7 +492,7 @@ func marshalParamsOverride(app *v1alpha1.Application, originalData []byte) ([]by
488492 if helmParamVersion == nil {
489493 return nil , fmt .Errorf ("%s parameter not found" , helmAnnotationParamVersion )
490494 }
491- err = setHelmValue (& helmNewValues , helmAnnotationParamVersion , helmParamVersion .Value )
495+ err = setHelmValue (helmNewValues , helmAnnotationParamVersion , helmParamVersion .Value )
492496 if err != nil {
493497 return nil , fmt .Errorf ("failed to set image parameter version value: %v" , err )
494498 }
@@ -499,13 +503,13 @@ func marshalParamsOverride(app *v1alpha1.Application, originalData []byte) ([]by
499503 return nil , fmt .Errorf ("%s parameter not found" , helmAnnotationParamName )
500504 }
501505
502- err = setHelmValue (& helmNewValues , helmAnnotationParamName , helmParamName .Value )
506+ err = setHelmValue (helmNewValues , helmAnnotationParamName , helmParamName .Value )
503507 if err != nil {
504508 return nil , fmt .Errorf ("failed to set image parameter name value: %v" , err )
505509 }
506510 }
507511
508- override , err = yaml . Marshal (helmNewValues )
512+ override = [] byte (helmNewValues . String () )
509513 } else {
510514 var params helmOverride
511515 newParams := helmOverride {
@@ -561,76 +565,113 @@ func mergeKustomizeOverride(t *kustomizeOverride, o *kustomizeOverride) {
561565 }
562566}
563567
564- // Check if a key exists in a MapSlice and return its index and value
565- func findHelmValuesKey (m yaml.MapSlice , key string ) (int , bool ) {
566- for i , item := range m {
567- if item .Key == key {
568- return i , true
568+ func splitKeys (input string ) []string {
569+ // Regular expression to match quoted and unquoted segments
570+ re := regexp .MustCompile (`'([^']*)'|"([^"]*)"|([^.".]+)` )
571+ matches := re .FindAllStringSubmatch (input , - 1 )
572+
573+ var result []string
574+ for _ , match := range matches {
575+ if match [1 ] != "" { // Single-quoted string
576+ result = append (result , match [1 ])
577+ } else if match [2 ] != "" { // Double-quoted string
578+ result = append (result , match [2 ])
579+ } else if match [3 ] != "" { // Unquoted segment
580+ result = append (result , match [3 ])
569581 }
570582 }
571- return - 1 , false
583+
584+ return result
572585}
573586
574587// set value of the parameter passed from the annotations.
575- func setHelmValue (currentValues * yaml.MapSlice , key string , value interface {}) error {
576- // Check if the full key exists
577- if idx , found := findHelmValuesKey (* currentValues , key ); found {
578- (* currentValues )[idx ].Value = value
579- return nil
588+ func setHelmValue (file * ast.File , keyPath , value string ) error {
589+ path := splitKeys (keyPath )
590+ if len (path ) == 0 {
591+ return fmt .Errorf ("empty key provided" )
580592 }
581593
582- var err error
583- keys := strings .Split (key , "." )
584- current := currentValues
585- var parent * yaml.MapSlice
586- parentIdx := - 1
587-
588- for i , k := range keys {
589- if idx , found := findHelmValuesKey (* current , k ); found {
590- if i == len (keys )- 1 {
591- // If we're at the final key, set the value and return
592- (* current )[idx ].Value = value
593- return nil
594- } else {
595- // Navigate deeper into the map
596- if nestedMap , ok := (* current )[idx ].Value .(yaml.MapSlice ); ok {
597- parent = current
598- parentIdx = idx
599- current = & nestedMap
600- } else {
601- return fmt .Errorf ("unexpected type %T for key %s" , (* current )[idx ].Value , k )
594+ mapping , ok := file .Docs [0 ].Body .(* ast.MappingNode )
595+ if ! ok {
596+ tk := token .New ("$" , "$" , & token.Position {})
597+ mapping = ast .Mapping (tk , false )
598+ file .Docs [0 ].Body = mapping
599+ }
600+
601+ // Traverse the path
602+ var lastNode * ast.MappingValueNode
603+ for index , key := range path {
604+ found := false
605+ var currentNode * ast.MappingValueNode
606+
607+ for _ , v := range mapping .Values {
608+ if v .Key .GetToken ().Value == key {
609+ currentNode = v
610+ if index == len (path )- 1 {
611+ lastNode = currentNode
612+ found = true
613+ break
614+ }
615+ // Move deeper into the structure
616+ if nextMapping , ok := v .Value .(* ast.MappingNode ); ok {
617+ mapping = nextMapping
618+ found = true
619+ break
602620 }
603621 }
604- } else {
605- newCurrent := yaml.MapSlice {}
606- var newParent yaml.MapSlice
622+ }
607623
608- if i == len (keys )- 1 {
609- newParent = append (* current , yaml.MapItem {Key : k , Value : value })
610- } else {
611- newParent = append (* current , yaml.MapItem {Key : k , Value : newCurrent })
612- }
624+ // If key does not exist, create it
625+ if ! found {
626+ // Create a token with proper position (assuming default line/column)
627+ keyToken := token .New (key , key , & token.Position {Column : index * 2 + 1 })
628+ newKey := ast .String (keyToken ) // Create key node
629+ mappingToken := token .New (key , key , & token.Position {})
630+ newMapping := ast .Mapping (mappingToken , false ) // Create empty mapping
613631
614- if parent == nil {
615- * currentValues = newParent
616- } else {
617- // if parentIdx has not been set (parent element is also new), set it to the last element
618- if parentIdx == - 1 {
619- parentIdx = len (* parent ) - 1
620- if parentIdx < 0 {
621- parentIdx = 0
622- }
632+ if currentNode != nil {
633+ comment := currentNode .Value .GetComment ()
634+ currentNode .Value = newMapping
635+ err := currentNode .Value .SetComment (comment )
636+ if err != nil {
637+ return err
623638 }
624- (* parent )[parentIdx ].Value = newParent
639+ lastNode = currentNode
640+ } else {
641+ // Add the new mapping to the parent mapping
642+ lastNode = ast .MappingValue (mappingToken , newKey , newMapping )
643+ mapping .Values = append (mapping .Values , lastNode )
625644 }
626-
627- parent = & newParent
628- current = & newCurrent
629- parentIdx = - 1
645+ mapping = newMapping
630646 }
631647 }
632648
633- return err
649+ if lastNode == nil {
650+ return fmt .Errorf ("key not found" )
651+ }
652+
653+ var valueToken * token.Token
654+ if token .IsNeedQuoted (value ) {
655+ valueToken = token .SingleQuote (value , value , & token.Position {})
656+ } else {
657+ valueToken = token .New (value , value , & token.Position {})
658+ }
659+ newValue := ast .String (valueToken )
660+ comment := lastNode .Value .GetComment ()
661+ if comment == nil {
662+ comment = lastNode .Key .GetComment ()
663+ }
664+ lastNode .Value = newValue
665+ err := lastNode .Key .SetComment (nil )
666+ if err != nil {
667+ return err
668+ }
669+ err = lastNode .Value .SetComment (comment )
670+ if err != nil {
671+ return err
672+ }
673+
674+ return nil
634675}
635676
636677func getWriteBackConfig (app * v1alpha1.Application , kubeClient * kube.ImageUpdaterKubernetesClient , argoClient ArgoCD ) (* WriteBackConfig , error ) {
0 commit comments