From 4d4c1d1e1652e5f97e0fd77ace9102bcd68f10b8 Mon Sep 17 00:00:00 2001
From: konstin <konstin@mailbox.org>
Date: Fri, 4 Oct 2024 12:54:42 +0200
Subject: [PATCH] Clear publish progress bar on retry

Previously, we would add broken progress bars.
---
 crates/uv-publish/src/lib.rs        | 26 +++++++++++++++-----------
 crates/uv/src/commands/reporters.rs |  5 ++---
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/crates/uv-publish/src/lib.rs b/crates/uv-publish/src/lib.rs
index 4276d5b7ed25..f85643b0339f 100644
--- a/crates/uv-publish/src/lib.rs
+++ b/crates/uv-publish/src/lib.rs
@@ -22,7 +22,7 @@ use std::{env, fmt, io};
 use thiserror::Error;
 use tokio::io::AsyncReadExt;
 use tokio_util::io::ReaderStream;
-use tracing::{debug, enabled, trace, warn, Level};
+use tracing::{debug, enabled, trace, Level};
 use url::Url;
 use uv_client::UvRetryableStrategy;
 use uv_configuration::{KeyringProviderType, TrustedPublishing};
@@ -30,7 +30,7 @@ use uv_distribution_filename::{DistFilename, SourceDistExtension, SourceDistFile
 use uv_fs::{ProgressReader, Simplified};
 use uv_metadata::read_metadata_async_seek;
 use uv_pypi_types::{Metadata23, MetadataError};
-use uv_warnings::warn_user_once;
+use uv_warnings::{warn_user, warn_user_once};
 
 pub use trusted_publishing::TrustedPublishingToken;
 
@@ -95,7 +95,7 @@ pub trait Reporter: Send + Sync + 'static {
     fn on_progress(&self, name: &str, id: usize);
     fn on_download_start(&self, name: &str, size: Option<u64>) -> usize;
     fn on_download_progress(&self, id: usize, inc: u64);
-    fn on_download_complete(&self);
+    fn on_download_complete(&self, id: usize);
 }
 
 impl PublishSendError {
@@ -298,7 +298,7 @@ pub async fn upload(
     let mut attempt = 0;
     loop {
         attempt += 1;
-        let request = build_request(
+        let (request, idx) = build_request(
             file,
             filename,
             registry,
@@ -312,8 +312,9 @@ pub async fn upload(
         .map_err(|err| PublishError::PublishPrepare(file.to_path_buf(), Box::new(err)))?;
 
         let result = request.send().await;
-        if attempt <= retries && UvRetryableStrategy.handle(&result) == Some(Retryable::Transient) {
-            warn!("Transient request failure for {}, retrying", registry);
+        if attempt < retries && UvRetryableStrategy.handle(&result) == Some(Retryable::Transient) {
+            reporter.on_download_complete(idx);
+            warn_user!("Transient request failure for {}, retrying", registry);
             continue;
         }
 
@@ -478,6 +479,9 @@ async fn form_metadata(
     Ok(form_metadata)
 }
 
+/// Build the upload request.
+///
+/// Returns the request and the reporter progress bar id.
 async fn build_request(
     file: &Path,
     filename: &DistFilename,
@@ -487,7 +491,7 @@ async fn build_request(
     password: Option<&str>,
     form_metadata: &[(&'static str, String)],
     reporter: Arc<impl Reporter>,
-) -> Result<RequestBuilder, PublishPrepareError> {
+) -> Result<(RequestBuilder, usize), PublishPrepareError> {
     let mut form = reqwest::multipart::Form::new();
     for (key, value) in form_metadata {
         form = form.text(*key, value.clone());
@@ -534,7 +538,7 @@ async fn build_request(
         let credentials = BASE64_STANDARD.encode(format!("{username}:{password}"));
         request = request.header(AUTHORIZATION, format!("Basic {credentials}"));
     }
-    Ok(request)
+    Ok((request, idx))
 }
 
 /// Returns `true` if the file was newly uploaded and `false` if it already existed.
@@ -636,7 +640,7 @@ mod tests {
             0
         }
         fn on_download_progress(&self, _id: usize, _inc: u64) {}
-        fn on_download_complete(&self) {}
+        fn on_download_complete(&self, _id: usize) {}
     }
 
     /// Snapshot the data we send for an upload request for a source distribution.
@@ -700,7 +704,7 @@ mod tests {
         project_urls: Source, https://github.com/unknown/tqdm
         "###);
 
-        let request = build_request(
+        let (request, _) = build_request(
             &file,
             &filename,
             &Url::parse("https://example.org/upload").unwrap(),
@@ -843,7 +847,7 @@ mod tests {
         project_urls: wiki, https://github.com/tqdm/tqdm/wiki
         "###);
 
-        let request = build_request(
+        let (request, _) = build_request(
             &file,
             &filename,
             &Url::parse("https://example.org/upload").unwrap(),
diff --git a/crates/uv/src/commands/reporters.rs b/crates/uv/src/commands/reporters.rs
index ea8dca37d205..284673b726f4 100644
--- a/crates/uv/src/commands/reporters.rs
+++ b/crates/uv/src/commands/reporters.rs
@@ -522,9 +522,8 @@ impl uv_publish::Reporter for PublishReporter {
         self.reporter.on_download_progress(id, inc);
     }
 
-    fn on_download_complete(&self) {
-        self.reporter.root.set_message("");
-        self.reporter.root.finish_and_clear();
+    fn on_download_complete(&self, id: usize) {
+        self.reporter.on_download_complete(id);
     }
 }