feat: introduct db context in graph and put status cache there

dejanb · dejanb · commit 3b7acd9edeca · 2025-02-06T14:00:53.000+01:00
diff --git a/modules/importer/src/runner/csaf/mod.rs b/modules/importer/src/runner/csaf/mod.rs
@@ -12,7 +12,6 @@ use crate::{
     },
     server::context::WalkerProgress,
 };
-use anyhow::anyhow;
 use csaf_walker::{
     metadata::MetadataRetriever,
     retrieve::RetrievingVisitor,
@@ -25,10 +24,7 @@ use reqwest::StatusCode;
 use std::collections::HashSet;
 use std::{sync::Arc, time::SystemTime};
 use tracing::instrument;
-use trustify_module_ingestor::{
-    graph::Graph,
-    service::{advisory::StatusCache, IngestorService},
-};
+use trustify_module_ingestor::{graph::Graph, service::IngestorService};
 use url::Url;
 use walker_common::fetcher::{Fetcher, FetcherOptions};
 
@@ -67,18 +63,10 @@ impl super::ImportRunner {
         };
 
         // storage (called by validator)
-
-        let mut status_cache = StatusCache::new();
-        status_cache
-            .load_statuses(&self.db)
-            .await
-            .map_err(|_| anyhow!("Failed to load statuses"))?;
-
         let ingestor = IngestorService {
             graph: Graph::new(self.db.clone()),
             storage: self.storage.clone(),
             analysis: self.analysis.clone(),
-            status_cache,
         };
 
         let storage = storage::StorageVisitor {
diff --git a/modules/ingestor/src/graph/db_context.rs b/modules/ingestor/src/graph/db_context.rs
@@ -0,0 +1,55 @@
+use std::collections::HashMap;
+
+use crate::graph::error::Error;
+use sea_orm::ConnectionTrait;
+use sea_orm::EntityTrait;
+use trustify_entity::status;
+use uuid::Uuid;
+
+#[derive(Debug, Clone)]
+pub struct DbContext {
+    pub status_cache: HashMap<String, Uuid>,
+}
+
+impl DbContext {
+    pub fn new() -> Self {
+        Self {
+            status_cache: HashMap::new(),
+        }
+    }
+
+    pub async fn load_statuses(&mut self, connection: &impl ConnectionTrait) -> Result<(), Error> {
+        self.status_cache.clear();
+        let statuses = status::Entity::find().all(connection).await?;
+        statuses
+            .iter()
+            .map(|s| self.status_cache.insert(s.slug.clone(), s.id))
+            .for_each(drop);
+
+        Ok(())
+    }
+
+    pub async fn get_status_id(
+        &mut self,
+        status: &str,
+        connection: &impl ConnectionTrait,
+    ) -> Result<Uuid, Error> {
+        if let Some(s) = self.status_cache.get(status) {
+            return Ok(*s);
+        }
+
+        // If not found, reload the cache and check again
+        self.load_statuses(connection).await?;
+
+        self.status_cache
+            .get(status)
+            .cloned()
+            .ok_or_else(|| crate::graph::error::Error::InvalidStatus(status.to_string()))
+    }
+}
+
+impl Default for DbContext {
+    fn default() -> Self {
+        Self::new()
+    }
+}
diff --git a/modules/ingestor/src/graph/mod.rs b/modules/ingestor/src/graph/mod.rs
@@ -1,19 +1,24 @@
 pub mod advisory;
 pub mod cpe;
+pub mod db_context;
 pub mod error;
 pub mod organization;
 pub mod product;
 pub mod purl;
 pub mod sbom;
 pub mod vulnerability;
 
+use db_context::DbContext;
 use sea_orm::DbErr;
 use std::fmt::Debug;
 use std::ops::{Deref, DerefMut};
+use std::sync::Arc;
+use tokio::sync::Mutex;
 
 #[derive(Debug, Clone)]
 pub struct Graph {
     pub(crate) db: trustify_common::db::Database,
+    pub(crate) db_context: Arc<Mutex<DbContext>>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -26,7 +31,10 @@ pub enum Error<E: Send> {
 
 impl Graph {
     pub fn new(db: trustify_common::db::Database) -> Self {
-        Self { db }
+        Self {
+            db,
+            db_context: Arc::new(Mutex::new(DbContext::new())),
+        }
     }
 }
 
diff --git a/modules/ingestor/src/service/advisory/csaf/creator.rs b/modules/ingestor/src/service/advisory/csaf/creator.rs
@@ -8,10 +8,7 @@ use crate::{
         Graph,
     },
     service::{
-        advisory::{
-            csaf::{product_status::ProductStatus, util::ResolveProductIdCache},
-            StatusCache,
-        },
+        advisory::csaf::{product_status::ProductStatus, util::ResolveProductIdCache},
         Error,
     },
 };
@@ -96,7 +93,6 @@ impl<'a> StatusCreator<'a> {
         &mut self,
         graph: &Graph,
         connection: &C,
-        mut status_cache: StatusCache,
     ) -> Result<(), Error> {
         let mut product_status_models = Vec::new();
         let mut purls = PurlCreator::new();
@@ -108,11 +104,10 @@ impl<'a> StatusCreator<'a> {
         let mut product_version_ranges = Vec::new();
 
         let product_statuses = self.products.clone();
+        let mut db_context = graph.db_context.lock().await;
 
         for product in product_statuses {
-            let status_id = status_cache
-                .get_status_id(product.status, connection)
-                .await?;
+            let status_id = db_context.get_status_id(product.status, connection).await?;
 
             // There should be only a few organizations per document,
             // so simple caching should work here.
@@ -248,7 +243,7 @@ impl<'a> StatusCreator<'a> {
                                     purl,
                                     scheme,
                                     spec,
-                                    status_cache
+                                    db_context
                                         .get_status_id(&Status::Affected.to_string(), connection)
                                         .await?,
                                 );
diff --git a/modules/ingestor/src/service/advisory/csaf/loader.rs b/modules/ingestor/src/service/advisory/csaf/loader.rs
@@ -8,10 +8,7 @@ use crate::{
     },
     model::IngestResult,
     service::{
-        advisory::{
-            csaf::{util::gen_identifier, StatusCreator},
-            StatusCache,
-        },
+        advisory::csaf::{util::gen_identifier, StatusCreator},
         Error, Warnings,
     },
 };
@@ -91,7 +88,6 @@ impl<'g> CsafLoader<'g> {
         labels: impl Into<Labels> + Debug,
         csaf: Csaf,
         digests: &Digests,
-        status_cache: StatusCache,
     ) -> Result<IngestResult, Error> {
         let warnings = Warnings::new();
 
@@ -116,7 +112,7 @@ impl<'g> CsafLoader<'g> {
             .await?;
 
         for vuln in csaf.vulnerabilities.iter().flatten() {
-            self.ingest_vulnerability(&csaf, &advisory, vuln, &warnings, &tx, status_cache.clone())
+            self.ingest_vulnerability(&csaf, &advisory, vuln, &warnings, &tx)
                 .await?;
         }
 
@@ -142,7 +138,6 @@ impl<'g> CsafLoader<'g> {
         vulnerability: &Vulnerability,
         report: &dyn ReportSink,
         connection: &C,
-        status_cache: StatusCache,
     ) -> Result<(), Error> {
         let Some(cve_id) = &vulnerability.cve else {
             return Ok(());
@@ -173,14 +168,8 @@ impl<'g> CsafLoader<'g> {
             .await?;
 
         if let Some(product_status) = &vulnerability.product_status {
-            self.ingest_product_statuses(
-                csaf,
-                &advisory_vulnerability,
-                product_status,
-                connection,
-                status_cache,
-            )
-            .await?;
+            self.ingest_product_statuses(csaf, &advisory_vulnerability, product_status, connection)
+                .await?;
         }
 
         for score in vulnerability.scores.iter().flatten() {
@@ -211,7 +200,6 @@ impl<'g> CsafLoader<'g> {
         advisory_vulnerability: &AdvisoryVulnerabilityContext<'_>,
         product_status: &ProductStatus,
         connection: &C,
-        status_cache: StatusCache,
     ) -> Result<(), Error> {
         let mut creator = StatusCreator::new(
             csaf,
@@ -226,7 +214,7 @@ impl<'g> CsafLoader<'g> {
         creator.add_all(&product_status.known_not_affected, "not_affected");
         creator.add_all(&product_status.known_affected, "affected");
 
-        creator.create(self.graph, connection, status_cache).await?;
+        creator.create(self.graph, connection).await?;
 
         Ok(())
     }
@@ -250,12 +238,7 @@ mod test {
         let (csaf, digests): (Csaf, _) = document("csaf/CVE-2023-20862.json").await?;
         let loader = CsafLoader::new(&graph);
         loader
-            .load(
-                ("file", "CVE-2023-20862.json"),
-                csaf,
-                &digests,
-                StatusCache::default(),
-            )
+            .load(("file", "CVE-2023-20862.json"), csaf, &digests)
             .await?;
 
         let loaded_vulnerability = graph.get_vulnerability("CVE-2023-20862", &ctx.db).await?;
@@ -329,9 +312,7 @@ mod test {
         let loader = CsafLoader::new(&graph);
 
         let (csaf, digests): (Csaf, _) = document("csaf/rhsa-2024_3666.json").await?;
-        loader
-            .load(("source", "test"), csaf, &digests, StatusCache::default())
-            .await?;
+        loader.load(("source", "test"), csaf, &digests).await?;
 
         let loaded_vulnerability = graph.get_vulnerability("CVE-2024-23672", &ctx.db).await?;
         assert!(loaded_vulnerability.is_some());
@@ -372,9 +353,7 @@ mod test {
         let loader = CsafLoader::new(&graph);
 
         let (csaf, digests): (Csaf, _) = document("csaf/cve-2023-0044.json").await?;
-        loader
-            .load(("source", "test"), csaf, &digests, StatusCache::default())
-            .await?;
+        loader.load(("source", "test"), csaf, &digests).await?;
 
         let loaded_vulnerability = graph.get_vulnerability("CVE-2023-0044", &ctx.db).await?;
         assert!(loaded_vulnerability.is_some());
diff --git a/modules/ingestor/src/service/advisory/mod.rs b/modules/ingestor/src/service/advisory/mod.rs
@@ -1,59 +1,3 @@
-use std::collections::HashMap;
-
-use crate::graph::error::Error;
-use sea_orm::ConnectionTrait;
-use sea_orm::EntityTrait;
-use trustify_entity::status;
-use uuid::Uuid;
-
 pub mod csaf;
 pub mod cve;
 pub mod osv;
-
-#[derive(Debug, Clone)]
-pub struct StatusCache {
-    pub cache: HashMap<String, Uuid>,
-}
-
-impl StatusCache {
-    pub fn new() -> Self {
-        Self {
-            cache: HashMap::new(),
-        }
-    }
-
-    pub async fn load_statuses(&mut self, connection: &impl ConnectionTrait) -> Result<(), Error> {
-        self.cache.clear();
-        let statuses = status::Entity::find().all(connection).await?;
-        statuses
-            .iter()
-            .map(|s| self.cache.insert(s.slug.clone(), s.id))
-            .for_each(drop);
-
-        Ok(())
-    }
-
-    pub async fn get_status_id(
-        &mut self,
-        status: &str,
-        connection: &impl ConnectionTrait,
-    ) -> Result<Uuid, Error> {
-        if let Some(s) = self.cache.get(status) {
-            return Ok(*s);
-        }
-
-        // If not found, reload the cache and check again
-        self.load_statuses(connection).await?;
-
-        self.cache
-            .get(status)
-            .cloned()
-            .ok_or_else(|| crate::graph::error::Error::InvalidStatus(status.to_string()))
-    }
-}
-
-impl Default for StatusCache {
-    fn default() -> Self {
-        Self::new()
-    }
-}
diff --git a/modules/ingestor/src/service/dataset/mod.rs b/modules/ingestor/src/service/dataset/mod.rs
@@ -21,8 +21,6 @@ use trustify_common::hashing::Digests;
 use trustify_entity::labels::Labels;
 use trustify_module_storage::{service::dispatch::DispatchBackend, service::StorageBackend};
 
-use super::advisory::StatusCache;
-
 pub struct DatasetLoader<'g> {
     graph: &'g Graph,
     storage: &'g DispatchBackend,
@@ -45,9 +43,6 @@ impl<'g> DatasetLoader<'g> {
 
         let mut zip = zip::ZipArchive::new(Cursor::new(buffer))?;
 
-        let mut status_cache = StatusCache::new();
-        status_cache.load_statuses(&self.graph.db).await?;
-
         for i in 0..zip.len() {
             let mut file = zip.by_index(i)?;
 
@@ -111,20 +106,11 @@ impl<'g> DatasetLoader<'g> {
                             .await
                             .map_err(|err| Error::Storage(anyhow!("{err}")))?;
 
-                        let sc = status_cache.clone();
-
                         // We need to box it, to work around async recursion limits
                         let result = Box::pin({
                             async move {
                                 format
-                                    .load(
-                                        self.graph,
-                                        labels,
-                                        None,
-                                        &Digests::digest(&data),
-                                        &data,
-                                        sc,
-                                    )
+                                    .load(self.graph, labels, None, &Digests::digest(&data), &data)
                                     .await
                             }
                         })
diff --git a/modules/ingestor/src/service/format.rs b/modules/ingestor/src/service/format.rs
diff --git a/modules/ingestor/src/service/mod.rs b/modules/ingestor/src/service/mod.rs
