do_memory_storage_redb/
heuristics.rs1use 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 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 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 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}