lumni 0.0.5

Portable file-utility for object-stores
Documentation
use std::collections::HashMap;

use async_trait::async_trait;

use super::get::get_object;
use super::head::head_object;
use super::list::list_files;
use crate::base::config::EnvironmentConfig;
use crate::s3::config::validate_config;
use crate::table::FileObjectTable;
use crate::{FileObjectFilter, LakestreamError, ObjectStoreTrait};

#[derive(Clone)]
pub struct S3Credentials {
    access_key: String,
    secret_key: String,
    session_token: Option<String>,
}

impl S3Credentials {
    pub fn new(
        access_key: String,
        secret_key: String,
        session_token: Option<String>,
    ) -> S3Credentials {
        S3Credentials {
            access_key,
            secret_key,
            session_token,
        }
    }

    pub fn access_key(&self) -> &str {
        &self.access_key
    }

    pub fn secret_key(&self) -> &str {
        &self.secret_key
    }

    pub fn session_token(&self) -> Option<&str> {
        self.session_token.as_deref()
    }
}

#[derive(Debug, Clone)]
pub struct S3Bucket {
    name: String,
    config: EnvironmentConfig,
}

impl S3Bucket {
    pub fn new(
        name: &str,
        mut config: EnvironmentConfig,
    ) -> Result<S3Bucket, LakestreamError> {
        validate_config(&mut config)?;

        Ok(S3Bucket {
            name: name.to_string(),
            config,
        })
    }

    pub fn config(&self) -> &EnvironmentConfig {
        &self.config
    }

    pub fn bucket_path(&self) -> String {
        let region = self.config.get("AWS_REGION").unwrap();
        let endpoint_url =
            self.config.get("S3_ENDPOINT_URL").map(String::as_str);
        let name = Some(self.name().to_string());

        configure_bucket_url(region, endpoint_url, name.as_deref())
    }
}

#[async_trait(?Send)]
impl ObjectStoreTrait for S3Bucket {
    fn name(&self) -> &str {
        &self.name
    }

    fn config(&self) -> &EnvironmentConfig {
        &self.config
    }
    async fn list_files(
        &self,
        prefix: Option<&str>,
        selected_columns: &Option<Vec<&str>>,
        recursive: bool,
        max_keys: Option<u32>,
        filter: &Option<FileObjectFilter>,
        table: &mut FileObjectTable,
    ) -> Result<(), LakestreamError> {
        if let Some(prefix) = prefix {
            // prefix should not exist as a file object
            let (status_code, _response_headers) =
                self.head_object(prefix.trim_end_matches("/")).await?;
            if status_code != 404 {
                return Err(LakestreamError::NoBucketInUri(prefix.to_string()));
            }
        }
        list_files(
            self,
            prefix,
            selected_columns,
            recursive,
            max_keys,
            filter,
            table,
        )
        .await
    }

    async fn get_object(
        &self,
        key: &str,
        data: &mut Vec<u8>,
    ) -> Result<(), LakestreamError> {
        get_object(self, key, data).await
    }

    async fn head_object(
        &self,
        key: &str,
    ) -> Result<(u16, HashMap<String, String>), LakestreamError> {
        head_object(self, key).await
    }
}

pub fn configure_bucket_url(
    region: &str,
    endpoint_url: Option<&str>,
    bucket_name: Option<&str>,
) -> String {
    match endpoint_url {
        Some(url) => match bucket_name {
            Some(name) => format!("{}/{}", url.trim_end_matches('/'), name),
            None => url.to_owned(),
        },
        None => match bucket_name {
            Some(name) => {
                format!("https://{}.s3.{}.amazonaws.com", name, region)
            }
            None => format!("https://s3.{}.amazonaws.com", region),
        },
    }
}