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 pub fn open(path: &Path) -> Result<Self, CodememError> {
34 let conn = Connection::open(path).map_err(|e| CodememError::Storage(e.to_string()))?;
35
36 conn.pragma_update(None, "journal_mode", "WAL")
38 .map_err(|e| CodememError::Storage(e.to_string()))?;
39 conn.pragma_update(None, "cache_size", -64000i64)
41 .map_err(|e| CodememError::Storage(e.to_string()))?;
42 conn.pragma_update(None, "foreign_keys", "ON")
44 .map_err(|e| CodememError::Storage(e.to_string()))?;
45 conn.pragma_update(None, "synchronous", "NORMAL")
47 .map_err(|e| CodememError::Storage(e.to_string()))?;
48 conn.pragma_update(None, "mmap_size", 268435456i64)
50 .map_err(|e| CodememError::Storage(e.to_string()))?;
51 conn.pragma_update(None, "temp_store", "MEMORY")
53 .map_err(|e| CodememError::Storage(e.to_string()))?;
54 conn.busy_timeout(std::time::Duration::from_secs(5))
56 .map_err(|e| CodememError::Storage(e.to_string()))?;
57
58 migrations::run_migrations(&conn)?;
60
61 Ok(Self {
62 conn: Mutex::new(conn),
63 })
64 }
65
66 pub fn open_in_memory() -> Result<Self, CodememError> {
68 let conn =
69 Connection::open_in_memory().map_err(|e| CodememError::Storage(e.to_string()))?;
70 conn.pragma_update(None, "foreign_keys", "ON")
71 .map_err(|e| CodememError::Storage(e.to_string()))?;
72 migrations::run_migrations(&conn)?;
73 Ok(Self {
74 conn: Mutex::new(conn),
75 })
76 }
77
78 pub fn content_hash(content: &str) -> String {
80 let mut hasher = Sha256::new();
81 hasher.update(content.as_bytes());
82 format!("{:x}", hasher.finalize())
83 }
84}
85
86pub(crate) struct MemoryRow {
88 pub(crate) id: String,
89 pub(crate) content: String,
90 pub(crate) memory_type: String,
91 pub(crate) importance: f64,
92 pub(crate) confidence: f64,
93 pub(crate) access_count: i64,
94 pub(crate) content_hash: String,
95 pub(crate) tags: String,
96 pub(crate) metadata: String,
97 pub(crate) namespace: Option<String>,
98 pub(crate) created_at: i64,
99 pub(crate) updated_at: i64,
100 pub(crate) last_accessed_at: i64,
101}
102
103impl MemoryRow {
104 pub(crate) fn into_memory_node(self) -> Result<MemoryNode, CodememError> {
105 let memory_type: MemoryType = self.memory_type.parse()?;
106 let tags: Vec<String> = serde_json::from_str(&self.tags).unwrap_or_default();
107 let metadata: HashMap<String, serde_json::Value> =
108 serde_json::from_str(&self.metadata).unwrap_or_default();
109
110 let created_at = chrono::DateTime::from_timestamp(self.created_at, 0)
111 .unwrap_or_default()
112 .with_timezone(&chrono::Utc);
113 let updated_at = chrono::DateTime::from_timestamp(self.updated_at, 0)
114 .unwrap_or_default()
115 .with_timezone(&chrono::Utc);
116 let last_accessed_at = chrono::DateTime::from_timestamp(self.last_accessed_at, 0)
117 .unwrap_or_default()
118 .with_timezone(&chrono::Utc);
119
120 Ok(MemoryNode {
121 id: self.id,
122 content: self.content,
123 memory_type,
124 importance: self.importance,
125 confidence: self.confidence,
126 access_count: self.access_count as u32,
127 content_hash: self.content_hash,
128 tags,
129 metadata,
130 namespace: self.namespace,
131 created_at,
132 updated_at,
133 last_accessed_at,
134 })
135 }
136}