1use crate::config::StorageConfig;
4use crate::drivers::{LocalDriver, MemoryDriver};
5use crate::storage::{FileMetadata, PutOptions, StorageDriver};
6use crate::Error;
7use bytes::Bytes;
8use dashmap::DashMap;
9use std::sync::Arc;
10use std::time::Duration;
11
12#[derive(Debug, Clone)]
14pub struct DiskConfig {
15 pub driver: DiskDriver,
17 pub root: Option<String>,
19 pub url: Option<String>,
21 #[cfg(feature = "s3")]
23 pub bucket: Option<String>,
24 #[cfg(feature = "s3")]
26 pub region: Option<String>,
27}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum DiskDriver {
32 Local,
34 Memory,
36 #[cfg(feature = "s3")]
38 S3,
39}
40
41impl Default for DiskConfig {
42 fn default() -> Self {
43 Self {
44 driver: DiskDriver::Local,
45 root: Some("storage".to_string()),
46 url: None,
47 #[cfg(feature = "s3")]
48 bucket: None,
49 #[cfg(feature = "s3")]
50 region: None,
51 }
52 }
53}
54
55impl DiskConfig {
56 pub fn local(root: impl Into<String>) -> Self {
58 Self {
59 driver: DiskDriver::Local,
60 root: Some(root.into()),
61 url: None,
62 #[cfg(feature = "s3")]
63 bucket: None,
64 #[cfg(feature = "s3")]
65 region: None,
66 }
67 }
68
69 pub fn memory() -> Self {
71 Self {
72 driver: DiskDriver::Memory,
73 root: None,
74 url: None,
75 #[cfg(feature = "s3")]
76 bucket: None,
77 #[cfg(feature = "s3")]
78 region: None,
79 }
80 }
81
82 pub fn with_url(mut self, url: impl Into<String>) -> Self {
84 self.url = Some(url.into());
85 self
86 }
87}
88
89#[derive(Clone)]
91pub struct Storage {
92 inner: Arc<StorageInner>,
93}
94
95struct StorageInner {
96 disks: DashMap<String, Arc<dyn StorageDriver>>,
97 default_disk: String,
98}
99
100impl Default for Storage {
101 fn default() -> Self {
102 Self::new()
103 }
104}
105
106impl Storage {
107 pub fn new() -> Self {
109 let inner = StorageInner {
110 disks: DashMap::new(),
111 default_disk: "local".to_string(),
112 };
113
114 let storage = Self {
115 inner: Arc::new(inner),
116 };
117
118 let local = LocalDriver::new("storage");
120 storage
121 .inner
122 .disks
123 .insert("local".to_string(), Arc::new(local));
124
125 storage
126 }
127
128 pub fn with_config(default_disk: &str, configs: Vec<(&str, DiskConfig)>) -> Self {
130 let inner = StorageInner {
131 disks: DashMap::new(),
132 default_disk: default_disk.to_string(),
133 };
134
135 let storage = Self {
136 inner: Arc::new(inner),
137 };
138
139 for (name, config) in configs {
140 let driver = Self::create_driver(&config);
141 storage.inner.disks.insert(name.to_string(), driver);
142 }
143
144 storage
145 }
146
147 pub fn with_storage_config(config: StorageConfig) -> Self {
165 let inner = StorageInner {
166 disks: DashMap::new(),
167 default_disk: config.default,
168 };
169
170 let storage = Self {
171 inner: Arc::new(inner),
172 };
173
174 for (name, disk_config) in config.disks {
175 let driver = Self::create_driver(&disk_config);
176 storage.inner.disks.insert(name, driver);
177 }
178
179 storage
180 }
181
182 fn create_driver(config: &DiskConfig) -> Arc<dyn StorageDriver> {
184 match config.driver {
185 DiskDriver::Local => {
186 let root = config.root.clone().unwrap_or_else(|| "storage".to_string());
187 let mut driver = LocalDriver::new(root);
188 if let Some(url) = &config.url {
189 driver = driver.with_url_base(url);
190 }
191 Arc::new(driver)
192 }
193 DiskDriver::Memory => {
194 let mut driver = MemoryDriver::new();
195 if let Some(url) = &config.url {
196 driver = driver.with_url_base(url);
197 }
198 Arc::new(driver)
199 }
200 #[cfg(feature = "s3")]
201 DiskDriver::S3 => {
202 let bucket = config.bucket.clone().unwrap_or_default();
203 let region = config
204 .region
205 .clone()
206 .unwrap_or_else(|| "us-east-1".to_string());
207 let url_base = config.url.clone();
208 let endpoint_url = std::env::var("AWS_URL").ok();
209 Arc::new(crate::drivers::S3Driver::new(
210 bucket,
211 region,
212 url_base,
213 endpoint_url,
214 ))
215 }
216 }
217 }
218
219 pub fn disk(&self, name: &str) -> Result<Disk, Error> {
221 let driver = self
222 .inner
223 .disks
224 .get(name)
225 .map(|d| d.clone())
226 .ok_or_else(|| Error::disk_not_configured(name))?;
227
228 Ok(Disk { driver })
229 }
230
231 pub fn default_disk(&self) -> Result<Disk, Error> {
233 self.disk(&self.inner.default_disk)
234 }
235
236 pub fn register_disk(&self, name: impl Into<String>, driver: Arc<dyn StorageDriver>) {
238 self.inner.disks.insert(name.into(), driver);
239 }
240
241 pub async fn exists(&self, path: &str) -> Result<bool, Error> {
245 self.default_disk()?.exists(path).await
246 }
247
248 pub async fn get(&self, path: &str) -> Result<Bytes, Error> {
250 self.default_disk()?.get(path).await
251 }
252
253 pub async fn get_string(&self, path: &str) -> Result<String, Error> {
255 self.default_disk()?.get_string(path).await
256 }
257
258 pub async fn put(&self, path: &str, contents: impl Into<Bytes>) -> Result<(), Error> {
260 self.default_disk()?.put(path, contents).await
261 }
262
263 pub async fn put_with_options(
265 &self,
266 path: &str,
267 contents: impl Into<Bytes>,
268 options: PutOptions,
269 ) -> Result<(), Error> {
270 self.default_disk()?
271 .put_with_options(path, contents, options)
272 .await
273 }
274
275 pub async fn delete(&self, path: &str) -> Result<(), Error> {
277 self.default_disk()?.delete(path).await
278 }
279
280 pub async fn copy(&self, from: &str, to: &str) -> Result<(), Error> {
282 self.default_disk()?.copy(from, to).await
283 }
284
285 pub async fn rename(&self, from: &str, to: &str) -> Result<(), Error> {
287 self.default_disk()?.rename(from, to).await
288 }
289
290 pub async fn url(&self, path: &str) -> Result<String, Error> {
292 self.default_disk()?.url(path).await
293 }
294}
295
296#[derive(Clone)]
298pub struct Disk {
299 driver: Arc<dyn StorageDriver>,
300}
301
302impl Disk {
303 pub fn new(driver: Arc<dyn StorageDriver>) -> Self {
305 Self { driver }
306 }
307
308 pub async fn exists(&self, path: &str) -> Result<bool, Error> {
310 self.driver.exists(path).await
311 }
312
313 pub async fn get(&self, path: &str) -> Result<Bytes, Error> {
315 self.driver.get(path).await
316 }
317
318 pub async fn get_string(&self, path: &str) -> Result<String, Error> {
320 self.driver.get_string(path).await
321 }
322
323 pub async fn put(&self, path: &str, contents: impl Into<Bytes>) -> Result<(), Error> {
325 self.driver
326 .put(path, contents.into(), PutOptions::new())
327 .await
328 }
329
330 pub async fn put_with_options(
332 &self,
333 path: &str,
334 contents: impl Into<Bytes>,
335 options: PutOptions,
336 ) -> Result<(), Error> {
337 self.driver.put(path, contents.into(), options).await
338 }
339
340 pub async fn delete(&self, path: &str) -> Result<(), Error> {
342 self.driver.delete(path).await
343 }
344
345 pub async fn copy(&self, from: &str, to: &str) -> Result<(), Error> {
347 self.driver.copy(from, to).await
348 }
349
350 pub async fn rename(&self, from: &str, to: &str) -> Result<(), Error> {
352 self.driver.rename(from, to).await
353 }
354
355 pub async fn size(&self, path: &str) -> Result<u64, Error> {
357 self.driver.size(path).await
358 }
359
360 pub async fn metadata(&self, path: &str) -> Result<FileMetadata, Error> {
362 self.driver.metadata(path).await
363 }
364
365 pub async fn url(&self, path: &str) -> Result<String, Error> {
367 self.driver.url(path).await
368 }
369
370 pub async fn temporary_url(&self, path: &str, expiration: Duration) -> Result<String, Error> {
372 self.driver.temporary_url(path, expiration).await
373 }
374
375 pub async fn files(&self, directory: &str) -> Result<Vec<String>, Error> {
377 self.driver.files(directory).await
378 }
379
380 pub async fn all_files(&self, directory: &str) -> Result<Vec<String>, Error> {
382 self.driver.all_files(directory).await
383 }
384
385 pub async fn directories(&self, directory: &str) -> Result<Vec<String>, Error> {
387 self.driver.directories(directory).await
388 }
389
390 pub async fn make_directory(&self, path: &str) -> Result<(), Error> {
392 self.driver.make_directory(path).await
393 }
394
395 pub async fn delete_directory(&self, path: &str) -> Result<(), Error> {
397 self.driver.delete_directory(path).await
398 }
399}
400
401#[cfg(test)]
402mod tests {
403 use super::*;
404
405 #[tokio::test]
406 async fn test_storage_default_disk() {
407 let storage = Storage::with_config("memory", vec![("memory", DiskConfig::memory())]);
408
409 storage.put("test.txt", "hello world").await.unwrap();
410 let contents = storage.get_string("test.txt").await.unwrap();
411 assert_eq!(contents, "hello world");
412 }
413
414 #[tokio::test]
415 async fn test_storage_multiple_disks() {
416 let storage = Storage::with_config(
417 "primary",
418 vec![
419 ("primary", DiskConfig::memory()),
420 ("backup", DiskConfig::memory()),
421 ],
422 );
423
424 storage
426 .disk("primary")
427 .unwrap()
428 .put("data.txt", "primary data")
429 .await
430 .unwrap();
431
432 storage
434 .disk("backup")
435 .unwrap()
436 .put("data.txt", "backup data")
437 .await
438 .unwrap();
439
440 let primary = storage
442 .disk("primary")
443 .unwrap()
444 .get_string("data.txt")
445 .await
446 .unwrap();
447 let backup = storage
448 .disk("backup")
449 .unwrap()
450 .get_string("data.txt")
451 .await
452 .unwrap();
453
454 assert_eq!(primary, "primary data");
455 assert_eq!(backup, "backup data");
456 }
457
458 #[tokio::test]
459 async fn test_disk_not_configured() {
460 let storage = Storage::with_config("memory", vec![("memory", DiskConfig::memory())]);
461 let result = storage.disk("nonexistent");
462 assert!(result.is_err());
463 }
464
465 #[tokio::test]
466 async fn test_register_disk() {
467 let storage = Storage::new();
468 let memory_driver = Arc::new(MemoryDriver::new());
469 storage.register_disk("test", memory_driver);
470
471 storage
472 .disk("test")
473 .unwrap()
474 .put("file.txt", "content")
475 .await
476 .unwrap();
477
478 assert!(storage
479 .disk("test")
480 .unwrap()
481 .exists("file.txt")
482 .await
483 .unwrap());
484 }
485}