datalab_cli/commands/

files.rs

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 a file to Datalab storage
    Upload(UploadArgs),
    /// List uploaded files
    List(ListArgs),
    /// Get file metadata
    Get(GetArgs),
    /// Download a file
    Download(DownloadArgs),
    /// Delete a file
    Delete(DeleteArgs),
}

#[derive(Args, Debug)]
pub struct UploadArgs {
    /// File to upload
    #[arg(value_name = "FILE")]
    pub file: PathBuf,

    /// Request timeout in seconds
    #[arg(long, default_value = "300", value_name = "SECS")]
    pub timeout: u64,
}

#[derive(Args, Debug)]
pub struct ListArgs {
    /// Results per page (max 100)
    #[arg(long, default_value = "50", value_name = "N")]
    pub limit: u32,

    /// Pagination offset
    #[arg(long, default_value = "0", value_name = "N")]
    pub offset: u32,

    /// Request timeout in seconds
    #[arg(long, default_value = "60", value_name = "SECS")]
    pub timeout: u64,
}

#[derive(Args, Debug)]
pub struct GetArgs {
    /// File ID
    #[arg(value_name = "FILE_ID")]
    pub file_id: String,

    /// Request timeout in seconds
    #[arg(long, default_value = "60", value_name = "SECS")]
    pub timeout: u64,
}

#[derive(Args, Debug)]
pub struct DownloadArgs {
    /// File ID
    #[arg(value_name = "FILE_ID")]
    pub file_id: String,

    /// Output file path
    #[arg(long, short, value_name = "FILE")]
    pub output: Option<PathBuf>,

    /// URL validity duration in seconds (max 86400)
    #[arg(long, default_value = "3600", value_name = "SECS")]
    pub expires_in: u32,

    /// Request timeout in seconds
    #[arg(long, default_value = "300", value_name = "SECS")]
    pub timeout: u64,
}

#[derive(Args, Debug)]
pub struct DeleteArgs {
    /// File ID
    #[arg(value_name = "FILE_ID")]
    pub file_id: String,

    /// Request timeout in seconds
    #[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(())
}