1use anyhow::{Context, Result};
2use rusqlite::{params, Connection, OptionalExtension};
3use serde::Serialize;
4
5pub struct TopicSegmentInput<'a> {
6 pub host_id: i64,
7 pub project_id: i64,
8 pub session_row_id: i64,
9 pub project: &'a str,
10 pub topic_key: &'a str,
11 pub title: &'a str,
12 pub summary: &'a str,
13 pub status: &'a str,
14 pub segment_index: i64,
15 pub covered_from_event_id: i64,
16 pub covered_to_event_id: i64,
17 pub evidence_event_ids: &'a str,
18 pub files: Option<&'a str>,
19 pub confidence: f64,
20}
21
22#[derive(Debug, Clone, Serialize)]
23pub struct TopicTraceEntry {
24 pub id: i64,
25 pub topic_key: String,
26 pub title: String,
27 pub summary: String,
28 pub status: String,
29 pub segment_index: i64,
30 pub covered_from_event_id: i64,
31 pub covered_to_event_id: i64,
32 pub evidence_event_ids: Vec<i64>,
33 #[serde(skip_serializing_if = "Option::is_none")]
34 pub files: Option<Vec<String>>,
35 pub created_at_epoch: i64,
36 pub updated_at_epoch: i64,
37}
38
39pub fn insert_topic_segment(conn: &Connection, seg: &TopicSegmentInput<'_>) -> Result<i64> {
40 let now = chrono::Utc::now().timestamp();
41 conn.execute(
42 "INSERT INTO topic_segments
43 (host_id, project_id, session_row_id, project, topic_key, title, summary,
44 status, segment_index, covered_from_event_id, covered_to_event_id,
45 evidence_event_ids, files, confidence, created_at_epoch, updated_at_epoch)
46 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?15)",
47 params![
48 seg.host_id,
49 seg.project_id,
50 seg.session_row_id,
51 seg.project,
52 seg.topic_key,
53 seg.title,
54 seg.summary,
55 seg.status,
56 seg.segment_index,
57 seg.covered_from_event_id,
58 seg.covered_to_event_id,
59 seg.evidence_event_ids,
60 seg.files,
61 seg.confidence,
62 now
63 ],
64 )?;
65 Ok(conn.last_insert_rowid())
66}
67
68pub fn topic_segment_exists(
69 conn: &Connection,
70 session_row_id: i64,
71 topic_key: &str,
72) -> Result<bool> {
73 let found: Option<i64> = conn
74 .query_row(
75 "SELECT id FROM topic_segments
76 WHERE session_row_id = ?1 AND topic_key = ?2 LIMIT 1",
77 params![session_row_id, topic_key],
78 |row| row.get(0),
79 )
80 .optional()?;
81 Ok(found.is_some())
82}
83
84pub fn load_trace_by_topic_key(
85 conn: &Connection,
86 project: &str,
87 topic_key: &str,
88 limit: i64,
89) -> Result<Vec<TopicTraceEntry>> {
90 let mut stmt = conn.prepare(
91 "SELECT id, topic_key, title, summary, status, segment_index,
92 covered_from_event_id, covered_to_event_id, evidence_event_ids,
93 files, created_at_epoch, updated_at_epoch
94 FROM (
95 SELECT id, topic_key, title, summary, status, segment_index,
96 covered_from_event_id, covered_to_event_id, evidence_event_ids,
97 files, created_at_epoch, updated_at_epoch
98 FROM topic_segments
99 WHERE project = ?1 AND topic_key = ?2
100 ORDER BY covered_from_event_id DESC, segment_index DESC, id DESC
101 LIMIT ?3
102 )
103 ORDER BY covered_from_event_id ASC, segment_index ASC, id ASC",
104 )?;
105 let rows = stmt.query_map(params![project, topic_key, limit.max(1)], |row| {
106 let evidence_json: String = row.get(8)?;
107 let files_json: Option<String> = row.get(9)?;
108 Ok(TopicTraceRaw {
109 id: row.get(0)?,
110 topic_key: row.get(1)?,
111 title: row.get(2)?,
112 summary: row.get(3)?,
113 status: row.get(4)?,
114 segment_index: row.get(5)?,
115 covered_from_event_id: row.get(6)?,
116 covered_to_event_id: row.get(7)?,
117 evidence_json,
118 files_json,
119 created_at_epoch: row.get(10)?,
120 updated_at_epoch: row.get(11)?,
121 })
122 })?;
123
124 let mut trace = Vec::new();
125 for row in rows {
126 trace.push(row?.try_into_entry()?);
127 }
128 Ok(trace)
129}
130
131struct TopicTraceRaw {
132 id: i64,
133 topic_key: String,
134 title: String,
135 summary: String,
136 status: String,
137 segment_index: i64,
138 covered_from_event_id: i64,
139 covered_to_event_id: i64,
140 evidence_json: String,
141 files_json: Option<String>,
142 created_at_epoch: i64,
143 updated_at_epoch: i64,
144}
145
146impl TopicTraceRaw {
147 fn try_into_entry(self) -> Result<TopicTraceEntry> {
148 let evidence_event_ids = serde_json::from_str::<Vec<i64>>(&self.evidence_json)
149 .with_context(|| format!("parse topic_segments evidence ids for id={}", self.id))?;
150 let files = match self.files_json {
151 Some(raw) => Some(
152 serde_json::from_str::<Vec<String>>(&raw)
153 .with_context(|| format!("parse topic_segments files for id={}", self.id))?,
154 ),
155 None => None,
156 };
157 Ok(TopicTraceEntry {
158 id: self.id,
159 topic_key: self.topic_key,
160 title: self.title,
161 summary: self.summary,
162 status: self.status,
163 segment_index: self.segment_index,
164 covered_from_event_id: self.covered_from_event_id,
165 covered_to_event_id: self.covered_to_event_id,
166 evidence_event_ids,
167 files,
168 created_at_epoch: self.created_at_epoch,
169 updated_at_epoch: self.updated_at_epoch,
170 })
171 }
172}
173
174#[cfg(test)]
175mod tests {
176 use super::*;
177
178 fn conn() -> Connection {
179 let conn = Connection::open_in_memory().expect("in-memory db");
180 crate::migrate::run_migrations(&conn).expect("migrations");
181 conn.execute_batch("PRAGMA foreign_keys=OFF;")
182 .expect("disable foreign keys");
183 conn
184 }
185
186 fn seg<'a>(
187 idx: i64,
188 topic_key: &'a str,
189 from: i64,
190 to: i64,
191 evidence: &'a str,
192 ) -> TopicSegmentInput<'a> {
193 TopicSegmentInput {
194 host_id: 1,
195 project_id: 1,
196 session_row_id: 7,
197 project: "/tmp/remem",
198 topic_key,
199 title: "title",
200 summary: "summary",
201 status: "resolved",
202 segment_index: idx,
203 covered_from_event_id: from,
204 covered_to_event_id: to,
205 evidence_event_ids: evidence,
206 files: None,
207 confidence: 0.75,
208 }
209 }
210
211 #[test]
212 fn insert_then_exists() -> Result<()> {
213 let conn = conn();
214 assert!(!topic_segment_exists(&conn, 7, "fts5-tokenizer")?);
215 let id = insert_topic_segment(&conn, &seg(0, "fts5-tokenizer", 100, 110, "[100,110]"))?;
216 assert!(id > 0);
217 assert!(topic_segment_exists(&conn, 7, "fts5-tokenizer")?);
218 assert!(!topic_segment_exists(&conn, 7, "other-topic")?);
219 assert!(!topic_segment_exists(&conn, 99, "fts5-tokenizer")?);
220 Ok(())
221 }
222
223 #[test]
224 fn overlapping_segments_coexist_and_trace_orders_by_range() -> Result<()> {
225 let conn = conn();
226 insert_topic_segment(
227 &conn,
228 &seg(1, "anti-bot-research", 3056, 3466, "[3056,3466]"),
229 )?;
230 insert_topic_segment(&conn, &seg(0, "anti-bot-research", 100, 120, "[100,120]"))?;
231 insert_topic_segment(&conn, &seg(2, "kexue-scraping", 3057, 3331, "[3057,3331]"))?;
232
233 let count: i64 = conn.query_row(
234 "SELECT COUNT(*) FROM topic_segments WHERE session_row_id = 7",
235 [],
236 |row| row.get(0),
237 )?;
238 assert_eq!(count, 3, "overlapping segments must coexist");
239
240 let trace = load_trace_by_topic_key(&conn, "/tmp/remem", "anti-bot-research", 10)?;
241 assert_eq!(trace.len(), 2);
242 assert_eq!(trace[0].covered_from_event_id, 100);
243 assert_eq!(trace[1].covered_from_event_id, 3056);
244 assert_eq!(trace[1].evidence_event_ids, vec![3056, 3466]);
245 Ok(())
246 }
247
248 #[test]
249 fn trace_limit_keeps_recent_segments_and_returns_chronologically() -> Result<()> {
250 let conn = conn();
251 for offset in 0..15 {
252 let event_id = 100 + offset;
253 let evidence = format!("[{event_id}]");
254 insert_topic_segment(
255 &conn,
256 &seg(offset, "release-plan", event_id, event_id, &evidence),
257 )?;
258 }
259
260 let trace = load_trace_by_topic_key(&conn, "/tmp/remem", "release-plan", 12)?;
261 let event_ids = trace
262 .iter()
263 .map(|entry| entry.covered_from_event_id)
264 .collect::<Vec<_>>();
265 assert_eq!(trace.len(), 12);
266 assert_eq!(event_ids[0], 103);
267 assert_eq!(event_ids[11], 114);
268 assert!(event_ids.windows(2).all(|pair| pair[0] < pair[1]));
269 Ok(())
270 }
271
272 #[test]
273 fn project_topic_trace_query_uses_project_trace_index() -> Result<()> {
274 let conn = conn();
275 let mut stmt = conn.prepare(
276 "EXPLAIN QUERY PLAN
277 SELECT id
278 FROM topic_segments
279 WHERE project = ?1 AND topic_key = ?2
280 ORDER BY covered_from_event_id DESC
281 LIMIT ?3",
282 )?;
283 let rows = stmt.query_map(params!["/tmp/remem", "release-plan", 12_i64], |row| {
284 row.get::<_, String>(3)
285 })?;
286 let plan = crate::db::query::collect_rows(rows)?.join("\n");
287
288 assert!(
289 plan.contains("idx_topic_segments_project_trace"),
290 "unexpected query plan: {plan}"
291 );
292 Ok(())
293 }
294}