use crate::constants::{content_types, defaults};
use crate::config::StorageConfig;
use crate::error::{BakerError, Result};
use aws_config::BehaviorVersion;
use aws_sdk_s3::primitives::ByteStream;
use aws_sdk_s3::Client;
use std::fs;
use std::path::Path;
use std::time::SystemTime;
#[derive(Debug, Clone)]
pub struct S3Object {
pub key: String,
pub size: u64,
pub last_modified: SystemTime,
pub etag: String,
}
pub struct S3Manager {
client: Client,
bucket_name: String,
bucket_prefix: Option<String>,
}
impl S3Manager {
pub async fn new(config: &StorageConfig) -> Result<Self> {
let region = config.region.as_deref().unwrap_or(defaults::AWS_REGION);
let aws_config = aws_config::defaults(BehaviorVersion::latest())
.profile_name(&config.profile)
.region(aws_sdk_s3::config::Region::new(region.to_string()))
.load()
.await;
let client = Client::new(&aws_config);
Ok(Self {
client,
bucket_name: config.bucket.name.clone(),
bucket_prefix: config.bucket.prefix.clone(),
})
}
pub fn with_client(client: Client, bucket_name: String, bucket_prefix: Option<String>) -> Self {
Self {
client,
bucket_name,
bucket_prefix,
}
}
pub fn get_client(&self) -> &Client {
&self.client
}
pub fn get_bucket_name(&self) -> &str {
&self.bucket_name
}
pub async fn upload_file<P: AsRef<Path>>(
&self,
local_path: P,
s3_key: &str,
verbose: bool,
) -> Result<()> {
let local_path = local_path.as_ref();
let full_key = self.get_full_key(s3_key);
if verbose {
println!(
" Uploading: {} -> s3://{}/{}",
local_path.display(),
self.bucket_name,
full_key
);
}
let file_content = fs::read(local_path)?;
let content_type = mime_guess::from_path(local_path)
.first_or_octet_stream()
.to_string();
self.client
.put_object()
.bucket(&self.bucket_name)
.key(&full_key)
.body(ByteStream::from(file_content))
.content_type(content_type)
.send()
.await
.map_err(|e| BakerError::S3(e.to_string()))?;
Ok(())
}
pub async fn upload_metadata(&self, s3_key: &str, content: &str) -> Result<()> {
let full_key = self.get_full_key(s3_key);
self.client
.put_object()
.bucket(&self.bucket_name)
.key(&full_key)
.body(ByteStream::from(content.as_bytes().to_vec()))
.content_type(content_types::YAML)
.send()
.await
.map_err(|e| BakerError::S3(e.to_string()))?;
Ok(())
}
pub async fn download_metadata(&self, s3_key: &str) -> Result<String> {
let full_key = self.get_full_key(s3_key);
let response = self
.client
.get_object()
.bucket(&self.bucket_name)
.key(&full_key)
.send()
.await
.map_err(|e| BakerError::S3(e.to_string()))?;
let body = response
.body
.collect()
.await
.map_err(|e| BakerError::S3(e.to_string()))?;
let bytes = body.into_bytes();
String::from_utf8(bytes.to_vec())
.map_err(|e| BakerError::S3(format!("Invalid UTF-8 in response: {}", e)))
}
pub async fn download_file<P: AsRef<Path>>(
&self,
s3_key: &str,
local_path: P,
verbose: bool,
) -> Result<()> {
let local_path = local_path.as_ref();
let full_key = self.get_full_key(s3_key);
if verbose {
println!(
" Downloading: s3://{}/{} -> {}",
self.bucket_name,
full_key,
local_path.display()
);
}
let response = self
.client
.get_object()
.bucket(&self.bucket_name)
.key(&full_key)
.send()
.await
.map_err(|e| BakerError::S3(e.to_string()))?;
if let Some(parent) = local_path.parent() {
fs::create_dir_all(parent)?;
}
let body = response
.body
.collect()
.await
.map_err(|e| BakerError::S3(e.to_string()))?;
fs::write(local_path, body.into_bytes())?;
Ok(())
}
pub async fn list_objects(&self, prefix: Option<&str>) -> Result<Vec<S3Object>> {
let full_prefix = self.get_full_key(prefix.unwrap_or(""));
let mut objects = Vec::new();
let mut continuation_token: Option<String> = None;
loop {
let mut request = self
.client
.list_objects_v2()
.bucket(&self.bucket_name)
.prefix(&full_prefix);
if let Some(token) = &continuation_token {
request = request.continuation_token(token);
}
let response = request
.send()
.await
.map_err(|e| BakerError::S3(e.to_string()))?;
if let Some(contents) = response.contents {
for obj in contents {
if let (Some(key), Some(size), Some(last_modified), Some(etag)) =
(obj.key, obj.size, obj.last_modified, obj.e_tag)
{
let system_time = SystemTime::UNIX_EPOCH
+ std::time::Duration::from_secs(last_modified.secs() as u64);
objects.push(S3Object {
key: self.remove_prefix(&key),
size: size as u64,
last_modified: system_time,
etag: etag.trim_matches('"').to_lowercase(),
});
}
}
}
continuation_token = response.next_continuation_token;
if continuation_token.is_none() {
break;
}
}
Ok(objects)
}
pub async fn head_object(&self, s3_key: &str) -> Result<Option<S3Object>> {
let full_key = self.get_full_key(s3_key);
let result = self
.client
.head_object()
.bucket(&self.bucket_name)
.key(&full_key)
.send()
.await;
match result {
Ok(response) => {
if let (Some(content_length), Some(last_modified), Some(etag)) =
(response.content_length, response.last_modified, response.e_tag)
{
let system_time = SystemTime::UNIX_EPOCH
+ std::time::Duration::from_secs(last_modified.secs() as u64);
Ok(Some(S3Object {
key: s3_key.to_string(),
size: content_length as u64,
last_modified: system_time,
etag: etag.trim_matches('"').to_lowercase(),
}))
} else {
Ok(None)
}
}
Err(e) => {
if e.to_string().contains("NotFound") || e.to_string().contains("404") {
Ok(None)
} else {
Err(BakerError::S3(e.to_string()))
}
}
}
}
pub async fn delete_object(&self, s3_key: &str, verbose: bool) -> Result<()> {
let full_key = self.get_full_key(s3_key);
if verbose {
println!(" Deleting: s3://{}/{}", self.bucket_name, full_key);
}
self.client
.delete_object()
.bucket(&self.bucket_name)
.key(&full_key)
.send()
.await
.map_err(|e| BakerError::S3(e.to_string()))?;
Ok(())
}
fn get_full_key(&self, key: &str) -> String {
if let Some(prefix) = &self.bucket_prefix {
if prefix.is_empty() {
key.to_string()
} else {
format!("{}/{}", prefix, key)
.replace("//", "/")
.trim_start_matches('/')
.to_string()
}
} else {
key.to_string()
}
}
fn remove_prefix(&self, full_key: &str) -> String {
if let Some(prefix) = &self.bucket_prefix {
if !prefix.is_empty() && full_key.starts_with(&format!("{}/", prefix)) {
return full_key[prefix.len() + 1..].to_string();
}
}
full_key.to_string()
}
}
pub async fn create_s3_client(config: &StorageConfig) -> Result<S3Manager> {
S3Manager::new(config).await
}