From 9074549fa42aba3e6654ebc173e8f3e39ffc83dc Mon Sep 17 00:00:00 2001
From: Jens Reimann <ctron@dentrassi.de>
Date: Fri, 8 Mar 2024 13:10:56 +0100
Subject: [PATCH] feat(import): add a progress meter for getting an ETA

---
 backend/importer/src/csaf/mod.rs | 23 ++---------------------
 backend/importer/src/lib.rs      |  1 +
 backend/importer/src/progress.rs | 26 ++++++++++++++++++++++++++
 backend/importer/src/sbom/mod.rs |  8 ++++++--
 4 files changed, 35 insertions(+), 23 deletions(-)
 create mode 100644 backend/importer/src/progress.rs

diff --git a/backend/importer/src/csaf/mod.rs b/backend/importer/src/csaf/mod.rs
index f5a648051..a7447800e 100644
--- a/backend/importer/src/csaf/mod.rs
+++ b/backend/importer/src/csaf/mod.rs
@@ -1,3 +1,4 @@
+use crate::progress::init_log_and_progress;
 use ::csaf::document::Category;
 use ::csaf::Csaf;
 use csaf_walker::retrieve::RetrievingVisitor;
@@ -47,7 +48,7 @@ pub struct ImportCsafCommand {
 
 impl ImportCsafCommand {
     pub async fn run(self) -> anyhow::Result<ExitCode> {
-        let progress = progress();
+        let progress = init_log_and_progress();
 
         let system = InnerSystem::with_config(&self.database).await?;
 
@@ -127,26 +128,6 @@ impl ImportCsafCommand {
     }
 }
 
-fn progress() -> Progress {
-    let mut builder = env_logger::builder();
-    let logger = builder.build();
-
-    match std::io::stdin().is_terminal() {
-        true => {
-            let max_level = logger.filter();
-            let multi = MultiProgress::new();
-
-            let log = LogWrapper::new(multi.clone(), logger);
-            // NOTE: LogWrapper::try_init is buggy and messes up the log levels
-            log::set_boxed_logger(Box::new(log)).unwrap();
-            log::set_max_level(max_level);
-
-            Progress::new(MultiIndicatif(multi))
-        }
-        false => Progress::new(NoProgress),
-    }
-}
-
 /// Process a single, validated advisory
 async fn process(system: &InnerSystem, doc: ValidatedAdvisory) -> anyhow::Result<()> {
     let csaf = serde_json::from_slice::<Csaf>(&doc.data)?;
diff --git a/backend/importer/src/lib.rs b/backend/importer/src/lib.rs
index c67fa5bbe..b2bef2006 100644
--- a/backend/importer/src/lib.rs
+++ b/backend/importer/src/lib.rs
@@ -4,6 +4,7 @@ use clap::Subcommand;
 use trustify_common::config::Database;
 
 mod csaf;
+mod progress;
 mod sbom;
 
 #[derive(Subcommand, Debug)]
diff --git a/backend/importer/src/progress.rs b/backend/importer/src/progress.rs
new file mode 100644
index 000000000..def86c682
--- /dev/null
+++ b/backend/importer/src/progress.rs
@@ -0,0 +1,26 @@
+use indicatif::MultiProgress;
+use indicatif_log_bridge::LogWrapper;
+use std::io::IsTerminal;
+use walker_common::progress::indicatif::MultiIndicatif;
+use walker_common::progress::{NoProgress, Progress};
+
+/// Set up the env_logger and attach a progress interface if we are running on a terminal.
+pub(crate) fn init_log_and_progress() -> Progress {
+    let mut builder = env_logger::builder();
+    let logger = builder.build();
+
+    match std::io::stdin().is_terminal() {
+        true => {
+            let max_level = logger.filter();
+            let multi = MultiProgress::new();
+
+            let log = LogWrapper::new(multi.clone(), logger);
+            // NOTE: LogWrapper::try_init is buggy and messes up the log levels
+            log::set_boxed_logger(Box::new(log)).unwrap();
+            log::set_max_level(max_level);
+
+            Progress::new(MultiIndicatif(multi))
+        }
+        false => Progress::new(NoProgress),
+    }
+}
diff --git a/backend/importer/src/sbom/mod.rs b/backend/importer/src/sbom/mod.rs
index 16bc9e168..85a8b4f6c 100644
--- a/backend/importer/src/sbom/mod.rs
+++ b/backend/importer/src/sbom/mod.rs
@@ -1,3 +1,4 @@
+use crate::progress::init_log_and_progress;
 use sbom_walker::{
     retrieve::RetrievingVisitor,
     source::{DispatchSource, FileSource, HttpSource},
@@ -26,7 +27,7 @@ pub struct ImportSbomCommand {
 
 impl ImportSbomCommand {
     pub async fn run(self) -> anyhow::Result<ExitCode> {
-        env_logger::init();
+        let progress = init_log_and_progress();
 
         log::info!("Ingesting SBOMs");
 
@@ -63,7 +64,10 @@ impl ImportSbomCommand {
 
         // walker
 
-        Walker::new(source).walk(visitor).await?;
+        Walker::new(source)
+            .with_progress(progress)
+            .walk(visitor)
+            .await?;
 
         Ok(ExitCode::SUCCESS)
     }