use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; use anyhow::{anyhow, Context, Result}; use async_trait::async_trait; use s3::bucket::Bucket; use crate::models::Tenant; #[async_trait] pub trait ObjectStorage: Send + Sync + 'static { async fn put_object( &self, key: &str, bytes: Vec, content_type: Option, content_disposition: Option, ) -> Result<()>; async fn presign_get_object( &self, key: &str, expires_in: Duration, response_content_disposition: Option<&str>, ) -> Result; async fn get_object(&self, key: &str) -> Result>; async fn get_object_range(&self, key: &str, start: u64, end: Option) -> Result>; async fn delete_object(&self, key: &str) -> Result<()>; } pub struct S3Storage { bucket: Bucket, } impl S3Storage { pub fn new(bucket: Bucket) -> Self { Self { bucket } } fn default_content_type(content_type: Option) -> String { content_type.unwrap_or_else(|| "application/octet-stream".to_string()) } } #[async_trait] impl ObjectStorage for S3Storage { async fn put_object( &self, key: &str, bytes: Vec, content_type: Option, content_disposition: Option, ) -> Result<()> { let mut builder = self .bucket .put_object_builder(key, &bytes) .with_content_type(Self::default_content_type(content_type)); if let Some(disposition) = content_disposition { builder = builder .with_content_disposition(disposition) .context("invalid content disposition header")?; } builder .execute() .await .context("failed to upload object to S3")?; Ok(()) } async fn presign_get_object( &self, key: &str, expires_in: Duration, response_content_disposition: Option<&str>, ) -> Result { let expiry_secs = u32::try_from(expires_in.as_secs()).context("presign expiry exceeds u32 range")?; let mut queries = HashMap::new(); if let Some(value) = response_content_disposition { queries.insert( "response-content-disposition".to_string(), value.to_string(), ); } self.bucket .presign_get(key, expiry_secs, (!queries.is_empty()).then_some(queries)) .await .context("failed to generate presigned download URL") } async fn get_object(&self, key: &str) -> Result> { let data = self .bucket .get_object(key) .await .context("failed to download object from S3")?; Ok(data.into_bytes().to_vec()) } async fn get_object_range(&self, key: &str, start: u64, end: Option) -> Result> { let data = self .bucket .get_object_range(key, start, end) .await .context("failed to download ranged object from S3")?; Ok(data.into_bytes().to_vec()) } async fn delete_object(&self, key: &str) -> Result<()> { self.bucket .delete_object(key) .await .context("failed to delete object from S3")?; Ok(()) } } #[derive(Clone)] pub struct TenantStorage { inner: Arc, root: String, } impl TenantStorage { pub fn new(inner: Arc, tenant: &Tenant) -> Result { let root = tenant .storage_root .as_ref() .ok_or_else(|| anyhow!("tenant {} missing storage_root", tenant.id))? .to_owned(); Ok(Self { inner, root }) } fn qualify(&self, key: &str) -> String { format!("{}{}", self.root, key) } pub fn root_prefix(&self) -> &str { &self.root } pub async fn put_object( &self, key: &str, bytes: Vec, content_type: Option, content_disposition: Option, ) -> Result<()> { let qualified = self.qualify(key); self.inner .put_object(&qualified, bytes, content_type, content_disposition) .await } pub async fn presign_get_object( &self, key: &str, expires_in: Duration, response_content_disposition: Option<&str>, ) -> Result { let qualified = self.qualify(key); self.inner .presign_get_object(&qualified, expires_in, response_content_disposition) .await } pub async fn get_object(&self, key: &str) -> Result> { let qualified = self.qualify(key); self.inner.get_object(&qualified).await } pub async fn get_object_range( &self, key: &str, start: u64, end: Option, ) -> Result> { let qualified = self.qualify(key); self.inner.get_object_range(&qualified, start, end).await } pub async fn delete_object(&self, key: &str) -> Result<()> { let qualified = self.qualify(key); self.inner.delete_object(&qualified).await } }