1use 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}