Merge pull request #125 from theseus-rs/add-archive-progress-bar-status

feat: add archive tracing progress bar status

Author: brianheineman

Cargo.lock

@@ -29,6 +29,7 @@ async-trait = "0.1.81"
 axum = "0.7.5"
 criterion = "0.5.1"
 flate2 = "1.0.31"
+futures-util = "0.3.30"
 hex = "0.4.3"
 home = "0.5.9"
 http = "1.1.0"
@@ -46,7 +47,7 @@ reqwest-middleware = "0.3.3"
 reqwest-retry = "0.6.1"
 reqwest-tracing = "0.5.3"
 semver = "1.0.23"
-serde = "1.0.206"
+serde = "1.0.207"
 serde_json = "1.0.124"
 sha1 = "0.10.6"
 sha2 = "0.10.8"

Cargo.toml

@@ -20,7 +20,8 @@ use tracing_subscriber::{fmt, Registry};
 /// Example of how to use postgresql embedded with axum.
 #[tokio::main]
 async fn main() -> Result<()> {
-    let progress_style = ProgressStyle::with_template("{span_child_prefix}{spinner} {span_name}")?;
+    let progress_style = ProgressStyle::with_template("{span_child_prefix}{spinner} {span_name} [{elapsed_precise}] [{wide_bar:.green.bold}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})")?
+        .progress_chars("=> ");
     let indicatif_layer = IndicatifLayer::new().with_progress_style(progress_style);
     let subscriber = Registry::default()
         .with(fmt::Layer::default().with_filter(LevelFilter::INFO))

examples/axum_embedded/src/main.rs

@@ -13,14 +13,15 @@ version.workspace = true
 anyhow = { workspace = true }
 async-trait = { workspace = true }
 flate2 = { workspace = true }
+futures-util = { workspace = true }
 hex = { workspace = true }
 http = { workspace = true }
 human_bytes = { workspace = true, default-features = false }
 md-5 = { workspace = true, optional = true }
 num-format = { workspace = true }
 quick-xml = { workspace = true, features = ["serialize"], optional = true }
 regex = { workspace = true }
-reqwest = { workspace = true, default-features = false, features = ["json"] }
+reqwest = { workspace = true, default-features = false, features = ["json", "stream"] }
 reqwest-middleware = { workspace = true }
 reqwest-retry = { workspace = true }
 reqwest-tracing = { workspace = true }
@@ -35,6 +36,7 @@ tempfile = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["full"], optional = true }
 tracing = { workspace = true, features = ["log"] }
+tracing-indicatif = { workspace = true }
 url = { workspace = true }
 xz2 = { workspace = true }
 zip = { workspace = true }

postgresql_archive/Cargo.toml

@@ -7,6 +7,7 @@ use crate::Error::{
 };
 use crate::{hasher, matcher, Result};
 use async_trait::async_trait;
+use futures_util::StreamExt;
 use http::{header, Extensions};
 use human_bytes::human_bytes;
 use regex::Regex;
@@ -17,9 +18,12 @@ use reqwest_retry::RetryTransientMiddleware;
 use reqwest_tracing::TracingMiddleware;
 use semver::{Version, VersionReq};
 use std::env;
+use std::io::Write;
 use std::str::FromStr;
 use std::sync::LazyLock;
-use tracing::{debug, instrument, warn};
+use tracing::{debug, instrument, warn, Span};
+use tracing_indicatif::span_ext::IndicatifSpanExt;
+
 use url::Url;
 
 const GITHUB_API_VERSION_HEADER: &str = "X-GitHub-Api-Version";
@@ -223,12 +227,19 @@ impl Repository for GitHub {
         debug!("Downloading archive {}", asset.browser_download_url);
         let request = client.get(&asset.browser_download_url);
         let response = request.send().await?.error_for_status()?;
-        let archive = response.bytes().await?;
-        let bytes = archive.to_vec();
+        let span = Span::current();
+        let content_length = response.content_length().unwrap_or_default();
+        let mut bytes = Vec::new();
+        let mut source = response.bytes_stream();
+        span.pb_set_length(content_length);
+        while let Some(chunk) = source.next().await {
+            bytes.write_all(&chunk?)?;
+            span.pb_set_position(bytes.len() as u64);
+        }
         debug!(
             "Archive {} downloaded: {}",
             asset.browser_download_url,
-            human_bytes(archive.len() as f64)
+            human_bytes(bytes.len() as f64)
         );
 
         if let Some(asset_hash) = asset_hash {

postgresql_archive/src/repository/github/repository.rs

@@ -4,6 +4,7 @@ use crate::repository::Archive;
 use crate::Error::{ArchiveHashMismatch, ParseError, RepositoryFailure, VersionNotFound};
 use crate::{hasher, Result};
 use async_trait::async_trait;
+use futures_util::StreamExt;
 use http::{header, Extensions};
 use human_bytes::human_bytes;
 use reqwest::{Request, Response};
@@ -13,8 +14,10 @@ use reqwest_retry::RetryTransientMiddleware;
 use reqwest_tracing::TracingMiddleware;
 use semver::{Version, VersionReq};
 use std::env;
+use std::io::Write;
 use std::sync::LazyLock;
-use tracing::{debug, instrument, warn};
+use tracing::{debug, instrument, warn, Span};
+use tracing_indicatif::span_ext::IndicatifSpanExt;
 
 static USER_AGENT: LazyLock<String> = LazyLock::new(|| {
     format!(
@@ -134,11 +137,18 @@ impl Repository for Maven {
         debug!("Downloading archive {archive_url}");
         let request = client.get(&archive_url);
         let response = request.send().await?.error_for_status()?;
-        let archive = response.bytes().await?;
-        let bytes = archive.to_vec();
+        let span = Span::current();
+        let content_length = response.content_length().unwrap_or_default();
+        let mut bytes = Vec::new();
+        let mut source = response.bytes_stream();
+        span.pb_set_length(content_length);
+        while let Some(chunk) = source.next().await {
+            bytes.write_all(&chunk?)?;
+            span.pb_set_position(bytes.len() as u64);
+        }
         debug!(
             "Archive {archive_url} downloaded: {}",
-            human_bytes(archive.len() as f64)
+            human_bytes(bytes.len() as f64)
         );
 
         let archive_hash = hasher_fn(&bytes)?;