From 940fbf04cc75e491739098af2f98d596ca0ad533 Mon Sep 17 00:00:00 2001
From: Dan Harris <1327726+thinkharderdev@users.noreply.github.com>
Date: Fri, 6 May 2022 15:49:06 -0400
Subject: [PATCH] Fix stage key extraction (#2472)

---
 .../scheduler/src/state/persistent_state.rs   | 46 ++++++++++++++++---
 1 file changed, 40 insertions(+), 6 deletions(-)

diff --git a/ballista/rust/scheduler/src/state/persistent_state.rs b/ballista/rust/scheduler/src/state/persistent_state.rs
index 92e94529ea84..4cce866762a5 100644
--- a/ballista/rust/scheduler/src/state/persistent_state.rs
+++ b/ballista/rust/scheduler/src/state/persistent_state.rs
@@ -318,16 +318,21 @@ fn extract_job_id_from_job_key(job_key: &str) -> Result<&str> {
 
 fn extract_stage_id_from_stage_key(stage_key: &str) -> Result<StageKey> {
     let splits: Vec<&str> = stage_key.split('/').collect();
-    if splits.len() < 4 {
+    if splits.len() > 4 {
+        Ok((
+            splits[splits.len() - 2].to_string(),
+            splits[splits.len() - 1].parse::<u32>().map_err(|e| {
+                BallistaError::Internal(format!(
+                    "Invalid stage ID in stage key: {}, {:?}",
+                    stage_key, e
+                ))
+            })?,
+        ))
+    } else {
         Err(BallistaError::Internal(format!(
             "Unexpected stage key: {}",
             stage_key
         )))
-    } else {
-        Ok((
-            splits.get(2).unwrap().to_string(),
-            splits.get(3).unwrap().parse::<u32>().unwrap(),
-        ))
     }
 }
 
@@ -352,3 +357,32 @@ fn encode_protobuf<T: Message + Default>(msg: &T) -> Result<Vec<u8>> {
     })?;
     Ok(value)
 }
+
+#[cfg(test)]
+mod test {
+    use super::extract_stage_id_from_stage_key;
+
+    #[test]
+    fn test_extract_stage_id_from_stage_key() {
+        let (job_id, stage_id) =
+            extract_stage_id_from_stage_key("/ballista/default/stages/2Yoyba8/1")
+                .expect("extracting stage key");
+
+        assert_eq!(job_id.as_str(), "2Yoyba8");
+        assert_eq!(stage_id, 1);
+
+        let (job_id, stage_id) =
+            extract_stage_id_from_stage_key("ballista/default/stages/2Yoyba8/1")
+                .expect("extracting stage key");
+
+        assert_eq!(job_id.as_str(), "2Yoyba8");
+        assert_eq!(stage_id, 1);
+
+        let (job_id, stage_id) =
+            extract_stage_id_from_stage_key("ballista//stages/2Yoyba8/1")
+                .expect("extracting stage key");
+
+        assert_eq!(job_id.as_str(), "2Yoyba8");
+        assert_eq!(stage_id, 1);
+    }
+}