Merge pull request #1172 from input-output-hk/ensemble/1095/client-snapshot-download-progress-with-json-output

Client snapshot download progress with json output

Author: Alenar

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-client"
-version = "0.3.34"
+version = "0.3.35"
 description = "A Mithril Client"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-client/Cargo.toml

@@ -1,12 +1,10 @@
-use std::{path::Path, sync::Arc};
-
 use async_recursion::async_recursion;
 use async_trait::async_trait;
 use futures::StreamExt;
-use indicatif::ProgressBar;
 use reqwest::{Client, Response, StatusCode};
 use semver::Version;
 use slog_scope::debug;
+use std::{path::Path, sync::Arc};
 use thiserror::Error;
 use tokio::{fs, io::AsyncWriteExt, sync::RwLock};
 
@@ -15,6 +13,8 @@ use mockall::automock;
 
 use mithril_common::{StdError, MITHRIL_API_VERSION_HEADER};
 
+use crate::utils::DownloadProgressReporter;
+
 /// Error tied with the Aggregator client
 #[derive(Error, Debug)]
 pub enum AggregatorHTTPClientError {
@@ -56,7 +56,7 @@ pub trait AggregatorClient: Sync + Send {
         &self,
         url: &str,
         filepath: &Path,
-        progress_bar: ProgressBar,
+        progress_reporter: DownloadProgressReporter,
     ) -> Result<(), AggregatorHTTPClientError>;
 
     /// Test if the given URL points to a valid location & existing content.
@@ -183,7 +183,7 @@ impl AggregatorClient for AggregatorHTTPClient {
         &self,
         url: &str,
         filepath: &Path,
-        progress_bar: ProgressBar,
+        progress_reporter: DownloadProgressReporter,
     ) -> Result<(), AggregatorHTTPClientError> {
         let response = self.get(url).await?;
         let mut local_file = fs::File::create(filepath).await.map_err(|e| {
@@ -215,7 +215,7 @@ impl AggregatorClient for AggregatorHTTPClient {
                 }
             })?;
             downloaded_bytes += chunk.len() as u64;
-            progress_bar.set_position(downloaded_bytes);
+            progress_reporter.report(downloaded_bytes);
         }
 
         Ok(())

mithril-client/src/aggregator_client/http_client.rs

@@ -1,20 +1,20 @@
 //! This module contains a struct to exchange snapshot information with the Aggregator
 
+use slog_scope::warn;
 use std::{
     path::{Path, PathBuf},
     sync::Arc,
 };
+use thiserror::Error;
 
-use indicatif::ProgressBar;
 use mithril_common::{
     entities::Snapshot,
     messages::{SnapshotListItemMessage, SnapshotListMessage, SnapshotMessage},
     StdResult,
 };
-use slog_scope::warn;
-use thiserror::Error;
 
-use super::AggregatorClient;
+use crate::aggregator_client::AggregatorClient;
+use crate::utils::DownloadProgressReporter;
 
 /// Error for the Snapshot client
 #[derive(Error, Debug)]
@@ -64,7 +64,7 @@ impl SnapshotClient {
         &self,
         snapshot: &Snapshot,
         download_dir: &Path,
-        progress_bar: ProgressBar,
+        progress_reporter: DownloadProgressReporter,
     ) -> StdResult<PathBuf> {
         let filepath = PathBuf::new()
             .join(download_dir)
@@ -74,7 +74,7 @@ impl SnapshotClient {
             if self.http_client.probe(url).await.is_ok() {
                 match self
                     .http_client
-                    .download(url, &filepath, progress_bar)
+                    .download(url, &filepath, progress_reporter)
                     .await
                 {
                     Ok(()) => return Ok(filepath),

mithril-client/src/aggregator_client/snapshot_client.rs

@@ -1,11 +1,12 @@
-use std::{path::PathBuf, sync::Arc};
-
 use clap::Parser;
 use config::{builder::DefaultState, Config, ConfigBuilder};
-use indicatif::ProgressDrawTarget;
+use std::{path::PathBuf, sync::Arc};
+
 use mithril_common::{messages::FromMessageAdapter, StdResult};
 
-use crate::{dependencies::DependenciesBuilder, FromSnapshotMessageAdapter};
+use crate::{
+    dependencies::DependenciesBuilder, utils::ProgressOutputType, FromSnapshotMessageAdapter,
+};
 
 /// Clap command to download the snapshot and verify the certificate.
 #[derive(Parser, Debug, Clone)]
@@ -36,17 +37,17 @@ impl SnapshotDownloadCommand {
         let snapshot_service = dependencies_builder.get_snapshot_service().await?;
         let snapshot_entity =
             FromSnapshotMessageAdapter::adapt(snapshot_service.show(&self.digest).await?);
-        let progress_target = if self.json {
-            ProgressDrawTarget::hidden()
+        let progress_output_type = if self.json {
+            ProgressOutputType::JsonReporter
         } else {
-            ProgressDrawTarget::stdout()
+            ProgressOutputType::TTY
         };
         let filepath = snapshot_service
             .download(
                 &snapshot_entity,
                 &self.download_dir,
                 &config.get_string("genesis_verification_key")?,
-                progress_target,
+                progress_output_type,
             )
             .await?;
 

mithril-client/src/commands/snapshot/download.rs

@@ -1,7 +1,7 @@
 use anyhow::Context;
 use async_trait::async_trait;
 use futures::Future;
-use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressState, ProgressStyle};
+use indicatif::{MultiProgress, ProgressBar, ProgressState, ProgressStyle};
 use slog_scope::{debug, warn};
 use std::{
     fmt::Write,
@@ -24,7 +24,10 @@ use mithril_common::{
 
 use crate::{
     aggregator_client::{AggregatorHTTPClientError, CertificateClient, SnapshotClient},
-    utils::{SnapshotUnpacker, SnapshotUnpackerError},
+    utils::{
+        DownloadProgressReporter, ProgressOutputType, ProgressPrinter, SnapshotUnpacker,
+        SnapshotUnpackerError,
+    },
 };
 
 /// [SnapshotService] related errors.
@@ -72,7 +75,7 @@ pub trait SnapshotService: Sync + Send {
         snapshot_entity: &SignedEntity<Snapshot>,
         pathdir: &Path,
         genesis_verification_key: &str,
-        progress_target: ProgressDrawTarget,
+        progress_output_type: ProgressOutputType,
     ) -> StdResult<PathBuf>;
 }
 
@@ -210,20 +213,20 @@ impl SnapshotService for MithrilClientSnapshotService {
         snapshot_entity: &SignedEntity<Snapshot>,
         download_dir: &Path,
         genesis_verification_key: &str,
-        progress_target: ProgressDrawTarget,
+        progress_output_type: ProgressOutputType,
     ) -> StdResult<PathBuf> {
         debug!("Snapshot service: download.");
 
         let db_dir = download_dir.join("db");
-        let progress_bar = MultiProgress::with_draw_target(progress_target);
-        progress_bar.println("1/7 - Checking local disk info…")?;
+        let progress_bar = ProgressPrinter::new(progress_output_type, 7);
+        progress_bar.report_step(1, "Checking local disk info…")?;
         let unpacker = SnapshotUnpacker;
 
         if let Err(e) = unpacker.check_prerequisites(&db_dir, snapshot_entity.artifact.size) {
             self.check_disk_space_error(e)?;
         }
 
-        progress_bar.println("2/7 - Fetching the certificate's information…")?;
+        progress_bar.report_step(2, "Fetching the certificate's information…")?;
         let certificate = self
             .certificate_client
             .get(&snapshot_entity.certificate_id)
@@ -234,23 +237,27 @@ impl SnapshotService for MithrilClientSnapshotService {
                 )
             })?;
 
-        progress_bar.println("3/7 - Verifying the certificate chain…")?;
+        progress_bar.report_step(3, "Verifying the certificate chain…")?;
         let verifier = self.verify_certificate_chain(genesis_verification_key, &certificate);
         self.wait_spinner(&progress_bar, verifier).await?;
 
-        progress_bar.println("4/7 - Downloading the snapshot…")?;
+        progress_bar.report_step(4, "Downloading the snapshot…")?;
         let pb = progress_bar.add(ProgressBar::new(snapshot_entity.artifact.size));
         pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
             .unwrap()
             .with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
             .progress_chars("#>-"));
         let snapshot_path = self
             .snapshot_client
-            .download(&snapshot_entity.artifact, download_dir, pb)
+            .download(
+                &snapshot_entity.artifact,
+                download_dir,
+                DownloadProgressReporter::new(pb, progress_output_type),
+            )
             .await
             .with_context(|| format!("Could not download file in '{}'", download_dir.display()))?;
 
-        progress_bar.println("5/7 - Unpacking the snapshot…")?;
+        progress_bar.report_step(5, "Unpacking the snapshot…")?;
         let unpacker = unpacker.unpack_snapshot(&snapshot_path, &db_dir);
         self.wait_spinner(&progress_bar, unpacker).await?;
 
@@ -262,14 +269,14 @@ impl SnapshotService for MithrilClientSnapshotService {
             );
         };
 
-        progress_bar.println("6/7 - Computing the snapshot digest…")?;
+        progress_bar.report_step(6, "Computing the snapshot digest…")?;
         let unpacked_snapshot_digest = self
             .immutable_digester
             .compute_digest(&db_dir, &certificate.beacon)
             .await
             .with_context(|| format!("Could not compute digest in '{}'", db_dir.display()))?;
 
-        progress_bar.println("7/7 - Verifying the snapshot signature…")?;
+        progress_bar.report_step(7, "Verifying the snapshot signature…")?;
         let expected_message = {
             let mut protocol_message = certificate.protocol_message.clone();
             protocol_message.set_message_part(
@@ -550,7 +557,7 @@ mod tests {
                 &snapshot,
                 &test_path,
                 &genesis_verification_key.to_json_hex().unwrap(),
-                ProgressDrawTarget::hidden(),
+                ProgressOutputType::Hidden,
             )
             .await
             .expect("Snapshot download should succeed.");
@@ -590,7 +597,7 @@ mod tests {
                 &snapshot,
                 &test_path,
                 &genesis_verification_key.to_json_hex().unwrap(),
-                ProgressDrawTarget::hidden(),
+                ProgressOutputType::Hidden,
             )
             .await
             .expect("Snapshot download should succeed.");
@@ -636,7 +643,7 @@ mod tests {
                 &signed_entity,
                 &test_path,
                 &genesis_verification_key.to_json_hex().unwrap(),
-                ProgressDrawTarget::hidden(),
+                ProgressOutputType::Hidden,
             )
             .await
             .expect_err("Snapshot digest comparison should fail.");
@@ -684,7 +691,7 @@ mod tests {
                 &snapshot,
                 &test_path,
                 &genesis_verification_key.to_json_hex().unwrap(),
-                ProgressDrawTarget::hidden(),
+                ProgressOutputType::Hidden,
             )
             .await
             .expect_err("Snapshot download should fail.");

mithril-client/src/services/snapshot.rs

@@ -1,6 +1,8 @@
 //! Utilities module
 //! This module contains tools needed mostly in services layers.
 
+mod progress_reporter;
 mod unpacker;
 
+pub use progress_reporter::*;
 pub use unpacker::*;