Fix a few bugs

Author: habedi

docs/README.md

@@ -8,7 +8,7 @@ The table below includes the information about all SQL functions exposed by Infe
 | 2  | `infera_unload_model(name VARCHAR)`                          | `BOOLEAN`        | Unloads a model, freeing its associated resources. Returns `true` on success.                                                               |
 | 3  | `infera_set_autoload_dir(path VARCHAR)`                      | `VARCHAR (JSON)` | Scans a directory for `.onnx` files, loads them automatically, and returns a JSON report of loaded models and any errors.                   |
 | 4  | `infera_get_loaded_models()`                                 | `VARCHAR (JSON)` | Returns a JSON array containing the names of all currently loaded models.                                                                   |
-| 5  | `infera_get_model_info(name VARCHAR)`                        | `VARCHAR (JSON)` | Returns a JSON object with metadata about a specific loaded model, including its name, input/output shapes, and status.                     |
+| 5  | `infera_get_model_info(name VARCHAR)`                        | `VARCHAR (JSON)` | Returns a JSON object with metadata about a specific loaded model (name, input/output shapes). If the model is not loaded, this function raises an error. |
 | 6  | `infera_predict(name VARCHAR, features... FLOAT)`            | `FLOAT`          | Performs inference on a batch of data, returning a single float value for each input row.                                                   |
 | 7  | `infera_predict_multi(name VARCHAR, features... FLOAT)`      | `VARCHAR (JSON)` | Performs inference and returns all outputs as a JSON-encoded array. This is useful for models that produce multiple predictions per sample. |
 | 8  | `infera_predict_multi_list(name VARCHAR, features... FLOAT)` | `LIST[FLOAT]`    | Performs inference and returns all outputs as a typed list of floats. Useful for multi-output models without JSON parsing.                  |
@@ -45,7 +45,7 @@ select infera_is_model_loaded('local_model');
 select infera_get_loaded_models();
 -- Output: ["local_model", "remote_model"]
 
--- Get information about a specific model
+-- Get information about a specific model (throws an error if the model is not loaded)
 select infera_get_model_info('local_model');
 -- Output: {"name":"local_model","input_shape":[-1,3],"output_shape":[-1,1],"loaded":true}
 
@@ -97,7 +97,7 @@ select infera_get_loaded_models();
 select infera_is_model_loaded('squeezenet');
 -- Output: true or false
 
--- Get detailed metadata for a specific model
+-- Get detailed metadata for a specific model (errors if the model is not loaded)
 select infera_get_model_info('squeezenet');
 /* Output:
 {
@@ -121,14 +121,14 @@ select infera_set_autoload_dir('path/to/your/models');
 select infera_clear_cache();
 -- Output: true
 
--- Get cache statistics
+-- Get cache statistics (field names as returned by the function)
 select infera_get_cache_info();
 /* Output:
 {
-  "path": "/path/to/cache",
-  "size_bytes": 204800,
+  "cache_dir": "/path/to/cache",
+  "total_size_bytes": 204800,
   "file_count": 10,
-  "size_limit": 10485760
+  "size_limit_bytes": 10485760
 }
 */
 ```

infera/bindings/infera_extension.cpp

@@ -227,6 +227,7 @@ static void Predict(DataChunk &args, ExpressionState &state, Vector &result) {
 
   infera::InferaInferenceResult res = infera::infera_predict(model_name_str.c_str(), features.data(), batch_size, feature_count);
   if (res.status != 0) {
+    infera::infera_free_result(res);
     throw InvalidInputException("Inference failed for model '" + model_name_str + "': " + GetInferaError());
   }
   if (res.rows != batch_size || res.cols != 1) {
@@ -434,16 +435,22 @@ static void GetModelInfo(DataChunk &args, ExpressionState &state, Vector &result
     throw InvalidInputException("Model name cannot be NULL");
   }
   std::string model_name_str = model_name.ToString();
-  char *json_meta = infera::infera_get_model_info(model_name_str.c_str());
+  char *json_meta_c = infera::infera_get_model_info(model_name_str.c_str());
+
+  // Convert to std::string and free the C string immediately to avoid leaks
+  std::string json_meta = json_meta_c ? std::string(json_meta_c) : std::string();
+  if (json_meta_c) {
+    infera::infera_free(json_meta_c);
+  }
 
-  if (json_meta == nullptr) {
-    throw InvalidInputException("Failed to get info for model '" + model_name_str + "': " + GetInferaError());
+  // If Rust returned an error JSON, surface it as a DuckDB error per contract/tests
+  if (json_meta.empty() || json_meta.find("\"error\"") != std::string::npos) {
+    throw InvalidInputException("Failed to get info for model '" + model_name_str + "'");
   }
 
   result.SetVectorType(VectorType::CONSTANT_VECTOR);
   ConstantVector::GetData<string_t>(result)[0] = StringVector::AddString(result, json_meta);
   ConstantVector::SetNull(result, false);
-  infera::infera_free(json_meta);
 }
 
 /**

infera/src/ffi_utils.rs

@@ -69,7 +69,45 @@ pub unsafe extern "C" fn infera_free(ptr: *mut c_char) {
 /// on the same result will lead to undefined behavior.
 #[no_mangle]
 pub unsafe extern "C" fn infera_free_result(res: InferaInferenceResult) {
-    if !res.data.is_null() && res.len > 0 {
-        let _ = Vec::from_raw_parts(res.data, res.len, res.len);
+    if !res.data.is_null() {
+        // SAFETY: `res.data` was allocated from a Box<[f32]> via `into_raw` with length `res.len`.
+        // Reconstruct the slice pointer and drop it to free the allocation correctly.
+        let slice_ptr: *mut [f32] = std::ptr::slice_from_raw_parts_mut(res.data, res.len);
+        let _ = Box::from_raw(slice_ptr);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_infera_free_result_zero_len_non_null() {
+        // Allocate an empty slice on heap and convert to raw
+        let empty: Box<[f32]> = Vec::<f32>::new().into_boxed_slice();
+        let ptr = Box::into_raw(empty) as *mut f32;
+        let res = InferaInferenceResult {
+            data: ptr,
+            len: 0,
+            rows: 0,
+            cols: 0,
+            status: 0,
+        };
+        unsafe { infera_free_result(res) }; // should not panic or leak
+    }
+
+    #[test]
+    fn test_infera_free_result_non_empty() {
+        let data: Vec<f32> = vec![1.0, 2.0, 3.0];
+        let len = data.len();
+        let ptr = Box::into_raw(data.into_boxed_slice()) as *mut f32;
+        let res = InferaInferenceResult {
+            data: ptr,
+            len,
+            rows: 1,
+            cols: len,
+            status: 0,
+        };
+        unsafe { infera_free_result(res) }; // should free without UB
     }
 }

infera/src/lib.rs

@@ -341,10 +341,7 @@ pub extern "C" fn infera_get_cache_info() -> *mut c_char {
             }
         }
 
-        let size_limit = env::var("INFERA_CACHE_SIZE_LIMIT")
-            .ok()
-            .and_then(|s| s.parse::<u64>().ok())
-            .unwrap_or(1024 * 1024 * 1024);
+        let size_limit = config::CONFIG.cache_size_limit;
 
         Ok(json!({
             "cache_dir": cache_dir_str,
@@ -625,4 +622,30 @@ mod tests {
             infera_unload_model(model_name.as_ptr());
         }
     }
+
+    #[test]
+    fn test_infera_get_model_info_nonexistent_returns_error_json() {
+        let name = CString::new("__missing_model__").unwrap();
+        let info_ptr = unsafe { infera_get_model_info(name.as_ptr()) };
+        let info_json = unsafe { CStr::from_ptr(info_ptr).to_str().unwrap() };
+        let value: serde_json::Value = serde_json::from_str(info_json).unwrap();
+        assert!(
+            value.get("error").is_some(),
+            "expected error field in JSON: {}",
+            info_json
+        );
+        unsafe { infera_free(info_ptr) };
+    }
+
+    #[test]
+    fn test_infera_get_cache_info_includes_configured_limit() {
+        let cache_info_ptr = infera_get_cache_info();
+        let cache_info_json = unsafe { CStr::from_ptr(cache_info_ptr).to_str().unwrap() };
+        let value: serde_json::Value = serde_json::from_str(cache_info_json).unwrap();
+        let size_limit = value["size_limit_bytes"]
+            .as_u64()
+            .expect("size_limit_bytes should be u64");
+        assert_eq!(size_limit, crate::config::CONFIG.cache_size_limit);
+        unsafe { infera_free(cache_info_ptr) };
+    }
 }

test/sql/test_edge_cases_more.test

@@ -31,8 +31,8 @@ select infera_unload_model('linear')
 ----
 true
 
-# model info after unload should contain model not found message
-query I
-select position('Model not found' in infera_get_model_info('linear')) > 0
+# model info after unload should error
+statement error
+select infera_get_model_info('linear')
 ----
-true
+Failed to get info for model 'linear'

test/sql/test_integration_and_errors.test

@@ -12,11 +12,11 @@ load 'build/release/extension/infera/infera.duckdb_extension'
 
 # test 1: error handling for non-existent models
 
-# infera_get_model_info for a missing model returns a non-null json containing an error field
-query I
-select infera_get_model_info('nonexistent_model') is null
+# infera_get_model_info for a missing model should error
+statement error
+select infera_get_model_info('nonexistent_model')
 ----
-0
+Failed to get info for model 'nonexistent_model'
 
 # unloading a non-existent model is benign (idempotent) and returns false
 statement ok