Add workers.celery.podDisruptionBudget

chart/newsfragments/61414.significant.rst

@@ -0,0 +1 @@
+``workers.podDisruptionBudget`` section is now deprecated in favor of ``workers.celery.podDisruptionBudget``. Please update your configuration accordingly.

chart/templates/NOTES.txt

@@ -375,6 +375,30 @@ DEPRECATION WARNING:
 
 {{- end }}
 
+{{- if .Values.workers.podDisruptionBudget.enabled }}
+
+ DEPRECATION WARNING:
+    `workers.podDisruptionBudget.enabled` has been renamed to `workers.celery.podDisruptionBudget.enabled`.
+    Please change your values as support for the old name will be dropped in a future release.
+
+{{- end }}
+
+{{- if ne (int .Values.workers.podDisruptionBudget.config.maxUnavailable) 1 }}
+
+ DEPRECATION WARNING:
+    `workers.podDisruptionBudget.config.maxUnavailable` has been renamed to `workers.celery.podDisruptionBudget.config.maxUnavailable`.
+    Please change your values as support for the old name will be dropped in a future release.
+
+{{- end }}
+
+{{- if hasKey .Values.workers.podDisruptionBudget.config "minUnavailable" }}
+
+ DEPRECATION WARNING:
+    `workers.podDisruptionBudget.config.minUnavailable` has been renamed to `workers.celery.podDisruptionBudget.config.minUnavailable`.
+    Please change your values as support for the old name will be dropped in a future release.
+
+{{- end }}
+
 {{- if not (empty .Values.webserver.defaultUser) }}
 
  DEPRECATION WARNING:

chart/values.schema.json

@@ -2229,30 +2229,30 @@
                     }
                 },
                 "podDisruptionBudget": {
-                    "description": "Worker pod disruption budget.",
+                    "description": "Worker pod disruption budget (deprecated, use `workers.celery.podDisruptionBudget` instead).",
                     "type": "object",
                     "additionalProperties": true,
                     "properties": {
                         "enabled": {
-                            "description": "Enable pod disruption budget.",
+                            "description": "Enable pod disruption budget (deprecated, use `workers.celery.podDisruptionBudget.enabled` instead).",
                             "type": "boolean",
                             "default": false
                         },
                         "config": {
-                            "description": "Disruption budget configuration.",
+                            "description": "Disruption budget configuration (deprecated, use `workers.celery.podDisruptionBudget.config` instead).",
                             "type": "object",
                             "additionalProperties": true,
                             "properties": {
                                 "maxUnavailable": {
-                                    "description": "Max unavailable pods for worker.",
+                                    "description": "Max unavailable pods for worker (deprecated, use `workers.celery.podDisruptionBudget.config.maxUnavailable` instead).",
                                     "type": [
                                         "integer",
                                         "string"
                                     ],
                                     "default": 1
                                 },
                                 "minAvailable": {
-                                    "description": "Min available pods for worker.",
+                                    "description": "Min available pods for worker (deprecated, use `workers.celery.podDisruptionBudget.config.minAvailable` instead).",
                                     "type": [
                                         "integer",
                                         "string"
@@ -2879,6 +2879,46 @@
                                 }
                             ]
                         },
+                        "podDisruptionBudget": {
+                            "description": "Airflow Celery worker pod disruption budget.",
+                            "type": "object",
+                            "additionalProperties": true,
+                            "properties": {
+                                "enabled": {
+                                    "description": "Enable pod disruption budget.",
+                                    "type": [
+                                        "boolean",
+                                        "null"
+                                    ],
+                                    "default": null
+                                },
+                                "config": {
+                                    "description": "Disruption budget configuration.",
+                                    "type": "object",
+                                    "additionalProperties": true,
+                                    "properties": {
+                                        "maxUnavailable": {
+                                            "description": "Max unavailable pods for worker.",
+                                            "type": [
+                                                "integer",
+                                                "string",
+                                                "null"
+                                            ],
+                                            "default": null
+                                        },
+                                        "minAvailable": {
+                                            "description": "Min available pods for worker.",
+                                            "type": [
+                                                "integer",
+                                                "string",
+                                                "null"
+                                            ],
+                                            "default": null
+                                        }
+                                    }
+                                }
+                            }
+                        },
                         "persistence": {
                             "description": "Persistence configuration for Airflow Celery workers.",
                             "type": "object",

chart/values.yaml

@@ -703,14 +703,19 @@ workers:
   # to separate value between Celery workers and pod-template-file.
   containerLifecycleHooks: {}
 
-  # Worker pod disruption budget
+  # Airflow Celery workers pod disruption budget
+  # (deprecated, use `workers.celery.podDisruptionBudget` instead)
   podDisruptionBudget:
+    # (deprecated, use `workers.celery.podDisruptionBudget.enabled` instead)
     enabled: false
 
     # PDB configuration
+    # (deprecated, use `workers.celery.podDisruptionBudget.config` instead)
     config:
       # minAvailable and maxUnavailable are mutually exclusive
+      # (deprecated, use `workers.celery.podDisruptionBudget.config.maxUnavailable` instead)
       maxUnavailable: 1
+      # (deprecated, use `workers.celery.podDisruptionBudget.config.minAvailable` instead)
       # minAvailable: 1
 
   # Create ServiceAccount for Airflow Celery workers and pods created with pod-template-file
@@ -1110,6 +1115,16 @@ workers:
     # Container level Lifecycle Hooks definition for Airflow Celery workers
     containerLifecycleHooks: {}
 
+    # Airflow Celery workers pod disruption budget
+    podDisruptionBudget:
+      enabled: ~
+
+      # PDB configuration
+      config:
+        # minAvailable and maxUnavailable are mutually exclusive
+        maxUnavailable: ~
+        # minAvailable: ~
+
     # Persistence volume configuration for Airflow Celery workers
     persistence:
       # Enable persistent volumes

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

@@ -17,41 +17,127 @@
 from __future__ import annotations
 
 import jmespath
+import pytest
 from chart_utils.helm_template_generator import render_chart
 
 
 class TestWorkerPdb:
     """Tests Worker PDB."""
 
-    def test_should_pass_validation_with_just_pdb_enabled(self):
-        render_chart(
-            values={"workers": {"podDisruptionBudget": {"enabled": True}}},
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"podDisruptionBudget": {"enabled": True}},
+            {"celery": {"podDisruptionBudget": {"enabled": True}}},
+        ],
+    )
+    def test_pod_disruption_budget_enabled(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
             show_only=["templates/workers/worker-poddisruptionbudget.yaml"],
         )
 
-    def test_should_add_component_specific_labels(self):
+        assert len(docs) == 1
+
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"podDisruptionBudget": {"enabled": True}},
+            {"celery": {"podDisruptionBudget": {"enabled": True}}},
+        ],
+    )
+    def test_pod_disruption_budget_name(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
+            show_only=["templates/workers/worker-poddisruptionbudget.yaml"],
+        )
+
+        assert jmespath.search("metadata.name", docs[0]) == "release-name-worker-pdb"
+
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"podDisruptionBudget": {"enabled": True}},
+            {"celery": {"podDisruptionBudget": {"enabled": True}}},
+        ],
+    )
+    def test_should_add_component_specific_labels(self, workers_values):
         docs = render_chart(
             values={
                 "workers": {
-                    "podDisruptionBudget": {"enabled": True},
+                    **workers_values,
                     "labels": {"test_label": "test_label_value"},
                 },
             },
             show_only=["templates/workers/worker-poddisruptionbudget.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_pass_validation_with_pdb_enabled_and_min_available_param(self):
-        render_chart(
-            values={
-                "workers": {
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"podDisruptionBudget": {"enabled": True}},
+            {"celery": {"podDisruptionBudget": {"enabled": True}}},
+        ],
+    )
+    def test_pod_disruption_budget_config_max_unavailable_default(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
+            show_only=["templates/workers/worker-poddisruptionbudget.yaml"],
+        )
+
+        assert jmespath.search("spec.maxUnavailable", docs[0]) == 1
+        assert jmespath.search("spec.minAvailable", docs[0]) is None
+
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"podDisruptionBudget": {"enabled": True, "config": {"maxUnavailable": 2}}},
+            {"celery": {"podDisruptionBudget": {"enabled": True, "config": {"maxUnavailable": 2}}}},
+            {
+                "podDisruptionBudget": {"enabled": True, "config": {"maxUnavailable": 3}},
+                "celery": {"podDisruptionBudget": {"enabled": True, "config": {"maxUnavailable": 2}}},
+            },
+        ],
+    )
+    def test_pod_disruption_budget_config_max_unavailable_overwrite(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
+            show_only=["templates/workers/worker-poddisruptionbudget.yaml"],
+        )
+
+        assert jmespath.search("spec.maxUnavailable", docs[0]) == 2
+        assert jmespath.search("spec.minAvailable", docs[0]) is None
+
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"podDisruptionBudget": {"enabled": True, "config": {"maxUnavailable": None, "minAvailable": 1}}},
+            {
+                "celery": {
                     "podDisruptionBudget": {
                         "enabled": True,
                         "config": {"maxUnavailable": None, "minAvailable": 1},
-                    },
+                    }
                 }
             },
+            {
+                "podDisruptionBudget": {"enabled": True, "config": {"maxUnavailable": 2, "minAvailable": 3}},
+                "celery": {
+                    "podDisruptionBudget": {
+                        "enabled": True,
+                        "config": {"maxUnavailable": None, "minAvailable": 1},
+                    }
+                },
+            },
+        ],
+    )
+    def test_pod_disruption_budget_config_min_available_set(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
             show_only=["templates/workers/worker-poddisruptionbudget.yaml"],
-        )  # checks that no validation exception is raised
+        )
+
+        assert jmespath.search("spec.maxUnavailable", docs[0]) is None
+        assert jmespath.search("spec.minAvailable", docs[0]) == 1