uqa_storage_sqlite/
backend.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Session factories and physical stores for the `SQLite` backend.
8
9use crate::{
10    Catalog, ManagedConnection, SQLiteBTreeIndexStore, SQLiteDocumentStore, SQLiteHNSWIndex,
11    SQLiteIVFIndex, SQLiteInvertedIndex, SQLiteVectorIndex,
12};
13use std::{collections::BTreeMap, sync::Arc};
14use uqa_analysis::Analyzer;
15use uqa_core::{DocId, Value};
16use uqa_storage::{
17    CatalogFacade, DocumentStore, InvertedIndex, PersistentStorageBackend,
18    PersistentStorageIdentity, PersistentStorageProvider, PersistentStorageSession,
19    StorageBackendError, StorageBackendResult, StorageSavepointId, VectorIndex,
20    VectorIndexOpenMode, VectorIndexSpec,
21};
22
23#[derive(Clone)]
24pub struct SQLiteStorageBackend {
25    conn: ManagedConnection,
26}
27
28impl SQLiteStorageBackend {
29    pub fn new(conn: ManagedConnection) -> Self {
30        Self { conn }
31    }
32
33    pub fn connection(&self) -> ManagedConnection {
34        self.conn.clone()
35    }
36
37    /// Create a backend whose stores use an independent transaction session
38    /// over the same physical `SQLite` pool.
39    #[must_use]
40    pub fn new_session(&self) -> Self {
41        Self::new(self.conn.new_session())
42    }
43}
44
45/// Database-level owner that creates isolated `SQLite` engine sessions.
46#[derive(Clone)]
47pub struct SQLiteStorageProvider {
48    connection: ManagedConnection,
49}
50
51impl SQLiteStorageProvider {
52    pub fn new(connection: ManagedConnection) -> Self {
53        Self { connection }
54    }
55}
56
57impl PersistentStorageProvider for SQLiteStorageProvider {
58    fn auxiliary_encryption_key(&self) -> Option<uqa_storage::StorageEncryptionKey> {
59        self.connection.auxiliary_encryption_key()
60    }
61
62    fn open_initial_session(&self) -> StorageBackendResult<PersistentStorageSession> {
63        let connection = self.connection.new_session();
64        let catalog: Arc<dyn CatalogFacade> =
65            Arc::new(Catalog::for_initial_restore(connection.clone()));
66        let backend: Arc<dyn PersistentStorageBackend> =
67            Arc::new(SQLiteStorageBackend::new(connection));
68        Ok(PersistentStorageSession::new(catalog, backend))
69    }
70
71    fn open_session(&self) -> StorageBackendResult<PersistentStorageSession> {
72        let connection = self.connection.new_session();
73        let catalog: Arc<dyn CatalogFacade> = Arc::new(Catalog::open(connection.clone())?);
74        let backend: Arc<dyn PersistentStorageBackend> =
75            Arc::new(SQLiteStorageBackend::new(connection));
76        Ok(PersistentStorageSession::new(catalog, backend))
77    }
78
79    fn storage_identity(&self) -> StorageBackendResult<Option<PersistentStorageIdentity>> {
80        let Some(path) = self.connection.database_path() else {
81            return Ok(None);
82        };
83        PersistentStorageIdentity::for_database_path(path)
84            .map(Some)
85            .map_err(|error| {
86                StorageBackendError::Other(format!(
87                    "resolve SQLite database identity `{}`: {error}",
88                    path.display()
89                ))
90            })
91    }
92}
93
94impl PersistentStorageBackend for SQLiteStorageBackend {
95    fn auxiliary_encryption_key(&self) -> Option<uqa_storage::StorageEncryptionKey> {
96        self.conn.auxiliary_encryption_key()
97    }
98
99    fn storage_identity(&self) -> StorageBackendResult<Option<PersistentStorageIdentity>> {
100        let Some(path) = self.conn.database_path() else {
101            return Ok(None);
102        };
103        PersistentStorageIdentity::for_database_path(path).map(Some)
104    }
105
106    fn open_session(&self) -> StorageBackendResult<PersistentStorageSession> {
107        let connection = self.conn.new_session();
108        let catalog: Arc<dyn CatalogFacade> = Arc::new(Catalog::open(connection.clone())?);
109        let backend: Arc<dyn PersistentStorageBackend> = Arc::new(Self::new(connection));
110        Ok(PersistentStorageSession::new(catalog, backend))
111    }
112
113    fn supports_concurrent_pinned_read_and_write(&self) -> bool {
114        self.conn.supports_concurrent_pinned_read_and_write()
115    }
116
117    fn document_store(&self, table: &str) -> Box<dyn DocumentStore> {
118        Box::new(SQLiteDocumentStore::new(self.conn.clone(), table))
119    }
120
121    fn inverted_index(&self, table: &str, analyzer: Analyzer) -> Box<dyn InvertedIndex> {
122        Box::new(SQLiteInvertedIndex::new(self.conn.clone(), table, analyzer))
123    }
124
125    fn vector_index(
126        &self,
127        table: &str,
128        field: &str,
129        dimensions: u32,
130        spec: VectorIndexSpec,
131        mode: VectorIndexOpenMode,
132    ) -> StorageBackendResult<Box<dyn VectorIndex>> {
133        let index: Box<dyn VectorIndex> = match spec {
134            VectorIndexSpec::BruteForce => Box::new(SQLiteVectorIndex::new(
135                self.conn.clone(),
136                table,
137                field,
138                dimensions,
139            )),
140            VectorIndexSpec::IVF(params) => {
141                params.validate()?;
142                match mode {
143                    VectorIndexOpenMode::Create => Box::new(SQLiteIVFIndex::with_params(
144                        self.conn.clone(),
145                        table,
146                        field,
147                        dimensions,
148                        params.nlist,
149                        params.nprobe,
150                        params.train_threshold,
151                    )),
152                    VectorIndexOpenMode::Restore => Box::new(SQLiteIVFIndex::open_existing(
153                        self.conn.clone(),
154                        table,
155                        field,
156                        dimensions,
157                        params.nlist,
158                        params.nprobe,
159                        params.train_threshold,
160                    )),
161                }
162            }
163            VectorIndexSpec::HNSW(params) => {
164                params.validate()?;
165                match mode {
166                    VectorIndexOpenMode::Create => Box::new(SQLiteHNSWIndex::with_params(
167                        self.conn.clone(),
168                        table,
169                        field,
170                        dimensions,
171                        params,
172                    )),
173                    VectorIndexOpenMode::Restore => {
174                        let index = SQLiteHNSWIndex::open_existing(
175                            self.conn.clone(),
176                            table,
177                            field,
178                            dimensions,
179                            params,
180                        );
181                        index.validate_existing()?;
182                        Box::new(index)
183                    }
184                }
185            }
186        };
187        Ok(index)
188    }
189
190    fn drop_vector_index_metadata(&self, table: &str, field: &str) -> StorageBackendResult<()> {
191        SQLiteIVFIndex::drop_metadata(&self.conn, table, field)?;
192        SQLiteHNSWIndex::drop_metadata(&self.conn, table, field)?;
193        Ok(())
194    }
195
196    fn persists_btree_indexes(&self) -> bool {
197        true
198    }
199
200    fn load_btree_index(
201        &self,
202        table: &str,
203        field: &uqa_storage::ValueIndexKey,
204    ) -> StorageBackendResult<Option<Vec<(DocId, Value)>>> {
205        Ok(SQLiteBTreeIndexStore::new(self.conn.clone()).load(table, field)?)
206    }
207
208    fn btree_index_fields(
209        &self,
210        table: &str,
211    ) -> StorageBackendResult<Vec<uqa_storage::ValueIndexKey>> {
212        Ok(SQLiteBTreeIndexStore::new(self.conn.clone()).fields(table)?)
213    }
214
215    fn btree_index_repairs(
216        &self,
217    ) -> StorageBackendResult<Vec<(String, uqa_storage::ValueIndexKey)>> {
218        Ok(SQLiteBTreeIndexStore::new(self.conn.clone()).repairs()?)
219    }
220
221    fn clear_btree_index_repair(
222        &self,
223        table: &str,
224        field: &uqa_storage::ValueIndexKey,
225    ) -> StorageBackendResult<()> {
226        SQLiteBTreeIndexStore::new(self.conn.clone()).clear_repair(table, field)?;
227        Ok(())
228    }
229
230    fn replace_btree_index(
231        &self,
232        table: &str,
233        field: &uqa_storage::ValueIndexKey,
234        values: &[(DocId, Value)],
235    ) -> StorageBackendResult<()> {
236        SQLiteBTreeIndexStore::new(self.conn.clone()).replace(table, field, values)?;
237        Ok(())
238    }
239
240    fn repair_btree_index(
241        &self,
242        table: &str,
243        field: &uqa_storage::ValueIndexKey,
244        _complete: &[(DocId, Value)],
245        stale_doc_ids: &[DocId],
246        missing: &[(DocId, Value)],
247    ) -> StorageBackendResult<()> {
248        SQLiteBTreeIndexStore::new(self.conn.clone()).repair(
249            table,
250            field,
251            stale_doc_ids,
252            missing,
253        )?;
254        Ok(())
255    }
256
257    fn replace_btree_indexes(
258        &self,
259        table: &str,
260        indexes: &[(&uqa_storage::ValueIndexKey, &[(DocId, Value)])],
261    ) -> StorageBackendResult<()> {
262        SQLiteBTreeIndexStore::new(self.conn.clone()).replace_many(table, indexes)?;
263        Ok(())
264    }
265
266    fn apply_btree_index_write(
267        &self,
268        table: &str,
269        doc_id: DocId,
270        values: Option<&BTreeMap<uqa_storage::ValueIndexKey, Value>>,
271    ) -> StorageBackendResult<()> {
272        SQLiteBTreeIndexStore::new(self.conn.clone()).apply_write(table, doc_id, values)?;
273        Ok(())
274    }
275
276    fn drop_btree_index(
277        &self,
278        table: &str,
279        field: &uqa_storage::ValueIndexKey,
280    ) -> StorageBackendResult<()> {
281        SQLiteBTreeIndexStore::new(self.conn.clone()).drop_index(table, field)?;
282        Ok(())
283    }
284
285    fn clear_btree_indexes(&self, table: &str) -> StorageBackendResult<()> {
286        SQLiteBTreeIndexStore::new(self.conn.clone()).clear_table(table)?;
287        Ok(())
288    }
289
290    fn vacuum(&self) -> StorageBackendResult<()> {
291        self.conn.vacuum()?;
292        Ok(())
293    }
294
295    fn begin_transaction(&self) -> StorageBackendResult<()> {
296        self.conn.begin_transaction()?;
297        Ok(())
298    }
299
300    fn begin_read_transaction(&self) -> StorageBackendResult<()> {
301        self.conn.begin_deferred_transaction()?;
302        Ok(())
303    }
304
305    fn begin_upgradeable_transaction(&self) -> StorageBackendResult<()> {
306        self.conn.begin_deferred_transaction()?;
307        Ok(())
308    }
309
310    fn in_transaction(&self) -> bool {
311        self.conn.in_transaction()
312    }
313
314    fn transaction_has_written(&self) -> StorageBackendResult<bool> {
315        Ok(self.conn.transaction_has_written()?)
316    }
317
318    fn change_version(&self) -> StorageBackendResult<Option<u64>> {
319        Ok(self.conn.data_version()?)
320    }
321
322    fn change_version_monitor_is_nonblocking(&self) -> StorageBackendResult<bool> {
323        Ok(self.conn.data_version_monitor_is_nonblocking()?)
324    }
325
326    fn pin_transaction_snapshot(&self) -> StorageBackendResult<()> {
327        self.conn.pin_transaction_snapshot()?;
328        Ok(())
329    }
330
331    fn commit_transaction(&self) -> StorageBackendResult<()> {
332        self.conn.commit_transaction()?;
333        Ok(())
334    }
335
336    fn rollback_transaction(&self) -> StorageBackendResult<()> {
337        self.conn.rollback_transaction()?;
338        Ok(())
339    }
340
341    fn savepoint(&self, id: StorageSavepointId) -> StorageBackendResult<()> {
342        self.conn.savepoint(&id.backend_name())?;
343        Ok(())
344    }
345
346    fn release_savepoint(&self, id: StorageSavepointId) -> StorageBackendResult<()> {
347        self.conn.release_savepoint(&id.backend_name())?;
348        Ok(())
349    }
350
351    fn rollback_to_savepoint(&self, id: StorageSavepointId) -> StorageBackendResult<()> {
352        self.conn.rollback_to_savepoint(&id.backend_name())?;
353        Ok(())
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use std::collections::BTreeMap;
360
361    use uqa_analysis::analyzer::standard_analyzer;
362    use uqa_core::Value;
363
364    use super::*;
365    use crate::{Catalog, SQLiteError};
366
367    #[test]
368    fn session_factory_reads_current_catalog_while_a_sibling_holds_a_writer_reservation() {
369        let directory = tempfile::tempdir().unwrap();
370        let connection = ManagedConnection::open_compressed(
371            &directory.path().join("session-writer-reservation.db"),
372            crate::SQLiteCompressionOptions::default(),
373        )
374        .unwrap();
375        let provider = SQLiteStorageProvider::new(connection);
376        let writer = provider.open_session().unwrap();
377        writer.backend.begin_transaction().unwrap();
378        writer
379            .catalog
380            .set_metadata("private-write", "uncommitted")
381            .unwrap();
382
383        let reader = provider.open_session().unwrap();
384        assert_eq!(reader.catalog.get_metadata("private-write").unwrap(), None);
385        writer.backend.commit_transaction().unwrap();
386        assert_eq!(
387            reader
388                .catalog
389                .get_metadata("private-write")
390                .unwrap()
391                .as_deref(),
392            Some("uncommitted")
393        );
394    }
395
396    #[test]
397    fn sqlite_backend_builds_document_index_and_vector_stores() {
398        let conn = ManagedConnection::open_in_memory().unwrap();
399        let _catalog = Catalog::open(conn.clone()).unwrap();
400        let backend = SQLiteStorageBackend::new(conn);
401
402        let mut doc = BTreeMap::new();
403        doc.insert("title".to_string(), Value::Str("rust storage".into()));
404        let mut docs = backend.document_store("articles");
405        docs.put(1, doc).unwrap();
406        assert_eq!(
407            docs.get_field(1, "title").unwrap(),
408            Some(Value::Str("rust storage".into()))
409        );
410
411        let mut inv = backend.inverted_index("articles", standard_analyzer("english"));
412        inv.add_document(
413            1,
414            BTreeMap::from([("title".to_string(), "rust storage".to_string())]),
415        )
416        .unwrap();
417        assert_eq!(inv.doc_freq("title", "rust").unwrap(), 1);
418
419        let mut vectors = backend
420            .vector_index(
421                "articles",
422                "embedding",
423                2,
424                VectorIndexSpec::IVF(uqa_storage::IVFIndexParams {
425                    nlist: 2,
426                    nprobe: 1,
427                    train_threshold: 2,
428                }),
429                VectorIndexOpenMode::Create,
430            )
431            .unwrap();
432        vectors.add(1, vec![1.0, 0.0]).unwrap();
433        let hits = vectors.search_knn(&[1.0, 0.0], 1).unwrap();
434        assert_eq!(hits.entries().len(), 1);
435        assert_eq!(hits.entries()[0].doc_id, 1);
436    }
437
438    #[test]
439    fn sqlite_backend_transaction_rolls_back_cross_store_writes() {
440        let conn = ManagedConnection::open_in_memory().unwrap();
441        let _catalog = Catalog::open(conn.clone()).unwrap();
442        let backend = SQLiteStorageBackend::new(conn);
443        let mut docs = backend.document_store("articles");
444        let mut inv = backend.inverted_index("articles", standard_analyzer("english"));
445
446        backend.begin_transaction().unwrap();
447        docs.put(
448            1,
449            BTreeMap::from([("title".to_string(), Value::Str("rollback".into()))]),
450        )
451        .unwrap();
452        inv.add_document(
453            1,
454            BTreeMap::from([("title".to_string(), "rollback".to_string())]),
455        )
456        .unwrap();
457        backend.rollback_transaction().unwrap();
458
459        assert_eq!(docs.len().unwrap(), 0);
460        assert_eq!(inv.doc_freq("title", "rollback").unwrap(), 0);
461    }
462
463    #[test]
464    fn sqlite_sessions_isolate_and_atomically_commit_cross_store_writes() {
465        let dir = tempfile::tempdir().unwrap();
466        let path = dir.path().join("cross-store-isolation.sqlite3");
467        let conn = ManagedConnection::open(&path).unwrap();
468        let catalog = Catalog::open(conn.clone()).unwrap();
469        let writer = SQLiteStorageBackend::new(conn.clone());
470        let observer_conn = conn.new_session();
471        let observer_catalog = Catalog::open(observer_conn.clone()).unwrap();
472        let observer = SQLiteStorageBackend::new(observer_conn);
473
474        let mut writer_docs = writer.document_store("articles");
475        let mut writer_inv = writer.inverted_index("articles", standard_analyzer("english"));
476        let mut writer_vectors = writer
477            .vector_index(
478                "articles",
479                "embedding",
480                2,
481                VectorIndexSpec::BruteForce,
482                VectorIndexOpenMode::Create,
483            )
484            .unwrap();
485        let observer_docs = observer.document_store("articles");
486        let observer_inv = observer.inverted_index("articles", standard_analyzer("english"));
487        let observer_vectors = observer
488            .vector_index(
489                "articles",
490                "embedding",
491                2,
492                VectorIndexSpec::BruteForce,
493                VectorIndexOpenMode::Restore,
494            )
495            .unwrap();
496
497        writer.begin_transaction().unwrap();
498        writer_docs
499            .put(
500                1,
501                BTreeMap::from([("title".to_string(), Value::Str("atomic rust".into()))]),
502            )
503            .unwrap();
504        writer_inv
505            .add_document(
506                1,
507                BTreeMap::from([("title".to_string(), "atomic rust".to_string())]),
508            )
509            .unwrap();
510        writer_vectors.add(1, vec![1.0, 0.0]).unwrap();
511        catalog
512            .save_scoring_params("transactional", r#"{"alpha":1.0}"#)
513            .unwrap();
514
515        assert_eq!(writer_docs.len().unwrap(), 1);
516        assert_eq!(writer_inv.doc_freq("title", "rust").unwrap(), 1);
517        assert_eq!(writer_vectors.count().unwrap(), 1);
518        assert!(catalog
519            .load_scoring_params("transactional")
520            .unwrap()
521            .is_some());
522
523        assert_eq!(observer_docs.len().unwrap(), 0);
524        assert_eq!(observer_inv.doc_freq("title", "rust").unwrap(), 0);
525        assert_eq!(observer_vectors.count().unwrap(), 0);
526        assert!(observer_catalog
527            .load_scoring_params("transactional")
528            .unwrap()
529            .is_none());
530
531        writer.commit_transaction().unwrap();
532        assert_eq!(observer_docs.len().unwrap(), 1);
533        assert_eq!(observer_inv.doc_freq("title", "rust").unwrap(), 1);
534        assert_eq!(observer_vectors.count().unwrap(), 1);
535        assert!(observer_catalog
536            .load_scoring_params("transactional")
537            .unwrap()
538            .is_some());
539    }
540
541    #[test]
542    fn ignored_legacy_index_error_cannot_commit_partial_document_write() {
543        let dir = tempfile::tempdir().unwrap();
544        let path = dir.path().join("ignored-index-error.sqlite3");
545        let conn = ManagedConnection::open(&path).unwrap();
546        let _catalog = Catalog::open(conn.clone()).unwrap();
547        let backend = SQLiteStorageBackend::new(conn.clone());
548        let observer = conn.new_session();
549        let mut docs = backend.document_store("articles");
550        let mut vectors = backend
551            .vector_index(
552                "articles",
553                "embedding",
554                2,
555                VectorIndexSpec::BruteForce,
556                VectorIndexOpenMode::Create,
557            )
558            .unwrap();
559
560        backend.begin_transaction().unwrap();
561        docs.put(
562            1,
563            BTreeMap::from([("title".to_string(), Value::Str("must roll back".into()))]),
564        )
565        .unwrap();
566        conn.with(|connection| {
567            connection.execute("DROP TABLE _vectors", [])?;
568            Ok(())
569        })
570        .unwrap();
571        // The vector write reports its error directly. Even if a caller
572        // ignores that Result, the managed transaction is poisoned and the
573        // partial document write cannot commit.
574        let ignored = vectors.add(1, vec![1.0, 0.0]);
575        assert!(ignored.is_err());
576        assert!(matches!(
577            backend.commit_transaction(),
578            Err(StorageBackendError::Backend { source, .. })
579                if matches!(source.downcast_ref::<SQLiteError>(), Some(SQLiteError::TransactionAborted(_)))
580        ));
581
582        let stored_docs: i64 = observer
583            .with(|connection| {
584                Ok(connection.query_row(
585                    "SELECT COUNT(*) FROM _documents WHERE table_name = 'articles'",
586                    [],
587                    |row| row.get(0),
588                )?)
589            })
590            .unwrap();
591        let vector_table_exists: i64 = observer
592            .with(|connection| {
593                Ok(connection.query_row(
594                    "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = '_vectors'",
595                    [],
596                    |row| row.get(0),
597                )?)
598            })
599            .unwrap();
600        assert_eq!(stored_docs, 0);
601        assert_eq!(vector_table_exists, 1);
602    }
603}
uqa_storage_sqlite/backend.rs

uqa_storage_sqlite/
backend.rs
