docbox-storage 0.8.2

#![forbid(unsafe_code)]
#![warn(missing_docs)]

//! # Storage
//!
//! Docbox storage backend abstraction, handles abstracting the task of working with file
//! storage to allow for multiple backends and easier testing.
//!
//! # Environment Variables
//!
//! See [s3] this is currently the only available backend for storage

use aws_config::SdkConfig;
use aws_sdk_s3::presigning::PresignedRequest;
use bytes::{Buf, Bytes};
use bytes_utils::SegmentedBuf;
use chrono::{DateTime, Utc};
use futures::{Stream, StreamExt};
use serde::{Deserialize, Serialize};
use std::{fmt::Debug, pin::Pin, time::Duration};
use thiserror::Error;

pub mod s3;

/// Configuration for a storage layer factory
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(tag = "provider", rename_all = "snake_case")]
pub enum StorageLayerFactoryConfig {
    /// Config for a S3 backend
    S3(s3::S3StorageLayerFactoryConfig),
}

impl Default for StorageLayerFactoryConfig {
    fn default() -> Self {
        Self::S3(Default::default())
    }
}

/// Errors that could occur when loading the storage layer factory
/// configuration from the environment
#[derive(Debug, Error)]
pub enum StorageLayerFactoryConfigError {
    /// Error from the S3 layer config
    #[error(transparent)]
    S3(#[from] s3::S3StorageLayerFactoryConfigError),
}

impl StorageLayerFactoryConfig {
    /// Load the configuration from the current environment variables
    pub fn from_env() -> Result<Self, StorageLayerFactoryConfigError> {
        s3::S3StorageLayerFactoryConfig::from_env()
            .map(Self::S3)
            .map_err(StorageLayerFactoryConfigError::S3)
    }
}

/// Storage layer factory for creating storage layer instances
/// with some underlying backend implementation
#[derive(Clone)]
pub enum StorageLayerFactory {
    /// S3 storage backend
    S3(s3::S3StorageLayerFactory),
}

/// Errors that can occur when using a storage layer
#[derive(Debug, Error)]
pub enum StorageLayerError {
    /// Error from the S3 layer
    #[error(transparent)]
    S3(Box<s3::S3StorageError>),

    /// Error collecting streamed response bytes
    #[error("failed to collect file contents")]
    CollectBytes,
}

impl From<s3::S3StorageError> for StorageLayerError {
    fn from(value: s3::S3StorageError) -> Self {
        Self::S3(Box::new(value))
    }
}

/// Options required to initialize a storage layer from a [StorageLayerFactory]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StorageLayerOptions {
    /// Name of the storage bucket
    pub bucket_name: String,
}

impl StorageLayerFactory {
    /// Create a [StorageLayerFactory] from the provided config
    pub fn from_config(aws_config: &SdkConfig, config: StorageLayerFactoryConfig) -> Self {
        match config {
            StorageLayerFactoryConfig::S3(config) => {
                Self::S3(s3::S3StorageLayerFactory::from_config(aws_config, config))
            }
        }
    }

    /// Create a simple layer for testing purposes
    #[cfg(debug_assertions)]
    pub fn create_test_layer(&self) -> StorageLayer {
        self.create_layer(StorageLayerOptions {
            bucket_name: "test".to_string(),
        })
    }

    /// Create a storage layer from the provided `options`
    pub fn create_layer(&self, options: StorageLayerOptions) -> StorageLayer {
        match self {
            StorageLayerFactory::S3(s3) => {
                let layer = s3.create_storage_layer(options.bucket_name);
                StorageLayer::S3(layer)
            }
        }
    }
}

/// Storage layer for a tenant with different underlying backends
#[derive(Clone)]
pub enum StorageLayer {
    /// Storage layer backed by S3
    S3(s3::S3StorageLayer),
}

/// Outcome from creating a bucket
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CreateBucketOutcome {
    /// Fresh bucket was created
    New,
    /// Bucket with the same name already exists
    Existing,
}

/// Options for properties of an uploaded file
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct UploadFileOptions {
    /// Content type of the uploaded file
    pub content_type: String,
    /// Tags to append to the file
    pub tags: Option<Vec<UploadFileTag>>,
}

/// Additional behavioral tags to use when uploading the file
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UploadFileTag {
    /// Tag that the file should expire after 1 day
    ExpireDays1,
    /// Tag that the file should expire after 30 days
    ExpireDays30,
}

impl StorageLayer {
    /// Get the name of the bucket
    pub fn bucket_name(&self) -> String {
        match self {
            StorageLayer::S3(layer) => layer.bucket_name(),
        }
    }

    /// Creates the tenant storage bucket
    ///
    /// In the event that the bucket already exists, this is treated as a
    /// [`Ok`] result rather than an error
    #[tracing::instrument(skip(self))]
    pub async fn create_bucket(&self) -> Result<CreateBucketOutcome, StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.create_bucket().await,
        }
    }

    /// Checks if the bucket exists
    #[tracing::instrument(skip(self))]
    pub async fn bucket_exists(&self) -> Result<bool, StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.bucket_exists().await,
        }
    }

    /// Deletes the tenant storage bucket
    ///
    /// In the event that the bucket did not exist before calling this
    /// function this is treated as an [`Ok`] result
    #[tracing::instrument(skip(self))]
    pub async fn delete_bucket(&self) -> Result<(), StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.delete_bucket().await,
        }
    }

    /// Create a presigned file upload URL
    #[tracing::instrument(skip(self))]
    pub async fn create_presigned(
        &self,
        key: &str,
        size: i64,
    ) -> Result<(PresignedRequest, DateTime<Utc>), StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.create_presigned(key, size).await,
        }
    }

    /// Create a presigned file download URL
    ///
    /// Presigned download creation will succeed even if the requested key
    /// is not present
    #[tracing::instrument(skip(self))]
    pub async fn create_presigned_download(
        &self,
        key: &str,
        expires_in: Duration,
    ) -> Result<(PresignedRequest, DateTime<Utc>), StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.create_presigned_download(key, expires_in).await,
        }
    }

    /// Uploads a file to the S3 bucket for the tenant
    #[tracing::instrument(skip(self, body), fields(body_length = body.len()))]
    pub async fn upload_file(
        &self,
        key: &str,
        body: Bytes,
        options: UploadFileOptions,
    ) -> Result<(), StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.upload_file(key, body, options).await,
        }
    }

    /// Add the SNS notification to a bucket
    #[tracing::instrument(skip(self))]
    pub async fn add_bucket_notifications(&self, sns_arn: &str) -> Result<(), StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.add_bucket_notifications(sns_arn).await,
        }
    }

    /// Sets the allowed CORS origins for accessing the storage from the frontend
    #[tracing::instrument(skip(self))]
    pub async fn set_bucket_cors_origins(
        &self,
        origins: Vec<String>,
    ) -> Result<(), StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.set_bucket_cors_origins(origins).await,
        }
    }

    /// Deletes the file with the provided `key`
    ///
    /// In the event that the file did not exist before calling this
    /// function this is treated as an [`Ok`] result
    #[tracing::instrument(skip(self))]
    pub async fn delete_file(&self, key: &str) -> Result<(), StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.delete_file(key).await,
        }
    }

    /// Gets a byte stream for a file from S3
    #[tracing::instrument(skip(self))]
    pub async fn get_file(&self, key: &str) -> Result<FileStream, StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.get_file(key).await,
        }
    }

    /// Get pending migrations for the storage layer based on the list of already applied
    /// migration names
    #[tracing::instrument(skip(self))]
    pub async fn get_pending_migrations(
        &self,
        applied_names: Vec<String>,
    ) -> Result<Vec<String>, StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.get_pending_migrations(applied_names).await,
        }
    }

    /// Apply a migration by name
    #[tracing::instrument(skip(self))]
    pub async fn apply_migration(&self, name: &str) -> Result<(), StorageLayerError> {
        match self {
            StorageLayer::S3(layer) => layer.apply_migration(name).await,
        }
    }
}

/// Internal trait defining required async implementations for a storage backend
pub(crate) trait StorageLayerImpl {
    fn bucket_name(&self) -> String;

    async fn create_bucket(&self) -> Result<CreateBucketOutcome, StorageLayerError>;

    async fn bucket_exists(&self) -> Result<bool, StorageLayerError>;

    async fn delete_bucket(&self) -> Result<(), StorageLayerError>;

    async fn create_presigned(
        &self,
        key: &str,
        size: i64,
    ) -> Result<(PresignedRequest, DateTime<Utc>), StorageLayerError>;

    async fn create_presigned_download(
        &self,
        key: &str,
        expires_in: Duration,
    ) -> Result<(PresignedRequest, DateTime<Utc>), StorageLayerError>;

    async fn upload_file(
        &self,
        key: &str,
        body: Bytes,
        options: UploadFileOptions,
    ) -> Result<(), StorageLayerError>;

    async fn add_bucket_notifications(&self, sns_arn: &str) -> Result<(), StorageLayerError>;

    async fn set_bucket_cors_origins(&self, origins: Vec<String>) -> Result<(), StorageLayerError>;

    async fn delete_file(&self, key: &str) -> Result<(), StorageLayerError>;

    async fn get_file(&self, key: &str) -> Result<FileStream, StorageLayerError>;

    async fn get_pending_migrations(
        &self,
        applied_names: Vec<String>,
    ) -> Result<Vec<String>, StorageLayerError>;

    async fn apply_migration(&self, name: &str) -> Result<(), StorageLayerError>;
}

/// Stream of bytes from a file
pub struct FileStream {
    /// Underlying stream
    pub stream: Pin<Box<dyn Stream<Item = std::io::Result<Bytes>> + Send>>,
}

impl Debug for FileStream {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FileStream").finish()
    }
}

impl Stream for FileStream {
    type Item = std::io::Result<Bytes>;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        self.stream.as_mut().poll_next(cx)
    }
}

impl FileStream {
    /// Collect the stream to completion as a single [Bytes] buffer
    pub async fn collect_bytes(mut self) -> Result<Bytes, StorageLayerError> {
        let mut output = SegmentedBuf::new();

        while let Some(result) = self.next().await {
            let chunk = result.map_err(|error| {
                tracing::error!(?error, "failed to collect file stream bytes");
                StorageLayerError::CollectBytes
            })?;

            output.push(chunk);
        }

        Ok(output.copy_to_bytes(output.remaining()))
    }
}