Possibility to set labels specific to the migrateDatabaseJob objects …

…and pods (apache#37490) --------- Co-authored-by: cyla <cyril@numberly.com>
abhishekbhakat · Mar 5, 2024 · b5da762 · b5da762
1 parent b61b157
commit b5da762
Show file tree

Hide file tree

Showing 5 changed files with 78 additions and 4 deletions.
diff --git a/chart/templates/jobs/migrate-database-job-serviceaccount.yaml b/chart/templates/jobs/migrate-database-job-serviceaccount.yaml
@@ -32,8 +32,8 @@ metadata:
     release: {{ .Release.Name }}
     chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
     heritage: {{ .Release.Service }}
-    {{- with .Values.labels }}
-      {{- toYaml . | nindent 4 }}
+    {{- if or (.Values.labels) (.Values.migrateDatabaseJob.labels) }}
+      {{- mustMerge .Values.migrateDatabaseJob.labels .Values.labels | toYaml | nindent 4 }}
     {{- end }}
   {{- with .Values.migrateDatabaseJob.serviceAccount.annotations }}
   annotations: {{- toYaml . | nindent 4 }}

diff --git a/chart/templates/jobs/migrate-database-job.yaml b/chart/templates/jobs/migrate-database-job.yaml
@@ -60,8 +60,8 @@ spec:
         tier: airflow
         component: run-airflow-migrations
         release: {{ .Release.Name }}
-        {{- with .Values.labels }}
-          {{- toYaml . | nindent 8 }}
+        {{- if or (.Values.labels) (.Values.migrateDatabaseJob.labels) }}
+          {{- mustMerge .Values.migrateDatabaseJob.labels .Values.labels | toYaml | nindent 8 }}
         {{- end }}
       {{- if or .Values.airflowPodAnnotations .Values.migrateDatabaseJob.annotations }}
       annotations:

diff --git a/chart/values.schema.json b/chart/values.schema.json
@@ -4005,6 +4005,14 @@
                         "type": "string"
                     }
                 },
+                "labels": {
+                    "description": "Labels to add to the migrate database job objects and pods.",
+                    "type": "object",
+                    "default": {},
+                    "additionalProperties": {
+                        "type": "string"
+                    }
+                },
                 "serviceAccount": {
                     "description": "Create ServiceAccount.",
                     "type": "object",

diff --git a/chart/values.yaml b/chart/values.yaml
@@ -1086,6 +1086,9 @@ migrateDatabaseJob:
   # jobAnnotations are annotations on the database migration job
   jobAnnotations: {}
 
+  # Labels specific to migrate database job objects and pods
+  labels: {}
+
   # When not set, the values defined in the global securityContext will be used
   securityContext: {}
   #  runAsUser: 50000

diff --git a/helm_tests/airflow_aux/test_migrate_database_job.py b/helm_tests/airflow_aux/test_migrate_database_job.py
@@ -60,6 +60,39 @@ def test_should_support_annotations(self):
         assert "fiz" in job_annotations
         assert "fuz" == job_annotations["fiz"]
 
+    def test_should_add_component_specific_labels(self):
+        docs = render_chart(
+            values={
+                "migrateDatabaseJob": {
+                    "labels": {"test_label": "test_label_value"},
+                },
+            },
+            show_only=["templates/jobs/migrate-database-job.yaml"],
+        )
+        assert "test_label" in jmespath.search("spec.template.metadata.labels", docs[0])
+        assert jmespath.search("spec.template.metadata.labels", docs[0])["test_label"] == "test_label_value"
+
+    def test_should_merge_common_labels_and_component_specific_labels(self):
+        docs = render_chart(
+            values={
+                "labels": {"test_common_label": "test_common_label_value"},
+                "migrateDatabaseJob": {
+                    "labels": {"test_specific_label": "test_specific_label_value"},
+                },
+            },
+            show_only=["templates/jobs/migrate-database-job.yaml"],
+        )
+        assert "test_common_label" in jmespath.search("spec.template.metadata.labels", docs[0])
+        assert (
+            jmespath.search("spec.template.metadata.labels", docs[0])["test_common_label"]
+            == "test_common_label_value"
+        )
+        assert "test_specific_label" in jmespath.search("spec.template.metadata.labels", docs[0])
+        assert (
+            jmespath.search("spec.template.metadata.labels", docs[0])["test_specific_label"]
+            == "test_specific_label_value"
+        )
+
     def test_should_create_valid_affinity_tolerations_and_node_selector(self):
         docs = render_chart(
             values={
@@ -340,6 +373,36 @@ def test_airflow_local_settings(self):
 class TestMigrateDatabaseJobServiceAccount:
     """Tests migrate database job service account."""
 
+    def test_should_add_component_specific_labels(self):
+        docs = render_chart(
+            values={
+                "migrateDatabaseJob": {
+                    "labels": {"test_label": "test_label_value"},
+                },
+            },
+            show_only=["templates/jobs/migrate-database-job-serviceaccount.yaml"],
+        )
+
+        assert "test_label" in jmespath.search("metadata.labels", docs[0])
+        assert jmespath.search("metadata.labels", docs[0])["test_label"] == "test_label_value"
+
+    def test_should_merge_common_labels_and_component_specific_labels(self):
+        docs = render_chart(
+            values={
+                "labels": {"test_common_label": "test_common_label_value"},
+                "migrateDatabaseJob": {
+                    "labels": {"test_specific_label": "test_specific_label_value"},
+                },
+            },
+            show_only=["templates/jobs/migrate-database-job-serviceaccount.yaml"],
+        )
+        assert "test_common_label" in jmespath.search("metadata.labels", docs[0])
+        assert jmespath.search("metadata.labels", docs[0])["test_common_label"] == "test_common_label_value"
+        assert "test_specific_label" in jmespath.search("metadata.labels", docs[0])
+        assert (
+            jmespath.search("metadata.labels", docs[0])["test_specific_label"] == "test_specific_label_value"
+        )
+
     def test_default_automount_service_account_token(self):
         docs = render_chart(
             values={