1use anyhow::{Context, Result};
20use rusqlite::{params, Connection, OptionalExtension};
21use std::path::PathBuf;
22use std::sync::Arc;
23use std::time::{SystemTime, UNIX_EPOCH};
24use tokio::sync::Mutex;
25
26fn now_millis() -> i64 {
27 SystemTime::now()
28 .duration_since(UNIX_EPOCH)
29 .map(|d| d.as_millis() as i64)
30 .unwrap_or(0)
31}
32
33#[derive(Clone)]
35pub struct PluginStorage {
36 conn: Arc<Mutex<Connection>>,
37}
38
39impl PluginStorage {
40 pub fn open(path: PathBuf) -> Result<Self> {
42 if let Some(parent) = path.parent() {
43 std::fs::create_dir_all(parent)
44 .with_context(|| format!("creating db dir {}", parent.display()))?;
45 }
46 let conn = Connection::open(&path)
47 .with_context(|| format!("opening plugin-storage db {}", path.display()))?;
48 Self::init_schema(&conn)?;
49 Ok(Self {
50 conn: Arc::new(Mutex::new(conn)),
51 })
52 }
53
54 pub fn in_memory() -> Result<Self> {
56 let conn = Connection::open_in_memory().context("opening in-memory plugin-storage db")?;
57 Self::init_schema(&conn)?;
58 Ok(Self {
59 conn: Arc::new(Mutex::new(conn)),
60 })
61 }
62
63 fn init_schema(conn: &Connection) -> Result<()> {
64 conn.execute_batch(
65 "PRAGMA journal_mode = WAL;
66 CREATE TABLE IF NOT EXISTS plugin_kv (
67 plugin_id TEXT NOT NULL,
68 namespace TEXT NOT NULL,
69 key TEXT NOT NULL,
70 value TEXT NOT NULL,
71 updated_at INTEGER NOT NULL,
72 PRIMARY KEY (plugin_id, namespace, key)
73 );",
74 )
75 .context("initializing plugin-storage schema")?;
76 Ok(())
77 }
78
79 pub async fn get(&self, plugin_id: &str, namespace: &str, key: &str) -> Result<Option<String>> {
81 let conn = self.conn.lock().await;
82 let v = conn
83 .query_row(
84 "SELECT value FROM plugin_kv WHERE plugin_id = ?1 AND namespace = ?2 AND key = ?3",
85 params![plugin_id, namespace, key],
86 |row| row.get::<_, String>(0),
87 )
88 .optional()
89 .context("reading plugin_kv")?;
90 Ok(v)
91 }
92
93 pub async fn set(
95 &self,
96 plugin_id: &str,
97 namespace: &str,
98 key: &str,
99 value: &str,
100 ) -> Result<()> {
101 let conn = self.conn.lock().await;
102 conn.execute(
103 "INSERT INTO plugin_kv (plugin_id, namespace, key, value, updated_at)
104 VALUES (?1, ?2, ?3, ?4, ?5)
105 ON CONFLICT(plugin_id, namespace, key)
106 DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at",
107 params![plugin_id, namespace, key, value, now_millis()],
108 )
109 .context("writing plugin_kv")?;
110 Ok(())
111 }
112
113 pub async fn delete(&self, plugin_id: &str, namespace: &str, key: &str) -> Result<()> {
115 let conn = self.conn.lock().await;
116 conn.execute(
117 "DELETE FROM plugin_kv WHERE plugin_id = ?1 AND namespace = ?2 AND key = ?3",
118 params![plugin_id, namespace, key],
119 )
120 .context("deleting plugin_kv")?;
121 Ok(())
122 }
123
124 pub async fn rekey_plugin(&self, from: &str, to: &str) -> Result<usize> {
144 let conn = self.conn.lock().await;
145 conn.execute(
146 "INSERT OR IGNORE INTO plugin_kv (plugin_id, namespace, key, value, updated_at)
147 SELECT ?2, namespace, key, value, updated_at FROM plugin_kv WHERE plugin_id = ?1",
148 rusqlite::params![from, to],
149 )?;
150 let removed = conn.execute(
151 "DELETE FROM plugin_kv WHERE plugin_id = ?1",
152 rusqlite::params![from],
153 )?;
154 Ok(removed)
155 }
156
157 pub async fn keys(&self, plugin_id: &str, namespace: &str) -> Result<Vec<String>> {
158 let conn = self.conn.lock().await;
159 let mut stmt = conn
160 .prepare(
161 "SELECT key FROM plugin_kv WHERE plugin_id = ?1 AND namespace = ?2
162 ORDER BY updated_at DESC",
163 )
164 .context("preparing plugin_kv keys query")?;
165 let rows = stmt
166 .query_map(params![plugin_id, namespace], |row| row.get::<_, String>(0))
167 .context("querying plugin_kv keys")?;
168 let mut out = Vec::new();
169 for r in rows {
170 out.push(r.context("reading plugin_kv key row")?);
171 }
172 Ok(out)
173 }
174}
175
176#[cfg(test)]
177mod tests {
178 use super::*;
179
180 #[tokio::test]
181 async fn set_get_delete_roundtrip() {
182 let s = PluginStorage::in_memory().unwrap();
183 assert_eq!(s.get("p", "ns", "k").await.unwrap(), None);
184 s.set("p", "ns", "k", "v1").await.unwrap();
185 assert_eq!(s.get("p", "ns", "k").await.unwrap().as_deref(), Some("v1"));
186 s.set("p", "ns", "k", "v2").await.unwrap();
188 assert_eq!(s.get("p", "ns", "k").await.unwrap().as_deref(), Some("v2"));
189 s.delete("p", "ns", "k").await.unwrap();
190 assert_eq!(s.get("p", "ns", "k").await.unwrap(), None);
191 }
192
193 #[tokio::test]
194 async fn plugins_are_isolated_by_id_and_namespace() {
195 let s = PluginStorage::in_memory().unwrap();
196 s.set("plugin-a", "default", "shared", "a").await.unwrap();
197 s.set("plugin-b", "default", "shared", "b").await.unwrap();
198 assert_eq!(
200 s.get("plugin-a", "default", "shared")
201 .await
202 .unwrap()
203 .as_deref(),
204 Some("a")
205 );
206 assert_eq!(
207 s.get("plugin-b", "default", "shared")
208 .await
209 .unwrap()
210 .as_deref(),
211 Some("b")
212 );
213 s.set("plugin-a", "other", "shared", "a2").await.unwrap();
215 assert_eq!(
216 s.get("plugin-a", "default", "shared")
217 .await
218 .unwrap()
219 .as_deref(),
220 Some("a")
221 );
222 }
223
224 #[tokio::test]
225 async fn keys_lists_namespaced_keys() {
226 let s = PluginStorage::in_memory().unwrap();
227 s.set("p", "goals", "conv-1", "x").await.unwrap();
228 s.set("p", "goals", "conv-2", "y").await.unwrap();
229 s.set("p", "other", "conv-3", "z").await.unwrap();
230 let mut keys = s.keys("p", "goals").await.unwrap();
231 keys.sort();
232 assert_eq!(keys, vec!["conv-1".to_string(), "conv-2".to_string()]);
233 }
234}