Skip to content

Commit

Permalink
refactor code to add progress bars
Browse files Browse the repository at this point in the history
  • Loading branch information
wolfv committed Dec 21, 2023
1 parent 9bf2685 commit 601e999
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 126 deletions.
24 changes: 15 additions & 9 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ struct RebuildOpts {
struct UploadOpts {
/// The package file to upload
#[arg(global = true, required = false)]
package_file: PathBuf,
package_files: Vec<PathBuf>,

/// The server type
#[clap(subcommand)]
Expand Down Expand Up @@ -604,11 +604,17 @@ async fn rebuild_from_args(args: RebuildOpts) -> miette::Result<()> {
}

async fn upload_from_args(args: UploadOpts) -> miette::Result<()> {
if ArchiveType::try_from(&args.package_file).is_none() {
return Err(miette::miette!(
"The file {} does not appear to be a conda package.",
args.package_file.to_string_lossy()
));
if args.package_files.is_empty() {
return Err(miette::miette!("No package files were provided."));
}

for package_file in &args.package_files {
if ArchiveType::try_from(&package_file).is_none() {

Check failure on line 612 in src/main.rs

View workflow job for this annotation

GitHub Actions / test (ubuntu-latest)

[clippy] reported by reviewdog 🐶 error: the borrowed expression implements the required traits --> src/main.rs:612:34 | 612 | if ArchiveType::try_from(&package_file).is_none() { | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/main.rs:612:34:e:error: the borrowed expression implements the required traits --> src/main.rs:612:34 | 612 | if ArchiveType::try_from(&package_file).is_none() { | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__

Check failure on line 612 in src/main.rs

View workflow job for this annotation

GitHub Actions / test (macos-latest)

[clippy] reported by reviewdog 🐶 error: the borrowed expression implements the required traits --> src/main.rs:612:34 | 612 | if ArchiveType::try_from(&package_file).is_none() { | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/main.rs:612:34:e:error: the borrowed expression implements the required traits --> src/main.rs:612:34 | 612 | if ArchiveType::try_from(&package_file).is_none() { | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__
return Err(miette::miette!(
"The file {} does not appear to be a conda package.",
package_file.to_string_lossy()
));
}
}

let store = get_auth_store(args.common.auth_file);
Expand All @@ -618,7 +624,7 @@ async fn upload_from_args(args: UploadOpts) -> miette::Result<()> {
upload::upload_package_to_quetz(
&store,
quetz_opts.api_key,
args.package_file,
&args.package_files,
quetz_opts.url,
quetz_opts.channel,
)
Expand All @@ -629,7 +635,7 @@ async fn upload_from_args(args: UploadOpts) -> miette::Result<()> {
&store,
artifactory_opts.username,
artifactory_opts.password,
args.package_file,
&args.package_files,
artifactory_opts.url,
artifactory_opts.channel,
)
Expand All @@ -639,7 +645,7 @@ async fn upload_from_args(args: UploadOpts) -> miette::Result<()> {
upload::upload_package_to_prefix(
&store,
prefix_opts.api_key,
args.package_file,
&args.package_files,
prefix_opts.url,
prefix_opts.channel,
)
Expand Down
235 changes: 118 additions & 117 deletions src/upload.rs
Original file line number Diff line number Diff line change
@@ -1,15 +1,17 @@
use std::{fmt::Write, path::PathBuf};

use futures::TryStreamExt;
use indicatif::{style::TemplateError, HumanBytes, ProgressState};
use std::{
fmt::Write,
path::{Path, PathBuf},
};
use tokio_util::io::ReaderStream;

use miette::{Context, IntoDiagnostic};
use rattler_conda_types::package::{IndexJson, PackageFile};
use rattler_digest::compute_file_digest;
use rattler_networking::{redact_known_secrets_from_error, Authentication, AuthenticationStorage};
use reqwest::Method;
use sha2::{Digest, Sha256};
use sha2::Sha256;
use tracing::info;
use url::Url;

Expand All @@ -32,10 +34,21 @@ fn default_bytes_style() -> Result<indicatif::ProgressStyle, TemplateError> {
))
}

fn get_client() -> Result<reqwest::Client, reqwest::Error> {
reqwest::Client::builder().no_gzip().build()
}

fn sha256_sum(package_file: &Path) -> Result<String, std::io::Error> {
Ok(format!(
"{:x}",
compute_file_digest::<Sha256>(&package_file)?
))
}

pub async fn upload_package_to_quetz(
storage: &AuthenticationStorage,
api_key: Option<String>,
package_file: PathBuf,
package_files: &Vec<PathBuf>,
url: Url,
channel: String,
) -> miette::Result<()> {
Expand All @@ -60,61 +73,40 @@ pub async fn upload_package_to_quetz(
},
};

let client = reqwest::Client::builder()
.no_gzip()
.build()
.expect("failed to create client");

let upload_url = url
.join(&format!(
"api/channels/{}/upload/{}",
channel,
package_file.file_name().unwrap().to_string_lossy()
))
.into_diagnostic()?;
let client = get_client().into_diagnostic()?;

let bytes = tokio::fs::read(package_file).await.into_diagnostic()?;
let upload_hash = sha2::Sha256::digest(&bytes);
for package_file in package_files {
let upload_url = url
.join(&format!(
"api/channels/{}/upload/{}",
channel,
package_file.file_name().unwrap().to_string_lossy()
))
.into_diagnostic()?;

let req = client
.request(Method::POST, upload_url)
.query(&[("force", "false"), ("sha256", &hex::encode(upload_hash))])
.body(bytes)
.header("X-API-Key", token)
.send()
.await
.map_err(redact_known_secrets_from_error)
.into_diagnostic()
.map_err(|e| miette::miette!("Sending package to Quetz server failed: {e}"))?;
let hash = sha256_sum(&package_file).into_diagnostic()?;

Check failure on line 87 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (ubuntu-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:87:31 | 87 | let hash = sha256_sum(&package_file).into_diagnostic()?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow = note: `-D clippy::needless-borrow` implied by `-D warnings` Raw Output: src/upload.rs:87:31:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:87:31 | 87 | let hash = sha256_sum(&package_file).into_diagnostic()?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow = note: `-D clippy::needless-borrow` implied by `-D warnings` __END__

Check failure on line 87 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (macos-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:87:31 | 87 | let hash = sha256_sum(&package_file).into_diagnostic()?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow = note: `-D clippy::needless-borrow` implied by `-D warnings` Raw Output: src/upload.rs:87:31:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:87:31 | 87 | let hash = sha256_sum(&package_file).into_diagnostic()?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow = note: `-D clippy::needless-borrow` implied by `-D warnings` __END__

req.error_for_status_ref()
.map_err(redact_known_secrets_from_error)
.into_diagnostic()
.map_err(|e| miette::miette!("Quetz server responded with error: {e}"))?;
let prepared_request = client
.request(Method::POST, upload_url)
.query(&[("force", "false"), ("sha256", &hash)])
.header("X-API-Key", token.clone());

send_request(prepared_request, &package_file).await?;

Check failure on line 94 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (ubuntu-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:94:40 | 94 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:94:40:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:94:40 | 94 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__

Check failure on line 94 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (macos-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:94:40 | 94 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:94:40:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:94:40 | 94 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__
}

info!("Packages successfully uploaded to Quetz server");

info!("Package was successfully uploaded to Quetz server");
Ok(())
}

pub async fn upload_package_to_artifactory(
storage: &AuthenticationStorage,
username: Option<String>,
password: Option<String>,
package_file: PathBuf,
package_files: &Vec<PathBuf>,
url: Url,
channel: String,
) -> miette::Result<()> {
let package_dir = tempfile::tempdir()
.into_diagnostic()
.wrap_err("Creating temporary directory failed")?;

rattler_package_streaming::fs::extract(&package_file, package_dir.path()).into_diagnostic()?;

let index_json = IndexJson::from_package_directory(package_dir.path()).into_diagnostic()?;
let subdir = index_json
.subdir
.ok_or_else(|| miette::miette!("index.json of the package has no subdirectory. Cannot determine which directory to upload to"))?;

let (username, password) = match (username, password) {
(Some(u), Some(p)) => (u, p),
(Some(_), _) | (_, Some(_)) => {
Expand All @@ -139,45 +131,46 @@ pub async fn upload_package_to_artifactory(
},
};

let client = reqwest::Client::builder()
.no_gzip()
.build()
.expect("failed to create client");
for package_file in package_files {
let package_dir = tempfile::tempdir()
.into_diagnostic()
.wrap_err("Creating temporary directory failed")?;

let package_name = package_file
.file_name()
.expect("no filename found")
.to_string_lossy();
rattler_package_streaming::fs::extract(&package_file, package_dir.path())

Check failure on line 139 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (ubuntu-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:139:48 | 139 | rattler_package_streaming::fs::extract(&package_file, package_dir.path()) | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:139:48:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:139:48 | 139 | rattler_package_streaming::fs::extract(&package_file, package_dir.path()) | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__

Check failure on line 139 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (macos-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:139:48 | 139 | rattler_package_streaming::fs::extract(&package_file, package_dir.path()) | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:139:48:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:139:48 | 139 | rattler_package_streaming::fs::extract(&package_file, package_dir.path()) | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__
.into_diagnostic()?;

let upload_url = url
.join(&format!("{}/{}/{}", channel, subdir, package_name))
.into_diagnostic()?;
let index_json = IndexJson::from_package_directory(package_dir.path()).into_diagnostic()?;
let subdir = index_json
.subdir
.ok_or_else(|| miette::miette!("index.json of the package has no subdirectory. Cannot determine which directory to upload to"))?;

let bytes = tokio::fs::read(package_file).await.into_diagnostic()?;
let client = get_client().into_diagnostic()?;

client
.request(Method::PUT, upload_url)
.body(bytes)
.basic_auth(username, Some(password))
.send()
.await
.map_err(redact_known_secrets_from_error)
.into_diagnostic()
.wrap_err("Sending package to artifactory server failed")?
.error_for_status_ref()
.map_err(redact_known_secrets_from_error)
.into_diagnostic()
.wrap_err("Artifactory responded with error")?;
let package_name = package_file
.file_name()
.expect("no filename found")
.to_string_lossy();

info!("Package was successfully uploaded to artifactory server");
let upload_url = url
.join(&format!("{}/{}/{}", channel, subdir, package_name))
.into_diagnostic()?;

let prepared_request = client
.request(Method::PUT, upload_url)
.basic_auth(username.clone(), Some(password.clone()));

send_request(prepared_request, &package_file).await?;

Check failure on line 162 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (ubuntu-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:162:40 | 162 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:162:40:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:162:40 | 162 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__

Check failure on line 162 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (macos-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:162:40 | 162 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:162:40:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:162:40 | 162 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__
}

info!("Packages successfully uploaded to Artifactory server");

Ok(())
}

pub async fn upload_package_to_prefix(
storage: &AuthenticationStorage,
api_key: Option<String>,
package_file: PathBuf,
package_files: &Vec<PathBuf>,
url: Url,
channel: String,
) -> miette::Result<()> {
Expand All @@ -202,78 +195,86 @@ pub async fn upload_package_to_prefix(
},
};

let filename = package_file
.file_name()
.unwrap()
.to_string_lossy()
.to_string();
for package_file in package_files {
let filename = package_file
.file_name()
.unwrap()
.to_string_lossy()
.to_string();

let file_size = package_file.metadata().into_diagnostic()?.len();
let file_size = package_file.metadata().into_diagnostic()?.len();

println!("Uploading package to: {}", url);
println!(
"Package file: {} ({})\n",
package_file.display(),
HumanBytes(file_size)
);
let url = url
.join(&format!("api/v1/upload/{}", channel))
.into_diagnostic()?;

let url = url
.join(&format!("api/v1/upload/{}", channel))
.into_diagnostic()?;
let client = get_client().into_diagnostic()?;

let client = reqwest::Client::builder()
.no_gzip()
.build()
.expect("failed to create client");
let hash = sha256_sum(&package_file).into_diagnostic()?;

Check failure on line 213 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (ubuntu-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:213:31 | 213 | let hash = sha256_sum(&package_file).into_diagnostic()?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:213:31:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:213:31 | 213 | let hash = sha256_sum(&package_file).into_diagnostic()?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__

Check failure on line 213 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (macos-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:213:31 | 213 | let hash = sha256_sum(&package_file).into_diagnostic()?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:213:31:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:213:31 | 213 | let hash = sha256_sum(&package_file).into_diagnostic()?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__

let sha256sum = format!(
"{:x}",
compute_file_digest::<Sha256>(&package_file).into_diagnostic()?
);
let prepared_request = client
.post(url.clone())
.header("X-File-Sha256", hash)
.header("X-File-Name", filename)
.header("Content-Length", file_size)
.header("Content-Type", "application/octet-stream")
.bearer_auth(token.clone());

send_request(prepared_request, &package_file).await?;

Check failure on line 223 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (ubuntu-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:223:40 | 223 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:223:40:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:223:40 | 223 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__

Check failure on line 223 in src/upload.rs

View workflow job for this annotation

GitHub Actions / test (macos-latest)

[clippy] reported by reviewdog 🐶 error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:223:40 | 223 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow Raw Output: src/upload.rs:223:40:e:error: this expression creates a reference which is immediately dereferenced by the compiler --> src/upload.rs:223:40 | 223 | send_request(prepared_request, &package_file).await?; | ^^^^^^^^^^^^^ help: change this to: `package_file` | = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_borrow __END__
}

info!("Packages successfully uploaded to prefix.dev server");

let file = tokio::fs::File::open(&package_file)
Ok(())
}

async fn send_request(
prepared_request: reqwest::RequestBuilder,
package_file: &Path,
) -> miette::Result<reqwest::Response> {
let file = tokio::fs::File::open(package_file)
.await
.into_diagnostic()?;

let file_size = file.metadata().await.into_diagnostic()?.len();
info!(
"Uploading package file: {} ({})\n",
package_file.file_name().unwrap().to_string_lossy(),
HumanBytes(file_size)
);
let progress_bar = indicatif::ProgressBar::new(file_size)
.with_prefix("Uploading")
.with_style(default_bytes_style().into_diagnostic()?);

let progress_bar_clone = progress_bar.clone();
let reader_stream = ReaderStream::new(file)
.inspect_ok(move |bytes| {
progress_bar.inc(bytes.len() as u64);
progress_bar_clone.inc(bytes.len() as u64);
})
.inspect_err(|e| {
println!("Error while uploading: {}", e);
});

let body = reqwest::Body::wrap_stream(reader_stream);

let response = client
.post(url.clone())
.header("X-File-Sha256", sha256sum)
.header("X-File-Name", filename)
.header("Content-Length", file_size)
.header("Content-Type", "application/octet-stream")
.bearer_auth(token)
let response = prepared_request
.body(body)
.send()
.await
.map_err(redact_known_secrets_from_error)
.into_diagnostic()?;

if response.status().is_success() {
println!("Upload successful!");
} else {
println!("Upload failed!");
if response.status() == 401 {
println!(
"Authentication failed! Did you set the correct API key for {}",
url
);
}
response
.error_for_status_ref()
.map_err(redact_known_secrets_from_error)
.into_diagnostic()
.wrap_err("Server responded with error")?;

std::process::exit(1);
}
progress_bar.finish();
info!(
"\nUpload complete for package file: {}",
package_file.file_name().unwrap().to_string_lossy()
);

Ok(())
Ok(response)
}

0 comments on commit 601e999

Please sign in to comment.