From 681d2ba0128e3e6cbb044970ba97af8e2fb000f9 Mon Sep 17 00:00:00 2001
From: Kj Justin <55255608+kajanajustin@users.noreply.github.com>
Date: Mon, 5 Feb 2024 08:09:10 -0500
Subject: [PATCH] Added ability to enable/disable scheduler and webserver 
 (#36991)

---------

Co-authored-by: Elad Kalif <45845474+eladkal@users.noreply.github.com>
Co-authored-by: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com>
---
 .../scheduler/scheduler-deployment.yaml       |  3 ++-
 .../scheduler/scheduler-networkpolicy.yaml    |  2 ++
 .../scheduler-poddisruptionbudget.yaml        |  2 ++
 .../scheduler/scheduler-service.yaml          |  2 ++
 .../scheduler/scheduler-serviceaccount.yaml   |  2 +-
 .../webserver/webserver-deployment.yaml       |  2 ++
 .../webserver/webserver-ingress.yaml          |  2 ++
 .../webserver/webserver-networkpolicy.yaml    |  2 ++
 .../webserver-poddisruptionbudget.yaml        |  2 ++
 .../webserver/webserver-service.yaml          | 12 ++++++----
 .../webserver/webserver-serviceaccount.yaml   |  2 +-
 chart/values.schema.json                      | 10 ++++++++
 chart/values.yaml                             |  2 ++
 docs/helm-chart/production-guide.rst          | 24 +++++++++++++++++++
 helm_tests/airflow_core/test_scheduler.py     | 17 +++++++++++++
 helm_tests/webserver/test_webserver.py        | 13 ++++++++++
 16 files changed, 92 insertions(+), 7 deletions(-)

diff --git a/chart/templates/scheduler/scheduler-deployment.yaml b/chart/templates/scheduler/scheduler-deployment.yaml
index ccee1609fbc5a..7ec2c46beddac 100644
--- a/chart/templates/scheduler/scheduler-deployment.yaml
+++ b/chart/templates/scheduler/scheduler-deployment.yaml
@@ -20,7 +20,7 @@
 ################################
 ## Airflow Scheduler Deployment/StatefulSet
 #################################
-
+{{- if .Values.scheduler.enabled }}
 # Are we using a local executor?
 {{- $local := contains "Local" .Values.executor }}
 # Is persistence enabled on the _workers_?
@@ -333,3 +333,4 @@ spec:
           requests:
             storage: {{ .Values.workers.persistence.size }}
   {{- end }}
+  {{- end }}
diff --git a/chart/templates/scheduler/scheduler-networkpolicy.yaml b/chart/templates/scheduler/scheduler-networkpolicy.yaml
index bbd985371e8d1..f1c04ad4bca7d 100644
--- a/chart/templates/scheduler/scheduler-networkpolicy.yaml
+++ b/chart/templates/scheduler/scheduler-networkpolicy.yaml
@@ -20,6 +20,7 @@
 ################################
 ## Airflow Scheduler NetworkPolicy
 #################################
+{{- if .Values.scheduler.enabled }}
 {{- if .Values.networkPolicies.enabled }}
 apiVersion: networking.k8s.io/v1
 kind: NetworkPolicy
@@ -55,3 +56,4 @@ spec:
       port: {{ .Values.ports.workerLogs }}
   {{- end }}
 {{- end }}
+{{- end }}
diff --git a/chart/templates/scheduler/scheduler-poddisruptionbudget.yaml b/chart/templates/scheduler/scheduler-poddisruptionbudget.yaml
index 3c6d7983b167e..c413b099f39e3 100644
--- a/chart/templates/scheduler/scheduler-poddisruptionbudget.yaml
+++ b/chart/templates/scheduler/scheduler-poddisruptionbudget.yaml
@@ -20,6 +20,7 @@
 ################################
 ## Airflow Scheduler PodDisruptionBudget
 #################################
+{{- if .Values.scheduler.enabled }}
 {{- if .Values.scheduler.podDisruptionBudget.enabled }}
 apiVersion: policy/v1
 kind: PodDisruptionBudget
@@ -42,3 +43,4 @@ spec:
       release: {{ .Release.Name }}
   {{- toYaml .Values.scheduler.podDisruptionBudget.config | nindent 2 }}
 {{- end }}
+{{- end }}
diff --git a/chart/templates/scheduler/scheduler-service.yaml b/chart/templates/scheduler/scheduler-service.yaml
index b2bf2a87db64f..aa28563deb8d1 100644
--- a/chart/templates/scheduler/scheduler-service.yaml
+++ b/chart/templates/scheduler/scheduler-service.yaml
@@ -20,6 +20,7 @@
 ################################
 ## Airflow Scheduler Service
 #################################
+{{- if .Values.scheduler.enabled }}
 {{- if or (eq .Values.executor "LocalExecutor") (eq .Values.executor "LocalKubernetesExecutor") }}
 apiVersion: v1
 kind: Service
@@ -46,3 +47,4 @@ spec:
       port: {{ .Values.ports.workerLogs }}
       targetPort: {{ .Values.ports.workerLogs }}
 {{- end }}
+{{- end }}
diff --git a/chart/templates/scheduler/scheduler-serviceaccount.yaml b/chart/templates/scheduler/scheduler-serviceaccount.yaml
index 1a359ab1829ef..310f168496766 100644
--- a/chart/templates/scheduler/scheduler-serviceaccount.yaml
+++ b/chart/templates/scheduler/scheduler-serviceaccount.yaml
@@ -20,7 +20,7 @@
 ################################
 ## Airflow Scheduler ServiceAccount
 #################################
-{{- if .Values.scheduler.serviceAccount.create }}
+{{- if and .Values.scheduler.enabled .Values.scheduler.serviceAccount.create }}
 apiVersion: v1
 kind: ServiceAccount
 automountServiceAccountToken: {{ .Values.scheduler.serviceAccount.automountServiceAccountToken }}
diff --git a/chart/templates/webserver/webserver-deployment.yaml b/chart/templates/webserver/webserver-deployment.yaml
index 1fa06b0724cea..bb6c9d4b53823 100644
--- a/chart/templates/webserver/webserver-deployment.yaml
+++ b/chart/templates/webserver/webserver-deployment.yaml
@@ -20,6 +20,7 @@
 ################################
 ## Airflow Webserver Deployment
 #################################
+{{- if .Values.webserver.enabled }}
 {{- $nodeSelector := or .Values.webserver.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.webserver.affinity .Values.affinity }}
 {{- $tolerations := or .Values.webserver.tolerations .Values.tolerations }}
@@ -291,3 +292,4 @@ spec:
         {{- if .Values.webserver.extraVolumes }}
           {{- tpl (toYaml .Values.webserver.extraVolumes) . | nindent 8 }}
         {{- end }}
+{{- end }}
diff --git a/chart/templates/webserver/webserver-ingress.yaml b/chart/templates/webserver/webserver-ingress.yaml
index 451a57cbc0ca9..6851be45ec1fd 100644
--- a/chart/templates/webserver/webserver-ingress.yaml
+++ b/chart/templates/webserver/webserver-ingress.yaml
@@ -20,6 +20,7 @@
 ################################
 ## Airflow Webserver Ingress
 #################################
+{{- if .Values.webserver.enabled }}
 {{- if or .Values.ingress.web.enabled .Values.ingress.enabled }}
 {{- $fullname := (include "airflow.fullname" .) }}
 apiVersion: networking.k8s.io/v1
@@ -109,3 +110,4 @@ spec:
   ingressClassName: {{ .Values.ingress.web.ingressClassName }}
   {{- end }}
 {{- end }}
+{{- end }}
diff --git a/chart/templates/webserver/webserver-networkpolicy.yaml b/chart/templates/webserver/webserver-networkpolicy.yaml
index d0be9e3283fa0..8945d9c2b62f5 100644
--- a/chart/templates/webserver/webserver-networkpolicy.yaml
+++ b/chart/templates/webserver/webserver-networkpolicy.yaml
@@ -20,6 +20,7 @@
 ################################
 ## Airflow Webserver NetworkPolicy
 #################################
+{{- if .Values.webserver.enabled }}
 {{- if .Values.networkPolicies.enabled }}
 {{- $from := or .Values.webserver.networkPolicy.ingress.from .Values.webserver.extraNetworkPolicies }}
 apiVersion: networking.k8s.io/v1
@@ -55,3 +56,4 @@ spec:
       {{- end }}
   {{- end }}
 {{- end }}
+{{- end }}
diff --git a/chart/templates/webserver/webserver-poddisruptionbudget.yaml b/chart/templates/webserver/webserver-poddisruptionbudget.yaml
index 840219a56ed88..f5670ca6e74f3 100644
--- a/chart/templates/webserver/webserver-poddisruptionbudget.yaml
+++ b/chart/templates/webserver/webserver-poddisruptionbudget.yaml
@@ -20,6 +20,7 @@
 ################################
 ## Airflow Webserver PodDisruptionBudget
 #################################
+{{- if .Values.webserver.enabled }}
 {{- if .Values.webserver.podDisruptionBudget.enabled }}
 apiVersion: policy/v1
 kind: PodDisruptionBudget
@@ -42,3 +43,4 @@ spec:
       release: {{ .Release.Name }}
   {{- toYaml .Values.webserver.podDisruptionBudget.config | nindent 2 }}
 {{- end }}
+{{- end }}
diff --git a/chart/templates/webserver/webserver-service.yaml b/chart/templates/webserver/webserver-service.yaml
index a3c49d32a5fb7..c2bdb6b0b5943 100644
--- a/chart/templates/webserver/webserver-service.yaml
+++ b/chart/templates/webserver/webserver-service.yaml
@@ -20,6 +20,7 @@
 ################################
 ## Airflow Webserver Service
 #################################
+{{- if .Values.webserver.enabled }}
 apiVersion: v1
 kind: Service
 metadata:
@@ -43,10 +44,12 @@ spec:
     component: webserver
     release: {{ .Release.Name }}
   ports:
-  {{- range .Values.webserver.service.ports }}
-    - {{ range $key, $val := . }}
-      {{- $key }}: {{ tpl (toString $val) $ }}
-      {{ end }}
+  ports:
+  {{ range .Values.webserver.service.ports }}
+    -
+      {{- range $key, $val := . }}
+      {{ $key }}: {{ tpl (toString $val) $ }}
+      {{- end }}
   {{- end }}
   {{- if .Values.webserver.service.loadBalancerIP }}
   loadBalancerIP: {{ .Values.webserver.service.loadBalancerIP }}
@@ -54,3 +57,4 @@ spec:
   {{- if .Values.webserver.service.loadBalancerSourceRanges }}
   loadBalancerSourceRanges: {{- toYaml .Values.webserver.service.loadBalancerSourceRanges | nindent 4 }}
   {{- end }}
+{{- end }}
diff --git a/chart/templates/webserver/webserver-serviceaccount.yaml b/chart/templates/webserver/webserver-serviceaccount.yaml
index 853e376a5daf8..a971da6e3275e 100644
--- a/chart/templates/webserver/webserver-serviceaccount.yaml
+++ b/chart/templates/webserver/webserver-serviceaccount.yaml
@@ -20,7 +20,7 @@
 ######################################
 ## Airflow Webserver ServiceAccount
 ######################################
-{{- if .Values.webserver.serviceAccount.create }}
+{{- if and .Values.webserver.enabled .Values.webserver.serviceAccount.create }}
 apiVersion: v1
 kind: ServiceAccount
 automountServiceAccountToken: {{ .Values.webserver.serviceAccount.automountServiceAccountToken }}
diff --git a/chart/values.schema.json b/chart/values.schema.json
index f6613adb96416..3f9e3fc8688ed 100644
--- a/chart/values.schema.json
+++ b/chart/values.schema.json
@@ -2193,6 +2193,11 @@
             "x-docsSection": "Scheduler",
             "additionalProperties": false,
             "properties": {
+                "enabled": {
+                    "description": "Enable scheduler",
+                    "type": "boolean",
+                    "default": true
+                },
                 "hostAliases": {
                     "description": "HostAliases for the scheduler pod.",
                     "items": {
@@ -4212,6 +4217,11 @@
             "x-docsSection": "Webserver",
             "additionalProperties": false,
             "properties": {
+                "enabled": {
+                    "description": "Enable webserver",
+                    "type": "boolean",
+                    "default": true
+                },
                 "configMapAnnotations": {
                     "description": "Extra annotations to apply to the webserver configmap.",
                     "type": "object",
diff --git a/chart/values.yaml b/chart/values.yaml
index fddcf67c872bb..9a7b7352c1754 100644
--- a/chart/values.yaml
+++ b/chart/values.yaml
@@ -790,6 +790,7 @@ workers:
 
 # Airflow scheduler settings
 scheduler:
+  enabled: true
   #  hostAliases for the scheduler pod
   hostAliases: []
   #  - ip: "127.0.0.1"
@@ -1150,6 +1151,7 @@ migrateDatabaseJob:
 
 # Airflow webserver settings
 webserver:
+  enabled: true
   # Add custom annotations to the webserver configmap
   configMapAnnotations: {}
   #  hostAliases for the webserver pod
diff --git a/docs/helm-chart/production-guide.rst b/docs/helm-chart/production-guide.rst
index 6451205611524..4261c9f9adfe3 100644
--- a/docs/helm-chart/production-guide.rst
+++ b/docs/helm-chart/production-guide.rst
@@ -208,11 +208,35 @@ They match, right? Good. Now, add the public key to your values. It'll look some
           github.com ssh-rsa AAAA...1/wsjk=
 
 
+External Scheduler
+^^^^^^^^^^^^^^^^^^
+
+To use an external Scheduler instance:
+
+.. code-block:: yaml
+
+  scheduler:
+    enabled: false
+
+Ensure that your external webserver/scheduler is connected to the same redis host. This will ensure the scheduler is aware of the workers deployed in the helm-chart.
+
 Accessing the Airflow UI
 ------------------------
 
 How you access the Airflow UI will depend on your environment; however, the chart does support various options:
 
+External Webserver
+^^^^^^^^^^^^^^^^^^
+
+To use an external Webserver:
+
+.. code-block:: yaml
+
+  webserver:
+    enabled: false
+
+Ensure that your external webserver/scheduler is connected to the same redis host. This will ensure the scheduler is aware of the workers deployed in the helm-chart.
+
 Ingress
 ^^^^^^^
 
diff --git a/helm_tests/airflow_core/test_scheduler.py b/helm_tests/airflow_core/test_scheduler.py
index addaaf896e7cf..42980b8865c92 100644
--- a/helm_tests/airflow_core/test_scheduler.py
+++ b/helm_tests/airflow_core/test_scheduler.py
@@ -901,3 +901,20 @@ def test_overridden_automount_service_account_token(self):
             show_only=["templates/scheduler/scheduler-serviceaccount.yaml"],
         )
         assert jmespath.search("automountServiceAccountToken", docs[0]) is False
+
+
+class TestSchedulerCreation:
+    """Tests scheduler deployment creation."""
+
+    def test_can_be_disabled(self):
+        """
+        Scheduler should be able to be disabled if the users desires.
+
+        For example, user may be disabled when using scheduler and having it deployed on another host.
+        """
+        docs = render_chart(
+            values={"scheduler": {"enabled": False}},
+            show_only=["templates/scheduler/scheduler-deployment.yaml"],
+        )
+
+        assert 0 == len(docs)
diff --git a/helm_tests/webserver/test_webserver.py b/helm_tests/webserver/test_webserver.py
index 7fd7e60257b9b..42a89bcc5c3fb 100644
--- a/helm_tests/webserver/test_webserver.py
+++ b/helm_tests/webserver/test_webserver.py
@@ -25,6 +25,19 @@
 class TestWebserverDeployment:
     """Tests webserver deployment."""
 
+    def test_can_be_disabled(self):
+        """
+        Webserver should be able to be disabled if the users desires.
+
+        For example, user may be disabled when using webserver and having it deployed on another host.
+        """
+        docs = render_chart(
+            values={"webserver": {"enabled": False}},
+            show_only=["templates/webserver/webserver-deployment.yaml"],
+        )
+
+        assert 0 == len(docs)
+
     def test_should_add_host_header_to_liveness_and_readiness_and_startup_probes(self):
         docs = render_chart(
             values={