antimatter 2.0.13

// src/stream.rs

use rusoto_core::Region;
use rusoto_s3::{GetObjectRequest, PutObjectRequest, S3Client, S3};
use std::fs::File;
use std::io::{self, Read, Write};
use tokio::io::AsyncReadExt;

/// Writes a `CapsuleBundle`'s bytes to file.
///
/// # Arguments
///
/// * `data` - The byte vector to write.
/// * `path` - The path to the `CapsuleBundle` file.
pub fn write_to_file(data: Vec<u8>, path: String) -> io::Result<()> {
    let mut out_file = File::create(path)?;
    out_file.write_all(&data)?;
    Ok(())
}

/// Reads a `CapsuleBundle` object from file.
///
/// # Arguments
///
/// * `path` - The path to the `CapsuleBundle` file.
///
/// # Returns
///
/// A `CapsuleBundle` bytes.
pub fn read_from_file(path: String) -> io::Result<Vec<u8>> {
    let mut file = File::open(path)?;
    let mut data = Vec::new();
    file.read_to_end(&mut data)?;
    Ok(data)
}

/// Writes a `CapsuleBundle`'s bytes to S3.
///
/// # Arguments
///
/// * `bucket` - The S3 bucket containing the `CapsuleBundle`'s resource.
/// * `key` - The S3 bucket key for the `CapsuleBundle`'s resource.
/// * `data` - The byte vector to write.
pub async fn upload_to_s3(
    bucket: String,
    key: String,
    data: Vec<u8>,
) -> Result<(), rusoto_core::RusotoError<rusoto_s3::PutObjectError>> {
    let s3_client = S3Client::new(Region::default());

    // Prepare S3 request.
    let req = PutObjectRequest {
        bucket,
        key,
        body: Some(data.into()),
        ..Default::default()
    };

    // Execute the S3 request.
    s3_client.put_object(req).await.map(|_| ())
}

/// Reads a `CapsuleBundle` object from S3.
///
/// # Arguments
///
/// * `bucket` - The S3 bucket containing the `CapsuleBundle`'s resource.
/// * `key` - The S3 bucket key for the `CapsuleBundle`'s resource.
///
/// # Returns
///
/// A `CapsuleBundle` bytes.
pub async fn download_from_s3(
    bucket: String,
    key: String,
) -> Result<Vec<u8>, rusoto_core::RusotoError<rusoto_s3::GetObjectError>> {
    let s3_client = S3Client::new(Region::default());

    // Prepare S3 request.
    let req = GetObjectRequest {
        bucket,
        key,
        ..Default::default()
    };

    // Execute the S3 request.
    let s3_obj = s3_client.get_object(req).await?;
    let mut content = Vec::new();
    if let Some(stream) = s3_obj.body {
        stream.into_async_read().read_to_end(&mut content).await?;
    }

    Ok(content)
}

#[cfg(test)]
mod tests {
    // NOTE: For these tests you must fill in your AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY.
    use super::*;

    #[tokio::test]
    #[ignore]
    async fn s3_upload() {
        let bucket = "rust-capsule-test".to_string();
        let key = "test-object".to_string();

        let content_to_upload = "This is some content".as_bytes().to_vec();

        // Upload content
        match upload_to_s3(bucket.clone(), key.clone(), content_to_upload).await {
            Ok(_) => println!("Successfully uploaded to S3"),
            Err(e) => println!("Error uploading to S3: {:?}", e),
        }
    }

    #[tokio::test]
    #[ignore]
    async fn s3_download() {
        let bucket = "rust-capsule-test".to_string();
        let key = "test-object".to_string();

        match download_from_s3(bucket, key).await {
            Ok(content) => {
                let text =
                    String::from_utf8(content).expect("Downloaded content was not valid UTF-8");
                println!("Downloaded content: {}", text);
            }
            Err(e) => println!("Error downloading from S3: {:?}", e),
        }
    }
}