Skip to main content

codemem_storage/
lib.rs

1//! codemem-storage: SQLite persistence layer for Codemem.
2//!
3//! Uses rusqlite with bundled SQLite, WAL mode, and embedded schema.
4
5use 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
18/// SQLite-backed storage for Codemem memories, embeddings, and graph data.
19///
20/// Wraps `rusqlite::Connection` in a `Mutex` to satisfy `Send + Sync` bounds
21/// required by the `StorageBackend` trait.
22pub struct Storage {
23    conn: Mutex<Connection>,
24}
25
26impl Storage {
27    /// Get a lock on the underlying connection.
28    pub(crate) fn conn(&self) -> std::sync::MutexGuard<'_, Connection> {
29        self.conn.lock().expect("Storage mutex poisoned")
30    }
31
32    /// Apply standard pragmas to a connection.
33    fn apply_pragmas(conn: &Connection) -> Result<(), CodememError> {
34        // WAL mode for concurrent reads
35        conn.pragma_update(None, "journal_mode", "WAL")
36            .map_err(|e| CodememError::Storage(e.to_string()))?;
37        // 64MB cache
38        conn.pragma_update(None, "cache_size", -64000i64)
39            .map_err(|e| CodememError::Storage(e.to_string()))?;
40        // Foreign keys ON
41        conn.pragma_update(None, "foreign_keys", "ON")
42            .map_err(|e| CodememError::Storage(e.to_string()))?;
43        // NORMAL sync (good balance of safety vs speed)
44        conn.pragma_update(None, "synchronous", "NORMAL")
45            .map_err(|e| CodememError::Storage(e.to_string()))?;
46        // 256MB mmap for faster reads
47        conn.pragma_update(None, "mmap_size", 268435456i64)
48            .map_err(|e| CodememError::Storage(e.to_string()))?;
49        // Temp tables in memory
50        conn.pragma_update(None, "temp_store", "MEMORY")
51            .map_err(|e| CodememError::Storage(e.to_string()))?;
52        // 5s busy timeout
53        conn.busy_timeout(std::time::Duration::from_secs(5))
54            .map_err(|e| CodememError::Storage(e.to_string()))?;
55        Ok(())
56    }
57
58    /// Open (or create) a Codemem database at the given path.
59    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    /// Open an existing database without running migrations.
69    ///
70    /// Use this in lifecycle hooks (context, prompt, summarize) where the
71    /// database has already been migrated by `codemem init` or `codemem serve`,
72    /// to avoid SQLITE_BUSY race conditions with the concurrent MCP server.
73    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    /// Open an in-memory database (for testing).
82    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    /// Compute SHA-256 hash of content for deduplication.
94    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
101/// Internal row struct for memory deserialization.
102pub(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}