uqa_storage/sqlite/document_store/
trait_impl.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Complete `DocumentStore` query, mutation, iteration, and snapshot contract.
8
9use super::{
10    allocation_error, blob_marker_info, chunk_bind_values, decode_legacy_document_body,
11    doc_id_in_placeholders, document_id_from_sqlite, hydrate_document_blobs,
12    load_marked_document_blob, params, read_doc_id, should_probe_doc_ids, sorted_unique_doc_ids,
13    sqlite_doc_id, take_requested_field, Arc, BTreeMap, DocId, Document, DocumentStore,
14    OptionalExtension, SQLiteDocumentStore, SQLiteError, SQLiteResult, StorageBackendResult, Value,
15    DOCUMENT_BLOBS_TABLE, DOC_ID_IN_CHUNK,
16};
17
18impl DocumentStore for SQLiteDocumentStore {
19    fn put(&mut self, doc_id: DocId, document: Document) -> StorageBackendResult<()> {
20        self.put_inner(doc_id, &document)?;
21        Ok(())
22    }
23
24    fn get(&self, doc_id: DocId) -> StorageBackendResult<Option<Document>> {
25        Ok(self.get_inner(doc_id)?)
26    }
27
28    fn contains_doc_id(&self, doc_id: DocId) -> StorageBackendResult<bool> {
29        let sqlite_doc_id = sqlite_doc_id(doc_id)?;
30        Ok(self.conn.with(|c| {
31            let found: Option<i64> = c
32                .prepare_cached(
33                    "SELECT 1 FROM _documents
34                         WHERE table_name = ?1 AND doc_id = ?2
35                         LIMIT 1",
36                )?
37                .query_row(params![self.table, sqlite_doc_id], |r| r.get(0))
38                .optional()?;
39            Ok(found.is_some())
40        })?)
41    }
42
43    fn get_field(
44        &self,
45        doc_id: DocId,
46        field: &str,
47    ) -> StorageBackendResult<Option<uqa_core::Value>> {
48        Ok(self.get_field_inner(doc_id, field)?)
49    }
50
51    fn find_doc_id_by_field(
52        &self,
53        field: &str,
54        value: &Value,
55    ) -> StorageBackendResult<Option<DocId>> {
56        Ok(self.find_doc_id_by_field_inner(field, value)?)
57    }
58
59    fn get_fields_bulk(
60        &self,
61        doc_ids: &[DocId],
62        field: &str,
63    ) -> StorageBackendResult<BTreeMap<DocId, Value>> {
64        let mut out: BTreeMap<DocId, Value> = doc_ids
65            .iter()
66            .copied()
67            .map(|doc_id| (doc_id, Value::Null))
68            .collect();
69        if doc_ids.is_empty() {
70            return Ok(out);
71        }
72        // Fetch the document body and extract the field in Rust: one
73        // JSON parse per row. Extracting through `json_type` +
74        // `json_extract` made `SQLite` parse the same body twice per
75        // requested field.
76        let mut decode_row = |c: &rusqlite::Connection,
77                              row: &rusqlite::Row<'_>|
78         -> SQLiteResult<()> {
79            let doc_id = read_doc_id(row, 0)?;
80            let body = row.get::<_, String>(1)?;
81            let mut document = decode_legacy_document_body(&body)?;
82            if let Some(value) = take_requested_field(c, &self.table, doc_id, &mut document, field)?
83            {
84                out.insert(doc_id, value);
85            }
86            Ok(())
87        };
88
89        // Selective requests probe by id; wide requests (half the
90        // table or more) sequential-scan once instead of issuing many
91        // B-tree probes.
92        let should_probe =
93            doc_ids.len() <= DOC_ID_IN_CHUNK || should_probe_doc_ids(doc_ids.len(), self.len()?);
94        if should_probe {
95            let leading = [rusqlite::types::Value::Text(self.table.clone())];
96            let sql = format!(
97                "SELECT doc_id, body FROM _documents
98                 WHERE table_name = ?1 AND doc_id IN ({})",
99                doc_id_in_placeholders(2, DOC_ID_IN_CHUNK)?
100            );
101            self.conn.with(|c| {
102                for chunk in doc_ids.chunks(DOC_ID_IN_CHUNK) {
103                    let mut stmt = c.prepare_cached(&sql)?;
104                    let bind = chunk_bind_values(&leading, chunk)?;
105                    let mut rows = stmt.query(rusqlite::params_from_iter(bind))?;
106                    while let Some(row) = rows.next()? {
107                        decode_row(c, row)?;
108                    }
109                }
110                Ok(())
111            })?;
112            return Ok(out);
113        }
114
115        let requested = sorted_unique_doc_ids(doc_ids)?;
116        self.conn.with(|c| {
117            let mut stmt = c.prepare_cached(
118                "SELECT doc_id, body FROM _documents
119                 WHERE table_name = ?1
120                 ORDER BY doc_id",
121            )?;
122            let mut rows = stmt.query(params![self.table])?;
123            while let Some(row) = rows.next()? {
124                let doc_id = read_doc_id(row, 0)?;
125                if requested.binary_search(&doc_id).is_err() {
126                    continue;
127                }
128                decode_row(c, row)?;
129            }
130            Ok(())
131        })?;
132        Ok(out)
133    }
134
135    fn get_fields_multi(
136        &self,
137        doc_ids: &[DocId],
138        fields: &[&str],
139    ) -> StorageBackendResult<BTreeMap<DocId, Vec<Value>>> {
140        let mut out: BTreeMap<DocId, Vec<Value>> = BTreeMap::new();
141        if doc_ids.is_empty() || fields.is_empty() {
142            return Ok(out);
143        }
144        // Fetch the document body and extract every requested field in
145        // Rust: one JSON parse per row, however many fields the caller
146        // asked for. The previous `json_type` + `json_extract` pair per
147        // field made `SQLite` parse the same body twice per field.
148        let decode_row = |c: &rusqlite::Connection,
149                          row: &rusqlite::Row<'_>|
150         -> SQLiteResult<(DocId, Vec<Value>)> {
151            let doc_id = read_doc_id(row, 0)?;
152            let body = row.get::<_, String>(1)?;
153            let document = decode_legacy_document_body(&body)?;
154            let mut values = Vec::new();
155            values
156                .try_reserve_exact(fields.len())
157                .map_err(|error| allocation_error("multi-field document values", error))?;
158            for field in fields {
159                let mut value = document.get(*field).cloned().unwrap_or(Value::Null);
160                if let Some(marker) = blob_marker_info(&value) {
161                    if let Some(decoded) =
162                        load_marked_document_blob(c, &self.table, doc_id, field, &marker)?
163                    {
164                        value = decoded;
165                    }
166                }
167                values.push(value);
168            }
169            Ok((doc_id, values))
170        };
171
172        let should_probe =
173            doc_ids.len() <= DOC_ID_IN_CHUNK || should_probe_doc_ids(doc_ids.len(), self.len()?);
174        if should_probe {
175            let leading = [rusqlite::types::Value::Text(self.table.clone())];
176            let sql = format!(
177                "SELECT doc_id, body FROM _documents
178                 WHERE table_name = ?1 AND doc_id IN ({})",
179                doc_id_in_placeholders(2, DOC_ID_IN_CHUNK)?
180            );
181            self.conn.with(|c| {
182                for chunk in doc_ids.chunks(DOC_ID_IN_CHUNK) {
183                    let mut stmt = c.prepare_cached(&sql)?;
184                    let bind = chunk_bind_values(&leading, chunk)?;
185                    let mut rows = stmt.query(rusqlite::params_from_iter(bind))?;
186                    while let Some(row) = rows.next()? {
187                        let (doc_id, values) = decode_row(c, row)?;
188                        out.insert(doc_id, values);
189                    }
190                }
191                Ok(())
192            })?;
193            return Ok(out);
194        }
195
196        let requested = sorted_unique_doc_ids(doc_ids)?;
197        self.conn.with(|c| {
198            let mut stmt = c.prepare_cached(
199                "SELECT doc_id, body FROM _documents
200                 WHERE table_name = ?1
201                 ORDER BY doc_id",
202            )?;
203            let mut rows = stmt.query(params![self.table])?;
204            while let Some(row) = rows.next()? {
205                let doc_id = read_doc_id(row, 0)?;
206                if requested.binary_search(&doc_id).is_err() {
207                    continue;
208                }
209                let (doc_id, values) = decode_row(c, row)?;
210                out.insert(doc_id, values);
211            }
212            Ok(())
213        })?;
214        Ok(out)
215    }
216
217    fn get_many(&self, doc_ids: &[DocId]) -> StorageBackendResult<BTreeMap<DocId, Document>> {
218        let mut out: BTreeMap<DocId, Document> = BTreeMap::new();
219        if doc_ids.is_empty() {
220            return Ok(out);
221        }
222        // Same probe-vs-scan split as `get_fields_bulk`.
223        let should_probe =
224            doc_ids.len() <= DOC_ID_IN_CHUNK || should_probe_doc_ids(doc_ids.len(), self.len()?);
225        if should_probe {
226            let leading = [rusqlite::types::Value::Text(self.table.clone())];
227            let sql = format!(
228                "SELECT doc_id, body FROM _documents
229                 WHERE table_name = ?1 AND doc_id IN ({})",
230                doc_id_in_placeholders(2, DOC_ID_IN_CHUNK)?
231            );
232            self.conn.with(|c| {
233                for chunk in doc_ids.chunks(DOC_ID_IN_CHUNK) {
234                    let mut stmt = c.prepare_cached(&sql)?;
235                    let bind = chunk_bind_values(&leading, chunk)?;
236                    let mut rows = stmt.query(rusqlite::params_from_iter(bind))?;
237                    while let Some(row) = rows.next()? {
238                        let doc_id = read_doc_id(row, 0)?;
239                        let body = row.get::<_, String>(1)?;
240                        let mut document = decode_legacy_document_body(&body)?;
241                        hydrate_document_blobs(c, &self.table, doc_id, &mut document)?;
242                        out.insert(doc_id, document);
243                    }
244                }
245                Ok(())
246            })?;
247            return Ok(out);
248        }
249
250        let requested = sorted_unique_doc_ids(doc_ids)?;
251        self.conn.with(|c| {
252            let mut stmt = c.prepare_cached(
253                "SELECT doc_id, body FROM _documents
254                 WHERE table_name = ?1
255                 ORDER BY doc_id",
256            )?;
257            let mut rows = stmt.query(params![self.table])?;
258            while let Some(row) = rows.next()? {
259                let doc_id = read_doc_id(row, 0)?;
260                if requested.binary_search(&doc_id).is_err() {
261                    continue;
262                }
263                let body = row.get::<_, String>(1)?;
264                let mut document = decode_legacy_document_body(&body)?;
265                hydrate_document_blobs(c, &self.table, doc_id, &mut document)?;
266                out.insert(doc_id, document);
267            }
268            Ok(())
269        })?;
270        Ok(out)
271    }
272
273    fn patch_fields(
274        &mut self,
275        doc_id: DocId,
276        updates: &BTreeMap<String, Value>,
277    ) -> StorageBackendResult<bool> {
278        Ok(self.patch_fields_inner(doc_id, updates)?)
279    }
280
281    fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
282        let sqlite_doc_id = sqlite_doc_id(doc_id)?;
283        self.conn.with(|c| {
284            c.prepare_cached(&format!(
285                "DELETE FROM {DOCUMENT_BLOBS_TABLE}
286                 WHERE table_name = ?1 AND doc_id = ?2"
287            ))?
288            .execute(params![self.table, sqlite_doc_id])?;
289            c.prepare_cached("DELETE FROM _documents WHERE table_name = ?1 AND doc_id = ?2")?
290                .execute(params![self.table, sqlite_doc_id])?;
291            Ok(())
292        })?;
293        Ok(())
294    }
295
296    fn clear(&mut self) -> StorageBackendResult<()> {
297        self.conn.with(|c| {
298            c.execute(
299                &format!("DELETE FROM {DOCUMENT_BLOBS_TABLE} WHERE table_name = ?1"),
300                params![self.table],
301            )?;
302            c.execute(
303                "DELETE FROM _documents WHERE table_name = ?1",
304                params![self.table],
305            )?;
306            Ok(())
307        })?;
308        Ok(())
309    }
310
311    fn doc_ids(&self) -> StorageBackendResult<Vec<DocId>> {
312        Ok(self.conn.with(|c| {
313            let mut stmt = c.prepare_cached(
314                "SELECT doc_id FROM _documents WHERE table_name = ?1 ORDER BY doc_id",
315            )?;
316            let rows = stmt.query_map(params![self.table], |r| r.get::<_, i64>(0))?;
317            let mut out = Vec::new();
318            for row in rows {
319                out.push(document_id_from_sqlite(row?)?);
320            }
321            Ok(out)
322        })?)
323    }
324
325    fn next_doc_id(&self, after: Option<DocId>) -> StorageBackendResult<Option<DocId>> {
326        let after = after.map(sqlite_doc_id).transpose()?;
327        Ok(self.conn.with(|connection| {
328            let doc_id: Option<i64> = match after {
329                Some(after) => connection
330                    .prepare_cached(
331                        "SELECT doc_id FROM _documents
332                         WHERE table_name = ?1 AND doc_id > ?2
333                         ORDER BY doc_id LIMIT 1",
334                    )?
335                    .query_row(params![self.table, after], |row| row.get::<_, i64>(0))
336                    .optional()?,
337                None => connection
338                    .prepare_cached(
339                        "SELECT doc_id FROM _documents
340                         WHERE table_name = ?1
341                         ORDER BY doc_id LIMIT 1",
342                    )?
343                    .query_row(params![self.table], |row| row.get::<_, i64>(0))
344                    .optional()?,
345            };
346            doc_id.map(document_id_from_sqlite).transpose()
347        })?)
348    }
349
350    fn next_doc_ids(&self, after: Option<DocId>, limit: usize) -> StorageBackendResult<Vec<DocId>> {
351        if limit == 0 {
352            return Ok(Vec::new());
353        }
354        let after = after.map(sqlite_doc_id).transpose()?;
355        let limit = i64::try_from(limit).map_err(|_| {
356            SQLiteError::StorageBackend(format!(
357                "document cursor limit {limit} is outside SQLite's integer range"
358            ))
359        })?;
360        Ok(self.conn.with(|connection| {
361            let mut out = Vec::new();
362            if let Some(after) = after {
363                let mut stmt = connection.prepare_cached(
364                    "SELECT doc_id FROM _documents
365                     WHERE table_name = ?1 AND doc_id > ?2
366                     ORDER BY doc_id LIMIT ?3",
367                )?;
368                let rows = stmt.query_map(params![self.table, after, limit], |row| {
369                    row.get::<_, i64>(0)
370                })?;
371                for row in rows {
372                    out.push(document_id_from_sqlite(row?)?);
373                }
374            } else {
375                let mut stmt = connection.prepare_cached(
376                    "SELECT doc_id FROM _documents
377                     WHERE table_name = ?1
378                     ORDER BY doc_id LIMIT ?2",
379                )?;
380                let rows =
381                    stmt.query_map(params![self.table, limit], |row| row.get::<_, i64>(0))?;
382                for row in rows {
383                    out.push(document_id_from_sqlite(row?)?);
384                }
385            }
386            Ok(out)
387        })?)
388    }
389
390    fn max_doc_id(&self) -> StorageBackendResult<DocId> {
391        SQLiteDocumentStore::max_doc_id(self)
392    }
393
394    fn len(&self) -> StorageBackendResult<usize> {
395        Ok(self.conn.with(|c| {
396            let n: i64 = c
397                .prepare_cached("SELECT COUNT(*) FROM _documents WHERE table_name = ?1")?
398                .query_row(params![self.table], |r| r.get(0))?;
399            usize::try_from(n).map_err(|_| {
400                SQLiteError::StorageBackend(format!(
401                    "document count {n} is outside the addressable range"
402                ))
403            })
404        })?)
405    }
406
407    fn snapshot(&self) -> StorageBackendResult<Arc<dyn DocumentStore>> {
408        Ok(Arc::new(self.clone()))
409    }
410}
uqa_storage/sqlite/document_store/trait_impl.rs

uqa_storage/sqlite/document_store/
trait_impl.rs
