recern_vector/
database.rs

1use std::collections::BTreeMap;
2use std::fs::{self, File};
3use std::io::Write;
4use std::path::{Path, PathBuf};
5
6use crate::collection::{Collection, CollectionConfig};
7use crate::error::{Error, Result};
8use crate::storage;
9
10/// A database file and the collections it holds.
11///
12/// The whole database is loaded into memory on open. Changes stay in memory
13/// until [`Database::save`], which atomically replaces the file.
14pub struct Database {
15    path: PathBuf,
16    collections: BTreeMap<String, Collection>,
17}
18
19impl Database {
20    /// Creates a new, empty database file. Fails if the file exists.
21    pub fn create(path: impl AsRef<Path>) -> Result<Self> {
22        let path = path.as_ref().to_path_buf();
23        if path.exists() {
24            return Err(Error::FileExists(path));
25        }
26        let db = Self {
27            path,
28            collections: BTreeMap::new(),
29        };
30        db.save()?;
31        Ok(db)
32    }
33
34    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
35        let path = path.as_ref().to_path_buf();
36        let collections = storage::decode(&fs::read(&path)?)?;
37        Ok(Self { path, collections })
38    }
39
40    pub fn open_or_create(path: impl AsRef<Path>) -> Result<Self> {
41        if path.as_ref().exists() {
42            Self::open(path)
43        } else {
44            Self::create(path)
45        }
46    }
47
48    pub fn path(&self) -> &Path {
49        &self.path
50    }
51
52    pub fn create_collection(
53        &mut self,
54        name: &str,
55        config: CollectionConfig,
56    ) -> Result<&mut Collection> {
57        if name.is_empty() || name.len() > 255 {
58            return Err(Error::InvalidArgument(
59                "collection name must be 1-255 bytes".into(),
60            ));
61        }
62        if self.collections.contains_key(name) {
63            return Err(Error::CollectionExists(name.to_owned()));
64        }
65        let collection = Collection::new(name, config)?;
66        Ok(self
67            .collections
68            .entry(name.to_owned())
69            .or_insert(collection))
70    }
71
72    pub fn collection(&self, name: &str) -> Result<&Collection> {
73        self.collections
74            .get(name)
75            .ok_or_else(|| Error::CollectionNotFound(name.to_owned()))
76    }
77
78    pub fn collection_mut(&mut self, name: &str) -> Result<&mut Collection> {
79        self.collections
80            .get_mut(name)
81            .ok_or_else(|| Error::CollectionNotFound(name.to_owned()))
82    }
83
84    pub fn drop_collection(&mut self, name: &str) -> Result<()> {
85        self.collections
86            .remove(name)
87            .map(|_| ())
88            .ok_or_else(|| Error::CollectionNotFound(name.to_owned()))
89    }
90
91    pub fn collections(&self) -> impl Iterator<Item = &Collection> {
92        self.collections.values()
93    }
94
95    /// Writes the database to a temporary file, syncs it, and renames it over
96    /// the original, so a crash leaves either the old or the new version.
97    pub fn save(&self) -> Result<()> {
98        let bytes = storage::encode(&self.collections);
99        let mut tmp = self.path.as_os_str().to_owned();
100        tmp.push(".tmp");
101        let tmp = PathBuf::from(tmp);
102        {
103            let mut file = File::create(&tmp)?;
104            file.write_all(&bytes)?;
105            file.sync_all()?;
106        }
107        fs::rename(&tmp, &self.path)?;
108        sync_parent_dir(&self.path);
109        Ok(())
110    }
111}
112
113/// Persists the rename itself. Not supported on every platform, so failures
114/// are ignored.
115fn sync_parent_dir(path: &Path) {
116    let dir = match path.parent() {
117        Some(dir) if !dir.as_os_str().is_empty() => dir,
118        _ => Path::new("."),
119    };
120    if let Ok(dir) = File::open(dir) {
121        let _ = dir.sync_all();
122    }
123}
recern_vector/database.rs

recern_vector/
database.rs
