1use anyhow::Context as _;
2use gpui::App;
3use sqlez_macros::sql;
4use util::ResultExt as _;
5
6use crate::{
7 query,
8 sqlez::{domain::Domain, thread_safe_connection::ThreadSafeConnection},
9 write_and_log,
10};
11
12pub struct KeyValueStore(crate::sqlez::thread_safe_connection::ThreadSafeConnection);
13
14impl Domain for KeyValueStore {
15 const NAME: &str = stringify!(KeyValueStore);
16
17 const MIGRATIONS: &[&str] = &[
18 sql!(
19 CREATE TABLE IF NOT EXISTS kv_store(
20 key TEXT PRIMARY KEY,
21 value TEXT NOT NULL
22 ) STRICT;
23 ),
24 sql!(
25 CREATE TABLE IF NOT EXISTS scoped_kv_store(
26 namespace TEXT NOT NULL,
27 key TEXT NOT NULL,
28 value TEXT NOT NULL,
29 PRIMARY KEY(namespace, key)
30 ) STRICT;
31 ),
32 ];
33}
34
35crate::static_connection!(KEY_VALUE_STORE, KeyValueStore, []);
36
37pub trait Dismissable {
38 const KEY: &'static str;
39
40 fn dismissed() -> bool {
41 KEY_VALUE_STORE
42 .read_kvp(Self::KEY)
43 .log_err()
44 .is_some_and(|s| s.is_some())
45 }
46
47 fn set_dismissed(is_dismissed: bool, cx: &mut App) {
48 write_and_log(cx, move || async move {
49 if is_dismissed {
50 KEY_VALUE_STORE
51 .write_kvp(Self::KEY.into(), "1".into())
52 .await
53 } else {
54 KEY_VALUE_STORE.delete_kvp(Self::KEY.into()).await
55 }
56 })
57 }
58}
59
60impl KeyValueStore {
61 query! {
62 pub fn read_kvp(key: &str) -> Result<Option<String>> {
63 SELECT value FROM kv_store WHERE key = (?)
64 }
65 }
66
67 pub async fn write_kvp(&self, key: String, value: String) -> anyhow::Result<()> {
68 log::debug!("Writing key-value pair for key {key}");
69 self.write_kvp_inner(key, value).await
70 }
71
72 query! {
73 async fn write_kvp_inner(key: String, value: String) -> Result<()> {
74 INSERT OR REPLACE INTO kv_store(key, value) VALUES ((?), (?))
75 }
76 }
77
78 query! {
79 pub async fn delete_kvp(key: String) -> Result<()> {
80 DELETE FROM kv_store WHERE key = (?)
81 }
82 }
83
84 pub fn scoped<'a>(&'a self, namespace: &'a str) -> ScopedKeyValueStore<'a> {
85 ScopedKeyValueStore {
86 store: self,
87 namespace,
88 }
89 }
90}
91
92pub struct ScopedKeyValueStore<'a> {
93 store: &'a KeyValueStore,
94 namespace: &'a str,
95}
96
97impl ScopedKeyValueStore<'_> {
98 pub fn read(&self, key: &str) -> anyhow::Result<Option<String>> {
99 self.store.select_row_bound::<(&str, &str), String>(
100 "SELECT value FROM scoped_kv_store WHERE namespace = (?) AND key = (?)",
101 )?((self.namespace, key))
102 .context("Failed to read from scoped_kv_store")
103 }
104
105 pub async fn write(&self, key: String, value: String) -> anyhow::Result<()> {
106 let namespace = self.namespace.to_owned();
107 self.store
108 .write(move |connection| {
109 connection.exec_bound::<(&str, &str, &str)>(
110 "INSERT OR REPLACE INTO scoped_kv_store(namespace, key, value) VALUES ((?), (?), (?))",
111 )?((&namespace, &key, &value))
112 .context("Failed to write to scoped_kv_store")
113 })
114 .await
115 }
116
117 pub async fn delete(&self, key: String) -> anyhow::Result<()> {
118 let namespace = self.namespace.to_owned();
119 self.store
120 .write(move |connection| {
121 connection.exec_bound::<(&str, &str)>(
122 "DELETE FROM scoped_kv_store WHERE namespace = (?) AND key = (?)",
123 )?((&namespace, &key))
124 .context("Failed to delete from scoped_kv_store")
125 })
126 .await
127 }
128
129 pub async fn delete_all(&self) -> anyhow::Result<()> {
130 let namespace = self.namespace.to_owned();
131 self.store
132 .write(move |connection| {
133 connection
134 .exec_bound::<&str>("DELETE FROM scoped_kv_store WHERE namespace = (?)")?(
135 &namespace,
136 )
137 .context("Failed to delete_all from scoped_kv_store")
138 })
139 .await
140 }
141}
142
143#[cfg(test)]
144mod tests {
145 use crate::kvp::KeyValueStore;
146
147 #[gpui::test]
148 async fn test_kvp() {
149 let db = KeyValueStore::open_test_db("test_kvp").await;
150
151 assert_eq!(db.read_kvp("key-1").unwrap(), None);
152
153 db.write_kvp("key-1".to_string(), "one".to_string())
154 .await
155 .unwrap();
156 assert_eq!(db.read_kvp("key-1").unwrap(), Some("one".to_string()));
157
158 db.write_kvp("key-1".to_string(), "one-2".to_string())
159 .await
160 .unwrap();
161 assert_eq!(db.read_kvp("key-1").unwrap(), Some("one-2".to_string()));
162
163 db.write_kvp("key-2".to_string(), "two".to_string())
164 .await
165 .unwrap();
166 assert_eq!(db.read_kvp("key-2").unwrap(), Some("two".to_string()));
167
168 db.delete_kvp("key-1".to_string()).await.unwrap();
169 assert_eq!(db.read_kvp("key-1").unwrap(), None);
170 }
171
172 #[gpui::test]
173 async fn test_scoped_kvp() {
174 let db = KeyValueStore::open_test_db("test_scoped_kvp").await;
175
176 let scope_a = db.scoped("namespace-a");
177 let scope_b = db.scoped("namespace-b");
178
179 // Reading a missing key returns None
180 assert_eq!(scope_a.read("key-1").unwrap(), None);
181
182 // Writing and reading back a key works
183 scope_a
184 .write("key-1".to_string(), "value-a1".to_string())
185 .await
186 .unwrap();
187 assert_eq!(scope_a.read("key-1").unwrap(), Some("value-a1".to_string()));
188
189 // Two namespaces with the same key don't collide
190 scope_b
191 .write("key-1".to_string(), "value-b1".to_string())
192 .await
193 .unwrap();
194 assert_eq!(scope_a.read("key-1").unwrap(), Some("value-a1".to_string()));
195 assert_eq!(scope_b.read("key-1").unwrap(), Some("value-b1".to_string()));
196
197 // delete removes a single key without affecting others in the namespace
198 scope_a
199 .write("key-2".to_string(), "value-a2".to_string())
200 .await
201 .unwrap();
202 scope_a.delete("key-1".to_string()).await.unwrap();
203 assert_eq!(scope_a.read("key-1").unwrap(), None);
204 assert_eq!(scope_a.read("key-2").unwrap(), Some("value-a2".to_string()));
205 assert_eq!(scope_b.read("key-1").unwrap(), Some("value-b1".to_string()));
206
207 // delete_all removes all keys in a namespace without affecting other namespaces
208 scope_a
209 .write("key-3".to_string(), "value-a3".to_string())
210 .await
211 .unwrap();
212 scope_a.delete_all().await.unwrap();
213 assert_eq!(scope_a.read("key-2").unwrap(), None);
214 assert_eq!(scope_a.read("key-3").unwrap(), None);
215 assert_eq!(scope_b.read("key-1").unwrap(), Some("value-b1".to_string()));
216 }
217}
218
219pub struct GlobalKeyValueStore(ThreadSafeConnection);
220
221impl Domain for GlobalKeyValueStore {
222 const NAME: &str = stringify!(GlobalKeyValueStore);
223 const MIGRATIONS: &[&str] = &[sql!(
224 CREATE TABLE IF NOT EXISTS kv_store(
225 key TEXT PRIMARY KEY,
226 value TEXT NOT NULL
227 ) STRICT;
228 )];
229}
230
231crate::static_connection!(GLOBAL_KEY_VALUE_STORE, GlobalKeyValueStore, [], global);
232
233impl GlobalKeyValueStore {
234 query! {
235 pub fn read_kvp(key: &str) -> Result<Option<String>> {
236 SELECT value FROM kv_store WHERE key = (?)
237 }
238 }
239
240 pub async fn write_kvp(&self, key: String, value: String) -> anyhow::Result<()> {
241 log::debug!("Writing global key-value pair for key {key}");
242 self.write_kvp_inner(key, value).await
243 }
244
245 query! {
246 async fn write_kvp_inner(key: String, value: String) -> Result<()> {
247 INSERT OR REPLACE INTO kv_store(key, value) VALUES ((?), (?))
248 }
249 }
250
251 query! {
252 pub async fn delete_kvp(key: String) -> Result<()> {
253 DELETE FROM kv_store WHERE key = (?)
254 }
255 }
256}