From fd6a45f50e479832fe02abc0349b776028c117cf Mon Sep 17 00:00:00 2001
From: HYBG <hybg1126@outlook.com>
Date: Fri, 20 Sep 2024 15:59:52 +0800
Subject: [PATCH] Gracefully close kyuubi & add customer labels

---
 .../kyuubi/templates/kyuubi-statefulset.yaml  | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/charts/kyuubi/templates/kyuubi-statefulset.yaml b/charts/kyuubi/templates/kyuubi-statefulset.yaml
index caea7d251b2..2984d0304db 100644
--- a/charts/kyuubi/templates/kyuubi-statefulset.yaml
+++ b/charts/kyuubi/templates/kyuubi-statefulset.yaml
@@ -21,6 +21,9 @@ metadata:
   name: {{ .Release.Name }}
   labels:
     {{- include "kyuubi.labels" . | nindent 4 }}
+    {{- with .Values.labels }}
+      {{- toYaml . | nindent 4 }}
+    {{- end }}
 spec:
   selector:
     matchLabels:
@@ -115,6 +118,27 @@ spec:
             {{- with .Values.volumeMounts }}
               {{- tpl (toYaml .) $ | nindent 12 }}
             {{- end }}
+          lifecycle:
+            {{- if .Values.lifecycle.postStart }}
+            postStart:
+              {{- if .Values.lifecycle.postStart.exec }}
+              exec:
+                {{- if .Values.lifecycle.postStart.exec.command }}
+                command:
+                  {{- toYaml .Values.lifecycle.postStart.exec.command | nindent 16 }}
+                {{- end }}
+              {{- end }}
+            {{- end }}
+            {{- if .Values.lifecycle.preStop }}
+            preStop:
+              {{- if .Values.lifecycle.preStop.exec }}
+              exec:
+                {{- if .Values.lifecycle.preStop.exec.command }}
+                command:
+                  {{- toYaml .Values.lifecycle.preStop.exec.command | nindent 16 }}
+                {{- end }}
+              {{- end }}
+            {{- end }}
         {{- with .Values.containers }}
           {{- tpl (toYaml .) $ | nindent 8 }}
         {{- end }}