From 44345f3a635d3aef3bf98d6a3134e8820564b105 Mon Sep 17 00:00:00 2001
From: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com>
Date: Tue, 25 May 2021 14:07:41 -0600
Subject: [PATCH] Support ``strategy``/``updateStrategy`` on scheduler (#16069)

Allow the schedulers strategy (for Deployment) or
updateStrategy (for StatefulSet) to be set via helm parameters.
---
 .../scheduler/scheduler-deployment.yaml       |  8 +++
 chart/tests/test_scheduler.py                 | 53 +++++++++++++++++++
 chart/values.schema.json                      | 16 ++++++
 chart/values.yaml                             |  6 +++
 4 files changed, 83 insertions(+)

diff --git a/chart/templates/scheduler/scheduler-deployment.yaml b/chart/templates/scheduler/scheduler-deployment.yaml
index da67975a0ee56..8137c846ccbbf 100644
--- a/chart/templates/scheduler/scheduler-deployment.yaml
+++ b/chart/templates/scheduler/scheduler-deployment.yaml
@@ -50,6 +50,14 @@ spec:
   serviceName: {{ .Release.Name }}-scheduler
 {{- end }}
   replicas: {{ .Values.scheduler.replicas }}
+  {{- if and $stateful .Values.scheduler.updateStrategy }}
+  updateStrategy:
+    {{- toYaml .Values.scheduler.updateStrategy | nindent 4 }}
+  {{- end }}
+  {{- if and (not $stateful) .Values.scheduler.strategy }}
+  strategy:
+    {{- toYaml .Values.scheduler.strategy | nindent 4 }}
+  {{- end }}
   selector:
     matchLabels:
       tier: airflow
diff --git a/chart/tests/test_scheduler.py b/chart/tests/test_scheduler.py
index 0e9cc60567c62..fd3e198518cbe 100644
--- a/chart/tests/test_scheduler.py
+++ b/chart/tests/test_scheduler.py
@@ -215,3 +215,56 @@ def test_airflow_local_settings(self):
             "subPath": "airflow_local_settings.py",
             "readOnly": True,
         } in jmespath.search("spec.template.spec.containers[0].volumeMounts", docs[0])
+
+    @parameterized.expand(
+        [
+            ("CeleryExecutor", False, {"rollingUpdate": {"partition": 0}}, None),
+            ("CeleryExecutor", True, {"rollingUpdate": {"partition": 0}}, None),
+            ("LocalExecutor", False, {"rollingUpdate": {"partition": 0}}, None),
+            ("LocalExecutor", True, {"rollingUpdate": {"partition": 0}}, {"rollingUpdate": {"partition": 0}}),
+            ("LocalExecutor", True, None, None),
+        ]
+    )
+    def test_scheduler_update_strategy(
+        self, executor, persistence, update_strategy, expected_update_strategy
+    ):
+        """updateStrategy should only be used when we have LocalExecutor and workers.persistence"""
+        docs = render_chart(
+            values={
+                "executor": executor,
+                "workers": {"persistence": {"enabled": persistence}},
+                "scheduler": {"updateStrategy": update_strategy},
+            },
+            show_only=["templates/scheduler/scheduler-deployment.yaml"],
+        )
+
+        assert expected_update_strategy == jmespath.search("spec.updateStrategy", docs[0])
+
+    @parameterized.expand(
+        [
+            ("LocalExecutor", False, None, None),
+            ("LocalExecutor", False, {"type": "Recreate"}, {"type": "Recreate"}),
+            ("LocalExecutor", True, {"type": "Recreate"}, None),
+            ("CeleryExecutor", True, None, None),
+            ("CeleryExecutor", False, None, None),
+            ("CeleryExecutor", True, {"type": "Recreate"}, {"type": "Recreate"}),
+            (
+                "CeleryExecutor",
+                False,
+                {"rollingUpdate": {"maxSurge": "100%", "maxUnavailable": "50%"}},
+                {"rollingUpdate": {"maxSurge": "100%", "maxUnavailable": "50%"}},
+            ),
+        ]
+    )
+    def test_scheduler_strategy(self, executor, persistence, strategy, expected_strategy):
+        """strategy should be used when we aren't using both LocalExecutor and workers.persistence"""
+        docs = render_chart(
+            values={
+                "executor": executor,
+                "workers": {"persistence": {"enabled": persistence}},
+                "scheduler": {"strategy": strategy},
+            },
+            show_only=["templates/scheduler/scheduler-deployment.yaml"],
+        )
+
+        assert expected_strategy == jmespath.search("spec.strategy", docs[0])
diff --git a/chart/values.schema.json b/chart/values.schema.json
index 88e64067422fd..9003d5ec18db1 100644
--- a/chart/values.schema.json
+++ b/chart/values.schema.json
@@ -1088,6 +1088,22 @@
                     "type": "integer",
                     "default": 1
                 },
+                "updateStrategy": {
+                    "description": "Specifies the strategy used to replace old Pods by new ones when deployed as a StatefulSet (when using LocalExecutor and workers.persistence).",
+                    "type": [
+                        "null",
+                        "object"
+                    ],
+                    "default": null
+                },
+                "strategy": {
+                    "description": "Specifies the strategy used to replace old Pods by new ones when deployed as a Deployment (when not using LocalExecutor and workers.persistence).",
+                    "type": [
+                        "null",
+                        "object"
+                    ],
+                    "default": null
+                },
                 "serviceAccount": {
                     "description": "Create ServiceAccount.",
                     "type": "object",
diff --git a/chart/values.yaml b/chart/values.yaml
index 89c8e5f91b54a..8cd3a75a1e4b1 100644
--- a/chart/values.yaml
+++ b/chart/values.yaml
@@ -429,6 +429,12 @@ scheduler:
   # Airflow 2.0 allows users to run multiple schedulers,
   # However this feature is only recommended for MySQL 8+ and Postgres
   replicas: 1
+  # Update Strategy when scheduler is deployed as a StatefulSet
+  # (when using LocalExecutor and workers.persistence)
+  updateStrategy: ~
+  # Update Strategy when scheduler is deployed as a Deployment
+  # (when not using LocalExecutor and workers.persistence)
+  strategy: ~
 
   # Create ServiceAccount
   serviceAccount: