Skip to main content

sz_orm_storage/
s3.rs

1use crate::error::StorageError;
2use crate::storage::Storage;
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8pub struct S3Storage {
9    pub bucket: String,
10    pub region: String,
11    store: Arc<RwLock<HashMap<String, Vec<u8>>>>,
12}
13
14impl S3Storage {
15    pub fn new(bucket: impl Into<String>, region: impl Into<String>) -> Self {
16        Self {
17            bucket: bucket.into(),
18            region: region.into(),
19            store: Arc::new(RwLock::new(HashMap::new())),
20        }
21    }
22
23    pub fn url_for(&self, key: &str) -> String {
24        format!("s3://{}.{}/{}", self.bucket, self.region, key)
25    }
26}
27
28#[async_trait]
29impl Storage for S3Storage {
30    async fn put(
31        &self,
32        key: &str,
33        data: &[u8],
34        _content_type: &str,
35    ) -> Result<String, StorageError> {
36        let mut store = self.store.write().await;
37        store.insert(key.to_string(), data.to_vec());
38        Ok(self.url_for(key))
39    }
40
41    async fn get(&self, key: &str) -> Result<Vec<u8>, StorageError> {
42        let store = self.store.read().await;
43        store
44            .get(key)
45            .cloned()
46            .ok_or_else(|| StorageError::NotFound(format!("s3://{}/{}", self.bucket, key)))
47    }
48
49    async fn delete(&self, key: &str) -> Result<(), StorageError> {
50        let mut store = self.store.write().await;
51        store.remove(key);
52        Ok(())
53    }
54
55    async fn exists(&self, key: &str) -> Result<bool, StorageError> {
56        let store = self.store.read().await;
57        Ok(store.contains_key(key))
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[tokio::test]
66    async fn test_s3_put_and_get() {
67        let storage = S3Storage::new("my-bucket", "us-east-1");
68        let url = storage
69            .put("file.txt", b"hello s3", "text/plain")
70            .await
71            .unwrap();
72        assert!(url.starts_with("s3://my-bucket.us-east-1/"));
73        assert!(url.contains("file.txt"));
74
75        let data = storage.get("file.txt").await.unwrap();
76        assert_eq!(data, b"hello s3");
77    }
78
79    #[tokio::test]
80    async fn test_s3_get_not_found() {
81        let storage = S3Storage::new("bucket", "us-west-2");
82        let result = storage.get("missing.txt").await;
83        assert!(result.is_err());
84        assert!(matches!(result.unwrap_err(), StorageError::NotFound(_)));
85    }
86
87    #[tokio::test]
88    async fn test_s3_delete() {
89        let storage = S3Storage::new("bucket", "us-east-1");
90        storage.put("del.txt", b"data", "text/plain").await.unwrap();
91        assert!(storage.exists("del.txt").await.unwrap());
92
93        storage.delete("del.txt").await.unwrap();
94        assert!(!storage.exists("del.txt").await.unwrap());
95    }
96
97    #[tokio::test]
98    async fn test_s3_exists_false_for_missing() {
99        let storage = S3Storage::new("bucket", "us-east-1");
100        assert!(!storage.exists("nope").await.unwrap());
101    }
102
103    #[tokio::test]
104    async fn test_s3_overwrite() {
105        let storage = S3Storage::new("bucket", "us-east-1");
106        storage.put("key", b"v1", "text/plain").await.unwrap();
107        storage.put("key", b"v2", "text/plain").await.unwrap();
108        let data = storage.get("key").await.unwrap();
109        assert_eq!(data, b"v2");
110    }
111
112    #[tokio::test]
113    async fn test_s3_multiple_keys() {
114        let storage = S3Storage::new("bucket", "us-east-1");
115        storage.put("a", b"a-data", "text/plain").await.unwrap();
116        storage.put("b", b"b-data", "text/plain").await.unwrap();
117        storage.put("c", b"c-data", "text/plain").await.unwrap();
118
119        assert_eq!(storage.get("a").await.unwrap(), b"a-data");
120        assert_eq!(storage.get("b").await.unwrap(), b"b-data");
121        assert_eq!(storage.get("c").await.unwrap(), b"c-data");
122
123        storage.delete("b").await.unwrap();
124        assert!(storage.exists("a").await.unwrap());
125        assert!(!storage.exists("b").await.unwrap());
126        assert!(storage.exists("c").await.unwrap());
127    }
128}