1use codemem_core::{CodememError, MemoryNode, MemoryType};
6use rusqlite::Connection;
7use sha2::{Digest, Sha256};
8use std::collections::HashMap;
9use std::path::Path;
10use std::sync::Mutex;
11
12mod backend;
13mod graph_persistence;
14mod memory;
15mod migrations;
16mod queries;
17
18pub struct Storage {
23 conn: Mutex<Connection>,
24}
25
26impl Storage {
27 pub(crate) fn conn(&self) -> std::sync::MutexGuard<'_, Connection> {
29 self.conn.lock().expect("Storage mutex poisoned")
30 }
31
32 fn apply_pragmas(conn: &Connection) -> Result<(), CodememError> {
34 conn.pragma_update(None, "journal_mode", "WAL")
36 .map_err(|e| CodememError::Storage(e.to_string()))?;
37 conn.pragma_update(None, "cache_size", -64000i64)
39 .map_err(|e| CodememError::Storage(e.to_string()))?;
40 conn.pragma_update(None, "foreign_keys", "ON")
42 .map_err(|e| CodememError::Storage(e.to_string()))?;
43 conn.pragma_update(None, "synchronous", "NORMAL")
45 .map_err(|e| CodememError::Storage(e.to_string()))?;
46 conn.pragma_update(None, "mmap_size", 268435456i64)
48 .map_err(|e| CodememError::Storage(e.to_string()))?;
49 conn.pragma_update(None, "temp_store", "MEMORY")
51 .map_err(|e| CodememError::Storage(e.to_string()))?;
52 conn.busy_timeout(std::time::Duration::from_secs(5))
54 .map_err(|e| CodememError::Storage(e.to_string()))?;
55 Ok(())
56 }
57
58 pub fn open(path: &Path) -> Result<Self, CodememError> {
60 let conn = Connection::open(path).map_err(|e| CodememError::Storage(e.to_string()))?;
61 Self::apply_pragmas(&conn)?;
62 migrations::run_migrations(&conn)?;
63 Ok(Self {
64 conn: Mutex::new(conn),
65 })
66 }
67
68 pub fn open_without_migrations(path: &Path) -> Result<Self, CodememError> {
74 let conn = Connection::open(path).map_err(|e| CodememError::Storage(e.to_string()))?;
75 Self::apply_pragmas(&conn)?;
76 Ok(Self {
77 conn: Mutex::new(conn),
78 })
79 }
80
81 pub fn open_in_memory() -> Result<Self, CodememError> {
83 let conn =
84 Connection::open_in_memory().map_err(|e| CodememError::Storage(e.to_string()))?;
85 conn.pragma_update(None, "foreign_keys", "ON")
86 .map_err(|e| CodememError::Storage(e.to_string()))?;
87 migrations::run_migrations(&conn)?;
88 Ok(Self {
89 conn: Mutex::new(conn),
90 })
91 }
92
93 pub fn content_hash(content: &str) -> String {
95 let mut hasher = Sha256::new();
96 hasher.update(content.as_bytes());
97 format!("{:x}", hasher.finalize())
98 }
99}
100
101pub(crate) struct MemoryRow {
103 pub(crate) id: String,
104 pub(crate) content: String,
105 pub(crate) memory_type: String,
106 pub(crate) importance: f64,
107 pub(crate) confidence: f64,
108 pub(crate) access_count: i64,
109 pub(crate) content_hash: String,
110 pub(crate) tags: String,
111 pub(crate) metadata: String,
112 pub(crate) namespace: Option<String>,
113 pub(crate) created_at: i64,
114 pub(crate) updated_at: i64,
115 pub(crate) last_accessed_at: i64,
116}
117
118impl MemoryRow {
119 pub(crate) fn into_memory_node(self) -> Result<MemoryNode, CodememError> {
120 let memory_type: MemoryType = self.memory_type.parse()?;
121 let tags: Vec<String> = serde_json::from_str(&self.tags).unwrap_or_default();
122 let metadata: HashMap<String, serde_json::Value> =
123 serde_json::from_str(&self.metadata).unwrap_or_default();
124
125 let created_at = chrono::DateTime::from_timestamp(self.created_at, 0)
126 .unwrap_or_default()
127 .with_timezone(&chrono::Utc);
128 let updated_at = chrono::DateTime::from_timestamp(self.updated_at, 0)
129 .unwrap_or_default()
130 .with_timezone(&chrono::Utc);
131 let last_accessed_at = chrono::DateTime::from_timestamp(self.last_accessed_at, 0)
132 .unwrap_or_default()
133 .with_timezone(&chrono::Utc);
134
135 Ok(MemoryNode {
136 id: self.id,
137 content: self.content,
138 memory_type,
139 importance: self.importance,
140 confidence: self.confidence,
141 access_count: self.access_count as u32,
142 content_hash: self.content_hash,
143 tags,
144 metadata,
145 namespace: self.namespace,
146 created_at,
147 updated_at,
148 last_accessed_at,
149 })
150 }
151}