iuna

chain_store.rs (5173B)

---

 use std::{
     fs,
     path::{Path, PathBuf},
     time::{SystemTime, UNIX_EPOCH},
 };
 
 use anyhow::{Context, Result};
 use rusqlite::{Connection, OptionalExtension, params};
 
 use crate::domain::ChainSnapshot;
 
 mod compact;
 use compact::{decode_compact_snapshot, encode_compact_snapshot};
 
 const SCHEMA: &str = r#"
 CREATE TABLE IF NOT EXISTS chain_snapshots (
     id INTEGER PRIMARY KEY CHECK (id = 1),
     height INTEGER NOT NULL,
     tip_hash TEXT NOT NULL,
     snapshot_blob BLOB NOT NULL,
     updated_at_ms INTEGER NOT NULL
 );
 "#;
 
 #[derive(Clone, Debug)]
 pub struct SqliteChainStore {
     path: PathBuf,
 }
 
 impl SqliteChainStore {
     pub fn open(path: impl AsRef<Path>) -> Result<Self> {
         let path = path.as_ref().to_path_buf();
         if let Some(parent) = path.parent() {
             fs::create_dir_all(parent).with_context(|| {
                 format!(
                     "failed to create chain database directory {}",
                     parent.display()
                 )
             })?;
         }
 
         let store = Self { path };
         store.with_connection_mut(|connection| {
             connection
                 .execute_batch(SCHEMA)
                 .context("failed to initialize chain database schema")?;
             Ok(())
         })?;
         Ok(store)
     }
 
     pub fn path(&self) -> &Path {
         &self.path
     }
 
     pub fn load(&self) -> Result<Option<ChainSnapshot>> {
         self.with_connection(|connection| {
             let snapshot_blob = connection
                 .query_row(
                     "SELECT snapshot_blob FROM chain_snapshots WHERE id = 1",
                     [],
                     |row| row.get::<_, Vec<u8>>(0),
                 )
                 .optional()
                 .context("failed to load chain snapshot from database")?;
 
             snapshot_blob
                 .map(|blob| {
                     decode_compact_snapshot(&blob)
                         .context("failed to parse compact chain snapshot from database")
                 })
                 .transpose()
         })
     }
 
     pub fn save(&self, snapshot: &ChainSnapshot) -> Result<()> {
         let (height, tip_hash) = snapshot_tip(snapshot).context("cannot persist empty chain")?;
         let snapshot_blob =
             encode_compact_snapshot(snapshot).context("failed to encode compact chain snapshot")?;
         let updated_at_ms = unix_ms();
 
         self.with_connection_mut(|connection| {
             let transaction = connection
                 .transaction()
                 .context("failed to start chain persistence transaction")?;
             transaction
                 .execute(
                     r#"
 INSERT INTO chain_snapshots (id, height, tip_hash, snapshot_blob, updated_at_ms)
 VALUES (1, ?1, ?2, ?3, ?4)
 ON CONFLICT(id) DO UPDATE SET
     height = excluded.height,
     tip_hash = excluded.tip_hash,
     snapshot_blob = excluded.snapshot_blob,
     updated_at_ms = excluded.updated_at_ms
 "#,
                     params![height, tip_hash, snapshot_blob, updated_at_ms],
                 )
                 .context("failed to persist chain snapshot")?;
             transaction
                 .commit()
                 .context("failed to commit chain persistence transaction")?;
             Ok(())
         })
     }
 
     pub fn clear_chain(&self) -> Result<()> {
         self.with_connection_mut(|connection| {
             let transaction = connection
                 .transaction()
                 .context("failed to start chain reset transaction")?;
             transaction
                 .execute("DELETE FROM chain_snapshots", [])
                 .context("failed to delete chain snapshot")?;
             transaction
                 .commit()
                 .context("failed to commit chain reset transaction")?;
             Ok(())
         })
     }
 
     fn with_connection<T>(&self, work: impl FnOnce(&Connection) -> Result<T>) -> Result<T> {
         let connection = self.open_connection()?;
         connection
             .execute_batch(
                 r#"
 PRAGMA busy_timeout = 5000;
 PRAGMA synchronous = NORMAL;
 "#,
             )
             .context("failed to configure chain database connection")?;
         work(&connection)
     }
 
     fn with_connection_mut<T>(&self, work: impl FnOnce(&mut Connection) -> Result<T>) -> Result<T> {
         let mut connection = self.open_connection()?;
         connection
             .execute_batch(
                 r#"
 PRAGMA journal_mode = WAL;
 PRAGMA busy_timeout = 5000;
 PRAGMA synchronous = NORMAL;
 "#,
             )
             .context("failed to configure chain database connection")?;
         work(&mut connection)
     }
 
     fn open_connection(&self) -> Result<Connection> {
         Connection::open(&self.path)
             .with_context(|| format!("failed to open chain database {}", self.path.display()))
     }
 }
 
 fn snapshot_tip(snapshot: &ChainSnapshot) -> Option<(u64, String)> {
     snapshot
         .blocks
         .last()
         .map(|block| (block.height, block.hash.clone()))
 }
 
 fn unix_ms() -> u64 {
     SystemTime::now()
         .duration_since(UNIX_EPOCH)
         .unwrap_or_default()
         .as_millis() as u64
 }
 
 #[cfg(test)]
 mod tests;