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 tracing::warn!(
203 "S3 driver is not yet implemented; all operations will return errors"
204 );
205 Arc::new(crate::drivers::S3Driver)
206 }
207 }
208 }
209
210 pub fn disk(&self, name: &str) -> Result<Disk, Error> {
212 let driver = self
213 .inner
214 .disks
215 .get(name)
216 .map(|d| d.clone())
217 .ok_or_else(|| Error::disk_not_configured(name))?;
218
219 Ok(Disk { driver })
220 }
221
222 pub fn default_disk(&self) -> Result<Disk, Error> {
224 self.disk(&self.inner.default_disk)
225 }
226
227 pub fn register_disk(&self, name: impl Into<String>, driver: Arc<dyn StorageDriver>) {
229 self.inner.disks.insert(name.into(), driver);
230 }
231
232 pub async fn exists(&self, path: &str) -> Result<bool, Error> {
236 self.default_disk()?.exists(path).await
237 }
238
239 pub async fn get(&self, path: &str) -> Result<Bytes, Error> {
241 self.default_disk()?.get(path).await
242 }
243
244 pub async fn get_string(&self, path: &str) -> Result<String, Error> {
246 self.default_disk()?.get_string(path).await
247 }
248
249 pub async fn put(&self, path: &str, contents: impl Into<Bytes>) -> Result<(), Error> {
251 self.default_disk()?.put(path, contents).await
252 }
253
254 pub async fn put_with_options(
256 &self,
257 path: &str,
258 contents: impl Into<Bytes>,
259 options: PutOptions,
260 ) -> Result<(), Error> {
261 self.default_disk()?
262 .put_with_options(path, contents, options)
263 .await
264 }
265
266 pub async fn delete(&self, path: &str) -> Result<(), Error> {
268 self.default_disk()?.delete(path).await
269 }
270
271 pub async fn copy(&self, from: &str, to: &str) -> Result<(), Error> {
273 self.default_disk()?.copy(from, to).await
274 }
275
276 pub async fn rename(&self, from: &str, to: &str) -> Result<(), Error> {
278 self.default_disk()?.rename(from, to).await
279 }
280
281 pub async fn url(&self, path: &str) -> Result<String, Error> {
283 self.default_disk()?.url(path).await
284 }
285}
286
287#[derive(Clone)]
289pub struct Disk {
290 driver: Arc<dyn StorageDriver>,
291}
292
293impl Disk {
294 pub fn new(driver: Arc<dyn StorageDriver>) -> Self {
296 Self { driver }
297 }
298
299 pub async fn exists(&self, path: &str) -> Result<bool, Error> {
301 self.driver.exists(path).await
302 }
303
304 pub async fn get(&self, path: &str) -> Result<Bytes, Error> {
306 self.driver.get(path).await
307 }
308
309 pub async fn get_string(&self, path: &str) -> Result<String, Error> {
311 self.driver.get_string(path).await
312 }
313
314 pub async fn put(&self, path: &str, contents: impl Into<Bytes>) -> Result<(), Error> {
316 self.driver
317 .put(path, contents.into(), PutOptions::new())
318 .await
319 }
320
321 pub async fn put_with_options(
323 &self,
324 path: &str,
325 contents: impl Into<Bytes>,
326 options: PutOptions,
327 ) -> Result<(), Error> {
328 self.driver.put(path, contents.into(), options).await
329 }
330
331 pub async fn delete(&self, path: &str) -> Result<(), Error> {
333 self.driver.delete(path).await
334 }
335
336 pub async fn copy(&self, from: &str, to: &str) -> Result<(), Error> {
338 self.driver.copy(from, to).await
339 }
340
341 pub async fn rename(&self, from: &str, to: &str) -> Result<(), Error> {
343 self.driver.rename(from, to).await
344 }
345
346 pub async fn size(&self, path: &str) -> Result<u64, Error> {
348 self.driver.size(path).await
349 }
350
351 pub async fn metadata(&self, path: &str) -> Result<FileMetadata, Error> {
353 self.driver.metadata(path).await
354 }
355
356 pub async fn url(&self, path: &str) -> Result<String, Error> {
358 self.driver.url(path).await
359 }
360
361 pub async fn temporary_url(&self, path: &str, expiration: Duration) -> Result<String, Error> {
363 self.driver.temporary_url(path, expiration).await
364 }
365
366 pub async fn files(&self, directory: &str) -> Result<Vec<String>, Error> {
368 self.driver.files(directory).await
369 }
370
371 pub async fn all_files(&self, directory: &str) -> Result<Vec<String>, Error> {
373 self.driver.all_files(directory).await
374 }
375
376 pub async fn directories(&self, directory: &str) -> Result<Vec<String>, Error> {
378 self.driver.directories(directory).await
379 }
380
381 pub async fn make_directory(&self, path: &str) -> Result<(), Error> {
383 self.driver.make_directory(path).await
384 }
385
386 pub async fn delete_directory(&self, path: &str) -> Result<(), Error> {
388 self.driver.delete_directory(path).await
389 }
390}
391
392#[cfg(test)]
393mod tests {
394 use super::*;
395
396 #[tokio::test]
397 async fn test_storage_default_disk() {
398 let storage = Storage::with_config("memory", vec![("memory", DiskConfig::memory())]);
399
400 storage.put("test.txt", "hello world").await.unwrap();
401 let contents = storage.get_string("test.txt").await.unwrap();
402 assert_eq!(contents, "hello world");
403 }
404
405 #[tokio::test]
406 async fn test_storage_multiple_disks() {
407 let storage = Storage::with_config(
408 "primary",
409 vec![
410 ("primary", DiskConfig::memory()),
411 ("backup", DiskConfig::memory()),
412 ],
413 );
414
415 storage
417 .disk("primary")
418 .unwrap()
419 .put("data.txt", "primary data")
420 .await
421 .unwrap();
422
423 storage
425 .disk("backup")
426 .unwrap()
427 .put("data.txt", "backup data")
428 .await
429 .unwrap();
430
431 let primary = storage
433 .disk("primary")
434 .unwrap()
435 .get_string("data.txt")
436 .await
437 .unwrap();
438 let backup = storage
439 .disk("backup")
440 .unwrap()
441 .get_string("data.txt")
442 .await
443 .unwrap();
444
445 assert_eq!(primary, "primary data");
446 assert_eq!(backup, "backup data");
447 }
448
449 #[tokio::test]
450 async fn test_disk_not_configured() {
451 let storage = Storage::with_config("memory", vec![("memory", DiskConfig::memory())]);
452 let result = storage.disk("nonexistent");
453 assert!(result.is_err());
454 }
455
456 #[tokio::test]
457 async fn test_register_disk() {
458 let storage = Storage::new();
459 let memory_driver = Arc::new(MemoryDriver::new());
460 storage.register_disk("test", memory_driver);
461
462 storage
463 .disk("test")
464 .unwrap()
465 .put("file.txt", "content")
466 .await
467 .unwrap();
468
469 assert!(storage
470 .disk("test")
471 .unwrap()
472 .exists("file.txt")
473 .await
474 .unwrap());
475 }
476}