context_forge/storage/
searcher.rs1use std::sync::Arc;
2
3use r2d2::Pool;
4use r2d2_sqlite::SqliteConnectionManager;
5
6use crate::engine::MATCH_ALL_QUERY;
7use crate::entry::ScoredEntry;
8use crate::storage::schema::row_to_entry;
9use crate::traits::Searcher;
10use crate::Result;
11
12pub struct SqliteSearcher {
14 pool: Arc<Pool<SqliteConnectionManager>>,
15}
16
17impl SqliteSearcher {
18 #[must_use]
20 pub fn new(pool: Arc<Pool<SqliteConnectionManager>>) -> Self {
21 Self { pool }
22 }
23
24 fn search_all(&self, scope: Option<&str>, limit: usize) -> Result<Vec<ScoredEntry>> {
32 let conn = self.pool.get()?;
33
34 let mut stmt = conn.prepare(
35 "SELECT id, content, timestamp, kind, scope, session_id, token_count, metadata \
36 FROM entries \
37 WHERE (?1 IS NULL OR scope = ?1) \
38 ORDER BY timestamp DESC \
39 LIMIT ?2",
40 )?;
41
42 let results = stmt
43 .query_map(
44 rusqlite::params![scope, i64::try_from(limit).unwrap_or(i64::MAX)],
45 |row| {
46 let entry = row_to_entry(row)?;
47 Ok(ScoredEntry { score: 1.0, entry })
48 },
49 )?
50 .collect::<std::result::Result<Vec<_>, _>>()?;
51
52 Ok(results)
53 }
54}
55
56impl Searcher for SqliteSearcher {
57 fn search(&self, query: &str, scope: Option<&str>, limit: usize) -> Result<Vec<ScoredEntry>> {
58 if query == MATCH_ALL_QUERY {
59 return self.search_all(scope, limit);
60 }
61
62 let conn = self.pool.get()?;
63
64 let mut stmt = conn.prepare(
65 "SELECT e.id, e.content, e.timestamp, e.kind, e.scope, e.session_id, \
66 e.token_count, e.metadata, bm25(entries_fts) AS score \
67 FROM entries_fts f \
68 JOIN entries e ON e.rowid = f.rowid \
69 WHERE entries_fts MATCH ?1 AND (?2 IS NULL OR e.scope = ?2) \
70 ORDER BY score \
71 LIMIT ?3",
72 )?;
73
74 let results = stmt
75 .query_map(
76 rusqlite::params![query, scope, i64::try_from(limit).unwrap_or(i64::MAX)],
77 |row| {
78 let entry = row_to_entry(row)?;
79 let raw_score: f64 = row.get("score")?;
80
81 Ok(ScoredEntry {
82 entry,
83 score: -raw_score,
85 })
86 },
87 )?
88 .collect::<std::result::Result<Vec<_>, _>>()?;
89
90 Ok(results)
91 }
92}