1use anyhow::Result;
2use std::path::Path;
3use std::sync::Arc;
4
5pub struct Db(DbStore);
6
7enum DbStore {
8 Null,
9 Real(rocksdb::DB),
10
11 #[cfg(any(test, feature = "test-support"))]
12 Fake {
13 data: parking_lot::Mutex<collections::HashMap<Vec<u8>, Vec<u8>>>,
14 },
15}
16
17impl Db {
18 /// Open or create a database at the given file path.
19 pub fn open(path: &Path) -> Result<Arc<Self>> {
20 let db = rocksdb::DB::open_default(path)?;
21 Ok(Arc::new(Self(DbStore::Real(db))))
22 }
23
24 /// Open a null database that stores no data, for use as a fallback
25 /// when there is an error opening the real database.
26 pub fn null() -> Arc<Self> {
27 Arc::new(Self(DbStore::Null))
28 }
29
30 /// Open a fake database for testing.
31 #[cfg(any(test, feature = "test-support"))]
32 pub fn open_fake() -> Arc<Self> {
33 Arc::new(Self(DbStore::Fake {
34 data: Default::default(),
35 }))
36 }
37
38 pub fn read<K, I>(&self, keys: I) -> Result<Vec<Option<Vec<u8>>>>
39 where
40 K: AsRef<[u8]>,
41 I: IntoIterator<Item = K>,
42 {
43 match &self.0 {
44 DbStore::Real(db) => db
45 .multi_get(keys)
46 .into_iter()
47 .map(|e| e.map_err(Into::into))
48 .collect(),
49
50 DbStore::Null => Ok(keys.into_iter().map(|_| None).collect()),
51
52 #[cfg(any(test, feature = "test-support"))]
53 DbStore::Fake { data: db } => {
54 let db = db.lock();
55 Ok(keys
56 .into_iter()
57 .map(|key| db.get(key.as_ref()).cloned())
58 .collect())
59 }
60 }
61 }
62
63 pub fn delete<K, I>(&self, keys: I) -> Result<()>
64 where
65 K: AsRef<[u8]>,
66 I: IntoIterator<Item = K>,
67 {
68 match &self.0 {
69 DbStore::Real(db) => {
70 let mut batch = rocksdb::WriteBatch::default();
71 for key in keys {
72 batch.delete(key);
73 }
74 db.write(batch)?;
75 }
76
77 DbStore::Null => {}
78
79 #[cfg(any(test, feature = "test-support"))]
80 DbStore::Fake { data: db } => {
81 let mut db = db.lock();
82 for key in keys {
83 db.remove(key.as_ref());
84 }
85 }
86 }
87 Ok(())
88 }
89
90 pub fn write<K, V, I>(&self, entries: I) -> Result<()>
91 where
92 K: AsRef<[u8]>,
93 V: AsRef<[u8]>,
94 I: IntoIterator<Item = (K, V)>,
95 {
96 match &self.0 {
97 DbStore::Real(db) => {
98 let mut batch = rocksdb::WriteBatch::default();
99 for (key, value) in entries {
100 batch.put(key, value);
101 }
102 db.write(batch)?;
103 }
104
105 DbStore::Null => {}
106
107 #[cfg(any(test, feature = "test-support"))]
108 DbStore::Fake { data: db } => {
109 let mut db = db.lock();
110 for (key, value) in entries {
111 db.insert(key.as_ref().into(), value.as_ref().into());
112 }
113 }
114 }
115 Ok(())
116 }
117}
118
119#[cfg(test)]
120mod tests {
121 use super::*;
122 use tempdir::TempDir;
123
124 #[gpui::test]
125 fn test_db() {
126 let dir = TempDir::new("db-test").unwrap();
127 let fake_db = Db::open_fake();
128 let real_db = Db::open(&dir.path().join("test.db")).unwrap();
129
130 for db in [&real_db, &fake_db] {
131 assert_eq!(
132 db.read(["key-1", "key-2", "key-3"]).unwrap(),
133 &[None, None, None]
134 );
135
136 db.write([("key-1", "one"), ("key-3", "three")]).unwrap();
137 assert_eq!(
138 db.read(["key-1", "key-2", "key-3"]).unwrap(),
139 &[
140 Some("one".as_bytes().to_vec()),
141 None,
142 Some("three".as_bytes().to_vec())
143 ]
144 );
145
146 db.delete(["key-3", "key-4"]).unwrap();
147 assert_eq!(
148 db.read(["key-1", "key-2", "key-3"]).unwrap(),
149 &[Some("one".as_bytes().to_vec()), None, None,]
150 );
151 }
152
153 drop(real_db);
154
155 let real_db = Db::open(&dir.path().join("test.db")).unwrap();
156 assert_eq!(
157 real_db.read(["key-1", "key-2", "key-3"]).unwrap(),
158 &[Some("one".as_bytes().to_vec()), None, None,]
159 );
160 }
161}