Skip to main content

do_memory_storage_redb/
heuristics.rs

1//! Heuristic storage operations for redb cache
2
3use crate::{Error, HEURISTICS_TABLE, RedbStorage};
4use do_memory_core::{Heuristic, Result};
5use redb::{ReadableDatabase, ReadableTable};
6use std::sync::Arc;
7use tracing::debug;
8use tracing::info;
9use uuid::Uuid;
10
11use crate::episodes::RedbQuery;
12
13impl RedbStorage {
14    /// Store a heuristic in cache
15    pub async fn store_heuristic(&self, heuristic: &Heuristic) -> Result<()> {
16        debug!("Storing heuristic in cache: {}", heuristic.heuristic_id);
17        let db = Arc::clone(&self.db);
18        let heuristic_id = heuristic.heuristic_id.to_string();
19        let heuristic_bytes = postcard::to_allocvec(heuristic)
20            .map_err(|e| Error::Storage(format!("Failed to serialize heuristic: {}", e)))?;
21
22        tokio::task::spawn_blocking(move || {
23            let write_txn = db
24                .begin_write()
25                .map_err(|e| Error::Storage(format!("Failed to begin write transaction: {}", e)))?;
26
27            {
28                let mut table = write_txn.open_table(HEURISTICS_TABLE).map_err(|e| {
29                    Error::Storage(format!("Failed to open heuristics table: {}", e))
30                })?;
31
32                table
33                    .insert(heuristic_id.as_str(), heuristic_bytes.as_slice())
34                    .map_err(|e| Error::Storage(format!("Failed to insert heuristic: {}", e)))?;
35            }
36
37            write_txn
38                .commit()
39                .map_err(|e| Error::Storage(format!("Failed to commit transaction: {}", e)))?;
40
41            Ok::<(), Error>(())
42        })
43        .await
44        .map_err(|e| Error::Storage(format!("Task join error: {}", e)))??;
45
46        info!("Successfully cached heuristic: {}", heuristic.heuristic_id);
47        Ok(())
48    }
49
50    /// Retrieve a heuristic from cache
51    pub async fn get_heuristic(&self, heuristic_id: Uuid) -> Result<Option<Heuristic>> {
52        debug!("Retrieving heuristic from cache: {}", heuristic_id);
53        let db = Arc::clone(&self.db);
54        let heuristic_id_str = heuristic_id.to_string();
55
56        tokio::task::spawn_blocking(move || {
57            let read_txn = db
58                .begin_read()
59                .map_err(|e| Error::Storage(format!("Failed to begin read transaction: {}", e)))?;
60
61            let table = read_txn
62                .open_table(HEURISTICS_TABLE)
63                .map_err(|e| Error::Storage(format!("Failed to open heuristics table: {}", e)))?;
64
65            match table
66                .get(heuristic_id_str.as_str())
67                .map_err(|e| Error::Storage(format!("Failed to get heuristic: {}", e)))?
68            {
69                Some(bytes_guard) => {
70                    let _bytes = bytes_guard.value();
71                    let heuristic: Heuristic =
72                        postcard::from_bytes(bytes_guard.value()).map_err(|e| {
73                            Error::Storage(format!("Failed to deserialize heuristic: {}", e))
74                        })?;
75                    Ok(Some(heuristic))
76                }
77                None => Ok(None),
78            }
79        })
80        .await
81        .map_err(|e| Error::Storage(format!("Task join error: {}", e)))?
82    }
83
84    /// Get all heuristics from cache (with optional limit)
85    pub async fn get_all_heuristics(&self, query: &RedbQuery) -> Result<Vec<Heuristic>> {
86        debug!("Retrieving all heuristics from cache");
87        let db = Arc::clone(&self.db);
88        let limit = query.limit;
89
90        tokio::task::spawn_blocking(move || {
91            let read_txn = db
92                .begin_read()
93                .map_err(|e| Error::Storage(format!("Failed to begin read transaction: {}", e)))?;
94
95            let table = read_txn
96                .open_table(HEURISTICS_TABLE)
97                .map_err(|e| Error::Storage(format!("Failed to open heuristics table: {}", e)))?;
98
99            let mut heuristics = Vec::new();
100            let iter = table
101                .iter()
102                .map_err(|e| Error::Storage(format!("Failed to iterate heuristics: {}", e)))?;
103
104            for (count, result) in iter.enumerate() {
105                if let Some(max) = limit {
106                    if count >= max {
107                        break;
108                    }
109                }
110
111                let (_, bytes_guard) = result.map_err(|e| {
112                    Error::Storage(format!("Failed to read heuristic entry: {}", e))
113                })?;
114
115                let heuristic: Heuristic =
116                    postcard::from_bytes(bytes_guard.value()).map_err(|e| {
117                        Error::Storage(format!("Failed to deserialize heuristic: {}", e))
118                    })?;
119
120                heuristics.push(heuristic);
121            }
122
123            Ok(heuristics)
124        })
125        .await
126        .map_err(|e| Error::Storage(format!("Task join error: {}", e)))?
127    }
128}