uqa_storage/sqlite/document_store/
store.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Single-document reads, writes, patching, and scalar lookup.
8
9use super::{
10    blob_marker, decode_json_field_value, decode_legacy_document_body, delete_document_blob,
11    document_id_from_sqlite, encode_document_blobs, encode_stored_value, hydrate_document_blobs,
12    params, sqlite_doc_id, upsert_document_blob, BTreeMap, DocId, Document, DocumentMetadata,
13    ManagedConnection, OptionalExtension, SQLiteDocumentStore, SQLiteResult, StorageBackendResult,
14    StoredDocument, Value, DOCUMENT_BLOBS_TABLE,
15};
16
17impl SQLiteDocumentStore {
18    pub fn new(conn: ManagedConnection, table: impl Into<String>) -> Self {
19        Self {
20            conn,
21            table: table.into(),
22        }
23    }
24
25    pub fn max_doc_id(&self) -> StorageBackendResult<DocId> {
26        Ok(self.conn.with(|c| {
27            let id: Option<i64> = c
28                .prepare_cached("SELECT MAX(doc_id) FROM _documents WHERE table_name = ?1")?
29                .query_row(params![self.table], |r| r.get(0))?;
30            id.map_or(Ok(0), document_id_from_sqlite)
31        })?)
32    }
33
34    pub(super) fn put_stored_inner(
35        &self,
36        doc_id: DocId,
37        document: &Document,
38        metadata: DocumentMetadata,
39    ) -> SQLiteResult<()> {
40        let sqlite_doc_id = sqlite_doc_id(doc_id)?;
41        let document: Document = document
42            .iter()
43            .filter(|(_, value)| !matches!(value, uqa_core::Value::Null))
44            .map(|(key, value)| (key.clone(), value.clone()))
45            .collect();
46        let (document, blobs) = encode_document_blobs(document)?;
47        let body = serde_json::to_string(&document)?;
48        self.conn.with(|c| {
49            c.prepare_cached(&format!(
50                "DELETE FROM {DOCUMENT_BLOBS_TABLE}
51                 WHERE table_name = ?1 AND doc_id = ?2"
52            ))?
53            .execute(params![self.table, sqlite_doc_id])?;
54            c.prepare_cached(
55                "INSERT INTO _documents (table_name, doc_id, body, tuple_xmin)
56                 VALUES (?1, ?2, ?3, ?4)
57                 ON CONFLICT (table_name, doc_id)
58                 DO UPDATE SET body = excluded.body, tuple_xmin = excluded.tuple_xmin",
59            )?
60            .execute(params![
61                self.table,
62                sqlite_doc_id,
63                body,
64                metadata.tuple_xmin().map(i64::from),
65            ])?;
66            for (field, bytes) in blobs {
67                c.prepare_cached(&format!(
68                    "INSERT OR REPLACE INTO {DOCUMENT_BLOBS_TABLE}
69                     (table_name, doc_id, field_name, bytes)
70                     VALUES (?1, ?2, ?3, ?4)"
71                ))?
72                .execute(params![self.table, sqlite_doc_id, field, bytes])?;
73            }
74            Ok(())
75        })
76    }
77
78    pub(super) fn get_inner(&self, doc_id: DocId) -> SQLiteResult<Option<Document>> {
79        self.get_stored_inner(doc_id)
80            .map(|document| document.map(StoredDocument::into_fields))
81    }
82
83    pub(super) fn get_stored_inner(&self, doc_id: DocId) -> SQLiteResult<Option<StoredDocument>> {
84        let sqlite_doc_id = sqlite_doc_id(doc_id)?;
85        self.conn.with(|c| {
86            let stored: Option<(String, Option<i64>)> = c
87                .prepare_cached(
88                    "SELECT body, tuple_xmin FROM _documents
89                     WHERE table_name = ?1 AND doc_id = ?2",
90                )?
91                .query_row(params![self.table, sqlite_doc_id], |row| {
92                    Ok((row.get(0)?, row.get(1)?))
93                })
94                .optional()?;
95            let Some((body, tuple_xmin)) = stored else {
96                return Ok(None);
97            };
98            let mut document = decode_legacy_document_body(&body)?;
99            hydrate_document_blobs(c, &self.table, doc_id, &mut document)?;
100            let metadata = tuple_xmin.map_or_else(
101                || Ok(DocumentMetadata::default()),
102                |tuple_xmin| {
103                    u32::try_from(tuple_xmin)
104                        .map(DocumentMetadata::with_tuple_xmin)
105                        .map_err(|_| {
106                            super::SQLiteError::StorageBackend(format!(
107                                "document `{}` row {doc_id} has an out-of-range tuple xmin",
108                                self.table
109                            ))
110                        })
111                },
112            )?;
113            Ok(Some(StoredDocument::with_metadata(document, metadata)))
114        })
115    }
116
117    pub(super) fn get_field_inner(
118        &self,
119        doc_id: DocId,
120        field: &str,
121    ) -> SQLiteResult<Option<Value>> {
122        let sqlite_doc_id = sqlite_doc_id(doc_id)?;
123        let path = sqlite_json_path(field);
124        self.conn.with(|c| {
125            let row: Option<(Option<String>, String)> = c
126                .prepare_cached(
127                    "SELECT json_type(body, ?3), json_quote(json_extract(body, ?3))
128                     FROM _documents
129                     WHERE table_name = ?1 AND doc_id = ?2",
130                )?
131                .query_row(params![self.table, sqlite_doc_id, path], |r| {
132                    Ok((r.get(0)?, r.get(1)?))
133                })
134                .optional()?;
135            let Some((json_type, json_text)) = row else {
136                return Ok(None);
137            };
138            decode_json_field_value(c, &self.table, doc_id, field, json_type, &json_text)
139        })
140    }
141
142    pub(super) fn find_doc_id_by_field_inner(
143        &self,
144        field: &str,
145        value: &Value,
146    ) -> SQLiteResult<Option<DocId>> {
147        let path = sqlite_json_path(field);
148        match value {
149            Value::Str(value) => self
150                .conn
151                .with(|c| find_doc_id_by_scalar(c, &self.table, &path, value)),
152            Value::Int(value) => self
153                .conn
154                .with(|c| find_doc_id_by_scalar(c, &self.table, &path, value)),
155            Value::Float(value) if value.is_finite() => self
156                .conn
157                .with(|c| find_doc_id_by_scalar(c, &self.table, &path, value)),
158            Value::Bool(value) => self.conn.with(|c| {
159                let json_type = if *value { "true" } else { "false" };
160                let doc_id: Option<i64> = c
161                    .query_row(
162                        "SELECT doc_id FROM _documents
163                         WHERE table_name = ?1 AND json_type(body, ?2) = ?3
164                         ORDER BY doc_id LIMIT 1",
165                        params![self.table, path, json_type],
166                        |r| r.get(0),
167                    )
168                    .optional()?;
169                doc_id.map(document_id_from_sqlite).transpose()
170            }),
171            _ => {
172                let doc_ids = self.conn.with(|c| {
173                    let mut stmt = c.prepare_cached(
174                        "SELECT doc_id FROM _documents
175                         WHERE table_name = ?1 ORDER BY doc_id",
176                    )?;
177                    let rows = stmt.query_map(params![self.table], |row| row.get::<_, i64>(0))?;
178                    let mut out = Vec::new();
179                    for row in rows {
180                        out.push(document_id_from_sqlite(row?)?);
181                    }
182                    Ok(out)
183                })?;
184                for doc_id in doc_ids {
185                    if self.get_field_inner(doc_id, field)?.as_ref() == Some(value) {
186                        return Ok(Some(doc_id));
187                    }
188                }
189                Ok(None)
190            }
191        }
192    }
193
194    pub(super) fn patch_fields_inner(
195        &self,
196        doc_id: DocId,
197        updates: &BTreeMap<String, Value>,
198    ) -> SQLiteResult<bool> {
199        if updates.is_empty() {
200            return Ok(true);
201        }
202        let sqlite_doc_id = sqlite_doc_id(doc_id)?;
203        self.conn.with(|c| {
204            let exists: Option<i64> = c
205                .prepare_cached(
206                    "SELECT 1 FROM _documents
207                     WHERE table_name = ?1 AND doc_id = ?2
208                     LIMIT 1",
209                )?
210                .query_row(params![self.table, sqlite_doc_id], |r| r.get(0))
211                .optional()?;
212            if exists.is_none() {
213                return Ok(false);
214            }
215
216            for (field, value) in updates {
217                let path = sqlite_json_path(field);
218                match value {
219                    Value::Null => {
220                        delete_document_blob(c, &self.table, doc_id, field)?;
221                        c.execute(
222                            "UPDATE _documents SET body = json_remove(body, ?3)
223                             WHERE table_name = ?1 AND doc_id = ?2",
224                            params![self.table, sqlite_doc_id, path],
225                        )?;
226                    }
227                    Value::Bytes(bytes) => {
228                        let marker = serde_json::to_string(&blob_marker(field.clone()))?;
229                        c.execute(
230                            "UPDATE _documents SET body = json_set(body, ?3, json(?4))
231                             WHERE table_name = ?1 AND doc_id = ?2",
232                            params![self.table, sqlite_doc_id, path, marker],
233                        )?;
234                        upsert_document_blob(c, &self.table, doc_id, field, bytes)?;
235                    }
236                    other => {
237                        let (stored, blob) = encode_stored_value(field, other.clone())?;
238                        let json = serde_json::to_string(&stored)?;
239                        c.execute(
240                            "UPDATE _documents SET body = json_set(body, ?3, json(?4))
241                             WHERE table_name = ?1 AND doc_id = ?2",
242                            params![self.table, sqlite_doc_id, path, json],
243                        )?;
244                        if let Some(bytes) = blob {
245                            upsert_document_blob(c, &self.table, doc_id, field, &bytes)?;
246                        } else {
247                            delete_document_blob(c, &self.table, doc_id, field)?;
248                        }
249                    }
250                }
251            }
252            Ok(true)
253        })
254    }
255}
256
257fn sqlite_json_path(field: &str) -> String {
258    if field.chars().all(|c| c == '_' || c.is_ascii_alphanumeric())
259        && field
260            .chars()
261            .next()
262            .is_some_and(|c| c == '_' || c.is_ascii_alphabetic())
263    {
264        format!("$.{field}")
265    } else {
266        format!("$.{}", serde_json::Value::String(field.to_string()))
267    }
268}
269
270fn find_doc_id_by_scalar<T: rusqlite::ToSql>(
271    conn: &rusqlite::Connection,
272    table: &str,
273    path: &str,
274    value: &T,
275) -> SQLiteResult<Option<DocId>> {
276    let doc_id: Option<i64> = conn
277        .query_row(
278            "SELECT doc_id FROM _documents
279             WHERE table_name = ?1 AND json_extract(body, ?2) = ?3
280             ORDER BY doc_id LIMIT 1",
281            (table, path, value),
282            |r| r.get(0),
283        )
284        .optional()?;
285    doc_id.map(document_id_from_sqlite).transpose()
286}
uqa_storage/sqlite/document_store/store.rs

uqa_storage/sqlite/document_store/
store.rs
