Add pdbs for triggerer and workers

chart/templates/triggerer/triggerer-poddisruptionbudget.yaml

@@ -0,0 +1,44 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements.  See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership.  The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License.  You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied.  See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+################################
+## Airflow Triggerer PodDisruptionBudget
+#################################
+{{- if and .Values.triggerer.enabled .Values.triggerer.podDisruptionBudget.enabled }}
+apiVersion: policy/v1
+kind: PodDisruptionBudget
+metadata:
+  name: {{ include "airflow.fullname" . }}-triggerer-pdb
+  labels:
+    tier: airflow
+    component: triggerer
+    release: {{ .Release.Name }}
+    chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+    heritage: {{ .Release.Service }}
+    {{- if or .Values.labels .Values.triggerer.labels }}
+      {{- mustMerge .Values.triggerer.labels .Values.labels | toYaml | nindent 4 }}
+    {{- end }}
+spec:
+  selector:
+    matchLabels:
+      tier: airflow
+      component: triggerer
+      release: {{ .Release.Name }}
+  {{- toYaml .Values.triggerer.podDisruptionBudget.config | nindent 2 }}
+{{- end }}

chart/templates/workers/worker-poddisruptionbudget.yaml

@@ -0,0 +1,44 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements.  See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership.  The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License.  You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied.  See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+################################
+## Airflow Worker PodDisruptionBudget
+#################################
+{{- if .Values.workers.podDisruptionBudget.enabled }}
+apiVersion: policy/v1
+kind: PodDisruptionBudget
+metadata:
+  name: {{ include "airflow.fullname" . }}-worker-pdb
+  labels:
+    tier: airflow
+    component: worker
+    release: {{ .Release.Name }}
+    chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+    heritage: {{ .Release.Service }}
+    {{- if or .Values.labels .Values.workers.labels }}
+      {{- mustMerge .Values.workers.labels .Values.labels | toYaml | nindent 4 }}
+    {{- end }}
+spec:
+  selector:
+    matchLabels:
+      tier: airflow
+      component: worker
+      release: {{ .Release.Name }}
+  {{- toYaml .Values.workers.podDisruptionBudget.config | nindent 2 }}
+{{- end }}

chart/values.schema.json

@@ -2160,6 +2160,41 @@
                         }
                     }
                 },
+                "podDisruptionBudget": {
+                    "description": "Worker pod disruption budget.",
+                    "type": "object",
+                    "additionalProperties": true,
+                    "properties": {
+                        "enabled": {
+                            "description": "Enable pod disruption budget.",
+                            "type": "boolean",
+                            "default": false
+                        },
+                        "config": {
+                            "description": "Disruption budget configuration.",
+                            "type": "object",
+                            "additionalProperties": true,
+                            "properties": {
+                                "maxUnavailable": {
+                                    "description": "Max unavailable pods for worker.",
+                                    "type": [
+                                        "integer",
+                                        "string"
+                                    ],
+                                    "default": 1
+                                },
+                                "minAvailable": {
+                                    "description": "Min available pods for worker.",
+                                    "type": [
+                                        "integer",
+                                        "string"
+                                    ],
+                                    "default": 1
+                                }
+                            }
+                        }
+                    }
+                },
                 "resources": {
                     "description": "Resource configuration for Airflow Celery workers and pods created with pod-template-file.",
                     "type": "object",
@@ -3454,6 +3489,41 @@
                         }
                     }
                 },
+                "podDisruptionBudget": {
+                    "description": "Triggerer pod disruption budget.",
+                    "type": "object",
+                    "additionalProperties": true,
+                    "properties": {
+                        "enabled": {
+                            "description": "Enable pod disruption budget.",
+                            "type": "boolean",
+                            "default": false
+                        },
+                        "config": {
+                            "description": "Disruption budget configuration.",
+                            "type": "object",
+                            "additionalProperties": true,
+                            "properties": {
+                                "maxUnavailable": {
+                                    "description": "Max unavailable pods for triggerer.",
+                                    "type": [
+                                        "integer",
+                                        "string"
+                                    ],
+                                    "default": 1
+                                },
+                                "minAvailable": {
+                                    "description": "Min available pods for triggerer.",
+                                    "type": [
+                                        "integer",
+                                        "string"
+                                    ],
+                                    "default": 1
+                                }
+                            }
+                        }
+                    }
+                },
                 "resources": {
                     "description": "Resources for triggerer pods.",
                     "type": "object",

chart/values.yaml

@@ -692,6 +692,16 @@ workers:
   # Airflow Celery workers and pods created with pod-template-file
   containerLifecycleHooks: {}
 
+  # Worker pod disruption budget
+  podDisruptionBudget:
+    enabled: false
+
+    # PDB configuration
+    config:
+      # minAvailable and maxUnavailable are mutually exclusive
+      maxUnavailable: 1
+      # minAvailable: 1
+
   # Create ServiceAccount for Airflow Celery workers and pods created with pod-template-file
   serviceAccount:
     # default value is true
@@ -1934,6 +1944,16 @@ triggerer:
     # Annotations to add to triggerer volumes
     annotations: {}
 
+  # Triggerer pod disruption budget
+  podDisruptionBudget:
+    enabled: false
+
+    # PDB configuration
+    config:
+      # minAvailable and maxUnavailable are mutually exclusive
+      maxUnavailable: 1
+      # minAvailable: 1
+
   resources: {}
   #  limits:
   #   cpu: 100m

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

@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import jmespath
+from chart_utils.helm_template_generator import render_chart
+
+
+class TestTriggererPdb:
+    """Tests Triggerer PDB."""
+
+    def test_should_pass_validation_with_just_pdb_enabled_v1(self):
+        render_chart(
+            values={"triggerer": {"podDisruptionBudget": {"enabled": True}}},
+            show_only=["templates/triggerer/triggerer-poddisruptionbudget.yaml"],
+        )  # checks that no validation exception is raised
+
+    def test_should_add_component_specific_labels(self):
+        docs = render_chart(
+            values={
+                "triggerer": {
+                    "podDisruptionBudget": {"enabled": True},
+                    "labels": {"test_label": "test_label_value"},
+                },
+            },
+            show_only=["templates/triggerer/triggerer-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={
+                "triggerer": {
+                    "podDisruptionBudget": {
+                        "enabled": True,
+                        "config": {"maxUnavailable": None, "minAvailable": 1},
+                    }
+                }
+            },
+            show_only=["templates/triggerer/triggerer-poddisruptionbudget.yaml"],
+        )  # checks that no validation exception is raised

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

@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import jmespath
+from chart_utils.helm_template_generator import render_chart
+
+
+class TestWorkerPdb:
+    """Tests Worker PDB."""
+
+    def test_should_pass_validation_with_just_pdb_enabled_v1(self):
+        render_chart(
+            values={"workers": {"podDisruptionBudget": {"enabled": True}}},
+            show_only=["templates/workers/worker-poddisruptionbudget.yaml"],
+        )  # checks that no validation exception is raised
+
+    def test_should_add_component_specific_labels(self):
+        docs = render_chart(
+            values={
+                "workers": {
+                    "podDisruptionBudget": {"enabled": True},
+                    "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": {
+                    "podDisruptionBudget": {
+                        "enabled": True,
+                        "config": {"maxUnavailable": None, "minAvailable": 1},
+                    },
+                }
+            },
+            show_only=["templates/workers/worker-poddisruptionbudget.yaml"],
+        )  # checks that no validation exception is raised