Skip to main content

uqa_storage_sqlite/
block_max_index.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Atomic `SQLite` persistence for the shared block-max score index.
8
9use rusqlite::params;
10use std::collections::BTreeMap;
11use uqa_storage::{BlockMaxIndex, StorageBackendError, TokenTermKey};
12
13pub trait SQLiteBlockMaxPersistence {
14    fn save_to_sqlite(&self, connection: &rusqlite::Connection) -> rusqlite::Result<()>;
15    fn load_from_sqlite(&mut self, connection: &rusqlite::Connection) -> rusqlite::Result<()>;
16}
17
18impl SQLiteBlockMaxPersistence for BlockMaxIndex {
19    fn save_to_sqlite(&self, conn: &rusqlite::Connection) -> rusqlite::Result<()> {
20        ensure_global_blockmax_shape(conn)?;
21        let transaction = conn.unchecked_transaction()?;
22        transaction.execute("DELETE FROM _global_blockmax", [])?;
23        for ((table, field, term), scores) in self.entries() {
24            for (block_idx, score) in scores.iter().enumerate() {
25                let block_idx = i64::try_from(block_idx)
26                    .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
27                transaction.execute(
28                    "INSERT INTO _global_blockmax
29                        (table_name, field, term, block_idx, max_score)
30                     VALUES (?1, ?2, ?3, ?4, ?5)",
31                    params![table, field, term.as_bytes(), block_idx, *score],
32                )?;
33            }
34        }
35        transaction.commit()
36    }
37
38    fn load_from_sqlite(&mut self, conn: &rusqlite::Connection) -> rusqlite::Result<()> {
39        ensure_global_blockmax_shape(conn)?;
40        let mut stmt = conn.prepare(
41            "SELECT table_name, field, term, block_idx, max_score
42             FROM _global_blockmax
43             ORDER BY table_name, field, term, block_idx",
44        )?;
45        let rows = stmt.query_map([], |row| {
46            Ok((
47                row.get::<_, String>(0)?,
48                row.get::<_, String>(1)?,
49                read_term_key(row.get_ref(2)?)?,
50                row.get::<_, i64>(3)?,
51                row.get::<_, f64>(4)?,
52            ))
53        })?;
54        let mut loaded = BTreeMap::<(String, String, TokenTermKey), Vec<f64>>::new();
55        for row in rows {
56            let (table, field, term, block_idx, score) = row?;
57            let idx = usize::try_from(block_idx)
58                .map_err(|_| rusqlite::Error::IntegralValueOutOfRange(3, block_idx))?;
59            let entry = loaded.entry((table, field, term)).or_default();
60            if idx != entry.len() {
61                return Err(rusqlite::Error::FromSqlConversionFailure(
62                    3,
63                    rusqlite::types::Type::Integer,
64                    Box::new(std::io::Error::new(
65                        std::io::ErrorKind::InvalidData,
66                        format!(
67                            "invalid block-max ordinal sequence: expected {}, found {idx}",
68                            entry.len()
69                        ),
70                    )),
71                ));
72            }
73            entry.push(score);
74        }
75        let mut replacement = Self::new(self.block_size()).map_err(storage_error_to_sqlite)?;
76        for ((table, field, term), scores) in loaded {
77            replacement
78                .set_block_maxes_key(&table, &field, &term, scores)
79                .map_err(storage_error_to_sqlite)?;
80        }
81        *self = replacement;
82        Ok(())
83    }
84}
85
86fn read_term_key(value: rusqlite::types::ValueRef<'_>) -> rusqlite::Result<TokenTermKey> {
87    match value {
88        rusqlite::types::ValueRef::Blob(bytes) => {
89            TokenTermKey::from_bytes(bytes.to_vec()).map_err(storage_error_to_sqlite)
90        }
91        rusqlite::types::ValueRef::Text(bytes) => std::str::from_utf8(bytes)
92            .map(TokenTermKey::from_text)
93            .map_err(|error| {
94                rusqlite::Error::FromSqlConversionFailure(
95                    2,
96                    rusqlite::types::Type::Text,
97                    Box::new(error),
98                )
99            }),
100        _ => Err(rusqlite::Error::InvalidColumnType(
101            2,
102            "term".into(),
103            value.data_type(),
104        )),
105    }
106}
107
108fn storage_error_to_sqlite(error: StorageBackendError) -> rusqlite::Error {
109    rusqlite::Error::ToSqlConversionFailure(Box::new(error))
110}
111
112fn ensure_global_blockmax_shape(conn: &rusqlite::Connection) -> rusqlite::Result<()> {
113    conn.execute(
114        "CREATE TABLE IF NOT EXISTS _global_blockmax (
115            table_name TEXT NOT NULL DEFAULT '',
116            field     TEXT NOT NULL,
117            term      BLOB NOT NULL,
118            block_idx INTEGER NOT NULL,
119            max_score REAL NOT NULL,
120            PRIMARY KEY (table_name, field, term, block_idx)
121        )",
122        [],
123    )?;
124    let mut stmt = conn.prepare("PRAGMA table_info(_global_blockmax)")?;
125    let cols = stmt
126        .query_map([], |row| row.get::<_, String>(1))?
127        .collect::<Result<Vec<_>, _>>()?;
128    drop(stmt);
129    if !cols.iter().any(|c| c == "table_name") {
130        conn.execute(
131            "ALTER TABLE _global_blockmax ADD COLUMN table_name TEXT NOT NULL DEFAULT ''",
132            [],
133        )?;
134    }
135    Ok(())
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    #[test]
142    fn corrupt_persisted_ordinal_does_not_replace_loaded_state() {
143        let connection = rusqlite::Connection::open_in_memory().unwrap();
144        ensure_global_blockmax_shape(&connection).unwrap();
145        connection
146            .execute(
147                "INSERT INTO _global_blockmax
148                    (table_name, field, term, block_idx, max_score)
149                 VALUES ('docs', 'body', 'bad', -1, 9.0)",
150                [],
151            )
152            .unwrap();
153        let mut index = BlockMaxIndex::default();
154        index
155            .set_block_maxes("old", "body", "term", vec![1.0])
156            .unwrap();
157
158        assert!(index.load_from_sqlite(&connection).is_err());
159        assert_eq!(index.block_max("old", "body", "term", 0), 1.0);
160    }
161
162    #[test]
163    fn failed_save_rolls_back_deleted_snapshot() {
164        let connection = rusqlite::Connection::open_in_memory().unwrap();
165        ensure_global_blockmax_shape(&connection).unwrap();
166        connection
167            .execute(
168                "INSERT INTO _global_blockmax
169                    (table_name, field, term, block_idx, max_score)
170                 VALUES ('old', 'body', 'term', 0, 1.0)",
171                [],
172            )
173            .unwrap();
174        connection
175            .execute_batch(
176                "CREATE TRIGGER fail_blockmax_insert
177                 BEFORE INSERT ON _global_blockmax
178                 BEGIN
179                     SELECT RAISE(ABORT, 'injected block-max failure');
180                 END;",
181            )
182            .unwrap();
183        let mut index = BlockMaxIndex::default();
184        index
185            .set_block_maxes("new", "body", "term", vec![2.0])
186            .unwrap();
187
188        assert!(index.save_to_sqlite(&connection).is_err());
189        let persisted: (String, f64) = connection
190            .query_row(
191                "SELECT table_name, max_score FROM _global_blockmax",
192                [],
193                |row| Ok((row.get(0)?, row.get(1)?)),
194            )
195            .unwrap();
196        assert_eq!(persisted, ("old".to_string(), 1.0));
197    }
198}