CLOUDP-186866 - Allow override of labels and annotations (#1337)

* Merge labels and annotations on created StatefulSets, Pods, and PersistentVolumeClaims


---------

Co-authored-by: cghislai <charlyghislain@gmail.com>

Author: nammn

api/v1/mongodbcommunity_types.go

@@ -343,6 +343,8 @@ type OverrideProcess struct {
 type StatefulSetConfiguration struct {
 	// +kubebuilder:pruning:PreserveUnknownFields
 	SpecWrapper StatefulSetSpecWrapper `json:"spec"`
+	// +optional
+	MetadataWrapper StatefulSetMetadataWrapper `json:"metadata"`
 }
 
 type LogLevel string
@@ -385,6 +387,21 @@ func (m *StatefulSetSpecWrapper) DeepCopy() *StatefulSetSpecWrapper {
 	}
 }
 
+// StatefulSetMetadataWrapper is a wrapper around Labels and Annotations
+type StatefulSetMetadataWrapper struct {
+	// +optional
+	Labels map[string]string `json:"labels,omitempty"`
+	// +optional
+	Annotations map[string]string `json:"annotations,omitempty"`
+}
+
+func (m *StatefulSetMetadataWrapper) DeepCopy() *StatefulSetMetadataWrapper {
+	return &StatefulSetMetadataWrapper{
+		Labels:      m.Labels,
+		Annotations: m.Annotations,
+	}
+}
+
 // MongodConfiguration holds the optional mongod configuration
 // that should be merged with the operator created one.
 type MongodConfiguration struct {

config/crd/bases/mongodbcommunity.mongodb.com_mongodbcommunity.yaml

@@ -337,6 +337,19 @@ spec:
                 description: StatefulSetConfiguration holds the optional custom StatefulSet
                   that should be merged into the operator created one.
                 properties:
+                  metadata:
+                    description: StatefulSetMetadataWrapper is a wrapper around Labels
+                      and Annotations
+                    properties:
+                      annotations:
+                        additionalProperties:
+                          type: string
+                        type: object
+                      labels:
+                        additionalProperties:
+                          type: string
+                        type: object
+                    type: object
                   spec:
                     type: object
                     x-kubernetes-preserve-unknown-fields: true

config/samples/arbitrary_statefulset_configuration/mongodb.com_v1_metadata.yaml

@@ -0,0 +1,59 @@
+apiVersion: mongodbcommunity.mongodb.com/v1
+kind: MongoDBCommunity
+metadata:
+  name: mdb0
+spec:
+  members: 3
+  type: ReplicaSet
+  version: "4.2.6"
+  security:
+    authentication:
+      modes: [ "SCRAM" ]
+  users:
+    - name: my-user
+      db: admin
+      passwordSecretRef: # a reference to the secret that will be used to generate the user's password
+        name: my-user-password
+      roles:
+        - name: clusterAdmin
+          db: admin
+        - name: userAdminAnyDatabase
+          db: admin
+      scramCredentialsSecretName: my-scram
+  additionalMongodConfig:
+    storage.wiredTiger.engineConfig.journalCompressor: zlib
+
+  statefulSet:
+    metadata:
+      annotations:
+        statefulSetAnnotationTest: testValue
+      labels:
+        statefulSetLabelTest: testValue
+    spec:
+      selector:
+        matchLabels:
+          podTemplateLabelTest: testValue
+
+      template:
+        metadata:
+          annotations:
+            podTemplateAnnotationTest: testValue
+          labels:
+            podTemplateLabelTest: testValue
+
+      volumeClaimTemplates:
+        - metadata:
+            name: data-volume
+            annotations:
+              pvcTemplateAnnotationTest: testValue
+            labels:
+              pvcTemplateLabelTest: testValue
+
+---
+apiVersion: v1
+kind: Secret
+metadata:
+  name: my-user-password
+type: Opaque
+stringData:
+  password: <your-password-here>

controllers/replica_set_controller.go

@@ -756,6 +756,10 @@ func buildStatefulSetModificationFunction(mdb mdbv1.MongoDBCommunity) statefulse
 		),
 
 		statefulset.WithCustomSpecs(mdb.Spec.StatefulSetConfiguration.SpecWrapper.Spec),
+		statefulset.WithObjectMetadata(
+			mdb.Spec.StatefulSetConfiguration.MetadataWrapper.Labels,
+			mdb.Spec.StatefulSetConfiguration.MetadataWrapper.Annotations,
+		),
 	)
 }
 

pkg/kube/statefulset/statefulset.go

@@ -306,6 +306,13 @@ func WithCustomSpecs(spec appsv1.StatefulSetSpec) Modification {
 	}
 }
 
+func WithObjectMetadata(labels map[string]string, annotations map[string]string) Modification {
+	return func(set *appsv1.StatefulSet) {
+		WithLabels(labels)(set)
+		WithAnnotations(annotations)(set)
+	}
+}
+
 func findVolumeClaimIndexByName(name string, pvcs []corev1.PersistentVolumeClaim) int {
 	for idx, pvc := range pvcs {
 		if pvc.Name == name {

pkg/kube/statefulset/statefulset_test.go

@@ -235,3 +235,35 @@ func TestWithAnnotations(t *testing.T) {
 	WithAnnotations(nil)(&sts)
 	assert.Len(t, sts.Annotations, 2)
 }
+
+func TestWithObjectMetadata(t *testing.T) {
+	sts, err := defaultStatefulSetBuilder().Build()
+	assert.NoError(t, err)
+	assert.Len(t, sts.Labels, 0)
+	assert.Len(t, sts.Annotations, 0)
+
+	// handles nil values gracefully
+	{
+		WithObjectMetadata(nil, nil)(&sts)
+	}
+
+	// Test that it works when there are no annotations
+	{
+		WithObjectMetadata(map[string]string{"label": "a"}, map[string]string{"annotation": "b"})(&sts)
+		assert.Equal(t, "b", sts.Annotations["annotation"])
+		assert.Equal(t, "a", sts.Labels["label"])
+	}
+
+	// test that WithObjectMetadata merges the maps
+	{
+		WithObjectMetadata(map[string]string{"label2": "a"}, map[string]string{"annotation2": "b"})(&sts)
+		assert.Equal(t, "b", sts.Annotations["annotation"])
+		assert.Equal(t, "b", sts.Annotations["annotation2"])
+	}
+
+	// Test that we can override a key
+	{
+		WithObjectMetadata(map[string]string{"label": "b"}, map[string]string{"annotation": "b"})(&sts)
+		assert.Equal(t, "b", sts.Annotations["annotation"])
+	}
+}

pkg/util/merge/merge_statefulset.go

@@ -170,6 +170,9 @@ func PersistentVolumeClaim(defaultPvc corev1.PersistentVolumeClaim, overridePvc
 		defaultPvc.Namespace = overridePvc.Namespace
 	}
 
+	defaultPvc.Labels = StringToStringMap(defaultPvc.Labels, overridePvc.Labels)
+	defaultPvc.Annotations = StringToStringMap(defaultPvc.Annotations, overridePvc.Annotations)
+
 	if overridePvc.Spec.VolumeMode != nil {
 		defaultPvc.Spec.VolumeMode = overridePvc.Spec.VolumeMode
 	}

requirements.txt

@@ -4,15 +4,14 @@ docker==4.3.1
 kubernetes==26.1.0
 jinja2==2.11.3
 MarkupSafe==2.0.1
-PyYAML==5.4.1
+PyYAML==6.0.1
 black==22.3.0
 mypy==0.961
 tqdm==v4.49.0
 boto3==1.16.21
 pymongo==3.11.4
 dnspython==2.0.0
 requests==2.31.0
-pyyaml==5.4.1
 ruamel.yaml==0.17.9
 semver==2.13.0
 rsa>=4.7 # not directly required, pinned by Snyk to avoid a vulnerability

test/e2e/e2eutil.go

@@ -28,6 +28,13 @@ func TestLabels() map[string]string {
 	}
 }
 
+// TestAnnotations create an annotations map
+func TestAnnotations() map[string]string {
+	return map[string]string{
+		"e2e-test-annotated": "true",
+	}
+}
+
 func TestDataDir() string {
 	return envvar.GetEnvOrDefault(testDataDirEnv, "/workspace/testdata")
 }

test/e2e/mongodbtests/mongodbtests.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"strings"
 	"testing"
 	"time"
 
@@ -217,6 +219,7 @@ func containsVolume(volumes []corev1.PersistentVolume, volumeName string) bool {
 	}
 	return false
 }
+
 func HasExpectedPersistentVolumes(volumes []corev1.PersistentVolume) func(t *testing.T) {
 	return func(t *testing.T) {
 		volumeList, err := getPersistentVolumesList()
@@ -230,6 +233,77 @@ func HasExpectedPersistentVolumes(volumes []corev1.PersistentVolume) func(t *tes
 		}
 	}
 }
+func HasExpectedMetadata(mdb *mdbv1.MongoDBCommunity, expectedLabels map[string]string, expectedAnnotations map[string]string) func(t *testing.T) {
+	return func(t *testing.T) {
+		namespace := mdb.Namespace
+
+		statefulSetList := appsv1.StatefulSetList{}
+		err := e2eutil.TestClient.Client.List(context.TODO(), &statefulSetList, client.InNamespace(namespace))
+		assert.NoError(t, err)
+		assert.NotEmpty(t, statefulSetList.Items)
+		for _, s := range statefulSetList.Items {
+			containsMetadata(t, &s.ObjectMeta, expectedLabels, expectedAnnotations, "statefulset "+s.Name)
+		}
+
+		volumeList := corev1.PersistentVolumeList{}
+		err = e2eutil.TestClient.Client.List(context.TODO(), &volumeList, client.InNamespace(namespace))
+		assert.NoError(t, err)
+		assert.NotEmpty(t, volumeList.Items)
+		for _, s := range volumeList.Items {
+			volName := s.Name
+			if strings.HasPrefix(volName, "data-volume-") || strings.HasPrefix(volName, "logs-volume-") {
+				containsMetadata(t, &s.ObjectMeta, expectedLabels, expectedAnnotations, "volume "+volName)
+			}
+		}
+
+		podList := corev1.PodList{}
+		err = e2eutil.TestClient.Client.List(context.TODO(), &podList, client.InNamespace(namespace))
+		assert.NoError(t, err)
+		assert.NotEmpty(t, podList.Items)
+
+		for _, s := range podList.Items {
+			// only consider stateful-sets (as opposite to the controller replica set)
+			for _, owner := range s.OwnerReferences {
+				if owner.Kind == "ReplicaSet" {
+					continue
+				}
+			}
+			// Ignore non-owned pods
+			if len(s.OwnerReferences) == 0 {
+				continue
+			}
+
+			// Ensure we are considering pods owned by a stateful set
+			hasStatefulSetOwner := false
+			for _, owner := range s.OwnerReferences {
+				if owner.Kind == "StatefulSet" {
+					hasStatefulSetOwner = true
+				}
+			}
+			if !hasStatefulSetOwner {
+				continue
+			}
+
+			containsMetadata(t, &s.ObjectMeta, expectedLabels, expectedAnnotations, "pod "+s.Name)
+		}
+	}
+}
+
+func containsMetadata(t *testing.T, metadata *metav1.ObjectMeta, expectedLabels map[string]string, expectedAnnotations map[string]string, msg string) {
+	labels := metadata.Labels
+	for k, v := range expectedLabels {
+		assert.Contains(t, labels, k, msg+" has label "+k)
+		value := labels[k]
+		assert.Equal(t, v, value, msg+" has label "+k+" with value "+v)
+	}
+
+	annotations := metadata.Annotations
+	for k, v := range expectedAnnotations {
+		assert.Contains(t, annotations, k, msg+" has annotation "+k)
+		value := annotations[k]
+		assert.Equal(t, v, value, msg+" has annotation "+k+" with value "+v)
+	}
+}
 
 // MongoDBReachesPendingPhase ensures the MongoDB resources gets to the Pending phase
 func MongoDBReachesPendingPhase(mdb *mdbv1.MongoDBCommunity) func(t *testing.T) {