add download and upload functionality on the cli

Author: Panaetius

Cargo.lock

@@ -70,7 +70,7 @@ docker_credential = "1.3.2"
 chrono = "0.4.42"
 openssl = { version = "0.10.75", features = ["vendored"] }
 tui-realm-treeview = "3.0.0"
-rust-s3 = "0.37.0"
+aws-sdk-s3 = "1.115.0"
 
 [build-dependencies]
 anyhow = "1.0.90"

coman/Cargo.toml

@@ -1,4 +1,3 @@
-use color_eyre::owo_colors::OwoColorize;
 use eyre::{Context, Report};
 use tokio::sync::mpsc;
 use tuirealm::{
@@ -19,7 +18,7 @@ use crate::{
     },
     components::{
         context_menu::ContextMenu, error_popup::ErrorPopup, info_popup::InfoPopup, login_popup::LoginPopup,
-        system_select_popup::SystemSelectPopup, toolbar::Toolbar, workload_list::WorkloadList,
+        system_select_popup::SystemSelectPopup, workload_list::WorkloadList,
         workload_log::WorkloadLog,
     },
     cscs::handlers::{cscs_login, cscs_system_set},

coman/src/app/model.rs

@@ -80,9 +80,19 @@ pub enum CscsFileCommands {
     #[clap(alias("ls"))]
     List { path: PathBuf },
     #[clap(alias("dl"))]
-    Download { remote: PathBuf, local: PathBuf },
+    Download {
+        #[clap(short, long, help = "account/project to use")]
+        account: Option<String>,
+        remote: PathBuf,
+        local: PathBuf,
+    },
     #[clap(alias("ul"))]
-    Upload { local: PathBuf, remote: PathBuf },
+    Upload {
+        #[clap(short, long, help = "account/project to use")]
+        account: Option<String>,
+        local: PathBuf,
+        remote: PathBuf,
+    },
 }
 
 #[derive(Subcommand, Debug)]

coman/src/cli.rs

@@ -8,7 +8,7 @@ use tuirealm::{
 
 use crate::app::{
     messages::{MenuMsg, Msg, View},
-    user_events::{FileEvent, UserEvent},
+    user_events::UserEvent,
 };
 
 #[derive(MockComponent)]

coman/src/components/context_menu.rs

@@ -81,7 +81,7 @@ impl FileTree {
 }
 impl Component<Msg, UserEvent> for FileTree {
     fn on(&mut self, ev: Event<UserEvent>) -> Option<Msg> {
-        let result = match ev {
+        match ev {
             Event::Keyboard(KeyEvent {
                 code: Key::Left,
                 modifiers: KeyModifiers::NONE,

coman/src/components/file_tree.rs

@@ -1,3 +1,4 @@
+pub(crate) mod context_menu;
 pub(crate) mod error_popup;
 pub(crate) mod file_tree;
 pub(crate) mod global_listener;
@@ -7,4 +8,3 @@ pub(crate) mod system_select_popup;
 pub(crate) mod toolbar;
 pub(crate) mod workload_list;
 pub(crate) mod workload_log;
-pub(crate) mod context_menu;

coman/src/components/mod.rs

@@ -8,7 +8,6 @@ use eyre::eyre;
 use lazy_static::lazy_static;
 use serde::{Deserialize, Serialize};
 
-use crate::trace_dbg;
 
 const DEFAULT_CONFIG_TOML: &str = include_str!("../.config/config.toml");
 

coman/src/config.rs

@@ -18,9 +18,10 @@ use firecrest_client::{
     types::{
         DownloadFileResponseTransferDirectives, File as CSCSFile, FileStat as CSCSFileStat,
         FileSystem as CSCSFileSystem, FileSystemDataType, HPCCluster, HealthCheckType, JobMetadataModel, JobModel,
-        SchedulerServiceHealth, UserInfoResponse,
+        S3TransferResponse, SchedulerServiceHealth, UserInfoResponse,
     },
 };
+use reqwest::Url;
 use strum::Display;
 
 use crate::trace_dbg;
@@ -121,7 +122,7 @@ where
 
 impl From<CSCSFile> for PathEntry {
     fn from(value: CSCSFile) -> Self {
-        let size = match usize::from_str_radix(&value.size, 10) {
+        let size = match value.size.parse::<usize>() {
             Ok(size) => size,
             Err(err) => panic!("Couldn't parse file size {}: {:?}", value.size, err),
         };
@@ -334,6 +335,31 @@ fn display_health(h: &Option<Vec<ServicesHealth>>) -> String {
         .unwrap_or("".to_string())
 }
 
+#[derive(Debug, Clone)]
+pub struct S3Upload {
+    pub parts_upload_urls: Vec<Url>,
+    pub complete_upload_url: Url,
+    pub part_size: u64,
+    pub num_parts: u64,
+}
+
+impl S3Upload {
+    fn convert(value: S3TransferResponse, size: u64) -> Result<Self> {
+        let complete_url = value.complete_upload_url.ok_or(eyre!("no upload completion url set"))?;
+        let part_urls = value.parts_upload_urls.ok_or(eyre!("no part upload urls set"))?;
+        let part_size = value.max_part_size.ok_or(eyre!("couldn't parse size"))? as u64;
+        Ok(Self {
+            parts_upload_urls: part_urls
+                .iter()
+                .map(|u| Url::parse(u).wrap_err("couldn't parse url"))
+                .collect::<Result<Vec<Url>>>()?,
+            part_size,
+            complete_upload_url: Url::parse(&complete_url)?,
+            num_parts: (size.div_ceil(part_size)),
+        })
+    }
+}
+
 pub struct CscsApi {
     client: FirecrestClient,
 }
@@ -442,16 +468,17 @@ impl CscsApi {
         account: &str,
         path: PathBuf,
         size: i64,
-    ) -> Result<(i64, Vec<String>, i64)> {
+    ) -> Result<(i64, S3Upload)> {
+        if account.is_empty() {
+            return Err(eyre!(
+                "An account is required, set it in the config file or with the --account flag"
+            ));
+        }
         let job = post_filesystem_transfer_upload(&self.client, system_name, account, path, size)
             .await
             .wrap_err("couldn't upload file")?;
         if let DownloadFileResponseTransferDirectives::S3(directives) = job.transfer_directives {
-            Ok((
-                job.transfer_job.job_id,
-                directives.parts_upload_urls.unwrap(),
-                directives.max_part_size.unwrap(),
-            ))
+            Ok((job.transfer_job.job_id, S3Upload::convert(directives, size as u64)?))
         } else {
             trace_dbg!(job);
             Err(eyre!("didn't get S3 transfer directive"))
@@ -463,12 +490,18 @@ impl CscsApi {
             .wrap_err("couldn't download file")?;
         Ok(content)
     }
-    pub async fn transfer_download(&self, system_name: &str, path: PathBuf) -> Result<(i64, String)> {
-        let job = post_filesystem_transfer_download(&self.client, system_name, path)
+    pub async fn transfer_download(&self, system_name: &str, account: &str, path: PathBuf) -> Result<(i64, Url)> {
+        if account.is_empty() {
+            return Err(eyre!(
+                "An account is required, set it in the config file or with the --account flag"
+            ));
+        }
+        let job = post_filesystem_transfer_download(&self.client, system_name, account, path)
             .await
             .wrap_err("couldn't transfer file")?;
         if let DownloadFileResponseTransferDirectives::S3(directives) = job.transfer_directives {
-            Ok((job.transfer_job.job_id, directives.download_url.unwrap()))
+            let download_url = Url::parse(&directives.download_url.unwrap())?;
+            Ok((job.transfer_job.job_id, download_url))
         } else {
             Err(eyre!("didn't get S3 transfer directive"))
         }

coman/src/cscs/api_client.rs

@@ -1,10 +1,19 @@
-use std::{path::PathBuf, time::Duration};
+use std::{
+    io::{SeekFrom, Write},
+    path::PathBuf,
+    time::{Duration, Instant},
+};
 
 use color_eyre::{Result, eyre::Context};
 use eyre::eyre;
+use futures::StreamExt;
 use inquire::{Password, Text};
 use itertools::Itertools;
 use reqwest::Url;
+use tokio::{
+    fs::File,
+    io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt, BufReader},
+};
 
 use crate::{
     cscs::{
@@ -14,7 +23,6 @@ use crate::{
             cscs_job_log, cscs_login, cscs_start_job, cscs_system_list, cscs_system_set,
         },
     },
-    trace_dbg,
     util::types::DockerImageUrl,
 };
 
@@ -123,8 +131,8 @@ pub(crate) async fn cli_cscs_file_list(path: PathBuf) -> Result<()> {
     }
 }
 
-pub(crate) async fn cli_cscs_file_download(remote: PathBuf, local: PathBuf) -> Result<()> {
-    match cscs_file_download(remote, local.clone()).await {
+pub(crate) async fn cli_cscs_file_download(remote: PathBuf, local: PathBuf, account: Option<String>) -> Result<()> {
+    match cscs_file_download(remote, local.clone(), account).await {
         Ok(None) => {
             println!("File successfully downloaded");
             Ok(())
@@ -148,44 +156,91 @@ pub(crate) async fn cli_cscs_file_download(remote: PathBuf, local: PathBuf) -> R
 
             // download from s3
             println!("Downloading file from s3, this might take a while");
-            let credentials = s3::creds::Credentials::default()?;
-            let url = Url::parse(&job_data.1)?;
-            let region = s3::region::Region::Custom {
-                region: "cscs-zonegroup".to_owned(),
-                endpoint: "https://rgw.cscs.ch".to_owned(),
-            };
-            let mut segments = url.path_segments().unwrap();
-            let bucket_name = segments.next().unwrap();
-            let path = segments.join("/");
-            let bucket = s3::bucket::Bucket::create_with_path_style(
-                bucket_name,
-                region,
-                credentials,
-                s3::BucketConfiguration::default(),
-            )
-            .await?
-            .bucket;
-            let mut async_output_file = tokio::fs::File::create(&local).await.expect("Unable to create file");
-            let status = bucket.get_object_to_writer(path, &mut async_output_file).await?;
-            println!("download finished with status {}", status);
+
+            let mut output = File::create(local).await?;
+            let mut stream = reqwest::get(job_data.1).await?.bytes_stream();
+            let mut progress = 0;
+            let mut start_time = Instant::now();
+            while let Some(chunk_result) = stream.next().await {
+                let chunk = chunk_result?;
+                output.write_all(&chunk).await?;
+                progress += chunk.len();
+
+                if start_time.elapsed() >= Duration::from_secs(1) {
+                    print!("\rDownloaded {}/{}Mb", progress / 1024 / 1024, job_data.2 / 1024 / 1024);
+                    std::io::stdout().flush()?;
+                    start_time = Instant::now();
+                }
+            }
+            output.flush().await?;
+            println!(); //force newline
+            println!("Download complete");
 
             Ok(())
         }
         Err(e) => Err(e),
     }
 }
 
-pub(crate) async fn cli_cscs_file_upload(local: PathBuf, remote: PathBuf) -> Result<()> {
-    match cscs_file_upload(local, remote).await {
+pub(crate) async fn cli_cscs_file_upload(local: PathBuf, remote: PathBuf, account: Option<String>) -> Result<()> {
+    match cscs_file_upload(local.clone(), remote, account).await {
         Ok(None) => {
             println!("File successfully uploaded");
             Ok(())
         }
         Ok(Some(transfer_data)) => {
             println!("starting file transfer, this might take a while");
-            let transfer_data = trace_dbg!(transfer_data);
-            Ok(())
+            let mut etags: Vec<String> = Vec::new();
+            let client = reqwest::Client::new();
+            let num_parts = transfer_data.1.num_parts;
+            for (chunk_id, transfer_url) in transfer_data.1.parts_upload_urls.into_iter().enumerate() {
+                println!(
+                    "Uploading part {}/{} ({}Mb)",
+                    chunk_id + 1,
+                    num_parts,
+                    transfer_data.1.part_size / 1024 / 1024
+                );
+                let etag = upload_chunk(
+                    local.clone(),
+                    (chunk_id as u64) * transfer_data.1.part_size,
+                    transfer_data.1.part_size,
+                    transfer_url,
+                )
+                .await?;
+                etags.push(etag);
+            }
+
+            let body = etags
+                .into_iter()
+                .enumerate()
+                .map(|(i, etag)| (i + 1, etag))
+                .map(|(i, etag)| format!("<Part><PartNumber>{}</PartNumber><ETag>{}</ETag></Part>", i, etag))
+                .join("");
+            let body = format!("<CompleteMultipartUpload>{}</CompleteMultipartUpload>", body);
+            let req = client.post(transfer_data.1.complete_upload_url).body(body).build()?;
+            let resp = client.execute(req).await?;
+            match resp.error_for_status() {
+                Ok(_) => {
+                    println!("done");
+                    Ok(())
+                }
+                Err(e) => Err(e).wrap_err("failed to complete upload"),
+            }
         }
         Err(e) => Err(e),
     }
 }
+
+async fn upload_chunk(path: PathBuf, offset: u64, size: u64, url: Url) -> Result<String> {
+    let client = reqwest::Client::new();
+
+    let source_file = File::open(path).await?;
+    let mut buf = vec![];
+    let mut reader = BufReader::new(source_file);
+    reader.seek(SeekFrom::Start(offset)).await?;
+    let mut chunk = reader.take(size);
+    chunk.read_to_end(&mut buf).await?;
+    let req = client.put(url).body(buf).build()?;
+    let resp = client.execute(req).await?;
+    Ok(resp.headers()["etag"].to_str()?.to_owned())
+}