dragonfly-client 1.5.5

/*
 *     Copyright 2025 The Dragonfly Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use clap::Parser;
use dragonfly_api::common::v2::ObjectStorage;
use dragonfly_api::dfdaemon::v2::{
    download_persistent_task_response, DownloadPersistentTaskRequest,
};
use dragonfly_api::errordetails::v2::Backend;
use dragonfly_client::terminal;
use dragonfly_client_core::{
    error::{ErrorType, OrErr},
    Error, Result,
};
use dragonfly_client_util::fs::fallocate;
use dragonfly_client_util::net::preferred_local_ip;
use indicatif::{ProgressBar, ProgressState, ProgressStyle};
use path_absolutize::*;
use std::path::{Path, PathBuf};
use std::time::Duration;
use std::{cmp::min, fmt::Write};
use tokio::fs::{self, OpenOptions};
use tokio::io::{AsyncSeekExt, AsyncWriteExt, BufWriter, SeekFrom};
use tracing::{debug, error, info};
use url::Url;

use super::*;

/// The subcommand of export.
#[derive(Debug, Clone, Parser)]
pub struct ExportCommand {
    #[arg(
        help = "Specify the URL to download. Format: scheme://<bucket>/<path>. Examples: s3://<bucket>/<path>, abs://<bucket>/<path>"
    )]
    url: Url,

    #[arg(
        long = "transfer-from-dfdaemon",
        default_value_t = false,
        env = "DFSTORE_EXPORT_TRANSFER_FROM_DFDAEMON",
        help = "Specify whether to transfer the content of downloading file from dfdaemon's unix domain socket. If it is true, dfstore will call dfdaemon to download the file, and dfdaemon will return the content of downloading file to dfstore via unix domain socket, and dfstore will copy the content to the output path. If it is false, dfdaemon will download the file and hardlink or copy the file to the output path."
    )]
    transfer_from_dfdaemon: bool,

    #[arg(
        long = "overwrite",
        default_value_t = false,
        env = "DFSTORE_EXPORT_OVERWRITE",
        help = "Specify whether to overwrite the output file if it already exists. If it is true, dfget will overwrite the output file. If it is false, dfget will return an error if the output file already exists. Cannot be used with `--force-hard-link=true`"
    )]
    overwrite: bool,

    #[arg(
        long = "force-hard-link",
        default_value_t = false,
        env = "DFSTORE_EXPORT_FORCE_HARD_LINK",
        help = "Specify whether the download file must be hard linked to the output path. If hard link is failed, download will be failed. If it is false, dfdaemon will copy the file to the output path if hard link is failed."
    )]
    force_hard_link: bool,

    #[arg(
        short = 'O',
        long = "output",
        help = "Specify the output path of exporting file"
    )]
    output: PathBuf,

    #[arg(
        long = "timeout",
        value_parser= humantime::parse_duration,
        default_value = "2h",
        env = "DFSTORE_EXPORT_TIMEOUT",
        help = "Specify the timeout for exporting a file"
    )]
    timeout: Duration,

    #[arg(
        long = "digest",
        required = false,
        help = "Verify the integrity of the downloaded file using the specified digest, support sha256, sha512, crc32. If the digest is not specified, the downloaded file will not be verified. Format: <algorithm>:<digest>, e.g. sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef, crc32:12345678"
    )]
    digest: Option<String>,

    #[arg(
        long,
        env = "DFSTORE_EXPORT_STORAGE_REGION",
        help = "Specify the region for the Object Storage Service (e.g., us-east-1)"
    )]
    storage_region: Option<String>,

    #[arg(
        long,
        env = "DFSTORE_EXPORT_STORAGE_ENDPOINT",
        help = "Specify the endpoint URL for the Object Storage Service (e.g., https://s3.amazonaws.com)"
    )]
    storage_endpoint: Option<String>,

    #[arg(
        long,
        env = "DFSTORE_EXPORT_STORAGE_ACCESS_KEY_ID",
        help = "Specify the access key ID for authenticating with the Object Storage Service"
    )]
    storage_access_key_id: Option<String>,

    #[arg(
        long,
        env = "DFSTORE_EXPORT_STORAGE_ACCESS_KEY_SECRET",
        help = "Specify the secret access key for authenticating with the Object Storage Service"
    )]
    storage_access_key_secret: Option<String>,

    #[arg(
        long,
        env = "DFSTORE_EXPORT_STORAGE_SECURITY_TOKEN",
        help = "Specify the security token for the Object Storage Service"
    )]
    storage_security_token: Option<String>,

    #[arg(
        long,
        env = "DFSTORE_EXPORT_STORAGE_INSECURE_SKIP_VERIFY",
        help = "Specify whether to skip verify TLS certification for object storage service"
    )]
    storage_insecure_skip_verify: Option<bool>,

    #[arg(
        long,
        env = "DFSTORE_EXPORT_STORAGE_SESSION_TOKEN",
        help = "Specify the session token for Amazon Simple Storage Service(S3)"
    )]
    storage_session_token: Option<String>,

    #[arg(
        long,
        env = "DFSTORE_EXPORT_STORAGE_CREDENTIAL_PATH",
        help = "Specify the local path to the credential file which is used for OAuth2 authentication for Google Cloud Storage Service(GCS)"
    )]
    storage_credential_path: Option<String>,

    #[arg(
        long,
        default_value = "publicRead",
        env = "DFSTORE_EXPORT_STORAGE_PREDEFINED_ACL",
        help = "Specify the predefined ACL for Google Cloud Storage Service(GCS)"
    )]
    storage_predefined_acl: Option<String>,

    #[arg(
        long,
        default_value_t = false,
        env = "DFSTORE_EXPORT_NO_PROGRESS",
        help = "Specify whether to disable the progress bar display"
    )]
    no_progress: bool,

    #[arg(
        short = 'e',
        long = "endpoint",
        default_value_os_t = dfdaemon::default_download_unix_socket_path(),
        help = "Endpoint of dfdaemon's GRPC server"
    )]
    endpoint: PathBuf,

    #[arg(
        short = 'l',
        long,
        default_value = "info",
        env = "DFSTORE_EXPORT_LOG_LEVEL",
        help = "Specify the logging level [trace, debug, info, warn, error]"
    )]
    log_level: Level,

    #[arg(
        long,
        default_value_t = false,
        env = "DFSTORE_EXPORT_CONSOLE",
        help = "Specify whether to print log"
    )]
    console: bool,
}

/// The buffer size used for transferring data from dfdaemon to dfget when
/// `--transfer-from-dfdaemon` is enabled.
const TRANSFER_WRITE_BUFFER_SIZE: usize = 8 * 1024 * 1024;

/// Implement the execute for ExportCommand.
impl ExportCommand {
    /// Executes the export command with comprehensive validation and advanced error handling.
    ///
    /// This function serves as the main entry point for the dfstore export command execution.
    /// It handles the complete workflow including argument parsing, validation, logging setup,
    /// dfdaemon client connection, and export operation execution. The function provides
    /// sophisticated error reporting with colored terminal output, including specialized
    /// handling for backend errors with HTTP status codes and headers.
    pub async fn execute(&self) -> Result<()> {
        // Parse command line arguments.
        Args::parse();

        // Initialize tracing.
        let _guards = init_command_tracing(self.log_level, self.console);

        // Validate the command line arguments.
        if let Err(err) = self.validate_args() {
            terminal::error("Validating Failed!");
            terminal::separator();
            terminal::field("Message:", err);
            terminal::separator();

            std::process::exit(1);
        }

        // Get dfdaemon download client.
        let dfdaemon_download_client =
            match get_dfdaemon_download_client(self.endpoint.to_path_buf()).await {
                Ok(client) => client,
                Err(err) => {
                    terminal::error("Connect Dfdaemon Failed!");
                    terminal::separator();
                    terminal::field(
                        "Message:",
                        format!(
                            "can not connect {}, please check the unix socket {}",
                            err,
                            self.endpoint.to_string_lossy()
                        ),
                    );
                    terminal::separator();

                    std::process::exit(1);
                }
            };

        // Run export command.
        if let Err(err) = self.run(dfdaemon_download_client).await {
            match err {
                Error::TonicStatus(status) => {
                    let details = status.details();
                    if let Ok(backend_err) = serde_json::from_slice::<Backend>(details) {
                        terminal::error("Exporting Failed!");
                        terminal::separator();

                        if let Some(status_code) = backend_err.status_code {
                            terminal::error_field("Bad Status Code:", status_code);
                        }

                        terminal::field("Message:", backend_err.message);

                        if !backend_err.header.is_empty() {
                            terminal::headers(
                                backend_err
                                    .header
                                    .iter()
                                    .map(|(key, value)| (key.as_str(), value.as_str())),
                            );
                        }

                        terminal::separator();
                    } else {
                        terminal::error("Exporting Failed!");
                        terminal::separator();
                        terminal::error_field("Bad Code:", status.code());
                        terminal::field("Message:", status.message());

                        if !status.details().is_empty() {
                            terminal::field(
                                "Details:",
                                std::str::from_utf8(status.details()).unwrap(),
                            );
                        }

                        terminal::separator();
                    }
                }
                Error::BackendError(err) => {
                    terminal::error("Exporting Failed!");
                    terminal::separator();
                    terminal::error_field("Message:", err.message);

                    if err.header.is_some() {
                        terminal::headers(
                            err.header
                                .unwrap_or_default()
                                .iter()
                                .map(|(key, value)| (key.as_str(), value.to_str().unwrap())),
                        );
                    }

                    terminal::separator();
                }
                err => {
                    terminal::error("Exporting Failed!");
                    terminal::separator();
                    terminal::error_field("Message:", err);
                    terminal::separator();
                }
            }

            std::process::exit(1);
        }

        Ok(())
    }

    /// Executes the export operation to retrieve stored files from the persistent system.
    ///
    /// This function handles the core export functionality by downloading a stored file from the
    /// dfdaemon persistent system. It supports two transfer modes: direct file transfer
    /// by dfdaemon (hardlink/copy) or streaming piece content through the client for manual
    /// file assembly. The operation provides real-time progress feedback and handles file
    /// creation, directory setup, and efficient piece-by-piece writing with sparse file allocation.
    async fn run(&self, dfdaemon_download_client: DfdaemonDownloadClient) -> Result<()> {
        // Dfstore needs to notify dfdaemon to transfer the piece content of downloading file via unix domain socket
        // when the `transfer_from_dfdaemon` is true. Otherwise, dfdaemon will download the file and hardlink or
        // copy the file to the output path.
        let (output_path, need_piece_content) = if self.transfer_from_dfdaemon {
            (None, true)
        } else {
            let absolute_path = Path::new(&self.output).absolutize()?;
            info!("export file to: {}", absolute_path.to_string_lossy());
            (Some(absolute_path.to_string_lossy().to_string()), false)
        };

        // Create dfdaemon client.
        let response = dfdaemon_download_client
            .download_persistent_task(DownloadPersistentTaskRequest {
                url: self.url.to_string(),
                object_storage: Some(ObjectStorage {
                    region: self.storage_region.clone(),
                    endpoint: self.storage_endpoint.clone(),
                    access_key_id: self.storage_access_key_id.clone(),
                    access_key_secret: self.storage_access_key_secret.clone(),
                    security_token: self.storage_security_token.clone(),
                    session_token: self.storage_session_token.clone(),
                    credential_path: self.storage_credential_path.clone(),
                    predefined_acl: self.storage_predefined_acl.clone(),
                    insecure_skip_verify: self.storage_insecure_skip_verify,
                }),
                // When scheduler triggers the export task, it will set true. If the export task is
                // triggered by the user, it will set false.
                persistent: false,
                output_path,
                timeout: Some(
                    prost_wkt_types::Duration::try_from(self.timeout)
                        .or_err(ErrorType::ParseError)?,
                ),
                need_piece_content,
                force_hard_link: self.force_hard_link,
                digest: self.digest.clone(),
                remote_ip: preferred_local_ip().map(|ip| ip.to_string()),
                overwrite: self.overwrite,
            })
            .await
            .inspect_err(|err| {
                error!("download persistent task failed: {}", err);
            })?;

        // If transfer_from_dfdaemon is true, then dfstore needs to create the output file and write the
        // piece content to the output file.
        let mut f = if self.transfer_from_dfdaemon {
            if let Some(parent) = self.output.parent() {
                if !parent.exists() {
                    fs::create_dir_all(parent).await.inspect_err(|err| {
                        error!("failed to create directory {:?}: {}", parent, err);
                    })?;
                }
            }

            let f = OpenOptions::new()
                .create(true)
                .truncate(true)
                .write(true)
                .mode(dfstore::DEFAULT_OUTPUT_FILE_MODE)
                .open(&self.output)
                .await
                .inspect_err(|err| {
                    error!("open file {:?} failed: {}", self.output, err);
                })?;

            Some(BufWriter::with_capacity(TRANSFER_WRITE_BUFFER_SIZE, f))
        } else {
            None
        };

        // Initialize progress bar.
        let progress_bar = if self.no_progress {
            ProgressBar::hidden()
        } else {
            ProgressBar::new(0)
        };

        progress_bar.set_style(
            ProgressStyle::with_template(
                "[{elapsed_precise}] [{wide_bar}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
            )
            .or_err(ErrorType::ParseError)?
            .with_key("eta", |state: &ProgressState, w: &mut dyn Write| {
                write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()
            })
            .progress_chars("#>-"),
        );

        //  Download file.
        let mut downloaded = 0;
        let mut content_length = None;
        let mut out_stream = response.into_inner();
        loop {
            match out_stream.message().await {
                Ok(Some(message)) => {
                    match message.response {
                        Some(download_persistent_task_response::Response::DownloadPersistentTaskStartedResponse(
                            response,
                        )) => {
                            if let Some(f) = &f {
                                if let Err(err) = fallocate(f.get_ref(), response.content_length).await {
                                    error!("fallocate {:?} failed: {}", self.output, err);
                                    fs::remove_file(&self.output).await.inspect_err(|err| {
                                        error!("remove file {:?} failed: {}", self.output, err);
                                    })?;

                                    return Err(err);
                                };
                            }

                            content_length = Some(response.content_length);
                            progress_bar.set_length(response.content_length);
                        }
                        Some(download_persistent_task_response::Response::DownloadPieceFinishedResponse(
                            response,
                        )) => {
                            let piece = match response.piece {
                                Some(piece) => piece,
                                None => {
                                    error!("response piece is missing");
                                    fs::remove_file(&self.output).await.inspect_err(|err| {
                                        error!("remove file {:?} failed: {}", self.output, err);
                                    })?;

                                    return Err(Error::InvalidParameter);
                                }
                            };

                            // Dfstore needs to write the piece content to the output file.
                            if let Some(f) = &mut f {
                                debug!("copy piece {} to {:?} started", piece.number, self.output);
                                if let Err(err) =f.seek(SeekFrom::Start(piece.offset)).await {
                                    error!("seek {:?} failed: {}", self.output, err);
                                    fs::remove_file(&self.output).await.inspect_err(|err| {
                                        error!("remove file {:?} failed: {}", self.output, err);
                                    })?;

                                    return Err(Error::IO(err));
                                };

                                let content = match piece.content {
                                    Some(content) => content,
                                    None => {
                                        error!("piece content is missing");
                                        fs::remove_file(&self.output).await.inspect_err(|err| {
                                            error!("remove file {:?} failed: {}", self.output, err);
                                        })?;

                                        return Err(Error::InvalidParameter);
                                    }
                                };

                                if let Err(err) =f.write_all(&content).await {
                                    error!("write {:?} failed: {}", self.output, err);
                                    fs::remove_file(&self.output).await.inspect_err(|err| {
                                        error!("remove file {:?} failed: {}", self.output, err);
                                    })?;

                                    return Err(Error::IO(err));
                                }

                                debug!("copy piece {} to {:?} success", piece.number, self.output);
                            };

                            downloaded += piece.length;
                            let position = min(downloaded + piece.length, progress_bar.length().unwrap_or(0));
                            progress_bar.set_position(position);
                        }
                        None => {
                            error!("response is missing");
                            fs::remove_file(&self.output).await.inspect_err(|err| {
                                error!("remove file {:?} failed: {}", self.output, err);
                            })?;

                            return Err(Error::UnexpectedResponse);
                        }
                    }
                }
                Ok(None) => break,
                Err(err) => {
                    error!("get message failed: {}", err);
                    fs::remove_file(&self.output).await.inspect_err(|err| {
                        error!("remove file {:?} failed: {}", self.output, err);
                    })?;

                    return Err(Error::TonicStatus(err));
                }
            }
        }

        // Abort if the stream ended before all pieces were received.
        if content_length != Some(downloaded) {
            error!(
                "download incomplete: received {} bytes, expected {} bytes",
                downloaded,
                content_length.unwrap_or_default()
            );
            fs::remove_file(&self.output).await.inspect_err(|err| {
                error!("remove file {:?} failed: {}", self.output, err);
            })?;

            return Err(Error::Unknown("download incomplete".to_string()));
        }

        if let Some(f) = &mut f {
            if let Err(err) = f.flush().await {
                error!("flush {:?} failed: {}", self.output, err);
                fs::remove_file(&self.output).await.inspect_err(|err| {
                    error!("remove file {:?} failed: {}", self.output, err);
                })?;

                return Err(Error::IO(err));
            }
        };
        info!("flush {:?} success", self.output);

        progress_bar.finish_with_message("downloaded");
        Ok(())
    }

    /// Validates command line arguments for the export operation to ensure safe file output.
    ///
    /// This function performs essential validation of the output path to prevent file conflicts
    /// and ensure the target location is suitable for export operations. It checks parent
    /// directory existence, prevents accidental file overwrites, and validates path accessibility
    /// before allowing the export operation to proceed.
    fn validate_args(&self) -> Result<()> {
        let absolute_path = Path::new(&self.output).absolutize()?;
        match absolute_path.parent() {
            Some(parent_path) => {
                if !parent_path.is_dir() {
                    return Err(Error::ValidationError(format!(
                        "output path {} is not a directory",
                        parent_path.to_string_lossy()
                    )));
                }
            }
            None => {
                return Err(Error::ValidationError(format!(
                    "output path {} is not exist",
                    self.output.to_string_lossy()
                )));
            }
        }

        if !self.overwrite && absolute_path.exists() {
            return Err(Error::ValidationError(format!(
                "output path {} is already exist",
                self.output.to_string_lossy()
            )));
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn validate_args_checks_output() {
        let tempdir = tempdir().unwrap();
        let dir = tempdir.path().to_str().unwrap();
        let new_file = format!("{dir}/new.txt");
        let existing_file = format!("{dir}/existing.txt");
        let missing_dir = format!("{dir}/missing");
        let file_in_missing_dir = format!("{missing_dir}/missing.txt");
        std::fs::File::create(&existing_file).unwrap();

        let test_cases: Vec<(Vec<&str>, Result<()>)> = vec![
            (
                vec!["s3://bucket/key", "--output", new_file.as_str()],
                Ok(()),
            ),
            (
                vec![
                    "s3://bucket/key",
                    "--output",
                    existing_file.as_str(),
                    "--overwrite",
                ],
                Ok(()),
            ),
            (
                vec!["s3://bucket/key", "--output", file_in_missing_dir.as_str()],
                Err(Error::ValidationError(format!(
                    "output path {missing_dir} is not a directory"
                ))),
            ),
            (
                vec!["s3://bucket/key", "--output", existing_file.as_str()],
                Err(Error::ValidationError(format!(
                    "output path {existing_file} is already exist"
                ))),
            ),
            (
                vec!["s3://bucket/key", "--output", "/"],
                Err(Error::ValidationError(
                    "output path / is not exist".to_string(),
                )),
            ),
        ];

        for (argv, expected) in test_cases {
            let command =
                ExportCommand::parse_from(std::iter::once("dfstore").chain(argv.iter().copied()));
            let result = command.validate_args();
            assert_eq!(
                result.map_err(|err| err.to_string()),
                expected.map_err(|err| err.to_string())
            );
        }
    }
}