diff --git a/charts/flink-job/templates/_helpers.tpl b/charts/flink-job/templates/_helpers.tpl
index 6fee07eb..ad5e2ee1 100644
--- a/charts/flink-job/templates/_helpers.tpl
+++ b/charts/flink-job/templates/_helpers.tpl
@@ -247,15 +247,6 @@ Validate the configuration
       {{- $configs = fromJson (include "flink-job._dictSet" (list $configs "state.savepoints.dir" $savepointsDir)) -}}
     {{- end -}}
   {{- end -}}
-  {{- if not (has .global.job.upgradeMode (list "stateless" "last-state" "savepoint")) -}}
-    {{- fail "job.upgradeMode must be either stateless, last-state, or savepoint" -}}
-  {{- end -}}
-  {{- if and (has .global.job.upgradeMode (list "last-state" "savepoint")) (not (hasKey $configs "state.checkpoints.dir")) -}}
-    {{- fail "storage.scheme and storage.baseDir or flinkConfiguration.'state.checkpoints.dir' is required when using job.upgradeMode=savepoint or last-state" -}}
-  {{- end -}}
-  {{- if and (eq .global.job.upgradeMode "savepoint") (not (hasKey $configs "state.savepoints.dir")) -}}
-    {{- fail "storage.scheme and storage.baseDir or flinkConfiguration.'state.savepoints.dir' is required when using job.upgradeMode=savepoint" -}}
-  {{- end -}}
   {{- $configs | toJson -}}
 {{- end -}}
 
diff --git a/charts/flink-job/values.schema.json b/charts/flink-job/values.schema.json
index 4a4285cd..8ae3e0b9 100644
--- a/charts/flink-job/values.schema.json
+++ b/charts/flink-job/values.schema.json
@@ -3,6 +3,68 @@
   "title": "JSON schema for flink-job/values.yaml",
   "type": "object",
   "allOf": [
+    {
+      "if": {
+        "properties": {
+          "job": {
+            "properties": {
+              "upgradeMode": {
+                "const": "savepoint"
+              }
+            }
+          }
+        }
+      },
+      "then": {
+        "anyOf": [
+          {
+            "properties": {
+              "flinkConfiguration": {
+                "required": ["state.savepoints.dir"]
+              }
+            }
+          },
+          {
+            "properties": {
+              "storage": {
+                "required": ["scheme", "baseDir"]
+              }
+            }
+          }
+        ]
+      }
+    },
+    {
+      "if": {
+        "properties": {
+          "job": {
+            "properties": {
+              "upgradeMode": {
+                "enum": ["savepoint", "last-state"]
+              }
+            }
+          }
+        }
+      },
+      "then": {
+        "anyOf": [
+          {
+            "properties": {
+              "flinkConfiguration": {
+                "required": ["state.checkpoints.dir"]
+              }
+            }
+          },
+          {
+            "properties": {
+              "storage": {
+                "required": ["scheme", "baseDir"]
+              }
+            }
+          }
+        ]
+      }
+    },
     {
       "if": {
         "properties": {