feat(update-collection-v3): order output keys based on input

src/go/cmd/update-collection-v3/main.go

@@ -1,12 +1,12 @@
 package main
 
 import (
-	"bytes"
 	"flag"
 	"fmt"
 	"io"
 	"log"
 	"os"
+	"sort"
 
 	disablethanos "github.com/SumoLogic/sumologic-kubernetes-collection/tools/cmd/update-collection-v3/migrations/disable-thanos"
 	"github.com/SumoLogic/sumologic-kubernetes-collection/tools/cmd/update-collection-v3/migrations/events"
@@ -156,26 +156,68 @@ func migrateYaml(input string) (string, error) {
 
 	// make the output consistently ordered
 	// without this logic, key ordering would depend on the final migration
-	// TODO: order keys the same as input
-	output, err = reorderYaml(output)
+	output, err = reorderYaml(output, input)
 
 	return output, err
 }
 
-// reorder yaml keys in the input
-// right now this just unmarshals into a map and marshals again
-// the result is alphabetical key ordering
-func reorderYaml(input string) (string, error) {
-	var outputMap map[string]interface{}
-	err := yaml.Unmarshal([]byte(input), &outputMap)
+// reorder yaml keys in the input based on their order in the output
+func reorderYaml(input string, original string) (string, error) {
+	var outputMapSlice, originalMapSlice yaml.MapSlice
+	var err error
+	err = yaml.UnmarshalWithOptions([]byte(input), &outputMapSlice, yaml.UseOrderedMap())
+	if err != nil {
+		return "", err
+	}
+	err = yaml.UnmarshalWithOptions([]byte(original), &originalMapSlice, yaml.UseOrderedMap())
 	if err != nil {
 		return "", err
 	}
-	buffer := bytes.Buffer{}
-	encoder := yaml.NewEncoder(&buffer, yaml.Indent(2))
-	err = encoder.Encode(outputMap)
 
-	return buffer.String(), err
+	sortByBlueprint(outputMapSlice, originalMapSlice)
+
+	outputBytes, err := yaml.MarshalWithOptions(outputMapSlice, yaml.Indent(2))
+
+	return string(outputBytes), err
+}
+
+// sortByBlueprint sorts the input based on the order of keys in the output
+// this is done recursively
+// keys not present in the blueprint go to the end and are sorted alphabetically
+func sortByBlueprint(input yaml.MapSlice, blueprint yaml.MapSlice) {
+	blueprintMap := blueprint.ToMap()
+	sort.Slice(input, func(i, j int) bool {
+		iKey := input[i].Key
+		jKey := input[j].Key
+		iPosition, jPosition := len(input), len(input)
+		for position, entry := range blueprint {
+			if entry.Key == iKey {
+				iPosition = position
+			}
+			if entry.Key == jKey {
+				jPosition = position
+			}
+		}
+		if iPosition == len(input) && jPosition == len(input) {
+			// if neither key are in the blueprint, sort alphabetically
+			return iKey.(string) < jKey.(string)
+		}
+		return iPosition < jPosition
+	})
+
+	// sort recursively for values which are also yaml.MapSlice and exist in the blueprint
+	var entryValueMapSlice yaml.MapSlice
+	var ok bool
+	for _, entry := range input {
+		if entryValueMapSlice, ok = (entry.Value).(yaml.MapSlice); !ok { // not a yaml.MapSlice
+			continue
+		}
+		if blueprintValue, ok := blueprintMap[entry.Key]; ok {
+			if blueprintValueMapSlice, ok := (blueprintValue).(yaml.MapSlice); ok {
+				sortByBlueprint(entryValueMapSlice, blueprintValueMapSlice)
+			}
+		}
+	}
 }
 
 type migrateFunc func(string) (string, error)

src/go/cmd/update-collection-v3/main_test.go

@@ -46,3 +46,86 @@ func testMigrationsInDirectory(t *testing.T, migrate migrateFunc, directory stri
 		})
 	}
 }
+
+func TestReorderYaml(t *testing.T) {
+	testCases := []struct {
+		inputYaml    string
+		originalYaml string
+		outputYaml   string
+		description  string
+	}{
+		{
+			inputYaml: `
+a: b
+c: d
+`,
+			originalYaml: `
+c: d
+a: b
+`,
+			outputYaml: `
+c: d
+a: b
+`,
+			description: "basic ordering",
+		},
+		{
+			inputYaml: `
+a: b
+c: d
+e: f
+`,
+			originalYaml: `
+c: d
+a: b
+`,
+			outputYaml: `
+c: d
+a: b
+e: f
+`,
+			description: "key not in original goes to the end",
+		},
+		{
+			inputYaml: `
+g: h
+e: f
+`,
+			originalYaml: `
+c: d
+a: b
+`,
+			outputYaml: `
+e: f
+g: h
+`,
+			description: "keys not in original are sorted alphabetically",
+		},
+		{
+			inputYaml: `
+nested: 
+  a: b
+  c: d
+`,
+			originalYaml: `
+nested:
+  c: d
+  a: b
+`,
+			outputYaml: `
+nested:
+  c: d
+  a: b
+`,
+			description: "basic nested ordering",
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.description, func(t *testing.T) {
+			actualOutput, err := reorderYaml(testCase.inputYaml, testCase.originalYaml)
+			require.NoError(t, err)
+			require.Equal(t, strings.Trim(testCase.outputYaml, "\n"), strings.Trim(actualOutput, "\n"))
+		})
+	}
+}

src/go/cmd/update-collection-v3/testdata/simple.input.yaml

@@ -1,3 +1,6 @@
+sumologic:
+  accessId: xxx
+  accessKey: yyy
 fluentd:
   events:
     enabled: false
@@ -14,6 +17,13 @@ kube-prometheus-stack:
         resources:
           limits:
             cpu: 20m
+metadata:
+  logs:
+    config:
+      keylogs: valuelogs
+  metrics:
+    config:
+      keymetrics: valuemetrics
 otelevents:
   config:
     override:
@@ -22,13 +32,3 @@ otellogs:
   config:
     override:
       key: value
-metadata:
-  logs:
-    config:
-      keylogs: valuelogs
-  metrics:
-    config:
-      keymetrics: valuemetrics
-sumologic:
-  accessId: xxx
-  accessKey: yyy

src/go/cmd/update-collection-v3/testdata/simple.output.yaml

@@ -1,3 +1,8 @@
+sumologic:
+  accessId: xxx
+  accessKey: yyy
+  events:
+    enabled: false
 fluentd:
   logs:
     autoscaling:
@@ -7,17 +12,17 @@ fluentd:
       enabled: true
 metadata:
   logs:
-    autoscaling:
-      enabled: true
     config:
       merge:
         keylogs: valuelogs
-  metrics:
     autoscaling:
       enabled: true
+  metrics:
     config:
       merge:
         keymetrics: valuemetrics
+    autoscaling:
+      enabled: true
 otelevents:
   config:
     merge:
@@ -26,8 +31,3 @@ otellogs:
   config:
     merge:
       key: value
-sumologic:
-  accessId: xxx
-  accessKey: yyy
-  events:
-    enabled: false

src/go/cmd/update-collection-v3/testdata/tracing.output.yaml

@@ -1,3 +1,6 @@
+sumologic:
+  traces:
+    enabled: true
 otelcolInstrumentation:
   config:
     processors:
@@ -17,9 +20,6 @@ otelcolInstrumentation:
         source_category_replace_dash: processors.source.category_replace_dash.replace
         source_host: "%{k8s.pod.hostname}"
         source_name: processors.source.name.replace
-sumologic:
-  traces:
-    enabled: true
 tracesSampler:
   config:
     processors: