add liveliness and readiness support for dag server

templates/_helpers.yaml

@@ -324,6 +324,12 @@ Create the name of the cleanup service account to use
       fieldRef:
         apiVersion: v1
         fieldPath: metadata.labels['workspace']
+{{- if .Values.loggingSidecar.livenessProbe  }}
+  livenessProbe: {{ tpl (toYaml .Values.loggingSidecar.livenessProbe) . | nindent 4 }}
+{{- end }}
+{{- if .Values.loggingSidecar.readinessProbe  }}
+  readinessProbe: {{ tpl (toYaml .Values.loggingSidecar.readinessProbe) . | nindent 4 }}
+{{- end }}
   volumeMounts:
   - mountPath: /etc/sidecar-logging
     name: config-volume

templates/dag-deploy/dag-server-statefulset.yaml

@@ -68,6 +68,9 @@ spec:
 {{- if .Values.dagDeploy.extraEnv }}
 {{ toYaml .Values.dagDeploy.extraEnv | indent 10 }}
 {{- end }}
+        {{- if .Values.dagDeploy.livenessProbe  }}
+          livenessProbe: {{ tpl (toYaml .Values.dagDeploy.livenessProbe) . | nindent 12 }}
+        {{- else }}
           livenessProbe:
             failureThreshold: 5
             httpGet:
@@ -78,12 +81,17 @@ spec:
             periodSeconds: 10
             successThreshold: 1
             timeoutSeconds: 5
+        {{- end }}
+        {{- if .Values.dagDeploy.readinessProbe  }}
+          readinessProbe: {{ tpl (toYaml .Values.dagDeploy.readinessProbe) . | nindent 12 }}
+        {{- else }}
           readinessProbe:
             httpGet:
               path: /healthz
               port: {{ .Values.dagDeploy.ports.dagServerHttp }}
             initialDelaySeconds: 10
             periodSeconds: 10
+        {{ end }}
           volumeMounts:
           - name: data
             mountPath: /data
@@ -106,20 +114,28 @@ spec:
           - containerPort: {{ .Values.authSidecar.port }}
             name: auth-proxy
             protocol: TCP
+        {{- if .Values.authSidecar.readinessProbe  }}
+          readinessProbe: {{ tpl (toYaml .Values.authSidecar.readinessProbe) . | nindent 12 }}
+        {{- else }}
           readinessProbe:
             httpGet:
               path: /healthz
               port: {{ .Values.authSidecar.port }}
               scheme: HTTP
             initialDelaySeconds: 10
             periodSeconds: 10
+        {{- end }}
+        {{- if .Values.authSidecar.livenessProbe  }}
+          livenessProbe: {{ tpl (toYaml .Values.authSidecar.livenessProbe) . | nindent 12 }}
+        {{- else }}
           livenessProbe:
             httpGet:
               path: /healthz
               port: {{ .Values.authSidecar.port }}
               scheme: HTTP
             initialDelaySeconds: 10
             periodSeconds: 10
+        {{- end }}
           volumeMounts:
           - mountPath: /etc/nginx/nginx.conf
             name: nginx-sidecar-conf

tests/chart/test_dag_server_statefulset.py

@@ -5,6 +5,9 @@
 
 from . import get_containers_by_name
 
+readinessProbe = {"httpGet": {"initialDelaySeconds": 20, "periodSeconds": 20, "path": "/rhealthz", "port": 8080, "scheme": "HTTP"}}
+livenessProbe = {"httpGet": {"initialDelaySeconds": 20, "periodSeconds": 20, "path": "/chealthz", "port": 8080, "scheme": "HTTP"}}
+
 
 def common_default_tests(doc):
     """Test cases for default dag-server sts enabled"""
@@ -232,3 +235,47 @@ def test_dag_server_statefulset_with_sidecar_and_authproxy_enabled(self, kube_ve
             "sanic dag_deploy.server.app -H 0.0.0.0 1> >( tee -a /var/log/sidecar-logging-consumer/out.log ) 2> >( tee -a /var/log/sidecar-logging-consumer/err.log >&2 )",
         ]
         assert "sidecar-log-consumer" in c_by_name
+
+    def test_dag_server_statefulset_liveliness_and_readiness_probes_with_dag_server_enabled(self, kube_version):
+        """Test that a valid statefulset is rendered when dag-server is enabled with readiness and liveliness probes."""
+        values = {"dagDeploy": {"enabled": True, "readinessProbe": readinessProbe, "livenessProbe": livenessProbe}}
+
+        docs = render_chart(
+            kube_version=kube_version,
+            show_only="templates/dag-deploy/dag-server-statefulset.yaml",
+            values=values,
+        )
+        assert len(docs) == 1
+        doc = docs[0]
+
+        common_default_tests(doc)
+
+        c_by_name = get_containers_by_name(doc)
+        assert readinessProbe == c_by_name["dag-server"]["readinessProbe"]
+        assert livenessProbe == c_by_name["dag-server"]["livenessProbe"]
+
+    def test_dag_server_statefulset_with_probes_and_authproxy_enabled(self, kube_version):
+        """Test that a valid statefulset is rendered when dag-server and authsidecar is enabled with readiness and liveliness probes."""
+        values = {
+            "dagDeploy": {"enabled": True},
+            "authSidecar": {"enabled": True, "readinessProbe": readinessProbe, "livenessProbe": livenessProbe},
+            "loggingSidecar": {"enabled": True, "readinessProbe": readinessProbe, "livenessProbe": livenessProbe},
+        }
+
+        docs = render_chart(
+            kube_version=kube_version,
+            show_only="templates/dag-deploy/dag-server-statefulset.yaml",
+            values=values,
+        )
+        assert len(docs) == 1
+        doc = docs[0]
+
+        c_by_name = get_containers_by_name(doc)
+        assert len(c_by_name) == 3
+        assert "dag-server" in c_by_name
+        assert "auth-proxy" in c_by_name
+        assert "sidecar-log-consumer" in c_by_name
+        assert readinessProbe == c_by_name["auth-proxy"]["readinessProbe"]
+        assert livenessProbe == c_by_name["auth-proxy"]["livenessProbe"]
+        assert readinessProbe == c_by_name["sidecar-log-consumer"]["readinessProbe"]
+        assert livenessProbe == c_by_name["sidecar-log-consumer"]["livenessProbe"]

values.yaml

@@ -495,6 +495,8 @@ authSidecar:
   port: 8084
   securityContext: {}
   resources: {}
+  livenessProbe: {}
+  readinessProbe: {}
 
 loggingSidecar:
   enabled: false
@@ -503,6 +505,8 @@ loggingSidecar:
   indexPattern: "%Y.%m.%d"
   indexNamePrefix: vector
   image: quay.io/astronomer/ap-vector:0.40.2-1
+  livenessProbe: {}
+  readinessProbe: {}
 
 # Ingress configuration
 ingress:
@@ -560,9 +564,8 @@ dagDeploy:
       repository: quay.io/astronomer/ap-dag-deploy
       tag: 0.5.1
       imagePullPolicy: IfNotPresent
-  livenessProbe:
-    initialDelaySeconds: 30
-    timeoutSeconds: 30
+  livenessProbe: {}
+  readinessProbe: {}
   terminationGracePeriodSeconds: 30
   podAnnotations: {}
   extraEnv: []