From 446de174043937b68192f5a26efbece394e5f205 Mon Sep 17 00:00:00 2001
From: HYBG <hybg1126@outlook.com>
Date: Wed, 15 Jan 2025 10:17:58 +0800
Subject: [PATCH] Gracefully close kyuubi & add customer labels

---
 .../kyuubi/templates/kyuubi-statefulset.yaml  | 24 +++++++++++++++++++
 charts/kyuubi/values.yaml                     | 12 ++++++++++
 2 files changed, 36 insertions(+)

diff --git a/charts/kyuubi/templates/kyuubi-statefulset.yaml b/charts/kyuubi/templates/kyuubi-statefulset.yaml
index 57601826f3a..5f909a770f9 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:
@@ -129,6 +132,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 }}
diff --git a/charts/kyuubi/values.yaml b/charts/kyuubi/values.yaml
index 4922a61779b..5aff250dd43 100644
--- a/charts/kyuubi/values.yaml
+++ b/charts/kyuubi/values.yaml
@@ -22,6 +22,9 @@
 # Kyuubi server numbers
 replicaCount: 2
 
+# Kyuubi labels
+labels: ~
+
 # controls how Kyuubi server pods are created during initial scale up,
 # when replacing pods on nodes, or when scaling down.
 # The default policy is `OrderedReady`, alternative policy is `Parallel`.
@@ -305,3 +308,12 @@ metrics:
     enabled: false
     # Content of Prometheus rule file
     groups: []
+
+# Kyuubi lifecycle
+lifecycle:
+  postStart:
+    exec:
+      command: ['sh', '-c', "echo \"ok!!!\""]
+  preStop:
+    exec:
+      command: ["/bin/sh","-c","/opt/kyuubi/bin/kyuubi-ctl delete server --host $(hostname -I) --port \"10009\""]
\ No newline at end of file