Skip to content

Commit

Permalink
feat(import): add a progress meter for getting an ETA
Browse files Browse the repository at this point in the history
  • Loading branch information
ctron committed Mar 8, 2024
1 parent 1062ccb commit 9074549
Show file tree
Hide file tree
Showing 4 changed files with 35 additions and 23 deletions.
23 changes: 2 additions & 21 deletions backend/importer/src/csaf/mod.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::progress::init_log_and_progress;
use ::csaf::document::Category;
use ::csaf::Csaf;
use csaf_walker::retrieve::RetrievingVisitor;
Expand Down Expand Up @@ -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?;

Expand Down Expand Up @@ -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)?;
Expand Down
1 change: 1 addition & 0 deletions backend/importer/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use clap::Subcommand;
use trustify_common::config::Database;

mod csaf;
mod progress;
mod sbom;

#[derive(Subcommand, Debug)]
Expand Down
26 changes: 26 additions & 0 deletions backend/importer/src/progress.rs
Original file line number Diff line number Diff line change
@@ -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),
}
}
8 changes: 6 additions & 2 deletions backend/importer/src/sbom/mod.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::progress::init_log_and_progress;
use sbom_walker::{
retrieve::RetrievingVisitor,
source::{DispatchSource, FileSource, HttpSource},
Expand Down Expand Up @@ -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");

Expand Down Expand Up @@ -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)
}
Expand Down

0 comments on commit 9074549

Please sign in to comment.