assetpack_core/sqlite/
rusqlite.rs

1use std::collections::{HashMap, HashSet};
2
3use rusqlite::{Connection, OptionalExtension, params, params_from_iter};
4
5use super::{
6  MerkleMode, MerkleProof, MerkleSummary, SqliteLayout, SqliteStoreOptions, StoredObjectRow, codec_to_str, hash_pair,
7  merkle_root_from_hashes, object_record_from_row, prepare_object, schema_statements, verified_object_from_row,
8};
9use crate::{Codec, Error, Hash32, ObjectKind, ObjectRecord, ObjectSource, Result, VerifiedObject};
10
11pub struct RusqliteStore<'connection> {
12  connection: &'connection Connection,
13  layout: SqliteLayout,
14  merkle: MerkleMode,
15}
16
17impl<'connection> RusqliteStore<'connection> {
18  pub fn from_connection(connection: &'connection Connection) -> Result<Self> {
19    Self::from_connection_with_options(connection, SqliteStoreOptions::default())
20  }
21
22  pub fn from_connection_with_namespace(connection: &'connection Connection, namespace: impl AsRef<str>) -> Result<Self> {
23    Self::from_connection_with_options(
24      connection,
25      SqliteStoreOptions {
26        namespace: Some(namespace.as_ref().to_string()),
27        ..Default::default()
28      },
29    )
30  }
31
32  pub fn from_connection_with_options(connection: &'connection Connection, options: SqliteStoreOptions) -> Result<Self> {
33    let layout = SqliteLayout::new(options.namespace.as_deref())?;
34    for statement in schema_statements(&layout, options.merkle) {
35      connection.execute_batch(&statement)?;
36    }
37    Ok(Self::view(connection, options, layout))
38  }
39
40  pub(in crate::sqlite) fn view(connection: &'connection Connection, options: SqliteStoreOptions, layout: SqliteLayout) -> Self {
41    Self {
42      connection,
43      layout,
44      merkle: options.merkle,
45    }
46  }
47
48  pub fn ensure_schema(connection: &Connection) -> Result<()> {
49    Self::ensure_schema_with_options(connection, &SqliteStoreOptions::default())
50  }
51
52  pub fn ensure_schema_with_options(connection: &Connection, options: &SqliteStoreOptions) -> Result<()> {
53    let layout = SqliteLayout::new(options.namespace.as_deref())?;
54    for statement in schema_statements(&layout, options.merkle) {
55      connection.execute_batch(&statement)?;
56    }
57    Ok(())
58  }
59
60  pub fn connection(&self) -> &Connection {
61    self.connection
62  }
63
64  pub fn namespace(&self) -> Option<&str> {
65    self.layout.namespace.as_deref()
66  }
67
68  pub fn merkle_mode(&self) -> MerkleMode {
69    self.merkle
70  }
71
72  pub fn put_objects_batch(&self, objects: &[ObjectRecord]) -> Result<()> {
73    if objects.is_empty() {
74      return Ok(());
75    }
76    let sql = self.layout.insert_object();
77    let mut statement = self.connection.prepare_cached(&sql)?;
78    for object in objects {
79      statement.execute(params![
80        object.hash.as_bytes().as_slice(),
81        object.kind as i64,
82        i64::try_from(object.decoded_len).map_err(|_| Error::Integrity("object size exceeds SQLite integer".into()))?,
83        i64::try_from(object.stored_bytes.len()).map_err(|_| Error::Integrity("stored size exceeds SQLite integer".into()))?,
84        codec_to_str(object.codec),
85        &object.stored_bytes,
86        unix_timestamp(),
87      ])?;
88    }
89    Ok(())
90  }
91
92  pub fn put_file_recipe_cache_batch(&self, entries: &[(Hash32, Hash32)]) -> Result<()> {
93    let sql = self.layout.upsert_recipe_cache();
94    let mut statement = self.connection.prepare_cached(&sql)?;
95    for (file_hash, recipe_hash) in entries {
96      statement.execute(params![
97        file_hash.as_bytes().as_slice(),
98        recipe_hash.as_bytes().as_slice(),
99        unix_timestamp()
100      ])?;
101    }
102    Ok(())
103  }
104
105  pub fn put_chunk(&self, hash: Hash32, content: &[u8], codec: Codec) -> Result<()> {
106    self.put_object(hash, ObjectKind::Chunk, content, codec)
107  }
108
109  pub fn put_recipe(&self, hash: Hash32, content: &[u8], codec: Codec) -> Result<()> {
110    self.put_object(hash, ObjectKind::Recipe, content, codec)
111  }
112
113  fn put_object(&self, hash: Hash32, kind: ObjectKind, content: &[u8], codec: Codec) -> Result<()> {
114    self.put_objects_batch(&[prepare_object(hash, kind, content, codec)?])
115  }
116
117  pub fn read_object_record(&self, hash: &Hash32) -> Result<Option<ObjectRecord>> {
118    let sql = self.layout.select_object();
119    self
120      .connection
121      .query_row(&sql, [hash.as_bytes().as_slice()], stored_row)
122      .optional()?
123      .map(object_record_from_row)
124      .transpose()
125  }
126
127  pub fn existing_hashes(&self, hashes: &[Hash32]) -> Result<HashSet<Hash32>> {
128    let mut found = HashSet::new();
129    for chunk in hashes.chunks(999) {
130      if chunk.is_empty() {
131        continue;
132      }
133      let placeholders = std::iter::repeat_n("?", chunk.len()).collect::<Vec<_>>().join(",");
134      let sql = format!("{}{placeholders})", self.layout.select_hashes_prefix());
135      let values = chunk.iter().map(|hash| hash.as_bytes().as_slice());
136      let mut statement = self.connection.prepare(&sql)?;
137      let rows = statement.query_map(params_from_iter(values), |row| row.get::<_, Vec<u8>>(0))?;
138      for row in rows {
139        found.insert(Hash32::from_bytes(&row?)?);
140      }
141    }
142    Ok(found)
143  }
144
145  pub fn chunk_stored_bytes(&self, hash: &Hash32) -> Result<Option<u64>> {
146    let sql = self.layout.select_chunk_stored_size();
147    let value = self
148      .connection
149      .query_row(&sql, [hash.as_bytes().as_slice()], |row| {
150        Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?))
151      })
152      .optional()?;
153    match value {
154      Some((kind, size)) if kind == ObjectKind::Chunk as i64 => {
155        Ok(Some(size.try_into().map_err(|_| Error::Integrity("negative stored size".into()))?))
156      }
157      _ => Ok(None),
158    }
159  }
160
161  pub fn chunk_sizes(&self) -> Result<Vec<(Hash32, u64)>> {
162    let sql = self.layout.select_chunk_sizes();
163    let mut statement = self.connection.prepare(&sql)?;
164    let rows = statement.query_map([ObjectKind::Chunk as i64], |row| {
165      Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, i64>(1)?))
166    })?;
167    rows
168      .map(|row| {
169        let (hash, size) = row?;
170        Ok((
171          Hash32::from_bytes(&hash)?,
172          size.try_into().map_err(|_| Error::Integrity("negative stored size".into()))?,
173        ))
174      })
175      .collect()
176  }
177
178  pub fn file_recipe_cache(&self) -> Result<HashMap<Hash32, Hash32>> {
179    let sql = self.layout.select_recipe_cache();
180    let mut statement = self.connection.prepare(&sql)?;
181    let rows = statement.query_map([], |row| Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, Vec<u8>>(1)?)))?;
182    let mut cache = HashMap::new();
183    for row in rows {
184      let (file, recipe) = row?;
185      cache.insert(Hash32::from_bytes(&file)?, Hash32::from_bytes(&recipe)?);
186    }
187    Ok(cache)
188  }
189
190  pub fn recipe_for_file_hash(&self, file_hash: &Hash32) -> Result<Option<Hash32>> {
191    let sql = self.layout.select_recipe_for_file();
192    self
193      .connection
194      .query_row(&sql, [file_hash.as_bytes().as_slice()], |row| row.get::<_, Vec<u8>>(0))
195      .optional()?
196      .map(|bytes| Hash32::from_bytes(&bytes))
197      .transpose()
198  }
199
200  pub fn all_objects(&self) -> Result<Vec<VerifiedObject>> {
201    let sql = self.layout.select_objects();
202    let mut statement = self.connection.prepare(&sql)?;
203    statement
204      .query_map([], stored_row)?
205      .map(|row| verified_object_from_row(row?))
206      .collect()
207  }
208
209  pub fn object_records(&self, offset: u64, limit: usize) -> Result<Vec<ObjectRecord>> {
210    let sql = self.layout.select_object_records();
211    let mut statement = self.connection.prepare(&sql)?;
212    statement
213      .query_map(params![limit as i64, offset as i64], stored_row)?
214      .map(|row| object_record_from_row(row?))
215      .collect()
216  }
217
218  pub fn integrity_check(&self) -> Result<()> {
219    let status: String = self
220      .connection
221      .query_row(super::statements::INTEGRITY_CHECK, [], |row| row.get(0))?;
222    if status.trim() == "ok" {
223      Ok(())
224    } else {
225      Err(Error::Integrity(status))
226    }
227  }
228
229  pub fn rebuild_merkle_index(&self) -> Result<MerkleSummary> {
230    self.ensure_merkle_enabled()?;
231    self.connection.execute(&self.layout.delete_merkle_leaves(), [])?;
232    self.connection.execute(&self.layout.delete_merkle_nodes(), [])?;
233    self.connection.execute(&self.layout.delete_merkle_meta(), [])?;
234    let hashes = self.object_hashes()?;
235    let leaf_sql = self.layout.insert_merkle_leaf();
236    let mut leaf_statement = self.connection.prepare_cached(&leaf_sql)?;
237    for (position, hash) in hashes.iter().enumerate() {
238      leaf_statement.execute(params![hash.as_bytes().as_slice(), position as i64])?;
239    }
240    let mut current = hashes.clone();
241    let node_sql = self.layout.insert_merkle_node();
242    let mut node_statement = self.connection.prepare_cached(&node_sql)?;
243    let mut level = 0i64;
244    while current.len() > 1 {
245      let next = current
246        .chunks(2)
247        .enumerate()
248        .map(|(position, pair)| {
249          let parent = hash_pair(&pair[0], pair.get(1).unwrap_or(&pair[0]));
250          node_statement.execute(params![level, position as i64, parent.as_bytes().as_slice()])?;
251          Ok(parent)
252        })
253        .collect::<Result<Vec<_>>>()?;
254      current = next;
255      level += 1;
256    }
257    let summary = merkle_root_from_hashes(&hashes);
258    let meta_sql = self.layout.upsert_merkle_meta();
259    let mut statement = self.connection.prepare_cached(&meta_sql)?;
260    for (key, value) in [
261      ("leaf_count", summary.leaf_count.to_string()),
262      ("root_hash", summary.root.to_hex()),
263      ("algo", "sha3-256".into()),
264      ("tree", "merkle-set-v1".into()),
265      ("sorted", "byte-lex".into()),
266    ] {
267      statement.execute(params![key, value])?;
268    }
269    Ok(summary)
270  }
271
272  pub fn merkle_summary(&self) -> Result<MerkleSummary> {
273    self.ensure_merkle_enabled()?;
274    let leaf_count = self
275      .meta("leaf_count")?
276      .parse()
277      .map_err(|_| Error::MissingMeta("leaf_count".into()))?;
278    let root = Hash32::from_hex(&self.meta("root_hash")?).map_err(|_| Error::MissingMeta("root_hash".into()))?;
279    Ok(MerkleSummary { root, leaf_count })
280  }
281
282  pub fn recompute_merkle_summary(&self) -> Result<MerkleSummary> {
283    self.ensure_merkle_enabled()?;
284    let sql = self.layout.select_merkle_hashes();
285    let mut statement = self.connection.prepare(&sql)?;
286    let rows = statement.query_map([], |row| row.get::<_, Vec<u8>>(0))?;
287    let hashes = rows.map(|row| Hash32::from_bytes(&row?)).collect::<Result<Vec<_>>>()?;
288    Ok(merkle_root_from_hashes(&hashes))
289  }
290
291  pub fn prove_membership(&self, hash: &Hash32) -> Result<Option<MerkleProof>> {
292    let summary = self.merkle_summary()?;
293    let position_sql = self.layout.select_merkle_position();
294    let Some(mut position) = self
295      .connection
296      .query_row(&position_sql, [hash.as_bytes().as_slice()], |row| row.get::<_, i64>(0))
297      .optional()?
298    else {
299      return Ok(None);
300    };
301    let leaf_pos = position as u64;
302    let mut count = summary.leaf_count;
303    let mut current = *hash;
304    let mut level = -1i64;
305    let mut siblings = Vec::new();
306    while count > 1 {
307      let sibling_pos = (position as u64) ^ 1;
308      let sibling = if sibling_pos >= count {
309        current
310      } else {
311        self.hash_at(level, sibling_pos)?
312      };
313      siblings.push(sibling);
314      current = if position % 2 == 0 {
315        hash_pair(&current, &sibling)
316      } else {
317        hash_pair(&sibling, &current)
318      };
319      position /= 2;
320      count = count.div_ceil(2);
321      level += 1;
322    }
323    Ok(Some(MerkleProof { leaf_pos, siblings }))
324  }
325
326  pub fn prove_memberships_batch(&self, hashes: &[Hash32]) -> Result<Vec<Option<MerkleProof>>> {
327    hashes.iter().map(|hash| self.prove_membership(hash)).collect()
328  }
329
330  pub fn verify_proof(target: &Hash32, proof: &MerkleProof, expected_root: &Hash32, leaf_count: u64) -> bool {
331    if leaf_count == 0 {
332      return false;
333    }
334    let mut current = *target;
335    let mut position = proof.leaf_pos;
336    for sibling in &proof.siblings {
337      current = if position.is_multiple_of(2) {
338        hash_pair(&current, sibling)
339      } else {
340        hash_pair(sibling, &current)
341      };
342      position /= 2;
343    }
344    current == *expected_root
345  }
346
347  pub fn verified_object_hashes(&self) -> Result<Vec<Hash32>> {
348    Ok(self.all_objects()?.into_iter().map(|object| object.hash).collect())
349  }
350
351  fn object_hashes(&self) -> Result<Vec<Hash32>> {
352    let sql = self.layout.select_object_hashes();
353    let mut statement = self.connection.prepare(&sql)?;
354    statement
355      .query_map([], |row| row.get::<_, Vec<u8>>(0))?
356      .map(|row| Hash32::from_bytes(&row?))
357      .collect()
358  }
359
360  fn meta(&self, key: &str) -> Result<String> {
361    let sql = self.layout.select_merkle_meta();
362    self
363      .connection
364      .query_row(&sql, [key], |row| row.get(0))
365      .optional()?
366      .ok_or_else(|| Error::MissingMeta(key.into()))
367  }
368
369  fn hash_at(&self, level: i64, position: u64) -> Result<Hash32> {
370    let (sql, values): (String, Vec<i64>) = if level < 0 {
371      (self.layout.select_leaf_hash_at(), vec![position as i64])
372    } else {
373      (self.layout.select_node_hash_at(), vec![level, position as i64])
374    };
375    let bytes: Vec<u8> = self.connection.query_row(&sql, params_from_iter(values), |row| row.get(0))?;
376    Hash32::from_bytes(&bytes)
377  }
378
379  fn ensure_merkle_enabled(&self) -> Result<()> {
380    if self.merkle == MerkleMode::Enabled {
381      Ok(())
382    } else {
383      Err(Error::Integrity("merkle disabled for sqlite store".into()))
384    }
385  }
386}
387
388impl ObjectSource for RusqliteStore<'_> {
389  fn read_object(&self, hash: &Hash32) -> Result<Option<VerifiedObject>> {
390    let sql = self.layout.select_object();
391    self
392      .connection
393      .query_row(&sql, [hash.as_bytes().as_slice()], stored_row)
394      .optional()?
395      .map(verified_object_from_row)
396      .transpose()
397  }
398}
399
400fn stored_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<StoredObjectRow> {
401  let hash = row.get::<_, Vec<u8>>(0)?;
402  Ok(StoredObjectRow {
403    hash: Hash32::from_bytes(&hash).map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?,
404    kind: row.get(1)?,
405    decoded_len: row.get(2)?,
406    codec: row.get(3)?,
407    stored_bytes: row.get(4)?,
408  })
409}
410
411fn unix_timestamp() -> i64 {
412  std::time::UNIX_EPOCH.elapsed().unwrap_or_default().as_secs() as i64
413}
assetpack_core/sqlite/rusqlite.rs

assetpack_core/sqlite/
rusqlite.rs
