sql_hummus/
kv.rs

1use sql_peas::StatementHandle;
2
3pub struct Kv {
4    inner: sql_peas::Connection,
5    contains_handle: StatementHandle,
6    get_handle: StatementHandle,
7    insert_handle: StatementHandle,
8    prefix_handle: StatementHandle,
9}
10
11pub struct KvCursor<'a> {
12    inner: sql_peas::Cursor<'a>,
13    prefix: &'a str,
14}
15
16#[derive(thiserror::Error, Debug)]
17pub enum Error {
18    #[error("SQLite error: {0}")]
19    Sqlite(#[from] sql_peas::Error),
20    #[error("Error: {0}")]
21    Other(&'static str),
22}
23
24type Result<T> = std::result::Result<T, Error>;
25
26impl Iterator for KvCursor<'_> {
27    type Item = Result<(String, String)>;
28
29    fn next(&mut self) -> Option<Self::Item> {
30        let row = match self.inner.next()? {
31            Err(e) => return Some(Err(e.into())),
32            Ok(x) => x,
33        };
34        let key = row.read::<&str, _>(0);
35        if !key.starts_with(self.prefix) {
36            return None;
37        }
38        let value = row.read::<&str, _>(1);
39        Some(Ok((key.into(), value.into())))
40    }
41}
42
43const KV_USER_VERSION: u32 = 1783402338;
44const KV_USER_VERSION_I64: i64 = KV_USER_VERSION as i64;
45const KV_TABLE_NAME: &str = "sql_hummus_0_kv";
46
47impl Kv {
48    pub fn new<P: AsRef<std::path::Path>>(p: P) -> Result<Self> {
49        let mut inner = sql_peas::Connection::open(p)?;
50
51        // Set some pragmas where new apps need different defaults than SQLite's defaults
52        inner.execute("PRAGMA trusted_schema = 0;")?;
53        inner.execute("PRAGMA foreign_keys = 1;")?;
54
55        let needs_setup = {
56            let handle = inner.prepare("PRAGMA user_version")?;
57            let stmt = inner.borrow_statement(handle)?;
58            let mut rows = stmt.iter();
59            let row = rows
60                .next()
61                .ok_or(Error::Other("Expected one row from PRAGMA user_version"))??;
62            let needs_setup = match row.read(0) {
63                0 => true,
64                KV_USER_VERSION_I64 => false,
65                _ => return Err(Error::Other("PRAGMA user_version looks like a non-KV file")),
66            };
67            inner.drop_statement(handle)?;
68            needs_setup
69        };
70
71        if needs_setup {
72            // Looks like a new file, so let's set everything up and then stamp it
73
74            inner.execute(format!("CREATE TABLE IF NOT EXISTS {KV_TABLE_NAME} (key TEXT PRIMARY KEY NOT NULL, value TEXT) WITHOUT ROWID"))?;
75
76            // Using `format!()` here because PRAGMA doesn't seem to like bound params.
77            // SECURITY: KV_USER_VERSION is not user-controlled input, it's a compile-time const that we control, so it's not a SQL injection vuln.
78            inner.execute(format!("PRAGMA user_version = {KV_USER_VERSION}"))?;
79        }
80
81        let contains_handle =
82            inner.prepare(format!("SELECT count() FROM {KV_TABLE_NAME} WHERE key = ?"))?;
83        let get_handle =
84            inner.prepare(format!("SELECT value FROM {KV_TABLE_NAME} WHERE key = ?"))?;
85        let insert_handle = inner.prepare(format!("INSERT INTO {KV_TABLE_NAME} (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value"))?;
86        let prefix_handle = inner.prepare(format!(
87            "SELECT key, value FROM {KV_TABLE_NAME} WHERE key >= ? ORDER BY key"
88        ))?;
89
90        Ok(Self {
91            inner,
92            contains_handle,
93            get_handle,
94            insert_handle,
95            prefix_handle,
96        })
97    }
98
99    pub fn clear(&self) -> Result<()> {
100        self.inner.execute(format!("DELETE FROM {KV_TABLE_NAME}"))?;
101        Ok(())
102    }
103
104    pub fn contains_key<S: AsRef<str>>(&self, key: S) -> Result<bool> {
105        let stmt = self.inner.borrow_statement(self.contains_handle)?;
106        stmt.bind((1, key.as_ref()))?;
107        let mut rows = stmt.iter();
108        let row = rows.next().ok_or(Error::Other("Expected one row here"))??;
109        match row.read::<i64, _>(0) {
110            0 => Ok(false),
111            1 => Ok(true),
112            _ => Err(Error::Other(
113                "Logic error - Two rows in KV table have the same key",
114            )),
115        }
116    }
117
118    pub fn get<S: AsRef<str>>(&self, key: S) -> Result<Option<String>> {
119        let stmt = self.inner.borrow_statement(self.get_handle)?;
120        stmt.bind((1, key.as_ref()))?;
121        let mut rows = stmt.iter();
122        match rows.next() {
123            None => Ok(None),
124            Some(Ok(row)) => Ok(Some(row.read::<&str, _>(0).into())),
125            Some(Err(e)) => Err(e)?,
126        }
127    }
128
129    pub fn insert<S: AsRef<str>, T: AsRef<str>>(&self, key: S, value: T) -> Result<Option<String>> {
130        let tx = self.inner.begin_immediate_transaction()?;
131        let old_value = self.get(&key)?;
132        {
133            let stmt = self.inner.borrow_statement(self.insert_handle)?;
134            stmt.bind((1, key.as_ref()))?;
135            stmt.bind((2, value.as_ref()))?;
136            if stmt.next()? != sql_peas::State::Done {
137                return Err(Error::Other("We didn't get State::Done during Kv::insert"));
138            }
139        }
140        tx.commit()?;
141        Ok(old_value)
142    }
143
144    pub fn with_prefix<'a>(&'a self, prefix: &'a str) -> Result<KvCursor<'a>> {
145        let stmt = self.inner.borrow_statement(self.prefix_handle)?;
146        stmt.bind((1, prefix))?;
147        let cursor = stmt.iter();
148        Ok(KvCursor {
149            inner: cursor,
150            prefix,
151        })
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn simple() -> Result<()> {
161        let path = "sql_hummus_test_temp_ZQ3KECAY.db";
162        const COLOR_MODE: &str = "/myapp/color_mode";
163
164        let cxn = Kv::new(path).unwrap();
165        cxn.clear()?;
166        assert!(!cxn.contains_key(COLOR_MODE).unwrap());
167
168        assert!(cxn.insert(COLOR_MODE, "light_mode").unwrap().is_none());
169        assert_eq!(cxn.get(COLOR_MODE).unwrap(), Some("light_mode".into()));
170        assert_eq!(
171            cxn.insert(COLOR_MODE, "dark_mode").unwrap(),
172            Some("light_mode".into())
173        );
174        assert_eq!(
175            cxn.insert(COLOR_MODE, "dark_mode").unwrap(),
176            Some("dark_mode".into())
177        );
178        assert_eq!(cxn.get(COLOR_MODE)?, Some("dark_mode".into()));
179
180        assert!(cxn.contains_key(COLOR_MODE)?);
181        assert!(!cxn.contains_key("Never used this key at all")?);
182
183        cxn.insert("/myapp/bookmarks/https://example.com/", "")?;
184        cxn.insert("/aaapppp/", "")?;
185        cxn.insert("/notmyapp/bookmarks/", "asdf")?;
186
187        assert_eq!(
188            cxn.with_prefix("/myapp/")?.collect::<Result<Vec<_>>>()?,
189            [
190                ("/myapp/bookmarks/https://example.com/", ""),
191                (COLOR_MODE, "dark_mode"),
192            ]
193            .into_iter()
194            .map(|(k, v)| (k.to_string(), v.to_string()))
195            .collect::<Vec<_>>(),
196        );
197
198        assert_eq!(
199            cxn.with_prefix("")?.collect::<Result<Vec<_>>>()?,
200            [
201                ("/aaapppp/", ""),
202                ("/myapp/bookmarks/https://example.com/", ""),
203                (COLOR_MODE, "dark_mode"),
204                ("/notmyapp/bookmarks/", "asdf"),
205            ]
206            .into_iter()
207            .map(|(k, v)| (k.to_string(), v.to_string()))
208            .collect::<Vec<_>>(),
209        );
210
211        Ok(())
212    }
213}
sql_hummus/kv.rs

sql_hummus/
kv.rs
