From 9a1da12776c10d94fa32827881b5b8b27f839920 Mon Sep 17 00:00:00 2001
From: Michal Belica <code@beli.sk>
Date: Thu, 30 Apr 2026 10:17:27 +0200
Subject: [PATCH] feat(helm): migration job annotation control

---
 .../templates/migrations-job.yaml             | 11 ++++---
 .../tests/migrations-job_tests.yaml           | 32 +++++++++++++++++++
 deploy/charts/litellm-helm/values.yaml        |  9 ++++++
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/deploy/charts/litellm-helm/templates/migrations-job.yaml b/deploy/charts/litellm-helm/templates/migrations-job.yaml
index c3f32fe32f3..71790114b5f 100644
--- a/deploy/charts/litellm-helm/templates/migrations-job.yaml
+++ b/deploy/charts/litellm-helm/templates/migrations-job.yaml
@@ -7,13 +7,16 @@ metadata:
   labels:
     {{- include "litellm.labels" . | nindent 4 }}
   annotations:
+    {{- with .Values.migrationJob.jobAnnotations }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
     {{- if .Values.migrationJob.hooks.argocd.enabled }}
-    argocd.argoproj.io/hook: PreSync
-    argocd.argoproj.io/hook-delete-policy: BeforeHookCreation
+    argocd.argoproj.io/hook: {{ .Values.migrationJob.hooks.argocd.hook | default "PreSync" | quote }}
+    argocd.argoproj.io/hook-delete-policy: {{ .Values.migrationJob.hooks.argocd.hookDeletePolicy | default "BeforeHookCreation" | quote }}
     {{- end }}
     {{- if .Values.migrationJob.hooks.helm.enabled }}
-    helm.sh/hook: "pre-install,pre-upgrade"
-    helm.sh/hook-delete-policy: "before-hook-creation"
+    helm.sh/hook: {{ .Values.migrationJob.hooks.helm.hook | default "pre-install,pre-upgrade" | quote }}
+    helm.sh/hook-delete-policy: {{ .Values.migrationJob.hooks.helm.hookDeletePolicy | default "before-hook-creation" | quote }}
     helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "1" | quote }}
     {{- end }}
     checksum/config: {{ toYaml .Values | sha256sum }}
diff --git a/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml b/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml
index 05dd37b4857..0cf9dbf47bb 100644
--- a/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml
+++ b/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml
@@ -254,3 +254,35 @@ tests:
           content:
             name: sidecar-tpl
             image: "ghcr.io/berriai/litellm-database:test"
+  - it: should set annotations on migrations job properly
+    template: migrations-job.yaml
+    set:
+      migrationJob:
+        enabled: true
+        jobAnnotations:
+          test: test1
+        hooks:
+          argocd:
+            enabled: true
+            hook: test2
+            hookDeletePolicy: test3
+          helm:
+            enabled: true
+            hook: test4
+            hookDeletePolicy: test5
+    asserts:
+      - equal:
+          path: metadata.annotations['test']
+          value: test1
+      - equal:
+          path: metadata.annotations['argocd.argoproj.io/hook']
+          value: test2
+      - equal:
+          path: metadata.annotations['argocd.argoproj.io/hook-delete-policy']
+          value: test3
+      - equal:
+          path: metadata.annotations['helm.sh/hook']
+          value: test4
+      - equal:
+          path: metadata.annotations['helm.sh/hook-delete-policy']
+          value: test5
diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml
index 690ca69e730..32b22b5e4a7 100644
--- a/deploy/charts/litellm-helm/values.yaml
+++ b/deploy/charts/litellm-helm/values.yaml
@@ -313,6 +313,9 @@ migrationJob:
   # Only used when migrationJob.hooks.helm.enabled=true and serviceAccount.create=true.
   # In that case, pre-install/pre-upgrade hooks run before normal resources, so this defaults to "default".
   serviceAccountName: ""
+  # extra annotations for the Job
+  jobAnnotations: {}
+  # extra annotations for pods created from the Job
   annotations: {}
   ttlSecondsAfterFinished: 120
   resources: {}
@@ -326,8 +329,14 @@ migrationJob:
   hooks:
     argocd:
       enabled: true
+      # Optional hook controls (defaults shown)
+      # hook: PreSync
+      # hookDeletePolicy: BeforeHookCreation
     helm:
       enabled: false
+      # Optional hook controls (defaults shown)
+      # hook: pre-install,pre-upgrade
+      # hookDeletePolicy: before-hook-creation
 
 # Additional environment variables to be added to the deployment as a map of key-value pairs
 envVars: {}