active_storage/drivers/
inmem.rs

1use std::{
2    collections::BTreeMap,
3    path::{Path, PathBuf},
4    sync::Mutex,
5    time::SystemTime,
6};
7
8use super::{Driver, DriverError};
9use crate::{contents::Contents, errors::DriverResult};
10
11#[derive(Debug, Clone)]
12pub struct File {
13    pub content: Vec<u8>,
14    pub last_modified: SystemTime,
15}
16
17#[derive(Debug, Default)]
18pub struct InMemoryDriver {
19    files: Mutex<BTreeMap<PathBuf, File>>,
20    directory: Mutex<BTreeMap<PathBuf, Vec<PathBuf>>>,
21}
22
23impl Clone for InMemoryDriver {
24    fn clone(&self) -> Self {
25        Self {
26            files: Mutex::new(self.files.lock().unwrap().clone()),
27            directory: Mutex::new(self.directory.lock().unwrap().clone()),
28        }
29    }
30}
31
32impl InMemoryDriver {
33    fn get_files(&self) -> BTreeMap<PathBuf, File> {
34        self.files
35            .lock()
36            .expect("inmem store failed getting a lock")
37            .clone()
38    }
39}
40
41#[async_trait::async_trait]
42impl Driver for InMemoryDriver {
43    async fn read(&self, path: &Path) -> DriverResult<Vec<u8>> {
44        let files = self.get_files();
45        let file = files.get(path).ok_or(DriverError::ResourceNotFound)?;
46
47        Ok(Contents::from(file.content.clone()).into())
48    }
49
50    async fn file_exists(&self, path: &Path) -> DriverResult<bool> {
51        Ok(self.get_files().contains_key(path))
52    }
53
54    async fn write(&self, path: &Path, content: Vec<u8>) -> DriverResult<()> {
55        self.files.lock().unwrap().insert(
56            path.to_path_buf(),
57            File {
58                last_modified: SystemTime::now(),
59                content,
60            },
61        );
62
63        if let Some(parent) = path.parent() {
64            self.directory
65                .lock()
66                .unwrap()
67                .entry(parent.to_path_buf())
68                .or_default()
69                .push(path.to_path_buf());
70        }
71
72        Ok(())
73    }
74
75    async fn delete(&self, path: &Path) -> DriverResult<()> {
76        if self.files.lock().unwrap().remove(path).is_none() {
77            return Err(DriverError::ResourceNotFound);
78        }
79
80        self.directory
81            .lock()
82            .unwrap()
83            .entry(path.parent().unwrap().to_path_buf())
84            .or_default()
85            .retain(|file_path| file_path != path);
86
87        Ok(())
88    }
89
90    async fn delete_directory(&self, path: &Path) -> DriverResult<()> {
91        if !self.directory.lock().unwrap().contains_key(path) {
92            return Err(DriverError::ResourceNotFound);
93        }
94
95        self.directory
96            .lock()
97            .unwrap()
98            .retain(|file_path, _| !file_path.starts_with(path));
99
100        self.files
101            .lock()
102            .unwrap()
103            .retain(|file_path, _| !file_path.starts_with(path));
104
105        Ok(())
106    }
107
108    async fn last_modified(&self, path: &Path) -> DriverResult<SystemTime> {
109        let file = self.get_files();
110        let file = file.get(path).ok_or(DriverError::ResourceNotFound)?;
111
112        Ok(file.last_modified)
113    }
114}
115
116#[cfg(test)]
117mod tests {
118
119    use insta::{assert_debug_snapshot, with_settings};
120    use lazy_static::lazy_static;
121
122    use super::*;
123
124    lazy_static! {
125        pub static ref CLEANUP_DATA: Vec<(&'static str, &'static str)> = vec![
126            (r"tv_sec: (\d+),", "tv_sec: TV_SEC"),
127            (r"tv_nsec: (\d+),", "tv_sec: TV_NSEC")
128        ];
129    }
130
131    #[tokio::test]
132    async fn validate_store() {
133        let driver = InMemoryDriver::default();
134
135        // cerate state
136        let _ = driver
137            .write(
138                PathBuf::from("foo").join("file-1.txt").as_path(),
139                b"".into(),
140            )
141            .await;
142
143        let _ = driver
144            .write(
145                PathBuf::from("foo").join("file-2.txt").as_path(),
146                b"".into(),
147            )
148            .await;
149
150        let _ = driver
151            .write(
152                PathBuf::from("bar").join("file-1.txt").as_path(),
153                b"".into(),
154            )
155            .await;
156        let _ = driver
157            .write(
158                PathBuf::from("bar").join("file-2.txt").as_path(),
159                b"".into(),
160            )
161            .await;
162
163        // snapshot the state
164        with_settings!({
165            filters => CLEANUP_DATA.to_vec()
166        }, {
167        assert_debug_snapshot!(driver);
168        });
169
170        // delete folder
171        assert!(driver
172            .delete_directory(PathBuf::from("foo").as_path())
173            .await
174            .is_ok());
175
176        // delete file
177        assert!(driver
178            .delete(PathBuf::from("bar").join("file-1.txt").as_path())
179            .await
180            .is_ok());
181
182        with_settings!({
183            filters => CLEANUP_DATA.to_vec()
184        }, {
185        assert_debug_snapshot!(driver);
186        });
187    }
188}