feat: add pod labels/annotations in the workload patcher (#1123)

KusionStack · May 16, 2024 · 78cbbbd · 78cbbbd
1 parent 5ca1c82
commit 78cbbbd
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 13 deletions.
diff --git a/pkg/apis/api.kusion.io/v1/types.go b/pkg/apis/api.kusion.io/v1/types.go
@@ -402,10 +402,14 @@ type AppConfiguration struct {
 type Patcher struct {
 	// Environments represent the environment variables patched to all containers in the workload.
 	Environments []v1.EnvVar `json:"environments" yaml:"environments"`
-	// Labels represent the labels patched to both the workload and pod.
+	// Labels represent the labels patched to the workload.
 	Labels map[string]string `json:"labels" yaml:"labels"`
-	// Annotations represent the annotations patched to both the workload and pod.
+	// PodLabels represent the labels patched to the pods.
+	PodLabels map[string]string `json:"podLabels" yaml:"podLabels"`
+	// Annotations represent the annotations patched to the workload.
 	Annotations map[string]string `json:"annotations" yaml:"annotations"`
+	// PodAnnotations represent the annotations patched to the pods.
+	PodAnnotations map[string]string `json:"podAnnotations" yaml:"podAnnotations"`
 }
 
 const ConfigBackends = "backends"

diff --git a/pkg/modules/generators/app_configurations_generator.go b/pkg/modules/generators/app_configurations_generator.go
@@ -163,6 +163,15 @@ func PatchWorkload(workload *v1.Resource, patcher *v1.Patcher) error {
 		if objLabels == nil {
 			objLabels = make(map[string]string)
 		}
+		// merge labels
+		for k, v := range patcher.Labels {
+			objLabels[k] = v
+		}
+		un.SetLabels(objLabels)
+	}
+
+	// patch pod labels
+	if patcher.PodLabels != nil {
 		podLabels, b, err := unstructured.NestedStringMap(un.Object, "spec", "template", "metadata", "labels")
 		if err != nil {
 			return fmt.Errorf("failed to get pod labels from workload:%s. %w", workload.ID, err)
@@ -171,11 +180,9 @@ func PatchWorkload(workload *v1.Resource, patcher *v1.Patcher) error {
 			podLabels = make(map[string]string)
 		}
 		// merge labels
-		for k, v := range patcher.Labels {
-			objLabels[k] = v
+		for k, v := range patcher.PodLabels {
 			podLabels[k] = v
 		}
-		un.SetLabels(objLabels)
 		err = unstructured.SetNestedStringMap(un.Object, podLabels, "spec", "template", "metadata", "labels")
 		if err != nil {
 			return err
@@ -188,6 +195,15 @@ func PatchWorkload(workload *v1.Resource, patcher *v1.Patcher) error {
 		if objAnnotations == nil {
 			objAnnotations = make(map[string]string)
 		}
+		// merge annotations
+		for k, v := range patcher.Annotations {
+			objAnnotations[k] = v
+		}
+		un.SetAnnotations(objAnnotations)
+	}
+
+	// patch pod annotations
+	if patcher.PodAnnotations != nil {
 		podAnnotations, b, err := unstructured.NestedStringMap(un.Object, "spec", "template", "metadata", "annotations")
 		if err != nil {
 			return fmt.Errorf("failed to get pod annotations from workload:%s. %w", workload.ID, err)
@@ -196,12 +212,9 @@ func PatchWorkload(workload *v1.Resource, patcher *v1.Patcher) error {
 			podAnnotations = make(map[string]string)
 		}
 		// merge annotations
-		for k, v := range patcher.Annotations {
-			objAnnotations[k] = v
+		for k, v := range patcher.PodAnnotations {
 			podAnnotations[k] = v
 		}
-
-		un.SetAnnotations(objAnnotations)
 		err = unstructured.SetNestedStringMap(un.Object, podAnnotations, "spec", "template", "metadata", "annotations")
 		if err != nil {
 			return err

diff --git a/pkg/modules/generators/app_configurations_generator_test.go b/pkg/modules/generators/app_configurations_generator_test.go
@@ -291,8 +291,10 @@ func Test_patchWorkload(t *testing.T) {
 
 	t.Run("Patch labels and annotations", func(t *testing.T) {
 		patcher := &v1.Patcher{
-			Labels:      map[string]string{"newLabel": "newValue"},
-			Annotations: map[string]string{"newAnnotation": "newValue"},
+			Labels:         map[string]string{"newLabel": "newValue"},
+			Annotations:    map[string]string{"newAnnotation": "newValue"},
+			PodLabels:      map[string]string{"newPodLabel": "newValue"},
+			PodAnnotations: map[string]string{"newPodAnnotation": "newValue"},
 		}
 
 		err := PatchWorkload(res, patcher)
@@ -305,7 +307,7 @@ func Test_patchWorkload(t *testing.T) {
 		assert.Equal(t, "newValue", workloadLabels["newLabel"])
 		assert.Equal(t, "oldValue", workloadLabels["oldLabel"])
 		// assert pod labels
-		assert.Equal(t, "newValue", podLabels["newLabel"])
+		assert.Equal(t, "newValue", podLabels["newPodLabel"])
 		assert.Equal(t, "oldValue", podLabels["oldLabel"])
 
 		annotations := res.Attributes["metadata"].(map[string]interface{})["annotations"].(map[string]interface{})
@@ -314,7 +316,7 @@ func Test_patchWorkload(t *testing.T) {
 		// assert deployment annotations
 		assert.Equal(t, "newValue", annotations["newAnnotation"])
 		// assert pod annotations
-		assert.Equal(t, "newValue", podAnnotations["newAnnotation"])
+		assert.Equal(t, "newValue", podAnnotations["newPodAnnotation"])
 		assert.Equal(t, "oldValue", podLabels["oldLabel"])
 	})