zenodo_rs/progress.rs
1//! Progress reporting hooks for uploads and downloads.
2//!
3//! The core client stays generic over any progress sink that implements
4//! [`TransferProgress`]. Under the optional `indicatif` feature,
5//! [`indicatif::ProgressBar`] implements this trait directly.
6//!
7//! # Examples
8//!
9//! ```rust
10//! #[cfg(feature = "indicatif")]
11//! #[tokio::main]
12//! async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
13//! use axum::{
14//! body::Body,
15//! extract::State,
16//! http::{header, HeaderValue},
17//! routing::get,
18//! Json, Router,
19//! };
20//! use indicatif::{ProgressBar, ProgressStyle};
21//! use serde_json::json;
22//! use std::sync::Arc;
23//! use zenodo_rs::{ArtifactSelector, Auth, Endpoint, RecordId, ZenodoClient};
24//!
25//! #[derive(Clone)]
26//! struct AppState {
27//! base: Arc<String>,
28//! }
29//!
30//! async fn record(State(state): State<AppState>) -> Json<serde_json::Value> {
31//! Json(json!({
32//! "id": 123,
33//! "recid": 123,
34//! "metadata": { "title": "Example" },
35//! "files": [{
36//! "id": "f-123",
37//! "key": "artifact.bin",
38//! "size": 5,
39//! "links": {
40//! "self": format!("{}download/123/artifact.bin", state.base),
41//! }
42//! }],
43//! "links": {}
44//! }))
45//! }
46//!
47//! async fn artifact() -> axum::response::Response {
48//! let mut response = axum::response::Response::new(Body::from("hello"));
49//! response.headers_mut().insert(
50//! header::CONTENT_TYPE,
51//! HeaderValue::from_static("application/octet-stream"),
52//! );
53//! response
54//! }
55//!
56//! let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
57//! let address = listener.local_addr()?;
58//! let base = Arc::new(format!("http://{address}/api/"));
59//! let app = Router::new()
60//! .route("/api/records/123", get(record))
61//! .route("/api/download/123/artifact.bin", get(artifact))
62//! .with_state(AppState { base: Arc::clone(&base) });
63//! let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
64//! let server = tokio::spawn(async move {
65//! axum::serve(listener, app)
66//! .with_graceful_shutdown(async {
67//! let _ = shutdown_rx.await;
68//! })
69//! .await
70//! });
71//!
72//! let client = ZenodoClient::builder(Auth::new("token"))
73//! .endpoint(Endpoint::Custom(base.parse()?))
74//! .build()?;
75//! let bar = ProgressBar::new(0);
76//! bar.set_style(ProgressStyle::with_template(
77//! "{bar:20.cyan/blue} {bytes}/{total_bytes}",
78//! )?);
79//! let temp_dir = tempfile::tempdir()?;
80//! let path = temp_dir.path().join("artifact.bin");
81//!
82//! let resolved = client
83//! .download_artifact_with_progress(
84//! &ArtifactSelector::latest_file(RecordId(123), "artifact.bin"),
85//! &path,
86//! bar.clone(),
87//! )
88//! .await?;
89//!
90//! assert_eq!(resolved.bytes_written, 5);
91//! assert_eq!(std::fs::read(&path)?, b"hello");
92//! assert_eq!(bar.position(), 5);
93//! let _ = shutdown_tx.send(());
94//! server.await??;
95//! Ok(())
96//! }
97//!
98//! #[cfg(not(feature = "indicatif"))]
99//! fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
100//! Ok(())
101//! }
102//! ```
103//!
104//! Pass `bar.clone()` into the progress-aware upload and download helpers when
105//! you want a real terminal progress bar during transfers.
106
107/// Progress sink for streaming uploads and downloads.
108///
109/// Implement this trait when you want upload and download helpers to report
110/// byte-level transfer progress into your own logging, UI, or terminal
111/// progress bar implementation.
112pub trait TransferProgress: Send + Sync {
113 /// Called once before the transfer starts.
114 ///
115 /// `total_bytes` is `Some(len)` when the total size is known up front and
116 /// `None` when the transfer length is unknown.
117 fn begin(&self, _total_bytes: Option<u64>) {}
118
119 /// Called after each successfully transferred chunk.
120 fn advance(&self, _delta: u64) {}
121
122 /// Called once after a transfer completes successfully.
123 fn finish(&self) {}
124}
125
126impl TransferProgress for () {}
127
128impl<P> TransferProgress for std::sync::Arc<P>
129where
130 P: TransferProgress + ?Sized,
131{
132 fn begin(&self, total_bytes: Option<u64>) {
133 self.as_ref().begin(total_bytes);
134 }
135
136 fn advance(&self, delta: u64) {
137 self.as_ref().advance(delta);
138 }
139
140 fn finish(&self) {
141 self.as_ref().finish();
142 }
143}
144
145#[cfg(feature = "indicatif")]
146impl TransferProgress for indicatif::ProgressBar {
147 fn begin(&self, total_bytes: Option<u64>) {
148 self.set_position(0);
149 if let Some(total_bytes) = total_bytes {
150 self.set_length(total_bytes);
151 }
152 }
153
154 fn advance(&self, delta: u64) {
155 self.inc(delta);
156 }
157
158 fn finish(&self) {
159 indicatif::ProgressBar::finish(self);
160 }
161}
162
163#[cfg(all(test, feature = "indicatif"))]
164mod tests {
165 use super::TransferProgress;
166
167 #[test]
168 fn indicatif_progress_bar_tracks_transfer_progress() {
169 let bar = indicatif::ProgressBar::new(0);
170
171 bar.begin(Some(5));
172 bar.advance(2);
173
174 assert_eq!(bar.length(), Some(5));
175 assert_eq!(bar.position(), 2);
176
177 TransferProgress::finish(&bar);
178 assert!(bar.is_finished());
179 }
180}