papercrate/backend/src/storage.rs

use std::time::Duration;

use std::sync::Arc;

use anyhow::{anyhow, Context, Result};
use async_trait::async_trait;
use aws_sdk_s3::presigning::PresigningConfig;
use aws_sdk_s3::primitives::ByteStream;
use aws_sdk_s3::Client as S3Client;

use crate::models::Tenant;

#[async_trait]
pub trait ObjectStorage: Send + Sync + 'static {
    async fn put_object(
        &self,
        key: &str,
        bytes: Vec<u8>,
        content_type: Option<String>,
        content_disposition: Option<String>,
    ) -> Result<()>;

    async fn presign_get_object(
        &self,
        key: &str,
        expires_in: Duration,
        response_content_disposition: Option<&str>,
    ) -> Result<String>;

    async fn get_object(&self, key: &str) -> Result<Vec<u8>>;

    async fn delete_object(&self, key: &str) -> Result<()>;
}

pub struct S3Storage {
    client: S3Client,
    bucket: String,
}

impl S3Storage {
    pub fn new(client: S3Client, bucket: impl Into<String>) -> Self {
        Self {
            client,
            bucket: bucket.into(),
        }
    }
}

#[async_trait]
impl ObjectStorage for S3Storage {
    async fn put_object(
        &self,
        key: &str,
        bytes: Vec<u8>,
        content_type: Option<String>,
        content_disposition: Option<String>,
    ) -> Result<()> {
        let mut request = self
            .client
            .put_object()
            .bucket(&self.bucket)
            .key(key)
            .body(ByteStream::from(bytes));

        if let Some(content_type) = content_type {
            request = request.content_type(content_type);
        }

        if let Some(content_disposition) = content_disposition {
            request = request.content_disposition(content_disposition);
        }

        request
            .send()
            .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<String> {
        let presign_config = PresigningConfig::builder()
            .expires_in(expires_in)
            .build()
            .context("failed to build S3 presigning config")?;

        let mut request = self.client.get_object().bucket(&self.bucket).key(key);
        if let Some(value) = response_content_disposition {
            request = request.response_content_disposition(value);
        }

        let presigned = request
            .presigned(presign_config)
            .await
            .context("failed to generate presigned download URL")?;

        Ok(presigned.uri().to_string())
    }

    async fn get_object(&self, key: &str) -> Result<Vec<u8>> {
        let response = self
            .client
            .get_object()
            .bucket(&self.bucket)
            .key(key)
            .send()
            .await
            .context("failed to download object from S3")?;

        let bytes = response
            .body
            .collect()
            .await
            .context("failed to read object stream")?
            .into_bytes()
            .to_vec();

        Ok(bytes)
    }

    async fn delete_object(&self, key: &str) -> Result<()> {
        self.client
            .delete_object()
            .bucket(&self.bucket)
            .key(key)
            .send()
            .await
            .context("failed to delete object from S3")?;
        Ok(())
    }
}
#[derive(Clone)]
pub struct TenantStorage {
    inner: Arc<dyn ObjectStorage>,
    root: String,
}

impl TenantStorage {
    pub fn new(inner: Arc<dyn ObjectStorage>, tenant: &Tenant) -> Result<Self> {
        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 async fn put_object(
        &self,
        key: &str,
        bytes: Vec<u8>,
        content_type: Option<String>,
        content_disposition: Option<String>,
    ) -> 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<String> {
        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<Vec<u8>> {
        let qualified = self.qualify(key);
        self.inner.get_object(&qualified).await
    }

    pub async fn delete_object(&self, key: &str) -> Result<()> {
        let qualified = self.qualify(key);
        self.inner.delete_object(&qualified).await
    }
}