cleaning up batching

gtfierro · gtfierro · commit f969a43f5130 · 2025-10-03T11:45:55.000-06:00
diff --git a/lib/src/api.rs b/lib/src/api.rs
@@ -99,8 +99,7 @@ pub struct OntoEnv {
     dependency_graph: DiGraph<GraphIdentifier, (), petgraph::Directed>,
     config: Config,
     failed_resolutions: HashSet<NamedNode>,
-    fetched_in_session: HashSet<OntologyLocation>,
-    io_batch_depth: usize,
+    batch_state: BatchState,
 }
 
 impl std::fmt::Debug for OntoEnv {
@@ -116,18 +115,45 @@ impl std::fmt::Debug for OntoEnv {
     }
 }
 
+#[derive(Default)]
+struct BatchState {
+    depth: usize,
+    seen_locations: HashSet<OntologyLocation>,
+}
+
+impl BatchState {
+    fn begin(&mut self) {
+        if self.depth == 0 {
+            self.seen_locations.clear();
+        }
+        self.depth += 1;
+    }
+
+    fn end(&mut self) {
+        self.depth = self.depth.saturating_sub(1);
+    }
+
+    fn has_seen(&self, location: &OntologyLocation) -> bool {
+        self.seen_locations.contains(location)
+    }
+
+    fn mark_seen(&mut self, location: &OntologyLocation) {
+        self.seen_locations.insert(location.clone());
+    }
+}
+
 struct BatchScope<'a> {
     env: &'a mut OntoEnv,
     completed: bool,
 }
 
 impl<'a> BatchScope<'a> {
     fn enter(env: &'a mut OntoEnv) -> Result<Self> {
-        if env.io_batch_depth == 0 {
-            env.fetched_in_session.clear();
+        env.batch_state.begin();
+        if let Err(err) = env.io.begin_batch() {
+            env.batch_state.end();
+            return Err(err);
         }
-        env.io_batch_depth += 1;
-        env.io.begin_batch()?;
         Ok(Self {
             env,
             completed: false,
@@ -137,7 +163,7 @@ impl<'a> BatchScope<'a> {
     fn run<T>(mut self, f: impl FnOnce(&mut OntoEnv) -> Result<T>) -> Result<T> {
         let result = f(self.env);
         let end_result = self.env.io.end_batch();
-        self.env.io_batch_depth = self.env.io_batch_depth.saturating_sub(1);
+        self.env.batch_state.end();
         self.completed = true;
         match (result, end_result) {
             (Ok(value), Ok(())) => Ok(value),
@@ -159,10 +185,15 @@ impl<'a> Drop for BatchScope<'a> {
         if let Err(err) = self.env.io.end_batch() {
             error!("Failed to finalize batched RDF write: {err}");
         }
-        self.env.io_batch_depth = self.env.io_batch_depth.saturating_sub(1);
+        self.env.batch_state.end();
     }
 }
 
+enum FetchOutcome {
+    Reused(GraphIdentifier),
+    Loaded(Ontology),
+}
+
 impl OntoEnv {
     fn new(env: Environment, io: Box<dyn GraphIO>, config: Config) -> Self {
         Self {
@@ -171,8 +202,7 @@ impl OntoEnv {
             config,
             dependency_graph: DiGraph::new(),
             failed_resolutions: HashSet::new(),
-            fetched_in_session: HashSet::new(),
-            io_batch_depth: 0,
+            batch_state: BatchState::default(),
         }
     }
 
@@ -437,8 +467,7 @@ impl OntoEnv {
             config,
             dependency_graph,
             failed_resolutions: HashSet::new(),
-            fetched_in_session: HashSet::new(),
-            io_batch_depth: 0,
+            batch_state: BatchState::default(),
         })
     }
 
@@ -538,8 +567,7 @@ impl OntoEnv {
             dependency_graph: DiGraph::new(),
             config,
             failed_resolutions: HashSet::new(),
-            fetched_in_session: HashSet::new(),
-            io_batch_depth: 0,
+            batch_state: BatchState::default(),
         };
 
         let _ = ontoenv.update_all(false)?;
@@ -599,6 +627,28 @@ impl OntoEnv {
         })
     }
 
+    fn fetch_location(
+        &mut self,
+        location: OntologyLocation,
+        overwrite: Overwrite,
+        refresh: RefreshStrategy,
+    ) -> Result<FetchOutcome> {
+        if let Some(existing_id) = self.try_reuse_cached(&location, refresh)? {
+            self.batch_state.mark_seen(&location);
+            return Ok(FetchOutcome::Reused(existing_id));
+        }
+
+        if !refresh.is_force() && self.batch_state.has_seen(&location) {
+            if let Some(existing) = self.env.get_ontology_by_location(&location) {
+                return Ok(FetchOutcome::Reused(existing.id().clone()));
+            }
+        }
+
+        let ontology = self.io.add(location.clone(), overwrite)?;
+        self.batch_state.mark_seen(&location);
+        Ok(FetchOutcome::Loaded(ontology))
+    }
+
     fn register_ontologies(
         &mut self,
         ontologies: Vec<Ontology>,
@@ -607,7 +657,7 @@ impl OntoEnv {
         let mut ids = Vec::with_capacity(ontologies.len());
         for ontology in ontologies {
             let id = ontology.id().clone();
-            self.env.add_ontology(ontology);
+            self.env.add_ontology(ontology)?;
             ids.push(id);
         }
 
@@ -628,33 +678,19 @@ impl OntoEnv {
     ) -> Result<GraphIdentifier> {
         self.failed_resolutions.clear();
 
-        if let Some(existing_id) = self.try_reuse_cached(&location, refresh)? {
-            debug!(
-                "Reusing cached ontology {} for location {}",
-                existing_id, location
-            );
-            self.fetched_in_session.insert(location);
-            return Ok(existing_id);
-        }
-
-        if !refresh.is_force() && self.fetched_in_session.contains(&location) {
-            if let Some(existing) = self.env.get_ontology_by_location(&location) {
-                debug!(
-                    "Skipping refetch for {} (already loaded this run)",
-                    location
-                );
-                return Ok(existing.id().clone());
+        match self.fetch_location(location.clone(), overwrite, refresh)? {
+            FetchOutcome::Reused(id) => {
+                debug!("Reusing cached ontology {} for location {}", id, location);
+                Ok(id)
+            }
+            FetchOutcome::Loaded(ont) => {
+                let ids = self.register_ontologies(vec![ont], update_dependencies)?;
+                Ok(ids
+                    .into_iter()
+                    .next()
+                    .expect("registered ontology list should contain new entry"))
             }
         }
-
-        let ont = self.io.add(location.clone(), overwrite)?;
-        let ids = self.register_ontologies(vec![ont], update_dependencies)?;
-        let id = ids
-            .into_iter()
-            .next()
-            .expect("registered ontology list should contain new entry");
-        self.fetched_in_session.insert(location);
-        Ok(id)
     }
 
     fn try_reuse_cached(
@@ -729,53 +765,10 @@ impl OntoEnv {
 
     fn update_all_inner(&mut self, all: bool) -> Result<Vec<GraphIdentifier>> {
         self.failed_resolutions.clear();
-        // remove ontologies which are no longer present in the search directories
-        // remove ontologies which are no longer present in the search directories
-        for graphid in self.missing_ontologies() {
-            self.io.remove(&graphid)?;
-            self.env.remove_ontology(&graphid);
-        }
-
-        // now, find all the new and updated ontologies in the search directories
-        // and add them to the environment
-        let updated_files: Vec<OntologyLocation> = if all {
-            let mut set: HashSet<OntologyLocation> = self
-                .env
-                .ontologies()
-                .values()
-                .filter_map(|o| o.location().cloned())
-                .collect();
-            for loc in self.find_files()? {
-                set.insert(loc);
-            }
-            set.into_iter().collect()
-        } else {
-            self.get_updated_locations()?
-        };
+        self.remove_missing_ontologies()?;
 
-        // load all of these files into the environment
-        let mut ontologies: Vec<Ontology> = vec![];
-        for location in updated_files {
-            // if 'strict' mode then fail on any errors when adding the ontology
-            // otherwise just warn
-
-            let result = self.io.add(location.clone(), Overwrite::Allow);
-            if result.is_err() {
-                if self.config.strict {
-                    return Err(result.unwrap_err());
-                } else {
-                    warn!(
-                        "Failed to read ontology file {}: {}",
-                        location,
-                        result.unwrap_err()
-                    );
-                    continue;
-                }
-            }
-
-            let new_ont = result.unwrap();
-            ontologies.push(new_ont);
-        }
+        let updated_files = self.collect_updated_files(all)?;
+        let ontologies = self.load_updated_ontologies(updated_files)?;
 
         let update_ids = self.register_ontologies(ontologies, true)?;
         Ok(update_ids)
@@ -842,6 +835,51 @@ impl OntoEnv {
             .collect()
     }
 
+    fn remove_missing_ontologies(&mut self) -> Result<()> {
+        for graphid in self.missing_ontologies() {
+            self.io.remove(&graphid)?;
+            self.env.remove_ontology(&graphid)?;
+        }
+        Ok(())
+    }
+
+    fn collect_updated_files(&mut self, all: bool) -> Result<Vec<OntologyLocation>> {
+        if all {
+            let mut set: HashSet<OntologyLocation> = self
+                .env
+                .ontologies()
+                .values()
+                .filter_map(|o| o.location().cloned())
+                .collect();
+            for loc in self.find_files()? {
+                set.insert(loc);
+            }
+            Ok(set.into_iter().collect())
+        } else {
+            self.get_updated_locations()
+        }
+    }
+
+    fn load_updated_ontologies(
+        &mut self,
+        updated_files: Vec<OntologyLocation>,
+    ) -> Result<Vec<Ontology>> {
+        let mut ontologies = Vec::with_capacity(updated_files.len());
+        for location in updated_files {
+            match self.io.add(location.clone(), Overwrite::Allow) {
+                Ok(ont) => ontologies.push(ont),
+                Err(err) => {
+                    if self.config.strict {
+                        return Err(err);
+                    } else {
+                        warn!("Failed to read ontology file {}: {}", location, err);
+                    }
+                }
+            }
+        }
+        Ok(ontologies)
+    }
+
     /// Returns a list of all files in the environment which have been updated (added or changed)
     /// Does not return files that have been removed
     pub fn get_updated_locations(&self) -> Result<Vec<OntologyLocation>> {
@@ -998,7 +1036,7 @@ impl OntoEnv {
                 match self.io.add(location, Overwrite::Preserve) {
                     Ok(new_ont) => {
                         let id = new_ont.id().clone();
-                        self.env.add_ontology(new_ont);
+                        self.env.add_ontology(new_ont)?;
                         stack.push_back(id);
                     }
                     Err(e) => {
diff --git a/lib/src/config.rs b/lib/src/config.rs
@@ -9,6 +9,8 @@ use serde::{Deserialize, Serialize};
 use std::io::{BufReader, Write};
 use std::path::{Path, PathBuf};
 
+const DEFAULT_INCLUDE_PATTERNS: &[&str] = &["*.ttl", "*.xml", "*.n3"];
+
 fn vec_pattern_ser<S>(patterns: &Vec<Pattern>, serializer: S) -> Result<S::Ok, S::Error>
 where
     S: serde::Serializer,
@@ -304,9 +306,12 @@ impl ConfigBuilder {
             self.locations.unwrap_or_else(|| vec![root.clone()])
         };
 
-        let includes_str = self
-            .includes
-            .unwrap_or_else(|| vec!["*.ttl".to_string(), "*.xml".to_string(), "*.n3".to_string()]);
+        let includes_str = self.includes.unwrap_or_else(|| {
+            DEFAULT_INCLUDE_PATTERNS
+                .iter()
+                .map(|s| s.to_string())
+                .collect()
+        });
         let excludes_str = self.excludes.unwrap_or_default();
 
         let includes = includes_str
diff --git a/lib/src/environment.rs b/lib/src/environment.rs
@@ -4,8 +4,9 @@
 use crate::io::GraphIO;
 use crate::ontology::{GraphIdentifier, Ontology, OntologyLocation};
 use crate::policy;
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use chrono::prelude::*;
+use log::warn;
 use oxigraph::model::{Graph, NamedNodeRef};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
@@ -84,17 +85,26 @@ impl Environment {
         &self.ontologies
     }
 
-    pub fn add_ontology(&mut self, mut ontology: Ontology) {
+    pub fn add_ontology(&mut self, mut ontology: Ontology) -> Result<()> {
         ontology.last_updated = Some(Utc::now());
-        self.locations
-            .insert(ontology.location().unwrap().clone(), ontology.id().clone());
+        let location = ontology
+            .location()
+            .cloned()
+            .ok_or_else(|| anyhow!("Cannot add ontology {} without a location", ontology.id()))?;
+        self.locations.insert(location, ontology.id().clone());
         self.ontologies.insert(ontology.id().clone(), ontology);
+        Ok(())
     }
 
-    pub fn remove_ontology(&mut self, id: &GraphIdentifier) -> Option<Ontology> {
-        self.locations
-            .remove(self.ontologies.get(id)?.location().unwrap());
-        self.ontologies.remove(id)
+    pub fn remove_ontology(&mut self, id: &GraphIdentifier) -> Result<Option<Ontology>> {
+        if let Some(existing) = self.ontologies.get(id) {
+            if let Some(location) = existing.location() {
+                self.locations.remove(location);
+            } else {
+                warn!("Removing ontology {} without recorded location", id);
+            }
+        }
+        Ok(self.ontologies.remove(id))
     }
 
     pub fn get_modified_time(&self, id: &GraphIdentifier) -> Option<DateTime<Utc>> {
@@ -145,9 +155,7 @@ impl Environment {
 
     /// Returns the first ontology with the given location
     pub fn get_ontology_by_location(&self, location: &OntologyLocation) -> Option<&Ontology> {
-        // choose the first ontology with the given location
-        self.ontologies
-            .values()
-            .find(|&ontology| ontology.location() == Some(location))
+        let id = self.locations.get(location)?;
+        self.ontologies.get(id)
     }
 }
