use crate::client::DatalabClient;
use crate::error::{DatalabError, Result};
use crate::output::Progress;
use clap::{Args, Subcommand};
use serde_json::json;
use std::fs;
use std::path::PathBuf;
#[derive(Subcommand, Debug)]
pub enum FilesCommand {
Upload(UploadArgs),
List(ListArgs),
Get(GetArgs),
Download(DownloadArgs),
Delete(DeleteArgs),
}
#[derive(Args, Debug)]
pub struct UploadArgs {
#[arg(value_name = "FILE")]
pub file: PathBuf,
#[arg(long, default_value = "300", value_name = "SECS")]
pub timeout: u64,
}
#[derive(Args, Debug)]
pub struct ListArgs {
#[arg(long, default_value = "50", value_name = "N")]
pub limit: u32,
#[arg(long, default_value = "0", value_name = "N")]
pub offset: u32,
#[arg(long, default_value = "60", value_name = "SECS")]
pub timeout: u64,
}
#[derive(Args, Debug)]
pub struct GetArgs {
#[arg(value_name = "FILE_ID")]
pub file_id: String,
#[arg(long, default_value = "60", value_name = "SECS")]
pub timeout: u64,
}
#[derive(Args, Debug)]
pub struct DownloadArgs {
#[arg(value_name = "FILE_ID")]
pub file_id: String,
#[arg(long, short, value_name = "FILE")]
pub output: Option<PathBuf>,
#[arg(long, default_value = "3600", value_name = "SECS")]
pub expires_in: u32,
#[arg(long, default_value = "300", value_name = "SECS")]
pub timeout: u64,
}
#[derive(Args, Debug)]
pub struct DeleteArgs {
#[arg(value_name = "FILE_ID")]
pub file_id: String,
#[arg(long, default_value = "60", value_name = "SECS")]
pub timeout: u64,
}
pub async fn execute(cmd: FilesCommand, progress: &Progress) -> Result<()> {
match cmd {
FilesCommand::Upload(args) => upload(args, progress).await,
FilesCommand::List(args) => list(args, progress).await,
FilesCommand::Get(args) => get(args, progress).await,
FilesCommand::Download(args) => download(args, progress).await,
FilesCommand::Delete(args) => delete(args, progress).await,
}
}
async fn upload(args: UploadArgs, progress: &Progress) -> Result<()> {
let client = DatalabClient::new(Some(args.timeout))?;
if !args.file.exists() {
return Err(DatalabError::FileNotFound(args.file.clone()));
}
progress.start("upload", Some(&args.file.to_string_lossy()));
let filename = args
.file
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("file")
.to_string();
let content_type = mime_guess::from_path(&args.file)
.first_or_octet_stream()
.to_string();
let request_body = json!({
"filename": filename,
"content_type": content_type,
});
let response = client.post_json("files/upload", &request_body).await?;
let upload_url = response
.get("upload_url")
.and_then(|v| v.as_str())
.ok_or_else(|| DatalabError::ProcessingFailed("No upload URL returned".to_string()))?;
let file_id = response
.get("file_id")
.and_then(|v| v.as_str())
.ok_or_else(|| DatalabError::ProcessingFailed("No file ID returned".to_string()))?;
client
.upload_file_to_presigned_url(upload_url, &args.file, &content_type, progress)
.await?;
let confirm_response = client.get(&format!("files/{}/confirm", file_id)).await?;
println!(
"{}",
serde_json::to_string_pretty(&json!({
"file_id": file_id,
"reference": response.get("reference"),
"upload_status": confirm_response.get("upload_status"),
"file_size": confirm_response.get("file_size"),
}))?
);
Ok(())
}
async fn list(args: ListArgs, progress: &Progress) -> Result<()> {
let client = DatalabClient::new(Some(args.timeout))?;
progress.start("list-files", None);
let path = format!("files?limit={}&offset={}", args.limit, args.offset);
let response = client.get(&path).await?;
println!("{}", serde_json::to_string_pretty(&response)?);
Ok(())
}
async fn get(args: GetArgs, progress: &Progress) -> Result<()> {
let client = DatalabClient::new(Some(args.timeout))?;
progress.start("get-file", Some(&args.file_id));
let path = format!("files/{}", args.file_id);
let response = client.get(&path).await?;
println!("{}", serde_json::to_string_pretty(&response)?);
Ok(())
}
async fn download(args: DownloadArgs, progress: &Progress) -> Result<()> {
let client = DatalabClient::new(Some(args.timeout))?;
progress.start("download", Some(&args.file_id));
let path = format!(
"files/{}/download?expires_in={}",
args.file_id, args.expires_in
);
let response = client.get(&path).await?;
if let Some(output_path) = args.output {
let download_url = response
.get("download_url")
.and_then(|v| v.as_str())
.ok_or_else(|| {
DatalabError::ProcessingFailed("No download URL returned".to_string())
})?;
let http_client = reqwest::Client::new();
let file_response = http_client.get(download_url).send().await?;
if !file_response.status().is_success() {
return Err(DatalabError::ApiError {
status: file_response.status().as_u16(),
message: "Failed to download file".to_string(),
});
}
let bytes = file_response.bytes().await?;
fs::write(&output_path, &bytes)?;
println!(
"{}",
serde_json::to_string_pretty(&json!({
"downloaded_to": output_path.to_string_lossy(),
"size": bytes.len(),
}))?
);
} else {
println!("{}", serde_json::to_string_pretty(&response)?);
}
Ok(())
}
async fn delete(args: DeleteArgs, progress: &Progress) -> Result<()> {
let client = DatalabClient::new(Some(args.timeout))?;
progress.start("delete", Some(&args.file_id));
let path = format!("files/{}", args.file_id);
client.delete(&path).await?;
println!(
"{}",
serde_json::to_string_pretty(&json!({
"deleted": true,
"file_id": args.file_id,
}))?
);
Ok(())
}