1use anyhow::Result;
2use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
3use sqlx::{Pool, Sqlite, SqlitePool};
4use std::path::Path;
5use std::str::FromStr;
6use std::sync::Arc;
7
8pub struct Db(DbStore);
9
10enum DbStore {
11 Null,
12 Live(Pool<Sqlite>),
13}
14
15// Things we need to think about:
16// Concurrency? - Needs some research
17// We need to configure or setup our database, create the tables and such
18
19// Write our first migration
20//
21
22// To make a migration:
23// Add to the migrations directory, a file with the name:
24// <NUMBER>_<DESCRIPTION>.sql. Migrations are executed in order of number
25
26impl Db {
27 /// Open or create a database at the given file path.
28 pub fn open(path: &Path) -> Result<Arc<Self>> {
29 let options = SqliteConnectOptions::from_str(path)?.create_if_missing(true);
30
31 Self::initialize(options)
32 }
33
34 /// Open a fake database for testing.
35 #[cfg(any(test, feature = "test-support"))]
36 pub fn open_fake() -> Arc<Self> {
37 let options = SqliteConnectOptions::from_str(":memory:")?;
38
39 Self::initialize(options)
40 }
41
42 fn initialize(options: SqliteConnectOptions) -> Result<Arc<Self>> {
43 let pool = Pool::<Sqlite>::connect_with(options)?;
44
45 sqlx::migrate!().run(&pool).await?;
46
47 Ok(Arc::new(Self(DbStore::Live(pool))))
48 }
49
50 /// Open a null database that stores no data, for use as a fallback
51 /// when there is an error opening the real database.
52 pub fn null() -> Arc<Self> {
53 Arc::new(Self(DbStore::Null))
54 }
55
56 pub fn read<K, I>(&self, keys: I) -> Result<Vec<Option<Vec<u8>>>>
57 where
58 K: AsRef<[u8]>,
59 I: IntoIterator<Item = K>,
60 {
61 match &self.0 {
62 DbStore::Real(db) => db
63 .multi_get(keys)
64 .into_iter()
65 .map(|e| e.map_err(Into::into))
66 .collect(),
67
68 DbStore::Null => Ok(keys.into_iter().map(|_| None).collect()),
69 }
70 }
71
72 pub fn delete<K, I>(&self, keys: I) -> Result<()>
73 where
74 K: AsRef<[u8]>,
75 I: IntoIterator<Item = K>,
76 {
77 match &self.0 {
78 DbStore::Real(db) => {
79 let mut batch = rocksdb::WriteBatch::default();
80 for key in keys {
81 batch.delete(key);
82 }
83 db.write(batch)?;
84 }
85
86 DbStore::Null => {}
87 }
88 Ok(())
89 }
90
91 pub fn write<K, V, I>(&self, entries: I) -> Result<()>
92 where
93 K: AsRef<[u8]>,
94 V: AsRef<[u8]>,
95 I: IntoIterator<Item = (K, V)>,
96 {
97 match &self.0 {
98 DbStore::Real(db) => {
99 let mut batch = rocksdb::WriteBatch::default();
100 for (key, value) in entries {
101 batch.put(key, value);
102 }
103 db.write(batch)?;
104 }
105
106 DbStore::Null => {}
107 }
108 Ok(())
109 }
110}
111
112#[cfg(test)]
113mod tests {
114 use super::*;
115 use tempdir::TempDir;
116
117 #[gpui::test]
118 fn test_db() {
119 let dir = TempDir::new("db-test").unwrap();
120 let fake_db = Db::open_fake();
121 let real_db = Db::open(&dir.path().join("test.db")).unwrap();
122
123 for db in [&real_db, &fake_db] {
124 assert_eq!(
125 db.read(["key-1", "key-2", "key-3"]).unwrap(),
126 &[None, None, None]
127 );
128
129 db.write([("key-1", "one"), ("key-3", "three")]).unwrap();
130 assert_eq!(
131 db.read(["key-1", "key-2", "key-3"]).unwrap(),
132 &[
133 Some("one".as_bytes().to_vec()),
134 None,
135 Some("three".as_bytes().to_vec())
136 ]
137 );
138
139 db.delete(["key-3", "key-4"]).unwrap();
140 assert_eq!(
141 db.read(["key-1", "key-2", "key-3"]).unwrap(),
142 &[Some("one".as_bytes().to_vec()), None, None,]
143 );
144 }
145
146 drop(real_db);
147
148 let real_db = Db::open(&dir.path().join("test.db")).unwrap();
149 assert_eq!(
150 real_db.read(["key-1", "key-2", "key-3"]).unwrap(),
151 &[Some("one".as_bytes().to_vec()), None, None,]
152 );
153 }
154}