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

uqa_storage_sqlite/
backend.rs
