Skip to main content

double_o/
store.rs

1use std::path::PathBuf;
2
3use rusqlite::Connection;
4use serde::{Deserialize, Serialize};
5
6use crate::error::Error;
7use crate::util;
8
9// ---------------------------------------------------------------------------
10// Types
11// ---------------------------------------------------------------------------
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct SessionMeta {
15    pub source: String,
16    pub session: String,
17    pub command: String,
18    pub timestamp: i64,
19}
20
21#[derive(Debug)]
22#[allow(dead_code)]
23pub struct SearchResult {
24    pub id: String,
25    pub content: String,
26    pub meta: Option<SessionMeta>,
27    pub similarity: Option<f64>,
28}
29
30// ---------------------------------------------------------------------------
31// Store trait
32// ---------------------------------------------------------------------------
33
34pub trait Store {
35    fn index(
36        &mut self,
37        project_id: &str,
38        content: &str,
39        meta: &SessionMeta,
40    ) -> Result<String, Error>;
41
42    fn search(
43        &mut self,
44        project_id: &str,
45        query: &str,
46        limit: usize,
47    ) -> Result<Vec<SearchResult>, Error>;
48
49    fn delete_by_session(&mut self, project_id: &str, session_id: &str) -> Result<usize, Error>;
50
51    fn cleanup_stale(&mut self, project_id: &str, max_age_secs: i64) -> Result<usize, Error>;
52}
53
54// ---------------------------------------------------------------------------
55// SqliteStore — default backend using FTS5 for text search
56// ---------------------------------------------------------------------------
57
58pub struct SqliteStore {
59    conn: Connection,
60}
61
62fn db_path() -> PathBuf {
63    dirs::data_dir()
64        .or_else(dirs::home_dir)
65        .unwrap_or_else(|| PathBuf::from("/tmp"))
66        .join(".oo")
67        .join("oo.db")
68}
69
70fn map_err(e: rusqlite::Error) -> Error {
71    Error::Store(e.to_string())
72}
73
74impl SqliteStore {
75    pub fn open() -> Result<Self, Error> {
76        Self::open_at(&db_path())
77    }
78
79    pub fn open_at(path: &std::path::Path) -> Result<Self, Error> {
80        if let Some(parent) = path.parent() {
81            std::fs::create_dir_all(parent).map_err(|e| Error::Store(e.to_string()))?;
82        }
83        let conn = Connection::open(path).map_err(map_err)?;
84        conn.execute_batch(
85            "CREATE TABLE IF NOT EXISTS entries (
86                id       TEXT PRIMARY KEY,
87                project  TEXT NOT NULL,
88                content  TEXT NOT NULL,
89                metadata TEXT,
90                created  INTEGER NOT NULL
91            );
92            CREATE VIRTUAL TABLE IF NOT EXISTS entries_fts USING fts5(
93                content,
94                content='entries',
95                content_rowid='rowid'
96            );
97            CREATE TRIGGER IF NOT EXISTS entries_ai AFTER INSERT ON entries BEGIN
98                INSERT INTO entries_fts(rowid, content)
99                VALUES (new.rowid, new.content);
100            END;
101            CREATE TRIGGER IF NOT EXISTS entries_ad AFTER DELETE ON entries BEGIN
102                INSERT INTO entries_fts(entries_fts, rowid, content)
103                VALUES ('delete', old.rowid, old.content);
104            END;
105            CREATE TRIGGER IF NOT EXISTS entries_au AFTER UPDATE ON entries BEGIN
106                INSERT INTO entries_fts(entries_fts, rowid, content)
107                VALUES ('delete', old.rowid, old.content);
108                INSERT INTO entries_fts(rowid, content)
109                VALUES (new.rowid, new.content);
110            END;",
111        )
112        .map_err(map_err)?;
113        Ok(Self { conn })
114    }
115}
116
117impl Store for SqliteStore {
118    fn index(
119        &mut self,
120        project_id: &str,
121        content: &str,
122        meta: &SessionMeta,
123    ) -> Result<String, Error> {
124        let id = uuid::Uuid::new_v4().to_string();
125        let meta_json = serde_json::to_string(meta).map_err(|e| Error::Store(e.to_string()))?;
126        self.conn
127            .execute(
128                "INSERT INTO entries (id, project, content, metadata, created)
129                 VALUES (?1, ?2, ?3, ?4, ?5)",
130                rusqlite::params![id, project_id, content, meta_json, meta.timestamp],
131            )
132            .map_err(map_err)?;
133        Ok(id)
134    }
135
136    fn search(
137        &mut self,
138        project_id: &str,
139        query: &str,
140        limit: usize,
141    ) -> Result<Vec<SearchResult>, Error> {
142        // Use FTS5 for full-text search, fall back to LIKE if query is too short
143        let results = if query.len() >= 2 {
144            let mut stmt = self
145                .conn
146                .prepare(
147                    "SELECT e.id, e.content, e.metadata, rank
148                     FROM entries_fts f
149                     JOIN entries e ON e.rowid = f.rowid
150                     WHERE entries_fts MATCH ?1 AND e.project = ?2
151                     ORDER BY rank
152                     LIMIT ?3",
153                )
154                .map_err(map_err)?;
155
156            // FTS5 query: wrap tokens for prefix matching
157            let fts_query = query
158                .split_whitespace()
159                .map(|w| format!("\"{w}\""))
160                .collect::<Vec<_>>()
161                .join(" ");
162
163            stmt.query_map(rusqlite::params![fts_query, project_id, limit], |row| {
164                let id: String = row.get(0)?;
165                let content: String = row.get(1)?;
166                let meta_json: Option<String> = row.get(2)?;
167                let rank: f64 = row.get(3)?;
168                Ok(SearchResult {
169                    id,
170                    content,
171                    meta: meta_json.as_deref().and_then(parse_meta),
172                    similarity: Some(-rank), // FTS5 rank is negative
173                })
174            })
175            .map_err(map_err)?
176            .filter_map(|r| r.ok())
177            .collect()
178        } else {
179            let mut stmt = self
180                .conn
181                .prepare(
182                    "SELECT id, content, metadata
183                     FROM entries
184                     WHERE project = ?1 AND content LIKE ?2
185                     ORDER BY created DESC
186                     LIMIT ?3",
187                )
188                .map_err(map_err)?;
189
190            let like = format!("%{query}%");
191            stmt.query_map(rusqlite::params![project_id, like, limit], |row| {
192                let id: String = row.get(0)?;
193                let content: String = row.get(1)?;
194                let meta_json: Option<String> = row.get(2)?;
195                Ok(SearchResult {
196                    id,
197                    content,
198                    meta: meta_json.as_deref().and_then(parse_meta),
199                    similarity: None,
200                })
201            })
202            .map_err(map_err)?
203            .filter_map(|r| r.ok())
204            .collect()
205        };
206
207        Ok(results)
208    }
209
210    fn delete_by_session(&mut self, project_id: &str, session_id: &str) -> Result<usize, Error> {
211        // Find entries matching this session
212        let ids: Vec<String> = {
213            let mut stmt = self
214                .conn
215                .prepare("SELECT id, metadata FROM entries WHERE project = ?1")
216                .map_err(map_err)?;
217            stmt.query_map(rusqlite::params![project_id], |row| {
218                let id: String = row.get(0)?;
219                let meta_json: Option<String> = row.get(1)?;
220                Ok((id, meta_json))
221            })
222            .map_err(map_err)?
223            .filter_map(|r| r.ok())
224            .filter(|(_, meta_json)| {
225                meta_json
226                    .as_deref()
227                    .and_then(parse_meta)
228                    .is_some_and(|m| m.source == "oo" && m.session == session_id)
229            })
230            .map(|(id, _)| id)
231            .collect()
232        };
233
234        let count = ids.len();
235        for id in &ids {
236            self.conn
237                .execute("DELETE FROM entries WHERE id = ?1", rusqlite::params![id])
238                .map_err(map_err)?;
239        }
240        Ok(count)
241    }
242
243    fn cleanup_stale(&mut self, project_id: &str, max_age_secs: i64) -> Result<usize, Error> {
244        let now = util::now_epoch();
245        let ids: Vec<String> = {
246            let mut stmt = self
247                .conn
248                .prepare("SELECT id, metadata FROM entries WHERE project = ?1")
249                .map_err(map_err)?;
250            stmt.query_map(rusqlite::params![project_id], |row| {
251                let id: String = row.get(0)?;
252                let meta_json: Option<String> = row.get(1)?;
253                Ok((id, meta_json))
254            })
255            .map_err(map_err)?
256            .filter_map(|r| r.ok())
257            .filter(|(_, meta_json)| {
258                meta_json
259                    .as_deref()
260                    .and_then(parse_meta)
261                    .is_some_and(|m| m.source == "oo" && (now - m.timestamp) > max_age_secs)
262            })
263            .map(|(id, _)| id)
264            .collect()
265        };
266
267        let count = ids.len();
268        for id in &ids {
269            self.conn
270                .execute("DELETE FROM entries WHERE id = ?1", rusqlite::params![id])
271                .map_err(map_err)?;
272        }
273        Ok(count)
274    }
275}
276
277// ---------------------------------------------------------------------------
278// VipuneStore — optional backend with semantic search
279// ---------------------------------------------------------------------------
280
281#[cfg(feature = "vipune-store")]
282pub struct VipuneStore {
283    store: vipune::MemoryStore,
284}
285
286#[cfg(feature = "vipune-store")]
287impl VipuneStore {
288    pub fn open() -> Result<Self, Error> {
289        let config = vipune::Config::load().map_err(|e| Error::Store(e.to_string()))?;
290        let store =
291            vipune::MemoryStore::new(&config.database_path, &config.embedding_model, config)
292                .map_err(|e| Error::Store(e.to_string()))?;
293        Ok(Self { store })
294    }
295}
296
297#[cfg(feature = "vipune-store")]
298impl Store for VipuneStore {
299    fn index(
300        &mut self,
301        project_id: &str,
302        content: &str,
303        meta: &SessionMeta,
304    ) -> Result<String, Error> {
305        let meta_json = serde_json::to_string(meta).map_err(|e| Error::Store(e.to_string()))?;
306        match self
307            .store
308            .add_with_conflict(project_id, content, Some(&meta_json), true)
309        {
310            Ok(vipune::AddResult::Added { id }) => Ok(id),
311            Ok(vipune::AddResult::Conflicts { .. }) => Ok(String::new()),
312            Err(e) => Err(Error::Store(e.to_string())),
313        }
314    }
315
316    fn search(
317        &mut self,
318        project_id: &str,
319        query: &str,
320        limit: usize,
321    ) -> Result<Vec<SearchResult>, Error> {
322        let memories = self
323            .store
324            .search_hybrid(project_id, query, limit, 0.3)
325            .map_err(|e| Error::Store(e.to_string()))?;
326        Ok(memories
327            .into_iter()
328            .map(|m| SearchResult {
329                id: m.id,
330                meta: m.metadata.as_deref().and_then(parse_meta),
331                content: m.content,
332                similarity: m.similarity,
333            })
334            .collect())
335    }
336
337    fn delete_by_session(&mut self, project_id: &str, session_id: &str) -> Result<usize, Error> {
338        let entries = self
339            .store
340            .list(project_id, 10_000)
341            .map_err(|e| Error::Store(e.to_string()))?;
342        let mut count = 0;
343        for entry in entries {
344            if let Some(meta) = entry.metadata.as_deref().and_then(parse_meta) {
345                if meta.source == "oo" && meta.session == session_id {
346                    self.store
347                        .delete(&entry.id)
348                        .map_err(|e| Error::Store(e.to_string()))?;
349                    count += 1;
350                }
351            }
352        }
353        Ok(count)
354    }
355
356    fn cleanup_stale(&mut self, project_id: &str, max_age_secs: i64) -> Result<usize, Error> {
357        let now = util::now_epoch();
358        let entries = self
359            .store
360            .list(project_id, 10_000)
361            .map_err(|e| Error::Store(e.to_string()))?;
362        let mut count = 0;
363        for entry in entries {
364            if let Some(meta) = entry.metadata.as_deref().and_then(parse_meta) {
365                if meta.source == "oo" && (now - meta.timestamp) > max_age_secs {
366                    self.store
367                        .delete(&entry.id)
368                        .map_err(|e| Error::Store(e.to_string()))?;
369                    count += 1;
370                }
371            }
372        }
373        Ok(count)
374    }
375}
376
377// ---------------------------------------------------------------------------
378// Helpers
379// ---------------------------------------------------------------------------
380
381fn parse_meta(json: &str) -> Option<SessionMeta> {
382    serde_json::from_str(json).ok()
383}
384
385/// Open the default store (SqliteStore, or VipuneStore if feature-enabled).
386pub fn open() -> Result<Box<dyn Store>, Error> {
387    #[cfg(feature = "vipune-store")]
388    {
389        return Ok(Box::new(VipuneStore::open()?));
390    }
391    #[cfg(not(feature = "vipune-store"))]
392    {
393        Ok(Box::new(SqliteStore::open()?))
394    }
395}
396
397// ---------------------------------------------------------------------------
398// Tests
399// ---------------------------------------------------------------------------
400
401#[cfg(test)]
402#[path = "store_tests.rs"]
403mod tests;