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    println!("Downloading query results from S3: {}", s3_url);
16
17    // Parse the S3 URL to extract bucket and key
18    let url = Url::parse(s3_url).context(format!("Failed to parse S3 URL: {}", s3_url))?;
19
20    // Log the URL components for debugging
21    println!(
22        "URL scheme: {}, host: {:?}, path: {}",
23        url.scheme(),
24        url.host_str(),
25        url.path()
26    );
27
28    let host = url
29        .host_str()
30        .ok_or_else(|| anyhow!("Invalid S3 URL: no host in {}", s3_url))?;
31
32    // Handle different S3 URL formats
33    let (bucket, key) = if let Some(stripped) = s3_url.strip_prefix("s3://") {
34        // s3://bucket-name/key format
35        let parts: Vec<&str> = stripped.splitn(2, '/').collect();
36
37        if parts.len() < 2 {
38            return Err(anyhow!("Invalid S3 URL format (s3://): {}", s3_url));
39        }
40
41        (parts[0].to_string(), parts[1].to_string())
42    } else if host.ends_with(".amazonaws.com") {
43        // https://bucket-name.s3.region.amazonaws.com/key format
44        let bucket_name = host
45            .split('.')
46            .next()
47            .ok_or_else(|| anyhow!("Invalid S3 URL: cannot extract bucket from host: {}", host))?;
48
49        // Remove leading slash from path
50        let object_key = url.path().strip_prefix('/').unwrap_or(url.path());
51
52        (bucket_name.to_string(), object_key.to_string())
53    } else {
54        // https://s3.region.amazonaws.com/bucket-name/key format
55        let path_segments = url
56            .path_segments()
57            .ok_or_else(|| anyhow!("Invalid S3 URL: no path in {}", s3_url))?
58            .collect::<Vec<_>>();
59
60        if path_segments.is_empty() {
61            return Err(anyhow!("Invalid S3 URL: empty path in {}", s3_url));
62        }
63
64        let bucket_name = path_segments[0];
65        let object_key = path_segments[1..].join("/");
66
67        (bucket_name.to_string(), object_key)
68    };
69
70    println!("Extracted bucket: {}, key: {}", bucket, key);
71
72    // Create output directory if it doesn't exist
73    fs::create_dir_all(output_dir)
74        .context(format!("Failed to create output directory: {}", output_dir))?;
75
76    // Extract filename from the key
77    let filename_from_key = Path::new(&key)
78        .file_name()
79        .ok_or_else(|| anyhow!("Could not extract filename from S3 key: {}", key))?
80        .to_string_lossy()
81        .to_string();
82
83    // Create output file path
84    let output_path = Path::new(output_dir).join(&filename_from_key);
85    println!("Will save to: {}", output_path.display());
86
87    // Get the object from S3
88    println!("Requesting object from S3...");
89    let resp = s3_client
90        .get_object()
91        .bucket(&bucket)
92        .key(&key)
93        .send()
94        .await
95        .context(format!(
96            "Failed to download file from S3 bucket: {}, key: {}",
97            bucket, key
98        ))?;
99
100    println!(
101        "S3 response received, content length: {:?}",
102        resp.content_length()
103    );
104
105    // Read the data
106    let data = resp
107        .body
108        .collect()
109        .await
110        .context("Failed to read S3 object data stream")?;
111    let bytes = data.into_bytes();
112
113    println!("Downloaded {} bytes from S3", bytes.len());
114
115    // Write to file
116    let mut file = File::create(&output_path).context(format!(
117        "Failed to create output file: {}",
118        output_path.display()
119    ))?;
120    file.write_all(&bytes).context(format!(
121        "Failed to write data to file: {}",
122        output_path.display()
123    ))?;
124
125    println!(
126        "Successfully downloaded {} bytes to {}",
127        bytes.len(),
128        output_path.display()
129    );
130
131    Ok(output_path)
132}
athena_cli/commands/inspect/download.rs

athena_cli/commands/inspect/
download.rs
