active_storage/drivers/
inmem.rs1use 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 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 with_settings!({
165 filters => CLEANUP_DATA.to_vec()
166 }, {
167 assert_debug_snapshot!(driver);
168 });
169
170 assert!(driver
172 .delete_directory(PathBuf::from("foo").as_path())
173 .await
174 .is_ok());
175
176 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}