From cf1fc5b0405a249509dd6a8ef3950cbca9e8d8c5 Mon Sep 17 00:00:00 2001
From: Ben Bellick <ben.bellick@datadoghq.com>
Date: Fri, 31 Oct 2025 16:32:01 -0400
Subject: [PATCH 1/3] chore: add some tests validating urn logic

---
 src/urn.rs | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/src/urn.rs b/src/urn.rs
index 4f47600f..e39e8b38 100644
--- a/src/urn.rs
+++ b/src/urn.rs
@@ -52,3 +52,34 @@ impl FromStr for Urn {
             .ok_or(InvalidUrn)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn can_parse_valid_urn() {
+        let valid_urn = "extension:io.substrait:fake_functions";
+        let result = valid_urn.parse::<Urn>();
+        assert_eq!(
+            result,
+            Ok(Urn {
+                owner: "io.substrait".to_string(),
+                id: "fake_functions".to_string()
+            })
+        );
+    }
+
+    #[test]
+    fn parse_fail_invalid_urns() {
+        let invalid_urns = vec![
+            "io.substrait:something",         // missing "extension:" at beginning
+            "other_type:substrait:something", // incorrect type at beginning
+            "extension:something",            // doesn't have both owner and id
+            "extension:one:two:three",        //has too many parts
+        ];
+        for invalid_urn in invalid_urns {
+            assert_eq!(invalid_urn.parse::<Urn>(), Err(InvalidUrn))
+        }
+    }
+}

From 1e3f24e3084f001e2d84ff95a299031f4b6ce1a4 Mon Sep 17 00:00:00 2001
From: Ben Bellick <ben.bellick@datadoghq.com>
Date: Fri, 31 Oct 2025 16:59:08 -0400
Subject: [PATCH 2/3] fix: do not allow extra ':' in urns

---
 src/urn.rs | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/urn.rs b/src/urn.rs
index e39e8b38..8f749987 100644
--- a/src/urn.rs
+++ b/src/urn.rs
@@ -40,16 +40,14 @@ impl FromStr for Urn {
     type Err = InvalidUrn;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        s.split_once(':')
-            .filter(|(extension, _)| *extension == "extension")
-            .map(|(_, segments)| segments)
-            .and_then(|segments| segments.split_once(':'))
-            .filter(|(owner, id)| !owner.is_empty() && !id.is_empty())
-            .map(|(owner, id)| Urn {
-                owner: owner.to_owned(),
-                id: id.to_owned(),
-            })
-            .ok_or(InvalidUrn)
+        let parts: Vec<&str> = s.split(':').collect();
+        match parts.as_slice() {
+            ["extension", owner, id] if !owner.is_empty() && !id.is_empty() => Ok(Urn {
+                owner: owner.to_string(),
+                id: id.to_string(),
+            }),
+            _ => Err(InvalidUrn),
+        }
     }
 }
 

From 97917918ce2dc8d3e03f15d3a6fa83a4eba56226 Mon Sep 17 00:00:00 2001
From: Ben Bellick <ben.bellick@datadoghq.com>
Date: Sat, 1 Nov 2025 12:15:26 -0400
Subject: [PATCH 3/3] tweak: drop allocation to validate urn

add a few more tests as well
---
 src/urn.rs | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/urn.rs b/src/urn.rs
index 8f749987..8e5c179c 100644
--- a/src/urn.rs
+++ b/src/urn.rs
@@ -40,12 +40,16 @@ impl FromStr for Urn {
     type Err = InvalidUrn;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        let parts: Vec<&str> = s.split(':').collect();
-        match parts.as_slice() {
-            ["extension", owner, id] if !owner.is_empty() && !id.is_empty() => Ok(Urn {
-                owner: owner.to_string(),
-                id: id.to_string(),
-            }),
+        let mut parts = s.split(':');
+        match (parts.next(), parts.next(), parts.next(), parts.next()) {
+            (Some("extension"), Some(owner), Some(id), None)
+                if !owner.is_empty() && !id.is_empty() =>
+            {
+                Ok(Urn {
+                    owner: owner.to_string(),
+                    id: id.to_string(),
+                })
+            }
             _ => Err(InvalidUrn),
         }
     }
@@ -75,6 +79,8 @@ mod tests {
             "other_type:substrait:something", // incorrect type at beginning
             "extension:something",            // doesn't have both owner and id
             "extension:one:two:three",        //has too many parts
+            "extension::id",                  //missing owner
+            "extension:owner:",               //missing id
         ];
         for invalid_urn in invalid_urns {
             assert_eq!(invalid_urn.parse::<Urn>(), Err(InvalidUrn))