1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use async_trait::async_trait;
use aws_config::meta::region::RegionProviderChain;
use aws_sdk_s3 as s3;
use s3::{
    error::SdkError,
    operation::{
        delete_object::DeleteObjectError, get_object::GetObjectError, put_object::PutObjectError,
    },
    primitives::ByteStream,
    Client,
};
use thiserror::Error;

use super::{StorageClientError, StorageClientInterface};

#[derive(Error, Debug)]
pub enum ClientError {
    #[error("GetObject Error: {0}")]
    GetObject(#[from] SdkError<GetObjectError>),
    #[error("PutObject Error: {0}")]
    PutObject(#[from] SdkError<PutObjectError>),
    #[error("DeleteObject Error: {0}")]
    DeleteObject(#[from] SdkError<DeleteObjectError>),
    #[error("S3 Client Error - {0}")]
    General(String),
}

impl From<ClientError> for StorageClientError {
    fn from(error: ClientError) -> Self {
        match error {
            ClientError::GetObject(err) => StorageClientError::GetObject(err.to_string()),
            ClientError::PutObject(err) => StorageClientError::PutObject(err.to_string()),
            ClientError::DeleteObject(err) => StorageClientError::DeleteObject(err.to_string()),
            ClientError::General(err) => StorageClientError::General(err),
        }
    }
}
#[derive(Clone, Debug)]
pub struct StorageClient {
    bucket: String,
    client: s3::Client,
}

impl StorageClient {
    #[allow(unused)]
    pub async fn new(bucket: String) -> Self {
        let region_provider = RegionProviderChain::default_provider().or_else("us-east-1");
        let config = aws_config::from_env().region(region_provider).load().await;
        let client = Client::new(&config);
        Self { bucket, client }
    }
}

#[async_trait]
impl StorageClientInterface for StorageClient {
    async fn get_object(&self, key: String) -> Result<Option<String>, StorageClientError> {
        let object_res = self
            .client
            .get_object()
            .bucket(&self.bucket)
            .key(key)
            .send()
            .await;

        let object = match object_res {
            Ok(object) => object,
            Err(err) => match err.into_service_error() {
                GetObjectError::NoSuchKey(_) => return Ok(None),
                err => {
                    return Err(StorageClientError::GetObject(err.to_string()));
                }
            },
        };

        let body_bytes = object
            .body
            .collect()
            .await
            .map_err(|err| ClientError::General(err.to_string()))?
            .to_vec();

        let body = String::from_utf8(body_bytes)
            .map_err(|err| ClientError::General(format!("Failed to parse object body: {err}")))?;

        Ok(Some(body))
    }

    async fn put_object(&self, key: String, body: String) -> Result<(), StorageClientError> {
        let body_bytes = body.as_bytes().to_vec();

        let _ = self
            .client
            .put_object()
            .bucket(&self.bucket)
            .key(key)
            .body(ByteStream::from(body_bytes))
            .send()
            .await
            .map_err(ClientError::PutObject)?;

        Ok(())
    }

    async fn delete_object(&self, key: String) -> Result<(), StorageClientError> {
        let _ = self
            .client
            .delete_object()
            .bucket(&self.bucket)
            .key(key)
            .send()
            .await
            .map_err(ClientError::DeleteObject)?;

        Ok(())
    }
}