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
56fn to_fts5_match(query: &str) -> Option<String> {
66 let terms: Vec<String> = query
67 .split(|c: char| !c.is_alphanumeric())
68 .filter(|term| !term.is_empty())
69 .map(|term| format!("\"{}\"", term.replace('"', "\"\"")))
70 .collect();
71
72 if terms.is_empty() {
73 None
74 } else {
75 Some(terms.join(" OR "))
76 }
77}
78
79impl Searcher for SqliteSearcher {
80 fn search(&self, query: &str, scope: Option<&str>, limit: usize) -> Result<Vec<ScoredEntry>> {
81 if query == MATCH_ALL_QUERY {
82 return self.search_all(scope, limit);
83 }
84
85 let Some(match_expr) = to_fts5_match(query) else {
86 return Ok(Vec::new());
87 };
88
89 let conn = self.pool.get()?;
90
91 let mut stmt = conn.prepare(
92 "SELECT e.id, e.content, e.timestamp, e.kind, e.scope, e.session_id, \
93 e.token_count, e.metadata, bm25(entries_fts) AS score \
94 FROM entries_fts f \
95 JOIN entries e ON e.rowid = f.rowid \
96 WHERE entries_fts MATCH ?1 AND (?2 IS NULL OR e.scope = ?2) \
97 ORDER BY score \
98 LIMIT ?3",
99 )?;
100
101 let results = stmt
102 .query_map(
103 rusqlite::params![match_expr, scope, i64::try_from(limit).unwrap_or(i64::MAX)],
104 |row| {
105 let entry = row_to_entry(row)?;
106 let raw_score: f64 = row.get("score")?;
107
108 Ok(ScoredEntry {
109 entry,
110 score: -raw_score,
112 })
113 },
114 )?
115 .collect::<std::result::Result<Vec<_>, _>>()?;
116
117 Ok(results)
118 }
119}
120
121#[cfg(test)]
122mod tests {
123 use super::to_fts5_match;
124
125 #[test]
126 fn single_term() {
127 assert_eq!(to_fts5_match("marco"), Some("\"marco\"".to_owned()));
128 }
129
130 #[test]
131 fn multi_word_with_punctuation() {
132 let input = r#"if I say "marco" you say "polo"."#;
133 let expected = "\"if\" OR \"I\" OR \"say\" OR \"marco\" OR \"you\" OR \"say\" OR \"polo\"";
134 assert_eq!(to_fts5_match(input), Some(expected.to_owned()));
135 }
136
137 #[test]
138 fn punctuation_only_returns_none() {
139 assert_eq!(to_fts5_match("...:-*"), None);
140 }
141
142 #[test]
143 fn empty_returns_none() {
144 assert_eq!(to_fts5_match(""), None);
145 }
146
147 #[test]
148 fn hyphen_and_colon_split_terms() {
149 assert_eq!(
150 to_fts5_match("well-known scope:test"),
151 Some("\"well\" OR \"known\" OR \"scope\" OR \"test\"".to_owned())
152 );
153 }
154
155 #[test]
156 fn unicode_term_survives_as_single_term() {
157 assert_eq!(to_fts5_match("café"), Some("\"café\"".to_owned()));
158 }
159}