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 unimplemented!("S3 driver requires async initialization")
204 }
205 }
206 }
207
208 pub fn disk(&self, name: &str) -> Result<Disk, Error> {
210 let driver = self
211 .inner
212 .disks
213 .get(name)
214 .map(|d| d.clone())
215 .ok_or_else(|| Error::disk_not_configured(name))?;
216
217 Ok(Disk { driver })
218 }
219
220 pub fn default_disk(&self) -> Result<Disk, Error> {
222 self.disk(&self.inner.default_disk)
223 }
224
225 pub fn register_disk(&self, name: impl Into<String>, driver: Arc<dyn StorageDriver>) {
227 self.inner.disks.insert(name.into(), driver);
228 }
229
230 pub async fn exists(&self, path: &str) -> Result<bool, Error> {
234 self.default_disk()?.exists(path).await
235 }
236
237 pub async fn get(&self, path: &str) -> Result<Bytes, Error> {
239 self.default_disk()?.get(path).await
240 }
241
242 pub async fn get_string(&self, path: &str) -> Result<String, Error> {
244 self.default_disk()?.get_string(path).await
245 }
246
247 pub async fn put(&self, path: &str, contents: impl Into<Bytes>) -> Result<(), Error> {
249 self.default_disk()?.put(path, contents).await
250 }
251
252 pub async fn put_with_options(
254 &self,
255 path: &str,
256 contents: impl Into<Bytes>,
257 options: PutOptions,
258 ) -> Result<(), Error> {
259 self.default_disk()?
260 .put_with_options(path, contents, options)
261 .await
262 }
263
264 pub async fn delete(&self, path: &str) -> Result<(), Error> {
266 self.default_disk()?.delete(path).await
267 }
268
269 pub async fn copy(&self, from: &str, to: &str) -> Result<(), Error> {
271 self.default_disk()?.copy(from, to).await
272 }
273
274 pub async fn rename(&self, from: &str, to: &str) -> Result<(), Error> {
276 self.default_disk()?.rename(from, to).await
277 }
278
279 pub async fn url(&self, path: &str) -> Result<String, Error> {
281 self.default_disk()?.url(path).await
282 }
283}
284
285#[derive(Clone)]
287pub struct Disk {
288 driver: Arc<dyn StorageDriver>,
289}
290
291impl Disk {
292 pub fn new(driver: Arc<dyn StorageDriver>) -> Self {
294 Self { driver }
295 }
296
297 pub async fn exists(&self, path: &str) -> Result<bool, Error> {
299 self.driver.exists(path).await
300 }
301
302 pub async fn get(&self, path: &str) -> Result<Bytes, Error> {
304 self.driver.get(path).await
305 }
306
307 pub async fn get_string(&self, path: &str) -> Result<String, Error> {
309 self.driver.get_string(path).await
310 }
311
312 pub async fn put(&self, path: &str, contents: impl Into<Bytes>) -> Result<(), Error> {
314 self.driver
315 .put(path, contents.into(), PutOptions::new())
316 .await
317 }
318
319 pub async fn put_with_options(
321 &self,
322 path: &str,
323 contents: impl Into<Bytes>,
324 options: PutOptions,
325 ) -> Result<(), Error> {
326 self.driver.put(path, contents.into(), options).await
327 }
328
329 pub async fn delete(&self, path: &str) -> Result<(), Error> {
331 self.driver.delete(path).await
332 }
333
334 pub async fn copy(&self, from: &str, to: &str) -> Result<(), Error> {
336 self.driver.copy(from, to).await
337 }
338
339 pub async fn rename(&self, from: &str, to: &str) -> Result<(), Error> {
341 self.driver.rename(from, to).await
342 }
343
344 pub async fn size(&self, path: &str) -> Result<u64, Error> {
346 self.driver.size(path).await
347 }
348
349 pub async fn metadata(&self, path: &str) -> Result<FileMetadata, Error> {
351 self.driver.metadata(path).await
352 }
353
354 pub async fn url(&self, path: &str) -> Result<String, Error> {
356 self.driver.url(path).await
357 }
358
359 pub async fn temporary_url(&self, path: &str, expiration: Duration) -> Result<String, Error> {
361 self.driver.temporary_url(path, expiration).await
362 }
363
364 pub async fn files(&self, directory: &str) -> Result<Vec<String>, Error> {
366 self.driver.files(directory).await
367 }
368
369 pub async fn all_files(&self, directory: &str) -> Result<Vec<String>, Error> {
371 self.driver.all_files(directory).await
372 }
373
374 pub async fn directories(&self, directory: &str) -> Result<Vec<String>, Error> {
376 self.driver.directories(directory).await
377 }
378
379 pub async fn make_directory(&self, path: &str) -> Result<(), Error> {
381 self.driver.make_directory(path).await
382 }
383
384 pub async fn delete_directory(&self, path: &str) -> Result<(), Error> {
386 self.driver.delete_directory(path).await
387 }
388}
389
390#[cfg(test)]
391mod tests {
392 use super::*;
393
394 #[tokio::test]
395 async fn test_storage_default_disk() {
396 let storage = Storage::with_config("memory", vec![("memory", DiskConfig::memory())]);
397
398 storage.put("test.txt", "hello world").await.unwrap();
399 let contents = storage.get_string("test.txt").await.unwrap();
400 assert_eq!(contents, "hello world");
401 }
402
403 #[tokio::test]
404 async fn test_storage_multiple_disks() {
405 let storage = Storage::with_config(
406 "primary",
407 vec![
408 ("primary", DiskConfig::memory()),
409 ("backup", DiskConfig::memory()),
410 ],
411 );
412
413 storage
415 .disk("primary")
416 .unwrap()
417 .put("data.txt", "primary data")
418 .await
419 .unwrap();
420
421 storage
423 .disk("backup")
424 .unwrap()
425 .put("data.txt", "backup data")
426 .await
427 .unwrap();
428
429 let primary = storage
431 .disk("primary")
432 .unwrap()
433 .get_string("data.txt")
434 .await
435 .unwrap();
436 let backup = storage
437 .disk("backup")
438 .unwrap()
439 .get_string("data.txt")
440 .await
441 .unwrap();
442
443 assert_eq!(primary, "primary data");
444 assert_eq!(backup, "backup data");
445 }
446
447 #[tokio::test]
448 async fn test_disk_not_configured() {
449 let storage = Storage::with_config("memory", vec![("memory", DiskConfig::memory())]);
450 let result = storage.disk("nonexistent");
451 assert!(result.is_err());
452 }
453
454 #[tokio::test]
455 async fn test_register_disk() {
456 let storage = Storage::new();
457 let memory_driver = Arc::new(MemoryDriver::new());
458 storage.register_disk("test", memory_driver);
459
460 storage
461 .disk("test")
462 .unwrap()
463 .put("file.txt", "content")
464 .await
465 .unwrap();
466
467 assert!(storage
468 .disk("test")
469 .unwrap()
470 .exists("file.txt")
471 .await
472 .unwrap());
473 }
474}