Skip to main content

prolly_store_sqlite/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::collections::{hash_map::Entry, HashMap};
4use std::path::Path;
5use std::sync::{Mutex, MutexGuard};
6use std::time::Duration;
7
8#[cfg(unix)]
9use rusqlite::OpenFlags;
10use rusqlite::{params, Connection, OptionalExtension};
11
12use prolly::{
13    BatchOp, Cid, Error, ManifestStore, ManifestStoreScan, ManifestUpdate, NamedRootManifest,
14    NodeStoreScan, RootCondition, RootManifest, RootWrite, Store, TransactionConflict,
15    TransactionNodeWrite, TransactionUpdate, TransactionalStore,
16};
17
18struct OrderedBatchReadPlan<'a> {
19    unique_keys: Vec<&'a [u8]>,
20    positions: Option<Vec<usize>>,
21}
22
23impl<'a> OrderedBatchReadPlan<'a> {
24    fn new(keys: &[&'a [u8]]) -> Self {
25        let mut unique_indexes = HashMap::with_capacity(keys.len());
26        let mut unique_keys = Vec::with_capacity(keys.len());
27        let mut positions = None;
28        for key in keys {
29            match unique_indexes.entry(*key) {
30                Entry::Occupied(entry) => positions
31                    .get_or_insert_with(|| (0..unique_keys.len()).collect::<Vec<_>>())
32                    .push(*entry.get()),
33                Entry::Vacant(entry) => {
34                    let index = unique_keys.len();
35                    unique_keys.push(*key);
36                    if let Some(positions) = positions.as_mut() {
37                        positions.push(index);
38                    }
39                    entry.insert(index);
40                }
41            }
42        }
43        Self {
44            unique_keys,
45            positions,
46        }
47    }
48
49    fn unique_keys(&self) -> &[&'a [u8]] {
50        &self.unique_keys
51    }
52
53    fn expand_owned<T: Clone>(&self, values: Vec<Option<T>>) -> Vec<Option<T>> {
54        match &self.positions {
55            Some(positions) => positions
56                .iter()
57                .map(|&index| values[index].clone())
58                .collect(),
59            None => values,
60        }
61    }
62}
63
64fn cid_from_store_key(key: &[u8], context: &str) -> Result<Cid, String> {
65    let bytes: [u8; 32] = key.try_into().map_err(|_| {
66        format!(
67            "{context} key has invalid CID length {}, expected 32",
68            key.len()
69        )
70    })?;
71    Ok(Cid(bytes))
72}
73
74fn sort_cids(cids: &mut [Cid]) {
75    cids.sort_by(|left, right| left.as_bytes().cmp(right.as_bytes()));
76}
77
78fn sort_named_root_manifests(roots: &mut [NamedRootManifest]) {
79    roots.sort_by(|left, right| left.name.cmp(&right.name));
80}
81
82const CREATE_TABLE_SQL: &str = "\
83CREATE TABLE IF NOT EXISTS prolly_nodes (
84    cid  BLOB PRIMARY KEY NOT NULL,
85    node BLOB NOT NULL
86) WITHOUT ROWID;";
87
88const CREATE_HINTS_TABLE_SQL: &str = "\
89CREATE TABLE IF NOT EXISTS prolly_hints (
90    namespace BLOB NOT NULL,
91    key       BLOB NOT NULL,
92    value     BLOB NOT NULL,
93    PRIMARY KEY (namespace, key)
94) WITHOUT ROWID;";
95
96const CREATE_ROOTS_TABLE_SQL: &str = "\
97CREATE TABLE IF NOT EXISTS prolly_roots (
98    name     BLOB PRIMARY KEY NOT NULL,
99    manifest BLOB NOT NULL
100) WITHOUT ROWID;";
101
102const SELECT_SQL: &str = "SELECT node FROM prolly_nodes WHERE cid = ?1";
103const SELECT_NODE_CIDS_SQL: &str = "SELECT cid FROM prolly_nodes ORDER BY cid";
104const UPSERT_SQL: &str = "\
105INSERT INTO prolly_nodes (cid, node)
106VALUES (?1, ?2)
107ON CONFLICT(cid) DO UPDATE SET node = excluded.node";
108const DELETE_SQL: &str = "DELETE FROM prolly_nodes WHERE cid = ?1";
109const SELECT_ROOT_SQL: &str = "SELECT manifest FROM prolly_roots WHERE name = ?1";
110const SELECT_ROOTS_SQL: &str = "SELECT name, manifest FROM prolly_roots ORDER BY name";
111const UPSERT_ROOT_SQL: &str = "\
112INSERT INTO prolly_roots (name, manifest)
113VALUES (?1, ?2)
114ON CONFLICT(name) DO UPDATE SET manifest = excluded.manifest";
115const DELETE_ROOT_SQL: &str = "DELETE FROM prolly_roots WHERE name = ?1";
116
117/// Configuration options for [`SqliteStore`].
118#[derive(Debug, Clone)]
119pub struct SqliteStoreConfig {
120    /// Busy timeout in milliseconds for contended SQLite locks.
121    pub busy_timeout_ms: u64,
122    /// Enable WAL journaling for file-backed databases.
123    pub enable_wal: bool,
124    /// Set SQLite synchronous mode to NORMAL when applying default pragmas.
125    pub synchronous_normal: bool,
126}
127
128impl Default for SqliteStoreConfig {
129    fn default() -> Self {
130        Self {
131            busy_timeout_ms: 5_000,
132            enable_wal: true,
133            synchronous_normal: true,
134        }
135    }
136}
137
138/// Error type for SQLite store operations.
139#[derive(Debug)]
140pub struct SqliteStoreError {
141    message: String,
142    source: Option<rusqlite::Error>,
143}
144
145impl SqliteStoreError {
146    /// Create a new error with a message.
147    pub fn new(message: impl Into<String>) -> Self {
148        Self {
149            message: message.into(),
150            source: None,
151        }
152    }
153
154    /// Create a new error from a rusqlite error.
155    pub fn from_sqlite(err: rusqlite::Error, context: impl Into<String>) -> Self {
156        Self {
157            message: format!("{}: {}", context.into(), err),
158            source: Some(err),
159        }
160    }
161}
162
163impl std::fmt::Display for SqliteStoreError {
164    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
165        write!(f, "SQLite error: {}", self.message)
166    }
167}
168
169impl std::error::Error for SqliteStoreError {
170    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
171        self.source
172            .as_ref()
173            .map(|e| e as &(dyn std::error::Error + 'static))
174    }
175}
176
177impl From<rusqlite::Error> for SqliteStoreError {
178    fn from(err: rusqlite::Error) -> Self {
179        Self {
180            message: err.to_string(),
181            source: Some(err),
182        }
183    }
184}
185
186/// SQLite-backed storage backend for Prolly Trees.
187///
188/// This store persists content-addressed nodes in a single SQLite table and
189/// supports atomic batch operations through transactions.
190pub struct SqliteStore {
191    conn: Mutex<Connection>,
192}
193
194/// Identity obtained from SQLite's actual open main-database file descriptor.
195#[derive(Clone, Copy, Debug, Eq, PartialEq)]
196pub struct SqliteMainFileIdentity {
197    /// Filesystem device containing the open database file.
198    pub device: u64,
199    /// Filesystem inode of the open database file.
200    pub inode: u64,
201    /// Length observed through the open descriptor.
202    pub length: u64,
203}
204
205impl SqliteStore {
206    /// Open or create a SQLite database at the given path with default config.
207    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, SqliteStoreError> {
208        Self::open_with_config(path, SqliteStoreConfig::default())
209    }
210
211    /// Open or create a SQLite database with custom configuration.
212    pub fn open_with_config<P: AsRef<Path>>(
213        path: P,
214        config: SqliteStoreConfig,
215    ) -> Result<Self, SqliteStoreError> {
216        let conn = Connection::open(path.as_ref()).map_err(|e| {
217            SqliteStoreError::from_sqlite(
218                e,
219                format!("Failed to open database at {:?}", path.as_ref()),
220            )
221        })?;
222        Self::from_connection(conn, config)
223    }
224
225    /// Open an existing SQLite database with default runtime configuration.
226    ///
227    /// Unlike [`Self::open`], this never creates the database file and does not
228    /// execute schema DDL. Callers must validate the required schema before
229    /// using this path.
230    pub fn open_existing<P: AsRef<Path>>(path: P) -> Result<Self, SqliteStoreError> {
231        Self::open_existing_verified(path, |_| Ok(()))
232    }
233
234    /// Open an existing database and verify SQLite's actual main-file handle
235    /// before executing any pragma, schema statement, or other SQL.
236    #[cfg(unix)]
237    pub fn open_existing_verified<P, F>(path: P, verifier: F) -> Result<Self, SqliteStoreError>
238    where
239        P: AsRef<Path>,
240        F: FnOnce(SqliteMainFileIdentity) -> Result<(), SqliteStoreError>,
241    {
242        let conn = Connection::open_with_flags(
243            path.as_ref(),
244            OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_NO_MUTEX,
245        )
246        .map_err(|error| {
247            SqliteStoreError::from_sqlite(
248                error,
249                format!("Failed to open existing database at {:?}", path.as_ref()),
250            )
251        })?;
252        verifier(sqlite_main_file_identity(&conn)?)?;
253        Self::from_existing_connection(conn, SqliteStoreConfig::default())
254    }
255
256    /// Non-Unix platforms cannot currently prove the SQLite VFS handle's
257    /// identity before SQL runs, so verified opens fail closed there.
258    #[cfg(not(unix))]
259    pub fn open_existing_verified<P, F>(_path: P, _verifier: F) -> Result<Self, SqliteStoreError>
260    where
261        P: AsRef<Path>,
262        F: FnOnce(SqliteMainFileIdentity) -> Result<(), SqliteStoreError>,
263    {
264        Err(SqliteStoreError::new(
265            "verified existing SQLite opens are unsupported on this platform",
266        ))
267    }
268
269    /// Create an in-memory SQLite store.
270    pub fn open_in_memory() -> Result<Self, SqliteStoreError> {
271        let conn = Connection::open_in_memory()
272            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to open in-memory database"))?;
273        Self::from_connection(conn, SqliteStoreConfig::default())
274    }
275
276    fn from_connection(
277        conn: Connection,
278        config: SqliteStoreConfig,
279    ) -> Result<Self, SqliteStoreError> {
280        Self::apply_runtime_config(&conn, &config)?;
281        conn.execute_batch(CREATE_TABLE_SQL)
282            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to initialize schema"))?;
283        conn.execute_batch(CREATE_HINTS_TABLE_SQL)
284            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to initialize hint schema"))?;
285        conn.execute_batch(CREATE_ROOTS_TABLE_SQL)
286            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to initialize root schema"))?;
287
288        Ok(Self {
289            conn: Mutex::new(conn),
290        })
291    }
292
293    fn from_existing_connection(
294        conn: Connection,
295        config: SqliteStoreConfig,
296    ) -> Result<Self, SqliteStoreError> {
297        Self::apply_runtime_config(&conn, &config)?;
298        Ok(Self {
299            conn: Mutex::new(conn),
300        })
301    }
302
303    fn apply_runtime_config(
304        conn: &Connection,
305        config: &SqliteStoreConfig,
306    ) -> Result<(), SqliteStoreError> {
307        conn.busy_timeout(Duration::from_millis(config.busy_timeout_ms))
308            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set busy timeout"))?;
309
310        if config.enable_wal {
311            conn.pragma_update(None, "journal_mode", "WAL")
312                .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to enable WAL mode"))?;
313        }
314        if config.synchronous_normal {
315            conn.pragma_update(None, "synchronous", "NORMAL")
316                .map_err(|e| {
317                    SqliteStoreError::from_sqlite(e, "Failed to set synchronous=NORMAL")
318                })?;
319        }
320        conn.pragma_update(None, "temp_store", "MEMORY")
321            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set temp_store=MEMORY"))?;
322        Ok(())
323    }
324
325    fn connection(&self) -> Result<MutexGuard<'_, Connection>, SqliteStoreError> {
326        self.conn
327            .lock()
328            .map_err(|e| SqliteStoreError::new(format!("lock poisoned: {}", e)))
329    }
330}
331
332#[cfg(unix)]
333/// Inspect SQLite's actual open main-database descriptor without executing SQL.
334pub fn sqlite_main_file_identity(
335    conn: &Connection,
336) -> Result<SqliteMainFileIdentity, SqliteStoreError> {
337    use std::ffi::{c_int, c_void};
338    use std::fs::File;
339    use std::os::fd::BorrowedFd;
340    use std::os::unix::fs::MetadataExt;
341
342    // Every bundled Unix SQLite VFS begins its concrete `unixFile` with this
343    // stable prefix. `SQLITE_FCNTL_FILE_POINTER` returns the actual main-file
344    // object owned by this connection, not a pathname-derived approximation.
345    #[repr(C)]
346    struct UnixFilePrefix {
347        methods: *const rusqlite::ffi::sqlite3_io_methods,
348        vfs: *mut rusqlite::ffi::sqlite3_vfs,
349        inode: *mut c_void,
350        fd: c_int,
351    }
352
353    let mut sqlite_file: *mut rusqlite::ffi::sqlite3_file = std::ptr::null_mut();
354    // SAFETY: `conn` remains alive for the call, `main` is NUL terminated, and
355    // SQLite writes one sqlite3_file pointer into `sqlite_file` for this opcode.
356    let rc = unsafe {
357        rusqlite::ffi::sqlite3_file_control(
358            conn.handle(),
359            c"main".as_ptr(),
360            rusqlite::ffi::SQLITE_FCNTL_FILE_POINTER,
361            (&mut sqlite_file as *mut *mut rusqlite::ffi::sqlite3_file).cast(),
362        )
363    };
364    if rc != rusqlite::ffi::SQLITE_OK || sqlite_file.is_null() {
365        return Err(SqliteStoreError::new(format!(
366            "SQLite did not expose its main-file handle (code {rc})"
367        )));
368    }
369    // SAFETY: the bundled Unix VFS concrete file begins with UnixFilePrefix,
370    // and SQLite retains the descriptor for the lifetime of this connection.
371    let fd = unsafe { (*(sqlite_file.cast::<UnixFilePrefix>())).fd };
372    // Duplicate the descriptor so the temporary File cannot close SQLite's
373    // owned descriptor when it is dropped.
374    let borrowed = unsafe { BorrowedFd::borrow_raw(fd) };
375    let owned = borrowed.try_clone_to_owned().map_err(|error| {
376        SqliteStoreError::new(format!(
377            "failed to duplicate SQLite main-file handle: {error}"
378        ))
379    })?;
380    let metadata = File::from(owned).metadata().map_err(|error| {
381        SqliteStoreError::new(format!("failed to stat SQLite main-file handle: {error}"))
382    })?;
383    Ok(SqliteMainFileIdentity {
384        device: metadata.dev(),
385        inode: metadata.ino(),
386        length: metadata.len(),
387    })
388}
389
390impl Store for SqliteStore {
391    type Error = SqliteStoreError;
392
393    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
394        let conn = self.connection()?;
395        conn.query_row(SELECT_SQL, params![key], |row| row.get(0))
396            .optional()
397            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read key"))
398    }
399
400    fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
401        let conn = self.connection()?;
402        conn.execute(UPSERT_SQL, params![key, value])
403            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to write key"))?;
404        Ok(())
405    }
406
407    fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
408        let conn = self.connection()?;
409        conn.execute(DELETE_SQL, params![key])
410            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to delete key"))?;
411        Ok(())
412    }
413
414    fn batch(&self, ops: &[BatchOp]) -> Result<(), Self::Error> {
415        let mut conn = self.connection()?;
416        let tx = conn
417            .transaction()
418            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to start transaction"))?;
419
420        {
421            let mut upsert = tx
422                .prepare_cached(UPSERT_SQL)
423                .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare batch write"))?;
424            let mut delete = tx
425                .prepare_cached(DELETE_SQL)
426                .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare batch delete"))?;
427
428            for op in ops {
429                match op {
430                    BatchOp::Upsert { key, value } => {
431                        upsert.execute(params![key, value]).map_err(|e| {
432                            SqliteStoreError::from_sqlite(e, "Failed to write key in batch")
433                        })?;
434                    }
435                    BatchOp::Delete { key } => {
436                        delete.execute(params![key]).map_err(|e| {
437                            SqliteStoreError::from_sqlite(e, "Failed to delete key in batch")
438                        })?;
439                    }
440                }
441            }
442        }
443
444        tx.commit()
445            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to commit transaction"))
446    }
447
448    fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
449        let conn = self.connection()?;
450        let mut stmt = conn
451            .prepare_cached(SELECT_SQL)
452            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare batch read"))?;
453        let plan = OrderedBatchReadPlan::new(keys);
454        let mut results = HashMap::with_capacity(plan.unique_keys().len());
455
456        for key in plan.unique_keys() {
457            if let Some(value) = stmt
458                .query_row(params![key], |row| row.get(0))
459                .optional()
460                .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read key in batch"))?
461            {
462                results.insert(key.to_vec(), value);
463            }
464        }
465
466        Ok(results)
467    }
468
469    fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
470        let conn = self.connection()?;
471        let mut stmt = conn.prepare_cached(SELECT_SQL).map_err(|e| {
472            SqliteStoreError::from_sqlite(e, "Failed to prepare ordered batch read")
473        })?;
474        let plan = OrderedBatchReadPlan::new(keys);
475        let mut unique_values = Vec::with_capacity(plan.unique_keys().len());
476
477        for key in plan.unique_keys() {
478            let value = stmt
479                .query_row(params![key], |row| row.get(0))
480                .optional()
481                .map_err(|e| {
482                    SqliteStoreError::from_sqlite(e, "Failed to read key in ordered batch")
483                })?;
484            unique_values.push(value);
485        }
486
487        Ok(plan.expand_owned(unique_values))
488    }
489
490    fn batch_get_ordered_unique(
491        &self,
492        keys: &[&[u8]],
493    ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
494        let conn = self.connection()?;
495        let mut stmt = conn.prepare_cached(SELECT_SQL).map_err(|e| {
496            SqliteStoreError::from_sqlite(e, "Failed to prepare unique ordered batch read")
497        })?;
498        let mut values = Vec::with_capacity(keys.len());
499
500        for key in keys {
501            let value = stmt
502                .query_row(params![key], |row| row.get(0))
503                .optional()
504                .map_err(|e| {
505                    SqliteStoreError::from_sqlite(e, "Failed to read key in unique ordered batch")
506                })?;
507            values.push(value);
508        }
509
510        Ok(values)
511    }
512
513    fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
514        let mut conn = self.connection()?;
515        let tx = conn
516            .transaction()
517            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to start transaction"))?;
518
519        {
520            let mut stmt = tx.prepare_cached(UPSERT_SQL).map_err(|e| {
521                SqliteStoreError::from_sqlite(e, "Failed to prepare batch_put write")
522            })?;
523            for (key, value) in entries {
524                stmt.execute(params![key, value]).map_err(|e| {
525                    SqliteStoreError::from_sqlite(e, "Failed to write key in batch_put")
526                })?;
527            }
528        }
529
530        tx.commit()
531            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to commit transaction"))
532    }
533
534    fn supports_hints(&self) -> bool {
535        true
536    }
537
538    fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
539        let conn = self.connection()?;
540        conn.query_row(
541            "SELECT value FROM prolly_hints WHERE namespace = ?1 AND key = ?2",
542            params![namespace, key],
543            |row| row.get(0),
544        )
545        .optional()
546        .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read hint"))
547    }
548
549    fn put_hint(&self, namespace: &[u8], key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
550        let conn = self.connection()?;
551        conn.execute(
552            "\
553            INSERT INTO prolly_hints (namespace, key, value) \
554            VALUES (?1, ?2, ?3) \
555            ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
556            params![namespace, key, value],
557        )
558        .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to write hint"))?;
559        Ok(())
560    }
561
562    fn batch_put_with_hint(
563        &self,
564        entries: &[(&[u8], &[u8])],
565        namespace: &[u8],
566        key: &[u8],
567        value: &[u8],
568    ) -> Result<(), Self::Error> {
569        let mut conn = self.connection()?;
570        let tx = conn
571            .transaction()
572            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to start transaction"))?;
573
574        {
575            let mut upsert_node = tx.prepare_cached(UPSERT_SQL).map_err(|e| {
576                SqliteStoreError::from_sqlite(e, "Failed to prepare batch_put write")
577            })?;
578            for (key, value) in entries {
579                upsert_node.execute(params![key, value]).map_err(|e| {
580                    SqliteStoreError::from_sqlite(e, "Failed to write key in batch_put")
581                })?;
582            }
583        }
584
585        tx.execute(
586            "\
587            INSERT INTO prolly_hints (namespace, key, value) \
588            VALUES (?1, ?2, ?3) \
589            ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
590            params![namespace, key, value],
591        )
592        .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to write hint in batch_put"))?;
593
594        tx.commit()
595            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to commit transaction"))
596    }
597}
598
599impl NodeStoreScan for SqliteStore {
600    type Error = SqliteStoreError;
601
602    fn list_node_cids(&self) -> Result<Vec<Cid>, Self::Error> {
603        let conn = self.connection()?;
604        let mut stmt = conn
605            .prepare_cached(SELECT_NODE_CIDS_SQL)
606            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare node CID listing"))?;
607        let rows = stmt
608            .query_map([], |row| row.get::<_, Vec<u8>>(0))
609            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to list node CIDs"))?;
610
611        let mut cids = Vec::new();
612        for row in rows {
613            let key = row
614                .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read listed node CID"))?;
615            cids.push(cid_from_store_key(&key, "SQLite node").map_err(SqliteStoreError::new)?);
616        }
617        sort_cids(&mut cids);
618        Ok(cids)
619    }
620}
621
622impl ManifestStore for SqliteStore {
623    type Error = SqliteStoreError;
624
625    fn get_root(&self, name: &[u8]) -> Result<Option<RootManifest>, Self::Error> {
626        let conn = self.connection()?;
627        let bytes = conn
628            .query_row(SELECT_ROOT_SQL, params![name], |row| row.get(0))
629            .optional()
630            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read root manifest"))?;
631        decode_root_manifest(bytes)
632    }
633
634    fn put_root(&self, name: &[u8], manifest: &RootManifest) -> Result<(), Self::Error> {
635        let conn = self.connection()?;
636        let bytes = encode_root_manifest(manifest)?;
637        conn.execute(UPSERT_ROOT_SQL, params![name, bytes])
638            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to write root manifest"))?;
639        Ok(())
640    }
641
642    fn delete_root(&self, name: &[u8]) -> Result<(), Self::Error> {
643        let conn = self.connection()?;
644        conn.execute(DELETE_ROOT_SQL, params![name])
645            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to delete root manifest"))?;
646        Ok(())
647    }
648
649    fn compare_and_swap_root(
650        &self,
651        name: &[u8],
652        expected: Option<&RootManifest>,
653        new: Option<&RootManifest>,
654    ) -> Result<ManifestUpdate, Self::Error> {
655        let expected_bytes = expected.map(encode_root_manifest).transpose()?;
656        let new_bytes = new.map(encode_root_manifest).transpose()?;
657
658        let mut conn = self.connection()?;
659        let tx = conn
660            .transaction()
661            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to start root transaction"))?;
662
663        let current_bytes = tx
664            .query_row(SELECT_ROOT_SQL, params![name], |row| row.get(0))
665            .optional()
666            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read root manifest"))?;
667
668        if current_bytes.as_deref() != expected_bytes.as_deref() {
669            return Ok(ManifestUpdate::Conflict {
670                current: decode_root_manifest(current_bytes)?,
671            });
672        }
673
674        match new_bytes {
675            Some(bytes) => {
676                tx.execute(UPSERT_ROOT_SQL, params![name, bytes])
677                    .map_err(|e| {
678                        SqliteStoreError::from_sqlite(e, "Failed to write root manifest")
679                    })?;
680            }
681            None => {
682                tx.execute(DELETE_ROOT_SQL, params![name]).map_err(|e| {
683                    SqliteStoreError::from_sqlite(e, "Failed to delete root manifest")
684                })?;
685            }
686        }
687
688        tx.commit()
689            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to commit root transaction"))?;
690        Ok(ManifestUpdate::Applied)
691    }
692}
693
694impl ManifestStoreScan for SqliteStore {
695    fn list_roots(&self) -> Result<Vec<NamedRootManifest>, Self::Error> {
696        let conn = self.connection()?;
697        let mut stmt = conn.prepare_cached(SELECT_ROOTS_SQL).map_err(|e| {
698            SqliteStoreError::from_sqlite(e, "Failed to prepare root manifest listing")
699        })?;
700        let rows = stmt
701            .query_map([], |row| {
702                Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, Vec<u8>>(1)?))
703            })
704            .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to list root manifests"))?;
705
706        let mut roots = Vec::new();
707        for row in rows {
708            let (name, bytes) = row.map_err(|e| {
709                SqliteStoreError::from_sqlite(e, "Failed to read listed root manifest")
710            })?;
711            let manifest = RootManifest::from_bytes(&bytes)
712                .map_err(|err| SqliteStoreError::new(err.to_string()))?;
713            roots.push(NamedRootManifest::new(name, manifest));
714        }
715        sort_named_root_manifests(&mut roots);
716        Ok(roots)
717    }
718}
719
720impl TransactionalStore for SqliteStore {
721    fn supports_transactions(&self) -> bool {
722        true
723    }
724
725    fn commit_transaction(
726        &self,
727        node_writes: &[TransactionNodeWrite],
728        root_conditions: &[RootCondition],
729        root_writes: &[RootWrite],
730    ) -> Result<TransactionUpdate, Error> {
731        let mut conn = self
732            .connection()
733            .map_err(|err| Error::Store(Box::new(err)))?;
734        let tx = conn.transaction().map_err(|err| {
735            Error::Store(Box::new(SqliteStoreError::from_sqlite(
736                err,
737                "Failed to start transaction commit",
738            )))
739        })?;
740
741        for condition in root_conditions {
742            let current_bytes = tx
743                .query_row(SELECT_ROOT_SQL, params![condition.name], |row| row.get(0))
744                .optional()
745                .map_err(|err| {
746                    Error::Store(Box::new(SqliteStoreError::from_sqlite(
747                        err,
748                        "Failed to read root manifest during transaction commit",
749                    )))
750                })?;
751            let current =
752                decode_root_manifest(current_bytes).map_err(|err| Error::Store(Box::new(err)))?;
753            if current != condition.expected {
754                return Ok(TransactionUpdate::Conflict(Box::new(
755                    TransactionConflict::new(
756                        condition.name.clone(),
757                        condition.expected.clone(),
758                        current,
759                    ),
760                )));
761            }
762        }
763
764        {
765            let mut upsert_node = tx.prepare_cached(UPSERT_SQL).map_err(|err| {
766                Error::Store(Box::new(SqliteStoreError::from_sqlite(
767                    err,
768                    "Failed to prepare transaction node write",
769                )))
770            })?;
771            let mut delete_node = tx.prepare_cached(DELETE_SQL).map_err(|err| {
772                Error::Store(Box::new(SqliteStoreError::from_sqlite(
773                    err,
774                    "Failed to prepare transaction node delete",
775                )))
776            })?;
777            for write in node_writes {
778                match write {
779                    TransactionNodeWrite::Upsert { key, value } => {
780                        upsert_node.execute(params![key, value]).map_err(|err| {
781                            Error::Store(Box::new(SqliteStoreError::from_sqlite(
782                                err,
783                                "Failed to write node during transaction commit",
784                            )))
785                        })?;
786                    }
787                    TransactionNodeWrite::Delete { key } => {
788                        delete_node.execute(params![key]).map_err(|err| {
789                            Error::Store(Box::new(SqliteStoreError::from_sqlite(
790                                err,
791                                "Failed to delete node during transaction commit",
792                            )))
793                        })?;
794                    }
795                }
796            }
797        }
798
799        {
800            let mut upsert_root = tx.prepare_cached(UPSERT_ROOT_SQL).map_err(|err| {
801                Error::Store(Box::new(SqliteStoreError::from_sqlite(
802                    err,
803                    "Failed to prepare transaction root write",
804                )))
805            })?;
806            let mut delete_root = tx.prepare_cached(DELETE_ROOT_SQL).map_err(|err| {
807                Error::Store(Box::new(SqliteStoreError::from_sqlite(
808                    err,
809                    "Failed to prepare transaction root delete",
810                )))
811            })?;
812            for write in root_writes {
813                match write {
814                    RootWrite::Put { name, manifest } => {
815                        let bytes = encode_root_manifest(manifest)
816                            .map_err(|err| Error::Store(Box::new(err)))?;
817                        upsert_root.execute(params![name, bytes]).map_err(|err| {
818                            Error::Store(Box::new(SqliteStoreError::from_sqlite(
819                                err,
820                                "Failed to write root during transaction commit",
821                            )))
822                        })?;
823                    }
824                    RootWrite::Delete { name } => {
825                        delete_root.execute(params![name]).map_err(|err| {
826                            Error::Store(Box::new(SqliteStoreError::from_sqlite(
827                                err,
828                                "Failed to delete root during transaction commit",
829                            )))
830                        })?;
831                    }
832                }
833            }
834        }
835
836        tx.commit().map_err(|err| {
837            Error::Store(Box::new(SqliteStoreError::from_sqlite(
838                err,
839                "Failed to commit transaction",
840            )))
841        })?;
842        Ok(TransactionUpdate::Applied {
843            nodes_written: node_writes.len(),
844            roots_written: root_writes.len(),
845        })
846    }
847}
848
849fn encode_root_manifest(manifest: &RootManifest) -> Result<Vec<u8>, SqliteStoreError> {
850    manifest
851        .to_bytes()
852        .map_err(|e| SqliteStoreError::new(format!("failed to encode root manifest: {e}")))
853}
854
855fn decode_root_manifest(bytes: Option<Vec<u8>>) -> Result<Option<RootManifest>, SqliteStoreError> {
856    bytes
857        .as_deref()
858        .map(RootManifest::from_bytes)
859        .transpose()
860        .map_err(|e| SqliteStoreError::new(format!("failed to decode root manifest: {e}")))
861}
862
863#[cfg(test)]
864mod tests {
865    use super::*;
866
867    #[test]
868    fn sqlite_store_put_get_delete() {
869        let store = SqliteStore::open_in_memory().unwrap();
870
871        store.put(b"key", b"value").unwrap();
872        assert_eq!(store.get(b"key").unwrap(), Some(b"value".to_vec()));
873
874        store.delete(b"key").unwrap();
875        assert_eq!(store.get(b"key").unwrap(), None);
876    }
877
878    #[test]
879    fn sqlite_store_batch_is_order_preserving_for_reads() {
880        let store = SqliteStore::open_in_memory().unwrap();
881        let ops = vec![
882            BatchOp::Upsert {
883                key: b"a",
884                value: b"1",
885            },
886            BatchOp::Upsert {
887                key: b"b",
888                value: b"2",
889            },
890            BatchOp::Upsert {
891                key: b"c",
892                value: b"3",
893            },
894        ];
895
896        store.batch(&ops).unwrap();
897
898        let keys: Vec<&[u8]> = vec![b"c", b"missing", b"a", b"c", b"missing", b"b"];
899        assert_eq!(
900            store.batch_get_ordered(&keys).unwrap(),
901            vec![
902                Some(b"3".to_vec()),
903                None,
904                Some(b"1".to_vec()),
905                Some(b"3".to_vec()),
906                None,
907                Some(b"2".to_vec())
908            ]
909        );
910    }
911
912    #[test]
913    fn sqlite_store_batch_put_updates_existing_keys() {
914        let store = SqliteStore::open_in_memory().unwrap();
915
916        store.put(b"a", b"old").unwrap();
917        store
918            .batch_put(&[(b"a".as_slice(), b"new".as_slice()), (b"b", b"2")])
919            .unwrap();
920
921        assert_eq!(store.get(b"a").unwrap(), Some(b"new".to_vec()));
922        assert_eq!(store.get(b"b").unwrap(), Some(b"2".to_vec()));
923    }
924
925    #[test]
926    fn sqlite_store_persists_hints_separately_from_nodes() {
927        let store = SqliteStore::open_in_memory().unwrap();
928
929        store.put_hint(b"rightmost", b"root", b"hint-v1").unwrap();
930        assert_eq!(
931            store.get_hint(b"rightmost", b"root").unwrap(),
932            Some(b"hint-v1".to_vec())
933        );
934        assert_eq!(store.get_hint(b"rightmost", b"missing").unwrap(), None);
935        assert_eq!(store.get(b"root").unwrap(), None);
936
937        store.put_hint(b"rightmost", b"root", b"hint-v2").unwrap();
938        assert_eq!(
939            store.get_hint(b"rightmost", b"root").unwrap(),
940            Some(b"hint-v2".to_vec())
941        );
942    }
943}