add progress bars

Cargo.toml

@@ -16,12 +16,14 @@ categories = ["development-tools::profiling"]
 anyhow = { version = "1.0.71" }
 axum = { version = "0.6.18" }
 clap = { version = "4.2.7", features = ["derive", "env"] }
+dialoguer = "0.10.4"
 directories = { version = "5.0.1" }
 flate2 = { version = "1.0.26" }
 futures-util = { version = "0.3.28", features = ["io"] }
 hex = "0.4.3"
 http = { version = "0.2.9" }
 include_dir = { version = "0.7.3" }
+indicatif = "0.17.5"
 once_cell = { version = "1.17.1" }
 remove_dir_all = { version = "0.8.2" }
 reqwest = { version = "0.11.18", default-features = false, features = ["json", "rustls-tls", "stream"] }

src/bin/am/commands/start.rs

@@ -1,3 +1,4 @@
+use crate::interactive;
 use anyhow::{anyhow, bail, Context, Result};
 use autometrics_am::prometheus;
 use axum::body::{self, Body};
@@ -11,6 +12,7 @@ use flate2::read::GzDecoder;
 use futures_util::future::join_all;
 use http::{StatusCode, Uri};
 use include_dir::{include_dir, Dir};
+use indicatif::{ProgressBar, ProgressState, ProgressStyle};
 use once_cell::sync::Lazy;
 use sha2::{Digest, Sha256};
 use std::fs::File;
@@ -35,8 +37,8 @@ static CLIENT: Lazy<reqwest::Client> = Lazy::new(|| {
 
 #[derive(Parser, Clone)]
 pub struct Arguments {
-    #[clap(value_parser = endpoint_parser)]
     /// The endpoint(s) that Prometheus will scrape.
+    #[clap(value_parser = endpoint_parser)]
     metrics_endpoints: Vec<Url>,
 
     /// The Prometheus version to use.
@@ -55,9 +57,10 @@ pub struct Arguments {
     enable_gateway: bool,
 }
 
-pub async fn handle_command(args: Arguments) -> Result<()> {
+pub async fn handle_command(mut args: Arguments) -> Result<()> {
     if args.metrics_endpoints.is_empty() && args.enable_gateway {
-        warn!("No metrics endpoints specified and gateway is not enabled");
+        let endpoint = interactive::user_input("Endpoint")?;
+        args.metrics_endpoints.push(Url::parse(&endpoint)?);
     }
 
     // First let's retrieve the directory for our application to store data in.
@@ -157,13 +160,32 @@ async fn download_prometheus(
         .await?
         .error_for_status()?;
 
+    let total_size = response
+        .content_length()
+        .ok_or_else(|| anyhow!("didn't receive content length"))?;
+    let mut downloaded = 0;
+
+    let pb = ProgressBar::new(total_size);
+
+    // https://github.com/console-rs/indicatif/blob/HEAD/examples/download.rs#L12
+    pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")?
+        .with_key("eta", |state: &ProgressState, w: &mut dyn std::fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap())
+        .progress_chars("=> "));
+
     let mut buffer = BufWriter::new(destination);
 
     while let Some(ref chunk) = response.chunk().await? {
         buffer.write_all(chunk)?;
         hasher.update(chunk);
+
+        let new_size = (downloaded + chunk.len() as u64).min(total_size);
+        downloaded = new_size;
+
+        pb.set_position(downloaded);
     }
 
+    pb.finish_and_clear();
+
     let checksum = hex::encode(hasher.finalize());
 
     Ok(checksum)
@@ -222,16 +244,25 @@ async fn unpack_prometheus(
     // This prefix will be removed from the files in the archive.
     let prefix = format!("prometheus-{prometheus_version}.{os}-{arch}/");
 
+    let pb = ProgressBar::new_spinner();
+    pb.set_style(ProgressStyle::default_spinner());
+    pb.enable_steady_tick(Duration::from_millis(120));
+    pb.set_message("Unpacking...");
+
     for entry in ar.entries()? {
         let mut entry = entry?;
+        let path = entry.path()?;
+
+        debug!("Unpacking {}", path.display());
 
         // Remove the prefix and join it with the base directory.
-        let path = entry.path()?.strip_prefix(&prefix)?.to_owned();
+        let path = path.strip_prefix(&prefix)?.to_owned();
         let path = prometheus_path.join(path);
 
         entry.unpack(&path)?;
     }
 
+    pb.finish_and_clear();
     Ok(())
 }
 

src/bin/am/commands/system/prune.rs

@@ -1,4 +1,5 @@
-use anyhow::{Context, Result};
+use crate::interactive;
+use anyhow::{bail, Context, Result};
 use clap::Parser;
 use directories::ProjectDirs;
 use std::io;
@@ -15,8 +16,8 @@ pub struct Arguments {
 pub async fn handle_command(args: Arguments) -> Result<()> {
     // If the users hasn't specified the `force` argument, then ask the user if
     // they want to continue.
-    if !args.force && !ask_continue_prune().await {
-        anyhow::bail!("Pruning cancelled")
+    if !args.force && !interactive::confirm("Prune all am program files?")? {
+        bail!("Pruning cancelled");
     }
 
     // Get local directory
@@ -39,10 +40,3 @@ pub async fn handle_command(args: Arguments) -> Result<()> {
     info!("Pruning complete");
     Ok(())
 }
-
-/// Ask the user if they want to continue pruning through the terminal. Returns
-/// `true` if the user wants to continue, `false` otherwise.
-async fn ask_continue_prune() -> bool {
-    // TODO: Make configurable
-    false
-}

src/bin/am/interactive.rs

@@ -0,0 +1,15 @@
+use dialoguer::theme::SimpleTheme;
+use dialoguer::{Confirm, Input};
+use std::io;
+
+pub fn user_input(prompt: impl Into<String>) -> io::Result<String> {
+    Ok(Input::with_theme(&SimpleTheme)
+        .with_prompt(prompt)
+        .interact()?)
+}
+
+pub fn confirm(prompt: impl Into<String>) -> io::Result<bool> {
+    Ok(Confirm::with_theme(&SimpleTheme)
+        .with_prompt(prompt)
+        .interact()?)
+}

src/bin/am/main.rs

@@ -2,37 +2,29 @@ use anyhow::{Context, Result};
 use clap::Parser;
 use commands::{handle_command, Application};
 use std::io;
-use tracing::metadata::LevelFilter;
-use tracing::{debug, error};
+use tracing::level_filters::LevelFilter;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 use tracing_subscriber::{EnvFilter, Registry};
 
 mod commands;
+mod interactive;
 
 #[tokio::main]
-async fn main() {
+async fn main() -> Result<()> {
     let app = Application::parse();
 
     if let Err(err) = init_logging() {
         eprintln!("Unable to initialize logging: {:#}", err);
         std::process::exit(1);
     }
 
-    let result = handle_command(app).await;
-
-    match result {
-        Ok(_) => debug!("Command completed successfully"),
-        Err(err) => {
-            error!("Command failed: {:?}", err);
-            std::process::exit(1);
-        }
-    }
+    handle_command(app).await
 }
 
 /// Initialize logging for the application.
 ///
-/// Currently we have a straight forward logging setup that will log everything
+/// Currently, we have a straight forward logging setup that will log everything
 /// that is level info and higher to stderr. Users are able to influence this by
 /// exporting the `RUST_LOG` environment variable.
 ///
@@ -41,12 +33,13 @@ async fn main() {
 /// within the `am` module, but will only show info for other modules.
 fn init_logging() -> Result<()> {
     // The filter layer controls which log levels to display.
-    let filter_layer = EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into());
+    let filter = EnvFilter::try_from_default_env()
+        .unwrap_or_else(|_| EnvFilter::default().add_directive(LevelFilter::INFO.into()));
 
     let log_layer = tracing_subscriber::fmt::layer().with_writer(io::stderr);
 
     Registry::default()
-        .with(filter_layer)
+        .with(filter)
         .with(log_layer)
         .try_init()
         .context("unable to initialize logger")?;