athena_cli/commands/inspect/
download.rs

1use anyhow::{anyhow, Context, Result};
2use aws_sdk_s3::Client;
3use std::fs::{self, File};
4use std::io::Write;
5use std::path::{Path, PathBuf};
6use url::Url;
7
8/// Downloads a query result file from S3 to the specified output directory
9pub async fn download_from_s3(
10    s3_client: &Client,
11    s3_url: &str,
12    output_dir: &str,
13    _query_id: &str,
14) -> Result<PathBuf> {
15    // Parse the S3 URL to extract bucket and key
16    let url = Url::parse(s3_url).context(format!("Failed to parse S3 URL: {}", s3_url))?;
17    let host = url
18        .host_str()
19        .ok_or_else(|| anyhow!("Invalid S3 URL: no host in {}", s3_url))?;
20
21    // Handle different S3 URL formats
22    let (bucket, key) = if let Some(stripped) = s3_url.strip_prefix("s3://") {
23        let parts: Vec<&str> = stripped.splitn(2, '/').collect();
24        if parts.len() < 2 {
25            return Err(anyhow!("Invalid S3 URL format (s3://): {}", s3_url));
26        }
27        (parts[0].to_string(), parts[1].to_string())
28    } else if host.ends_with(".amazonaws.com") {
29        let bucket_name = host
30            .split('.')
31            .next()
32            .ok_or_else(|| anyhow!("Invalid S3 URL: cannot extract bucket from host: {}", host))?;
33        let object_key = url.path().strip_prefix('/').unwrap_or(url.path());
34        (bucket_name.to_string(), object_key.to_string())
35    } else {
36        let path_segments = url
37            .path_segments()
38            .ok_or_else(|| anyhow!("Invalid S3 URL: no path in {}", s3_url))?
39            .collect::<Vec<_>>();
40        if path_segments.is_empty() {
41            return Err(anyhow!("Invalid S3 URL: empty path in {}", s3_url));
42        }
43        let bucket_name = path_segments[0];
44        let object_key = path_segments[1..].join("/");
45        (bucket_name.to_string(), object_key)
46    };
47
48    // Create output directory if it doesn't exist
49    fs::create_dir_all(output_dir)
50        .context(format!("Failed to create output directory: {}", output_dir))?;
51
52    // Extract filename from the key
53    let filename_from_key = Path::new(&key)
54        .file_name()
55        .ok_or_else(|| anyhow!("Could not extract filename from S3 key: {}", key))?
56        .to_string_lossy()
57        .to_string();
58
59    // Create output file path
60    let output_path = Path::new(output_dir).join(&filename_from_key);
61
62    // Get the object from S3
63    let resp = s3_client
64        .get_object()
65        .bucket(&bucket)
66        .key(&key)
67        .send()
68        .await
69        .context(format!(
70            "Failed to download file from S3 bucket: {}, key: {}",
71            bucket, key
72        ))?;
73
74    // Read the data
75    let data = resp
76        .body
77        .collect()
78        .await
79        .context("Failed to read S3 object data stream")?;
80    let bytes = data.into_bytes();
81
82    // Write to file
83    let mut file = File::create(&output_path).context(format!(
84        "Failed to create output file: {}",
85        output_path.display()
86    ))?;
87    file.write_all(&bytes).context(format!(
88        "Failed to write data to file: {}",
89        output_path.display()
90    ))?;
91
92    Ok(output_path)
93}
athena_cli/commands/inspect/download.rs

athena_cli/commands/inspect/
download.rs
