gitdataai/lib/storage/lib.rs

pub mod error;
pub mod local;
pub mod s3;

use std::time::Duration;

use async_trait::async_trait;
pub use aws_sdk_s3::primitives::ByteStream;
use aws_sdk_s3::primitives::ByteStreamError;
pub use error::{StorageError, StorageResult};
pub use local::{LocalStorage, LocalStorageConfig};
pub use s3::{S3Storage, S3StorageConfig};

#[derive(Clone, Debug)]
pub enum AppStorageConfig {
    Local(LocalStorageConfig),
    S3(S3StorageConfig),
}

#[derive(Clone)]
pub enum AppStorage {
    Local(LocalStorage),
    S3(S3Storage),
}

#[derive(Clone, Debug, Default)]
pub struct PutObjectOptions {
    pub content_type: Option<String>,
    pub content_length: Option<i64>,
    pub cache_control: Option<String>,
}

#[derive(Clone, Debug)]
pub struct StoredObject {
    pub key: String,
    pub url: String,
    pub e_tag: Option<String>,
    pub version_id: Option<String>,
}

#[derive(Debug)]
pub struct StorageObjectStream {
    pub body: ByteStream,
    pub content_length: Option<i64>,
    pub content_type: Option<String>,
    pub e_tag: Option<String>,
}

#[derive(Clone, Debug)]
pub struct StorageObject {
    pub bytes: Vec<u8>,
    pub content_length: Option<i64>,
    pub content_type: Option<String>,
    pub e_tag: Option<String>,
}

#[async_trait]
pub trait ObjectStorage: Send + Sync {
    async fn put_stream(
        &self,
        key: &str,
        body: ByteStream,
        options: PutObjectOptions,
    ) -> StorageResult<StoredObject>;

    async fn put_bytes(
        &self,
        key: &str,
        bytes: Vec<u8>,
        options: PutObjectOptions,
    ) -> StorageResult<StoredObject>;

    async fn get_stream(&self, key: &str)
    -> StorageResult<StorageObjectStream>;

    async fn get_bytes(&self, key: &str) -> StorageResult<StorageObject>;

    async fn delete(&self, key: &str) -> StorageResult<()>;

    fn public_url(&self, key: &str) -> StorageResult<Option<String>>;

    async fn presigned_get_url(
        &self,
        key: &str,
        expires_in: Duration,
    ) -> StorageResult<String>;
}

impl AppStorage {
    pub async fn init(config: AppStorageConfig) -> StorageResult<Self> {
        match config {
            AppStorageConfig::Local(config) => {
                Ok(Self::Local(LocalStorage::connect(config).await?))
            }
            AppStorageConfig::S3(config) => {
                Ok(Self::S3(S3Storage::connect(config).await?))
            }
        }
    }
}

#[async_trait]
impl ObjectStorage for AppStorage {
    async fn put_stream(
        &self,
        key: &str,
        body: ByteStream,
        options: PutObjectOptions,
    ) -> StorageResult<StoredObject> {
        match self {
            Self::Local(storage) => {
                storage.put_stream(key, body, options).await
            }
            Self::S3(storage) => storage.put_stream(key, body, options).await,
        }
    }

    async fn put_bytes(
        &self,
        key: &str,
        bytes: Vec<u8>,
        options: PutObjectOptions,
    ) -> StorageResult<StoredObject> {
        match self {
            Self::Local(storage) => {
                storage.put_bytes(key, bytes, options).await
            }
            Self::S3(storage) => storage.put_bytes(key, bytes, options).await,
        }
    }

    async fn get_stream(
        &self,
        key: &str,
    ) -> StorageResult<StorageObjectStream> {
        match self {
            Self::Local(storage) => storage.get_stream(key).await,
            Self::S3(storage) => storage.get_stream(key).await,
        }
    }

    async fn get_bytes(&self, key: &str) -> StorageResult<StorageObject> {
        match self {
            Self::Local(storage) => storage.get_bytes(key).await,
            Self::S3(storage) => storage.get_bytes(key).await,
        }
    }

    async fn delete(&self, key: &str) -> StorageResult<()> {
        match self {
            Self::Local(storage) => storage.delete(key).await,
            Self::S3(storage) => storage.delete(key).await,
        }
    }

    fn public_url(&self, key: &str) -> StorageResult<Option<String>> {
        match self {
            Self::Local(storage) => storage.public_url(key),
            Self::S3(storage) => storage.public_url(key),
        }
    }

    async fn presigned_get_url(
        &self,
        key: &str,
        expires_in: Duration,
    ) -> StorageResult<String> {
        match self {
            Self::Local(storage) => {
                storage.presigned_get_url(key, expires_in).await
            }
            Self::S3(storage) => {
                storage.presigned_get_url(key, expires_in).await
            }
        }
    }
}

pub async fn collect_byte_stream(
    body: ByteStream,
) -> Result<Vec<u8>, ByteStreamError> {
    body.collect().await.map(|data| data.to_vec())
}

impl TryFrom<&config::AppConfig> for AppStorageConfig {
    type Error = StorageError;

    fn try_from(config: &config::AppConfig) -> Result<Self, Self::Error> {
        let backend = config
            .storage_backend()
            .map_err(|error| StorageError::Config(error.to_string()))?;
        match backend.as_str() {
            "local" | "fs" | "filesystem" => {
                Ok(Self::Local(LocalStorageConfig::try_from(config)?))
            }
            "s3" => Ok(Self::S3(S3StorageConfig::try_from(config)?)),
            backend => Err(StorageError::Config(format!(
                "unsupported storage backend: {backend}"
            ))),
        }
    }
}