feat(helm): allow custom labels, annotations, command and args on migrationJob (#42242)

Co-authored-by: ryan-crabbe-berri <ryan@berri.ai>

helm/litellm-helm/templates/migrations-job.yaml

@@ -6,6 +6,9 @@ metadata:
   name: {{ include "litellm.fullname" . }}-migrations
   labels:
     {{- include "litellm.labels" . | nindent 4 }}
+    {{- with .Values.migrationJob.jobLabels }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
   annotations:
     {{- if .Values.migrationJob.hooks.argocd.enabled }}
     argocd.argoproj.io/hook: PreSync
@@ -17,6 +20,9 @@ metadata:
     helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "1" | quote }}
     {{- end }}
     checksum/config: {{ toYaml .Values | sha256sum }}
+    {{- with .Values.migrationJob.jobAnnotations }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
 spec:
   template:
     metadata:
@@ -25,6 +31,9 @@ spec:
         {{- with .Values.podLabels }}
         {{- toYaml . | nindent 8 }}
         {{- end }}
+        {{- with .Values.migrationJob.podLabels }}
+        {{- toYaml . | nindent 8 }}
+        {{- end }}
       annotations:
         {{- with .Values.migrationJob.annotations }}
         {{- toYaml . | nindent 8 }}
@@ -47,7 +56,16 @@ spec:
           imagePullPolicy: {{ .Values.image.pullPolicy }}
           securityContext:
             {{- toYaml .Values.securityContext | nindent 12 }}
+          {{- if .Values.migrationJob.command }}
+          command: {{ toYaml .Values.migrationJob.command | nindent 12 }}
+          {{- else }}
           command: ["python", "litellm/proxy/prisma_migration.py"]
+          {{- end }}
+          {{- if .Values.migrationJob.args }}
+          args: {{ toYaml .Values.migrationJob.args | nindent 12 }}
+          {{- else if .Values.migrationJob.command }}
+          args: []
+          {{- end }}
           workingDir: "/app"
           env:
             {{- if .Values.db.useExisting }}

helm/litellm-helm/tests/migrations-job_tests.yaml

@@ -360,3 +360,73 @@ tests:
     asserts:
       - notExists:
           path: spec.activeDeadlineSeconds
+
+  - it: should set custom jobLabels and jobAnnotations on Job metadata
+    template: migrations-job.yaml
+    set:
+      migrationJob:
+        enabled: true
+        jobLabels:
+          environment: production
+          team: platform
+        jobAnnotations:
+          example.com/cost-center: "1234"
+    asserts:
+      - equal:
+          path: metadata.labels.environment
+          value: production
+      - equal:
+          path: metadata.labels.team
+          value: platform
+      - equal:
+          path: metadata.annotations['example.com/cost-center']
+          value: "1234"
+
+  - it: should set custom podLabels on Pod template
+    template: migrations-job.yaml
+    set:
+      migrationJob:
+        enabled: true
+        podLabels:
+          custom.io/pod-role: migration
+    asserts:
+      - equal:
+          path: spec.template.metadata.labels['custom.io/pod-role']
+          value: migration
+
+  - it: should override container command and args
+    template: migrations-job.yaml
+    set:
+      migrationJob:
+        enabled: true
+        command:
+          - sh
+        args:
+          - -c
+          - echo migrating
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].command
+          value:
+            - sh
+      - equal:
+          path: spec.template.spec.containers[0].args
+          value:
+            - -c
+            - echo migrating
+
+  - it: should clear container args when only command is specified
+    template: migrations-job.yaml
+    set:
+      migrationJob:
+        enabled: true
+        command:
+          - sh
+    asserts:
+      - equal:
+          path: spec.template.spec.containers[0].command
+          value:
+            - sh
+      - equal:
+          path: spec.template.spec.containers[0].args
+          value: []

helm/litellm-helm/values.yaml

@@ -572,6 +572,11 @@ migrationJob:
   # In that case, pre-install/pre-upgrade hooks run before normal resources, so this defaults to "default".
   serviceAccountName: ""
   annotations: {}
+  jobLabels: {} # Custom labels for the Job metadata
+  jobAnnotations: {} # Custom annotations for the Job metadata
+  podLabels: {} # Custom labels for the Job pod template
+  command: [] # Override container command (defaults to ["python", "litellm/proxy/prisma_migration.py"])
+  args: [] # Override container args
   ttlSecondsAfterFinished: 120
   resources: {}
   #  Unset by default. This job runs the database migration and exits, so it does not