sqlite_wasm_rs/vfs/
memory.rs

1//! Memory VFS, used as the default VFS
2
3use crate::libsqlite3::*;
4use crate::utils::LazyCell;
5use crate::vfs::utils::{
6    check_import_db, ImportDbError, MemChunksFile, SQLiteIoMethods, SQLiteVfs, SQLiteVfsFile,
7    VfsAppData, VfsError, VfsFile, VfsResult, VfsStore,
8};
9
10use std::cell::RefCell;
11use std::collections::HashMap;
12
13type Result<T> = std::result::Result<T, MemVfsError>;
14
15enum MemFile {
16    Main(MemChunksFile),
17    Temp(MemChunksFile),
18}
19
20impl MemFile {
21    fn new(flags: i32) -> Self {
22        if flags & SQLITE_OPEN_MAIN_DB == 0 {
23            Self::Temp(MemChunksFile::default())
24        } else {
25            Self::Main(MemChunksFile::waiting_for_write())
26        }
27    }
28
29    fn file(&self) -> &MemChunksFile {
30        let (MemFile::Main(file) | MemFile::Temp(file)) = self;
31        file
32    }
33
34    fn file_mut(&mut self) -> &mut MemChunksFile {
35        let (MemFile::Main(file) | MemFile::Temp(file)) = self;
36        file
37    }
38}
39
40impl VfsFile for MemFile {
41    fn read(&self, buf: &mut [u8], offset: usize) -> VfsResult<bool> {
42        self.file().read(buf, offset)
43    }
44
45    fn write(&mut self, buf: &[u8], offset: usize) -> VfsResult<()> {
46        self.file_mut().write(buf, offset)
47    }
48
49    fn truncate(&mut self, size: usize) -> VfsResult<()> {
50        self.file_mut().truncate(size)
51    }
52
53    fn flush(&mut self) -> VfsResult<()> {
54        self.file_mut().flush()
55    }
56
57    fn size(&self) -> VfsResult<usize> {
58        self.file().size()
59    }
60}
61
62type MemAppData = RefCell<HashMap<String, MemFile>>;
63
64struct MemStore;
65
66impl VfsStore<MemFile, MemAppData> for MemStore {
67    fn add_file(vfs: *mut sqlite3_vfs, file: &str, flags: i32) -> VfsResult<()> {
68        let app_data = unsafe { Self::app_data(vfs) };
69        app_data
70            .borrow_mut()
71            .insert(file.into(), MemFile::new(flags));
72        Ok(())
73    }
74
75    fn contains_file(vfs: *mut sqlite3_vfs, file: &str) -> VfsResult<bool> {
76        let app_data = unsafe { Self::app_data(vfs) };
77        Ok(app_data.borrow().contains_key(file))
78    }
79
80    fn delete_file(vfs: *mut sqlite3_vfs, file: &str) -> VfsResult<()> {
81        let app_data = unsafe { Self::app_data(vfs) };
82        if app_data.borrow_mut().remove(file).is_none() {
83            return Err(VfsError::new(
84                SQLITE_IOERR_DELETE,
85                format!("{file} not found"),
86            ));
87        }
88        Ok(())
89    }
90
91    fn with_file<F: Fn(&MemFile) -> VfsResult<i32>>(
92        vfs_file: &SQLiteVfsFile,
93        f: F,
94    ) -> VfsResult<i32> {
95        let name = unsafe { vfs_file.name() };
96        let app_data = unsafe { Self::app_data(vfs_file.vfs) };
97        match app_data.borrow().get(name) {
98            Some(file) => f(file),
99            None => Err(VfsError::new(SQLITE_IOERR, format!("{name} not found"))),
100        }
101    }
102
103    fn with_file_mut<F: Fn(&mut MemFile) -> VfsResult<i32>>(
104        vfs_file: &SQLiteVfsFile,
105        f: F,
106    ) -> VfsResult<i32> {
107        let name = unsafe { vfs_file.name() };
108        let app_data = unsafe { Self::app_data(vfs_file.vfs) };
109        match app_data.borrow_mut().get_mut(name) {
110            Some(file) => f(file),
111            None => Err(VfsError::new(SQLITE_IOERR, format!("{name} not found"))),
112        }
113    }
114}
115
116struct MemIoMethods;
117
118impl SQLiteIoMethods for MemIoMethods {
119    type File = MemFile;
120    type AppData = MemAppData;
121    type Store = MemStore;
122
123    const VERSION: ::std::os::raw::c_int = 1;
124}
125
126struct MemVfs;
127
128impl SQLiteVfs<MemIoMethods> for MemVfs {
129    const VERSION: ::std::os::raw::c_int = 1;
130}
131
132#[cfg_attr(target_feature = "atomics", thread_local)]
133static APP_DATA: LazyCell<&'static VfsAppData<MemAppData>> =
134    LazyCell::new(|| unsafe { &*VfsAppData::new(MemAppData::default()).leak() });
135
136fn app_data() -> &'static VfsAppData<MemAppData> {
137    &APP_DATA
138}
139
140#[derive(thiserror::Error, Debug)]
141pub enum MemVfsError {
142    #[error(transparent)]
143    ImportDb(#[from] ImportDbError),
144    #[error("Generic error: {0}")]
145    Generic(String),
146}
147
148/// MemVfs management tools exposed to clients.
149pub struct MemVfsUtil(&'static VfsAppData<MemAppData>);
150
151unsafe impl Send for MemVfsUtil {}
152
153unsafe impl Sync for MemVfsUtil {}
154
155impl Default for MemVfsUtil {
156    fn default() -> Self {
157        MemVfsUtil::new()
158    }
159}
160
161impl MemVfsUtil {
162    /// Get management tool
163    pub fn new() -> Self {
164        MemVfsUtil(app_data())
165    }
166}
167
168impl MemVfsUtil {
169    fn import_db_unchecked_impl(
170        &self,
171        filename: &str,
172        bytes: &[u8],
173        page_size: usize,
174        clear_wal: bool,
175    ) -> Result<()> {
176        if self.exists(filename) {
177            return Err(MemVfsError::Generic(format!(
178                "{filename} file already exists"
179            )));
180        }
181
182        self.0.borrow_mut().insert(filename.into(), {
183            let mut file = MemFile::Main(MemChunksFile::new(page_size));
184            file.write(bytes, 0).unwrap();
185            if clear_wal {
186                file.write(&[1, 1], 18).unwrap();
187            }
188            file
189        });
190
191        Ok(())
192    }
193
194    /// Import the database.
195    ///
196    /// If the database is imported with WAL mode enabled,
197    /// it will be forced to write back to legacy mode, see
198    /// <https://sqlite.org/forum/forumpost/67882c5b04>
199    ///
200    /// If the imported database is encrypted, use `import_db_unchecked` instead.
201    pub fn import_db(&self, filename: &str, bytes: &[u8]) -> Result<()> {
202        let page_size = check_import_db(bytes)?;
203        self.import_db_unchecked_impl(filename, bytes, page_size, true)
204    }
205
206    /// `import_db` without checking, can be used to import encrypted database.
207    pub fn import_db_unchecked(
208        &self,
209        filename: &str,
210        bytes: &[u8],
211        page_size: usize,
212    ) -> Result<()> {
213        self.import_db_unchecked_impl(filename, bytes, page_size, false)
214    }
215
216    /// Export the database.
217    pub fn export_db(&self, filename: &str) -> Result<Vec<u8>> {
218        let name2file = self.0.borrow();
219
220        if let Some(file) = name2file.get(filename) {
221            let file_size = file.size().unwrap();
222            let mut ret = vec![0; file_size];
223            file.read(&mut ret, 0).unwrap();
224            Ok(ret)
225        } else {
226            Err(MemVfsError::Generic(
227                "The file to be exported does not exist".into(),
228            ))
229        }
230    }
231
232    /// Delete the specified database, please make sure that the database is closed.
233    pub fn delete_db(&self, filename: &str) {
234        self.0.borrow_mut().remove(filename);
235    }
236
237    /// Delete all database, please make sure that all database is closed.
238    pub fn clear_all(&self) {
239        std::mem::take(&mut *self.0.borrow_mut());
240    }
241
242    /// Does the database exists.
243    pub fn exists(&self, filename: &str) -> bool {
244        self.0.borrow().contains_key(filename)
245    }
246
247    /// List all files.
248    pub fn list(&self) -> Vec<String> {
249        self.0.borrow().keys().cloned().collect()
250    }
251
252    /// Number of files.
253    pub fn count(&self) -> usize {
254        self.0.borrow().len()
255    }
256}
257
258pub(crate) fn install() -> ::std::os::raw::c_int {
259    let app_data = app_data();
260    let vfs = Box::leak(Box::new(MemVfs::vfs(
261        c"memvfs".as_ptr().cast(),
262        app_data as *const _ as *mut _,
263    )));
264    unsafe { sqlite3_vfs_register(vfs, 1) }
265}
266
267#[cfg(test)]
268mod tests {
269    use crate::{
270        mem_vfs::{MemAppData, MemFile, MemStore},
271        utils::{test_suite::test_vfs_store, VfsAppData},
272    };
273    use wasm_bindgen_test::wasm_bindgen_test;
274
275    #[wasm_bindgen_test]
276    fn test_memory_vfs_store() {
277        test_vfs_store::<MemAppData, MemFile, MemStore>(VfsAppData::new(MemAppData::default()))
278            .unwrap();
279    }
280}