fix(query): handle malformed UDF script metadata

Author: KKould

src/query/script_udf_support/src/transform_udf_script.rs

@@ -609,7 +609,7 @@ impl TransformUdfScript {
                     imports_stage_info,
                     ..
                 }) => {
-                    let mut dependencies = Self::extract_deps(&code_str);
+                    let mut dependencies = Self::extract_deps(&code_str)?;
                     dependencies.extend_from_slice(packages.as_slice());
 
                     let stage_fingerprints = Self::collect_stage_fingerprints(imports_stage_info)?;
@@ -677,7 +677,7 @@ impl TransformUdfScript {
         Ok(script_runtimes)
     }
 
-    fn extract_deps(script: &str) -> Vec<String> {
+    fn extract_deps(script: &str) -> Result<Vec<String>> {
         let mut ss = String::new();
         let mut meta_start = false;
         for line in script.lines() {
@@ -693,19 +693,23 @@ impl TransformUdfScript {
             }
         }
 
-        let parsed = ss.parse::<toml::Value>().unwrap();
+        let parsed = ss.parse::<toml::Value>().map_err(|err| {
+            ErrorCode::SemanticError(format!("Failed to parse UDF script metadata as TOML: {err}"))
+        })?;
 
         if parsed.get("dependencies").is_none() {
-            return Vec::new();
+            return Ok(Vec::new());
         }
 
-        if let Some(deps) = parsed["dependencies"].as_array() {
-            deps.iter()
+        let deps = if let Some(deps) = parsed["dependencies"].as_array() {
+            deps
+                .iter()
                 .filter_map(|v| v.as_str().map(|s| s.to_string()))
                 .collect()
         } else {
             Vec::new()
-        }
+        };
+        Ok(deps)
     }
 
     fn prepare_block_entries(
@@ -858,6 +862,25 @@ impl TransformUdfScript {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_deps_returns_error_for_malformed_toml() {
+        let err = TransformUdfScript::extract_deps(
+            r#"# /// script
+# dependencies = [
+# ///
+"#,
+        )
+        .expect_err("malformed UDF script metadata should return an error");
+
+        assert_eq!(err.code(), ErrorCode::SEMANTIC_ERROR);
+        assert!(err.message().contains("Failed to parse UDF script metadata as TOML"));
+    }
+}
+
 mod venv {
     use std::fs;
     use std::fs::File;

src/query/sql/src/planner/semantic/type_check/udf.rs

@@ -264,7 +264,7 @@ fn escape_python_double_quoted(value: &str) -> String {
     value.replace('\\', "\\\\").replace('"', "\\\"")
 }
 
-fn extract_script_metadata_deps(script: &str) -> Vec<String> {
+fn extract_script_metadata_deps(script: &str) -> Result<Vec<String>> {
     let mut ss = String::new();
     let mut meta_start = false;
     for line in script.lines() {
@@ -280,18 +280,45 @@ fn extract_script_metadata_deps(script: &str) -> Vec<String> {
         }
     }
 
-    let parsed = ss.parse::<toml::Value>().unwrap();
+    let parsed = ss.parse::<toml::Value>().map_err(|err| {
+        ErrorCode::SemanticError(format!(
+            "Failed to parse UDF script metadata as TOML: {err}"
+        ))
+    })?;
 
     if parsed.get("dependencies").is_none() {
-        return Vec::new();
+        return Ok(Vec::new());
     }
 
-    if let Some(deps) = parsed["dependencies"].as_array() {
+    let deps = if let Some(deps) = parsed["dependencies"].as_array() {
         deps.iter()
             .filter_map(|value| value.as_str().map(|item| item.to_string()))
             .collect()
     } else {
         Vec::new()
+    };
+    Ok(deps)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_script_metadata_deps_returns_error_for_malformed_toml() {
+        let err = extract_script_metadata_deps(
+            r#"# /// script
+# dependencies = [
+# ///
+"#,
+        )
+        .expect_err("malformed UDF script metadata should return an error");
+
+        assert_eq!(err.code(), ErrorCode::SEMANTIC_ERROR);
+        assert!(
+            err.message()
+                .contains("Failed to parse UDF script metadata as TOML")
+        );
     }
 }
 
@@ -607,7 +634,7 @@ impl UdfAdapter for FullTypeCheckAdapter {
             ErrorCode::SemanticError(format!("Failed to parse UDF code as utf-8: {err}"))
         })?;
         let import_assets = self.build_udf_cloud_imports(&imports)?;
-        let mut merged_packages = extract_script_metadata_deps(&resolved_code);
+        let mut merged_packages = extract_script_metadata_deps(&resolved_code)?;
         merged_packages.extend_from_slice(&packages);
         let input_types = arg_types.iter().map(udf_type_string).collect::<Vec<_>>();
         let result_type = udf_type_string(&return_type);