[chart/v1-2x-test] Add workers.celery.extraInitContainers & workers.kubernetes.extraInitContainers (#64741) (#64979)

* Refactor workers extraInitContainers-related tests

* Add workers.celery.extraInitContainers & workers.kubernetes.extraInitContainers

* Add newsfragment
(cherry picked from commit 52c4d32)

Co-authored-by: Przemysław Mirowski <17602603+Miretpl@users.noreply.github.com>

Author: jscheffl

chart/docs/using-additional-containers.rst

@@ -52,7 +52,7 @@ Init Containers
 ---------------
 
 You can also deploy extra init containers through the ``extraInitContainers`` parameter.
-You can define different containers for the scheduler, webserver, api server, worker, triggerer, dag processor, create user job and migrate database job pods.
+You can define different containers for the scheduler, webserver/api-server, Celery/Kubernetes workers, triggerer, dag processor, create user job and migrate database job pods.
 
 For example, an init container that just says hello:
 

chart/files/pod-template-file.kubernetes-helm-yaml

@@ -52,8 +52,8 @@ spec:
     {{- if and .Values.dags.gitSync.enabled (not .Values.dags.persistence.enabled) }}
       {{- include "git_sync_container" (dict "Values" .Values "is_init" "true" "Template" .Template) | nindent 4 }}
     {{- end }}
-    {{- if .Values.workers.extraInitContainers }}
-      {{- tpl (toYaml .Values.workers.extraInitContainers) . | nindent 4 }}
+    {{- if or .Values.workers.kubernetes.extraInitContainers .Values.workers.extraInitContainers }}
+      {{- tpl (toYaml (.Values.workers.kubernetes.extraInitContainers | default .Values.workers.extraInitContainers)) . | nindent 4 }}
     {{- end }}
     {{- if or .Values.workers.kubernetes.kerberosInitContainer.enabled .Values.workers.kerberosInitContainer.enabled }}
     - name: kerberos-init

chart/newsfragments/64741.significant.rst

@@ -0,0 +1 @@
+``workers.extraInitContainers`` field is now deprecated in favor of ``workers.celery.extraInitContainers`` and ``workers.kubernetes.extraInitContainers``. Please update your configuration accordingly.

chart/templates/NOTES.txt

@@ -709,6 +709,14 @@ DEPRECATION WARNING:
 
 {{- end }}
 
+{{- if not (empty .Values.workers.extraInitContainers) }}
+
+ DEPRECATION WARNING:
+    `workers.extraInitContainers` has been renamed to `workers.celery.extraInitContainers`/`workers.kubernetes.extraInitContainers`.
+    Please change your values as support for the old name will be dropped in a future release.
+
+{{- end }}
+
 {{- if not (empty .Values.workers.runtimeClassName) }}
 
  DEPRECATION WARNING:

chart/values.schema.json

@@ -2301,7 +2301,7 @@
                     }
                 },
                 "extraInitContainers": {
-                    "description": "Add additional init containers into Airflow Celery workers and pods created with pod-template-file (templated).",
+                    "description": "Add additional init containers into Airflow Celery workers and pods created with pod-template-file (templated) (deprecated, use ``workers.celery.extraInitContainers`` and/or ``workers.kubernetes.extraInitContainers`` instead).",
                     "type": "array",
                     "default": [],
                     "items": {
@@ -3388,6 +3388,14 @@
                                 "$ref": "#/definitions/io.k8s.api.core.v1.Container"
                             }
                         },
+                        "extraInitContainers": {
+                            "description": "Add additional init containers into Airflow Celery workers (templated).",
+                            "type": "array",
+                            "default": [],
+                            "items": {
+                                "$ref": "#/definitions/io.k8s.api.core.v1.Container"
+                            }
+                        },
                         "extraPorts": {
                             "description": "Expose additional ports of Airflow Celery worker container.",
                             "type": "array",
@@ -3886,6 +3894,14 @@
                                 "$ref": "#/definitions/io.k8s.api.core.v1.Container"
                             }
                         },
+                        "extraInitContainers": {
+                            "description": "Add additional init containers into pods created with pod-template-file (templated).",
+                            "type": "array",
+                            "default": [],
+                            "items": {
+                                "$ref": "#/definitions/io.k8s.api.core.v1.Container"
+                            }
+                        },
                         "nodeSelector": {
                             "description": "Select certain nodes for pods created with pod-template-file.",
                             "type": "object",

chart/values.yaml

@@ -1052,6 +1052,10 @@ workers:
 
   # Add additional init containers into Airflow Celery workers
   # and pods created with pod-template-file (templated).
+  # (deprecated, use
+  #   `workers.celery.extraInitContainers` and/or
+  #   `workers.kubernetes.extraInitContainers`
+  # instead)
   extraInitContainers: []
 
   # Additional volumes and volume mounts attached to the
@@ -1450,6 +1454,9 @@ workers:
     # Launch additional containers into Airflow Celery worker (templated)
     extraContainers: []
 
+    # Add additional init containers into Airflow Celery workers (templated)
+    extraInitContainers: []
+
     # Expose additional ports of Airflow Celery workers. These can be used for additional metric collection.
     extraPorts: []
 
@@ -1617,6 +1624,9 @@ workers:
     # container finishes so Airflow can continue the worker shutdown process!
     extraContainers: []
 
+    # Add additional init containers into pods created with pod-template-file (templated)
+    extraInitContainers: []
+
     # Select certain nodes for pods created with pod-template-file
     nodeSelector: {}
 

helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py

@@ -1067,38 +1067,68 @@ def test_airflow_and_workers_pod_annotations(self):
         assert "my_annotation" in annotations
         assert "workerPodAnnotations" in annotations["my_annotation"]
 
-    def test_should_add_extra_init_containers(self):
-        docs = render_chart(
-            values={
-                "workers": {
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {
+                "extraInitContainers": [
+                    {"name": "test-init-container", "image": "test-registry/test-repo:test-tag"}
+                ]
+            },
+            {
+                "kubernetes": {
                     "extraInitContainers": [
                         {"name": "test-init-container", "image": "test-registry/test-repo:test-tag"}
-                    ],
+                    ]
+                }
+            },
+            {
+                "extraInitContainers": [{"name": "test", "image": "repo:tag"}],
+                "kubernetes": {
+                    "extraInitContainers": [
+                        {"name": "test-init-container", "image": "test-registry/test-repo:test-tag"}
+                    ]
                 },
             },
+        ],
+    )
+    def test_should_add_extra_init_containers(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
             show_only=["templates/pod-template-file.yaml"],
             chart_dir=self.temp_chart_dir,
         )
 
-        assert jmespath.search("spec.initContainers[-1]", docs[0]) == {
-            "name": "test-init-container",
-            "image": "test-registry/test-repo:test-tag",
-        }
+        assert jmespath.search("spec.initContainers", docs[0]) == [
+            {
+                "name": "test-init-container",
+                "image": "test-registry/test-repo:test-tag",
+            }
+        ]
 
-    def test_should_template_extra_init_containers(self):
-        docs = render_chart(
-            values={
-                "workers": {
-                    "extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}],
-                },
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}]},
+            {"kubernetes": {"extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}]}},
+            {
+                "extraInitContainers": [{"name": "container"}],
+                "kubernetes": {"extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}]},
             },
+        ],
+    )
+    def test_should_template_extra_init_containers(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
             show_only=["templates/pod-template-file.yaml"],
             chart_dir=self.temp_chart_dir,
         )
 
-        assert jmespath.search("spec.initContainers[-1]", docs[0]) == {
-            "name": "release-name-test-init-container",
-        }
+        assert jmespath.search("spec.initContainers", docs[0]) == [
+            {
+                "name": "release-name-test-init-container",
+            }
+        ]
 
     @pytest.mark.parametrize(
         "workers_values",

helm-tests/tests/helm_tests/airflow_core/test_worker.py

@@ -279,36 +279,66 @@ def test_logs_mount_on_wait_for_migrations_initcontainer(self, logs_values, expe
                 for m in mounts
             )
 
-    def test_should_add_extra_init_containers(self):
-        docs = render_chart(
-            values={
-                "workers": {
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {
+                "extraInitContainers": [
+                    {"name": "test-init-container", "image": "test-registry/test-repo:test-tag"}
+                ]
+            },
+            {
+                "celery": {
                     "extraInitContainers": [
                         {"name": "test-init-container", "image": "test-registry/test-repo:test-tag"}
-                    ],
+                    ]
+                }
+            },
+            {
+                "extraInitContainers": [{"name": "container", "image": "repo:tag"}],
+                "celery": {
+                    "extraInitContainers": [
+                        {"name": "test-init-container", "image": "test-registry/test-repo:test-tag"}
+                    ]
                 },
             },
+        ],
+    )
+    def test_should_add_extra_init_containers(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
             show_only=["templates/workers/worker-deployment.yaml"],
         )
 
-        assert jmespath.search("spec.template.spec.initContainers[-1]", docs[0]) == {
-            "name": "test-init-container",
-            "image": "test-registry/test-repo:test-tag",
-        }
+        # [1:] -> Skipping wait-for-airflow-migrations init container
+        assert jmespath.search("spec.template.spec.initContainers[1:]", docs[0]) == [
+            {
+                "name": "test-init-container",
+                "image": "test-registry/test-repo:test-tag",
+            }
+        ]
 
-    def test_should_template_extra_init_containers(self):
-        docs = render_chart(
-            values={
-                "workers": {
-                    "extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}],
-                },
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}]},
+            {"celery": {"extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}]}},
+            {
+                "extraInitContainers": [{"name": "container"}],
+                "celery": {"extraInitContainers": [{"name": "{{ .Release.Name }}-test-init-container"}]},
             },
+        ],
+    )
+    def test_should_template_extra_init_containers(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
             show_only=["templates/workers/worker-deployment.yaml"],
         )
 
-        assert jmespath.search("spec.template.spec.initContainers[-1]", docs[0]) == {
-            "name": "release-name-test-init-container"
-        }
+        # [1:] -> Skipping wait-for-airflow-migrations init container
+        assert jmespath.search("spec.template.spec.initContainers[1:]", docs[0]) == [
+            {"name": "release-name-test-init-container"}
+        ]
 
     def test_should_add_extra_volume_and_extra_volume_mount(self):
         docs = render_chart(
@@ -940,19 +970,24 @@ def test_extra_init_container_restart_policy_is_configurable(self):
         docs = render_chart(
             values={
                 "workers": {
-                    "extraInitContainers": [
-                        {
-                            "name": "test-init-container",
-                            "image": "test-registry/test-repo:test-tag",
-                            "restartPolicy": "Always",
-                        }
-                    ]
+                    "celery": {
+                        "extraInitContainers": [
+                            {
+                                "name": "test-init-container",
+                                "image": "test-registry/test-repo:test-tag",
+                                "restartPolicy": "Always",
+                            }
+                        ]
+                    }
                 },
             },
             show_only=["templates/workers/worker-deployment.yaml"],
         )
 
-        assert jmespath.search("spec.template.spec.initContainers[1].restartPolicy", docs[0]) == "Always"
+        # [1:] -> Skipping wait-for-airflow-migrations init container
+        assert jmespath.search("spec.template.spec.initContainers[1:] | [*].restartPolicy", docs[0]) == [
+            "Always"
+        ]
 
     @pytest.mark.parametrize(
         ("log_values", "expected_volume"),

helm-tests/tests/helm_tests/airflow_core/test_worker_sets.py

@@ -2513,6 +2513,20 @@ def test_overwrite_extra_containers(self, workers_values):
                     ],
                 },
             },
+            {
+                "celery": {
+                    "enableDefault": False,
+                    "extraInitContainers": [{"name": "test", "image": "test"}],
+                    "sets": [
+                        {
+                            "name": "set1",
+                            "extraInitContainers": [
+                                {"name": "{{ .Chart.Name }}", "image": "test-registry/test-repo:test-tag"}
+                            ],
+                        }
+                    ],
+                },
+            },
         ],
     )
     def test_overwrite_extra_init_containers(self, workers_values):