Skip to main content

codemem_storage/
queries.rs

1//! Stats, consolidation, pattern queries, and session management on Storage.
2
3use crate::Storage;
4use codemem_core::{CodememError, ConsolidationLogEntry, Session, StorageStats};
5use rusqlite::params;
6
7impl Storage {
8    // ── Health / Diagnostics ────────────────────────────────────────────
9
10    /// Run SQLite `PRAGMA integrity_check`. Returns `true` if the database is OK.
11    pub fn integrity_check(&self) -> Result<bool, CodememError> {
12        let conn = self.conn()?;
13        let result: String = conn
14            .query_row("PRAGMA integrity_check", [], |row| row.get(0))
15            .map_err(|e| CodememError::Storage(e.to_string()))?;
16        Ok(result == "ok")
17    }
18
19    /// Return the current schema version (max applied migration number).
20    pub fn schema_version(&self) -> Result<u32, CodememError> {
21        let conn = self.conn()?;
22        let version: u32 = conn
23            .query_row(
24                "SELECT COALESCE(MAX(version), 0) FROM schema_version",
25                [],
26                |row| row.get(0),
27            )
28            .map_err(|e| CodememError::Storage(e.to_string()))?;
29        Ok(version)
30    }
31
32    // ── Stats ───────────────────────────────────────────────────────────
33
34    /// Get database statistics in a single query.
35    pub fn stats(&self) -> Result<StorageStats, CodememError> {
36        let conn = self.conn()?;
37
38        let (memory_count, embedding_count, node_count, edge_count) = conn
39            .query_row(
40                "SELECT
41                    (SELECT COUNT(*) FROM memories) AS memory_count,
42                    (SELECT COUNT(*) FROM memory_embeddings) AS embedding_count,
43                    (SELECT COUNT(*) FROM graph_nodes) AS node_count,
44                    (SELECT COUNT(*) FROM graph_edges) AS edge_count",
45                [],
46                |row| {
47                    Ok((
48                        row.get::<_, i64>(0)?,
49                        row.get::<_, i64>(1)?,
50                        row.get::<_, i64>(2)?,
51                        row.get::<_, i64>(3)?,
52                    ))
53                },
54            )
55            .map_err(|e| CodememError::Storage(e.to_string()))?;
56
57        Ok(StorageStats {
58            memory_count: memory_count as usize,
59            embedding_count: embedding_count as usize,
60            node_count: node_count as usize,
61            edge_count: edge_count as usize,
62        })
63    }
64
65    // ── Consolidation Log ──────────────────────────────────────────────
66
67    /// Record a consolidation run.
68    pub fn insert_consolidation_log(
69        &self,
70        cycle_type: &str,
71        affected_count: usize,
72    ) -> Result<(), CodememError> {
73        let conn = self.conn()?;
74        let now = chrono::Utc::now().timestamp();
75        conn.execute(
76            "INSERT INTO consolidation_log (cycle_type, run_at, affected_count) VALUES (?1, ?2, ?3)",
77            params![cycle_type, now, affected_count as i64],
78        )
79        .map_err(|e| CodememError::Storage(e.to_string()))?;
80        Ok(())
81    }
82
83    /// Get the last consolidation run for each cycle type.
84    pub fn last_consolidation_runs(&self) -> Result<Vec<ConsolidationLogEntry>, CodememError> {
85        let conn = self.conn()?;
86        let mut stmt = conn
87            .prepare(
88                "SELECT cycle_type, run_at, affected_count FROM consolidation_log
89                 WHERE (cycle_type, run_at) IN (
90                     SELECT cycle_type, MAX(run_at) FROM consolidation_log GROUP BY cycle_type
91                 )
92                 ORDER BY cycle_type",
93            )
94            .map_err(|e| CodememError::Storage(e.to_string()))?;
95
96        let entries = stmt
97            .query_map([], |row| {
98                Ok(ConsolidationLogEntry {
99                    cycle_type: row.get(0)?,
100                    run_at: row.get(1)?,
101                    affected_count: row.get::<_, i64>(2)? as usize,
102                })
103            })
104            .map_err(|e| CodememError::Storage(e.to_string()))?
105            .collect::<Result<Vec<_>, _>>()
106            .map_err(|e| CodememError::Storage(e.to_string()))?;
107
108        Ok(entries)
109    }
110
111    // ── Pattern Detection Queries ───────────────────────────────────────
112
113    /// Find repeated search patterns (Grep/Glob) by extracting the "pattern" field
114    /// from memory metadata JSON. Returns (pattern, count, memory_ids) tuples where
115    /// count >= min_count, ordered by count descending.
116    pub fn get_repeated_searches(
117        &self,
118        min_count: usize,
119        namespace: Option<&str>,
120    ) -> Result<Vec<(String, usize, Vec<String>)>, CodememError> {
121        let conn = self.conn()?;
122        let sql = if namespace.is_some() {
123            "SELECT json_extract(metadata, '$.pattern') AS pat,
124                    COUNT(*) AS cnt,
125                    GROUP_CONCAT(id, ',') AS ids
126             FROM memories
127             WHERE json_extract(metadata, '$.tool') IN ('Grep', 'Glob')
128               AND pat IS NOT NULL
129               AND namespace = ?1
130             GROUP BY pat
131             HAVING cnt >= ?2
132             ORDER BY cnt DESC"
133        } else {
134            "SELECT json_extract(metadata, '$.pattern') AS pat,
135                    COUNT(*) AS cnt,
136                    GROUP_CONCAT(id, ',') AS ids
137             FROM memories
138             WHERE json_extract(metadata, '$.tool') IN ('Grep', 'Glob')
139               AND pat IS NOT NULL
140             GROUP BY pat
141             HAVING cnt >= ?1
142             ORDER BY cnt DESC"
143        };
144
145        let mut stmt = conn
146            .prepare(sql)
147            .map_err(|e| CodememError::Storage(e.to_string()))?;
148
149        let rows = if let Some(ns) = namespace {
150            stmt.query_map(params![ns, min_count as i64], |row| {
151                Ok((
152                    row.get::<_, String>(0)?,
153                    row.get::<_, i64>(1)?,
154                    row.get::<_, String>(2)?,
155                ))
156            })
157            .map_err(|e| CodememError::Storage(e.to_string()))?
158            .collect::<Result<Vec<_>, _>>()
159            .map_err(|e| CodememError::Storage(e.to_string()))?
160        } else {
161            stmt.query_map(params![min_count as i64], |row| {
162                Ok((
163                    row.get::<_, String>(0)?,
164                    row.get::<_, i64>(1)?,
165                    row.get::<_, String>(2)?,
166                ))
167            })
168            .map_err(|e| CodememError::Storage(e.to_string()))?
169            .collect::<Result<Vec<_>, _>>()
170            .map_err(|e| CodememError::Storage(e.to_string()))?
171        };
172
173        Ok(rows
174            .into_iter()
175            .map(|(pat, cnt, ids_str)| {
176                let ids: Vec<String> = ids_str.split(',').map(String::from).collect();
177                (pat, cnt as usize, ids)
178            })
179            .collect())
180    }
181
182    /// Find file hotspots by extracting the "file_path" field from memory metadata.
183    pub fn get_file_hotspots(
184        &self,
185        min_count: usize,
186        namespace: Option<&str>,
187    ) -> Result<Vec<(String, usize, Vec<String>)>, CodememError> {
188        let conn = self.conn()?;
189        let sql = if namespace.is_some() {
190            "SELECT json_extract(metadata, '$.file_path') AS fp,
191                    COUNT(*) AS cnt,
192                    GROUP_CONCAT(id, ',') AS ids
193             FROM memories
194             WHERE fp IS NOT NULL
195               AND namespace = ?1
196             GROUP BY fp
197             HAVING cnt >= ?2
198             ORDER BY cnt DESC"
199        } else {
200            "SELECT json_extract(metadata, '$.file_path') AS fp,
201                    COUNT(*) AS cnt,
202                    GROUP_CONCAT(id, ',') AS ids
203             FROM memories
204             WHERE fp IS NOT NULL
205             GROUP BY fp
206             HAVING cnt >= ?1
207             ORDER BY cnt DESC"
208        };
209
210        let mut stmt = conn
211            .prepare(sql)
212            .map_err(|e| CodememError::Storage(e.to_string()))?;
213
214        let rows = if let Some(ns) = namespace {
215            stmt.query_map(params![ns, min_count as i64], |row| {
216                Ok((
217                    row.get::<_, String>(0)?,
218                    row.get::<_, i64>(1)?,
219                    row.get::<_, String>(2)?,
220                ))
221            })
222            .map_err(|e| CodememError::Storage(e.to_string()))?
223            .collect::<Result<Vec<_>, _>>()
224            .map_err(|e| CodememError::Storage(e.to_string()))?
225        } else {
226            stmt.query_map(params![min_count as i64], |row| {
227                Ok((
228                    row.get::<_, String>(0)?,
229                    row.get::<_, i64>(1)?,
230                    row.get::<_, String>(2)?,
231                ))
232            })
233            .map_err(|e| CodememError::Storage(e.to_string()))?
234            .collect::<Result<Vec<_>, _>>()
235            .map_err(|e| CodememError::Storage(e.to_string()))?
236        };
237
238        Ok(rows
239            .into_iter()
240            .map(|(fp, cnt, ids_str)| {
241                let ids: Vec<String> = ids_str.split(',').map(String::from).collect();
242                (fp, cnt as usize, ids)
243            })
244            .collect())
245    }
246
247    /// Get tool usage statistics from memory metadata.
248    /// Returns (tool_name, count) pairs sorted by count descending.
249    pub fn get_tool_usage_stats(
250        &self,
251        namespace: Option<&str>,
252    ) -> Result<Vec<(String, usize)>, CodememError> {
253        let conn = self.conn()?;
254        let sql = if namespace.is_some() {
255            "SELECT json_extract(metadata, '$.tool') AS tool,
256                    COUNT(*) AS cnt
257             FROM memories
258             WHERE tool IS NOT NULL
259               AND namespace = ?1
260             GROUP BY tool
261             ORDER BY cnt DESC"
262        } else {
263            "SELECT json_extract(metadata, '$.tool') AS tool,
264                    COUNT(*) AS cnt
265             FROM memories
266             WHERE tool IS NOT NULL
267             GROUP BY tool
268             ORDER BY cnt DESC"
269        };
270
271        let mut stmt = conn
272            .prepare(sql)
273            .map_err(|e| CodememError::Storage(e.to_string()))?;
274
275        let rows = if let Some(ns) = namespace {
276            stmt.query_map(params![ns], |row| {
277                Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
278            })
279            .map_err(|e| CodememError::Storage(e.to_string()))?
280            .collect::<Result<Vec<_>, _>>()
281            .map_err(|e| CodememError::Storage(e.to_string()))?
282        } else {
283            stmt.query_map([], |row| {
284                Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
285            })
286            .map_err(|e| CodememError::Storage(e.to_string()))?
287            .collect::<Result<Vec<_>, _>>()
288            .map_err(|e| CodememError::Storage(e.to_string()))?
289        };
290
291        Ok(rows
292            .into_iter()
293            .map(|(tool, cnt)| (tool, cnt as usize))
294            .collect())
295    }
296
297    /// Find decision chains: files with multiple Edit/Write memories over time.
298    pub fn get_decision_chains(
299        &self,
300        min_count: usize,
301        namespace: Option<&str>,
302    ) -> Result<Vec<(String, usize, Vec<String>)>, CodememError> {
303        let conn = self.conn()?;
304        let sql = if namespace.is_some() {
305            "SELECT json_extract(metadata, '$.file_path') AS fp,
306                    COUNT(*) AS cnt,
307                    GROUP_CONCAT(id, ',') AS ids
308             FROM memories
309             WHERE json_extract(metadata, '$.tool') IN ('Edit', 'Write')
310               AND fp IS NOT NULL
311               AND namespace = ?1
312             GROUP BY fp
313             HAVING cnt >= ?2
314             ORDER BY cnt DESC"
315        } else {
316            "SELECT json_extract(metadata, '$.file_path') AS fp,
317                    COUNT(*) AS cnt,
318                    GROUP_CONCAT(id, ',') AS ids
319             FROM memories
320             WHERE json_extract(metadata, '$.tool') IN ('Edit', 'Write')
321               AND fp IS NOT NULL
322             GROUP BY fp
323             HAVING cnt >= ?1
324             ORDER BY cnt DESC"
325        };
326
327        let mut stmt = conn
328            .prepare(sql)
329            .map_err(|e| CodememError::Storage(e.to_string()))?;
330
331        let rows = if let Some(ns) = namespace {
332            stmt.query_map(params![ns, min_count as i64], |row| {
333                Ok((
334                    row.get::<_, String>(0)?,
335                    row.get::<_, i64>(1)?,
336                    row.get::<_, String>(2)?,
337                ))
338            })
339            .map_err(|e| CodememError::Storage(e.to_string()))?
340            .collect::<Result<Vec<_>, _>>()
341            .map_err(|e| CodememError::Storage(e.to_string()))?
342        } else {
343            stmt.query_map(params![min_count as i64], |row| {
344                Ok((
345                    row.get::<_, String>(0)?,
346                    row.get::<_, i64>(1)?,
347                    row.get::<_, String>(2)?,
348                ))
349            })
350            .map_err(|e| CodememError::Storage(e.to_string()))?
351            .collect::<Result<Vec<_>, _>>()
352            .map_err(|e| CodememError::Storage(e.to_string()))?
353        };
354
355        Ok(rows
356            .into_iter()
357            .map(|(fp, cnt, ids_str)| {
358                let ids: Vec<String> = ids_str.split(',').map(String::from).collect();
359                (fp, cnt as usize, ids)
360            })
361            .collect())
362    }
363
364    // ── Insight / Tag Queries ──────────────────────────────────────────
365
366    /// List memories that contain a specific tag, optionally scoped to a namespace.
367    /// Uses `json_each` for proper JSON array querying instead of LIKE patterns.
368    pub fn list_memories_by_tag(
369        &self,
370        tag: &str,
371        namespace: Option<&str>,
372        limit: usize,
373    ) -> Result<Vec<codemem_core::MemoryNode>, CodememError> {
374        let conn = self.conn()?;
375
376        let (sql, params_vec): (String, Vec<Box<dyn rusqlite::types::ToSql>>) = if let Some(ns) =
377            namespace
378        {
379            (
380                "SELECT m.id, m.content, m.memory_type, m.importance, m.confidence, m.access_count, \
381                 m.content_hash, m.tags, m.metadata, m.namespace, m.created_at, m.updated_at, m.last_accessed_at \
382                 FROM memories m, json_each(m.tags) AS jt \
383                 WHERE jt.value = ?1 AND m.namespace = ?2 \
384                 ORDER BY m.created_at DESC LIMIT ?3"
385                    .to_string(),
386                vec![
387                    Box::new(tag.to_string()) as Box<dyn rusqlite::types::ToSql>,
388                    Box::new(ns.to_string()),
389                    Box::new(limit as i64),
390                ],
391            )
392        } else {
393            (
394                "SELECT m.id, m.content, m.memory_type, m.importance, m.confidence, m.access_count, \
395                 m.content_hash, m.tags, m.metadata, m.namespace, m.created_at, m.updated_at, m.last_accessed_at \
396                 FROM memories m, json_each(m.tags) AS jt \
397                 WHERE jt.value = ?1 \
398                 ORDER BY m.created_at DESC LIMIT ?2"
399                    .to_string(),
400                vec![
401                    Box::new(tag.to_string()) as Box<dyn rusqlite::types::ToSql>,
402                    Box::new(limit as i64),
403                ],
404            )
405        };
406
407        let params_refs: Vec<&dyn rusqlite::types::ToSql> =
408            params_vec.iter().map(|b| &**b).collect();
409
410        let mut stmt = conn
411            .prepare(&sql)
412            .map_err(|e| CodememError::Storage(e.to_string()))?;
413
414        let rows = stmt
415            .query_map(params_refs.as_slice(), |row| {
416                let created_ts: i64 = row.get(10)?;
417                let updated_ts: i64 = row.get(11)?;
418                let accessed_ts: i64 = row.get(12)?;
419                let tags_json: String = row.get(7)?;
420                let metadata_json: String = row.get(8)?;
421                let memory_type_str: String = row.get(2)?;
422
423                Ok(codemem_core::MemoryNode {
424                    id: row.get(0)?,
425                    content: row.get(1)?,
426                    memory_type: memory_type_str
427                        .parse()
428                        .unwrap_or(codemem_core::MemoryType::Context),
429                    importance: row.get(3)?,
430                    confidence: row.get(4)?,
431                    access_count: row.get::<_, i64>(5).unwrap_or(0) as u32,
432                    content_hash: row.get(6)?,
433                    tags: serde_json::from_str(&tags_json).unwrap_or_default(),
434                    metadata: serde_json::from_str(&metadata_json).unwrap_or_default(),
435                    namespace: row.get(9)?,
436                    created_at: chrono::DateTime::from_timestamp(created_ts, 0)
437                        .unwrap_or_default()
438                        .with_timezone(&chrono::Utc),
439                    updated_at: chrono::DateTime::from_timestamp(updated_ts, 0)
440                        .unwrap_or_default()
441                        .with_timezone(&chrono::Utc),
442                    last_accessed_at: chrono::DateTime::from_timestamp(accessed_ts, 0)
443                        .unwrap_or_default()
444                        .with_timezone(&chrono::Utc),
445                })
446            })
447            .map_err(|e| CodememError::Storage(e.to_string()))?
448            .collect::<Result<Vec<_>, _>>()
449            .map_err(|e| CodememError::Storage(e.to_string()))?;
450
451        Ok(rows)
452    }
453
454    // ── Session Management ─────────────────────────────────────────────
455
456    /// Start a new session.
457    pub fn start_session(&self, id: &str, namespace: Option<&str>) -> Result<(), CodememError> {
458        let conn = self.conn()?;
459        let now = chrono::Utc::now().timestamp();
460        conn.execute(
461            "INSERT OR IGNORE INTO sessions (id, namespace, started_at) VALUES (?1, ?2, ?3)",
462            params![id, namespace, now],
463        )
464        .map_err(|e| CodememError::Storage(e.to_string()))?;
465        Ok(())
466    }
467
468    /// End a session by setting ended_at and optionally a summary.
469    pub fn end_session(&self, id: &str, summary: Option<&str>) -> Result<(), CodememError> {
470        let conn = self.conn()?;
471        let now = chrono::Utc::now().timestamp();
472        conn.execute(
473            "UPDATE sessions SET ended_at = ?1, summary = ?2 WHERE id = ?3",
474            params![now, summary, id],
475        )
476        .map_err(|e| CodememError::Storage(e.to_string()))?;
477        Ok(())
478    }
479
480    /// List sessions, optionally filtered by namespace.
481    pub fn list_sessions(&self, namespace: Option<&str>) -> Result<Vec<Session>, CodememError> {
482        self.list_sessions_with_limit(namespace, usize::MAX)
483    }
484
485    // ── Session Activity Tracking ─────────────────────────────────
486
487    /// Record a session activity event.
488    pub fn record_session_activity(
489        &self,
490        session_id: &str,
491        tool_name: &str,
492        file_path: Option<&str>,
493        directory: Option<&str>,
494        pattern: Option<&str>,
495    ) -> Result<(), CodememError> {
496        let conn = self.conn()?;
497        let now = chrono::Utc::now().timestamp();
498        conn.execute(
499            "INSERT INTO session_activity (session_id, tool_name, file_path, directory, pattern, created_at)
500             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
501            params![session_id, tool_name, file_path, directory, pattern, now],
502        )
503        .map_err(|e| CodememError::Storage(e.to_string()))?;
504        Ok(())
505    }
506
507    /// Get a summary of session activity counts using a single query with conditional aggregation.
508    pub fn get_session_activity_summary(
509        &self,
510        session_id: &str,
511    ) -> Result<codemem_core::SessionActivitySummary, CodememError> {
512        let conn = self.conn()?;
513
514        let (files_read, files_edited, searches, total_actions) = conn
515            .query_row(
516                "SELECT
517                     COUNT(DISTINCT CASE WHEN tool_name = 'Read' AND file_path IS NOT NULL THEN file_path END),
518                     COUNT(DISTINCT CASE WHEN tool_name IN ('Edit', 'Write') AND file_path IS NOT NULL THEN file_path END),
519                     SUM(CASE WHEN tool_name IN ('Grep', 'Glob') THEN 1 ELSE 0 END),
520                     COUNT(*)
521                 FROM session_activity
522                 WHERE session_id = ?1",
523                params![session_id],
524                |row| {
525                    Ok((
526                        row.get::<_, i64>(0)?,
527                        row.get::<_, i64>(1)?,
528                        row.get::<_, i64>(2)?,
529                        row.get::<_, i64>(3)?,
530                    ))
531                },
532            )
533            .map_err(|e| CodememError::Storage(e.to_string()))?;
534
535        Ok(codemem_core::SessionActivitySummary {
536            files_read: files_read as usize,
537            files_edited: files_edited as usize,
538            searches: searches as usize,
539            total_actions: total_actions as usize,
540        })
541    }
542
543    /// Get the most active directories in a session.
544    pub fn get_session_hot_directories(
545        &self,
546        session_id: &str,
547        limit: usize,
548    ) -> Result<Vec<(String, usize)>, CodememError> {
549        let conn = self.conn()?;
550        let mut stmt = conn
551            .prepare(
552                "SELECT directory, COUNT(*) AS cnt FROM session_activity
553                 WHERE session_id = ?1 AND directory IS NOT NULL
554                 GROUP BY directory ORDER BY cnt DESC LIMIT ?2",
555            )
556            .map_err(|e| CodememError::Storage(e.to_string()))?;
557
558        let rows = stmt
559            .query_map(params![session_id, limit as i64], |row| {
560                Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
561            })
562            .map_err(|e| CodememError::Storage(e.to_string()))?
563            .collect::<Result<Vec<_>, _>>()
564            .map_err(|e| CodememError::Storage(e.to_string()))?;
565
566        Ok(rows
567            .into_iter()
568            .map(|(dir, cnt)| (dir, cnt as usize))
569            .collect())
570    }
571
572    /// Check whether an auto-insight dedup tag exists for a session.
573    /// Uses `json_extract` with proper parameter binding instead of LIKE on JSON.
574    pub fn has_auto_insight(
575        &self,
576        session_id: &str,
577        dedup_tag: &str,
578    ) -> Result<bool, CodememError> {
579        let conn = self.conn()?;
580        let count: i64 = conn
581            .query_row(
582                "SELECT COUNT(*) FROM memories
583                 WHERE json_extract(metadata, '$.session_id') = ?1
584                   AND json_extract(metadata, '$.auto_insight_tag') = ?2",
585                params![session_id, dedup_tag],
586                |row| row.get(0),
587            )
588            .map_err(|e| CodememError::Storage(e.to_string()))?;
589        Ok(count > 0)
590    }
591
592    /// Count Read events in a directory during a session.
593    pub fn count_directory_reads(
594        &self,
595        session_id: &str,
596        directory: &str,
597    ) -> Result<usize, CodememError> {
598        let conn = self.conn()?;
599        let count: i64 = conn
600            .query_row(
601                "SELECT COUNT(*) FROM session_activity
602                 WHERE session_id = ?1 AND tool_name = 'Read' AND directory = ?2",
603                params![session_id, directory],
604                |row| row.get(0),
605            )
606            .map_err(|e| CodememError::Storage(e.to_string()))?;
607        Ok(count as usize)
608    }
609
610    /// Check if a file was read in the current session.
611    pub fn was_file_read_in_session(
612        &self,
613        session_id: &str,
614        file_path: &str,
615    ) -> Result<bool, CodememError> {
616        let conn = self.conn()?;
617        let count: i64 = conn
618            .query_row(
619                "SELECT COUNT(*) FROM session_activity
620                 WHERE session_id = ?1 AND tool_name = 'Read' AND file_path = ?2",
621                params![session_id, file_path],
622                |row| row.get(0),
623            )
624            .map_err(|e| CodememError::Storage(e.to_string()))?;
625        Ok(count > 0)
626    }
627
628    /// Count how many times a search pattern was used in a session.
629    pub fn count_search_pattern_in_session(
630        &self,
631        session_id: &str,
632        pattern: &str,
633    ) -> Result<usize, CodememError> {
634        let conn = self.conn()?;
635        let count: i64 = conn
636            .query_row(
637                "SELECT COUNT(*) FROM session_activity
638                 WHERE session_id = ?1 AND tool_name IN ('Grep', 'Glob') AND pattern = ?2",
639                params![session_id, pattern],
640                |row| row.get(0),
641            )
642            .map_err(|e| CodememError::Storage(e.to_string()))?;
643        Ok(count as usize)
644    }
645
646    /// List sessions with a limit.
647    pub(crate) fn list_sessions_with_limit(
648        &self,
649        namespace: Option<&str>,
650        limit: usize,
651    ) -> Result<Vec<Session>, CodememError> {
652        let conn = self.conn()?;
653        let sql_with_ns = "SELECT id, namespace, started_at, ended_at, memory_count, summary FROM sessions WHERE namespace = ?1 ORDER BY started_at DESC LIMIT ?2";
654        let sql_all = "SELECT id, namespace, started_at, ended_at, memory_count, summary FROM sessions ORDER BY started_at DESC LIMIT ?1";
655
656        let map_row = |row: &rusqlite::Row<'_>| -> rusqlite::Result<Session> {
657            let started_ts: i64 = row.get(2)?;
658            let ended_ts: Option<i64> = row.get(3)?;
659            Ok(Session {
660                id: row.get(0)?,
661                namespace: row.get(1)?,
662                started_at: chrono::DateTime::from_timestamp(started_ts, 0)
663                    .unwrap_or_default()
664                    .with_timezone(&chrono::Utc),
665                ended_at: ended_ts.and_then(|ts| {
666                    chrono::DateTime::from_timestamp(ts, 0).map(|dt| dt.with_timezone(&chrono::Utc))
667                }),
668                memory_count: row.get::<_, i64>(4).unwrap_or(0) as u32,
669                summary: row.get(5)?,
670            })
671        };
672
673        if let Some(ns) = namespace {
674            let mut stmt = conn
675                .prepare(sql_with_ns)
676                .map_err(|e| CodememError::Storage(e.to_string()))?;
677            let rows = stmt
678                .query_map(params![ns, limit as i64], map_row)
679                .map_err(|e| CodememError::Storage(e.to_string()))?;
680            rows.collect::<Result<Vec<_>, _>>()
681                .map_err(|e| CodememError::Storage(e.to_string()))
682        } else {
683            let mut stmt = conn
684                .prepare(sql_all)
685                .map_err(|e| CodememError::Storage(e.to_string()))?;
686            let rows = stmt
687                .query_map(params![limit as i64], map_row)
688                .map_err(|e| CodememError::Storage(e.to_string()))?;
689            rows.collect::<Result<Vec<_>, _>>()
690                .map_err(|e| CodememError::Storage(e.to_string()))
691        }
692    }
693
694    // ── Tag-based Queries ─────────────────────────────────────────
695
696    /// Find memory IDs whose tags JSON array contains the given tag value.
697    /// Optionally scoped to a namespace. Excludes the given `exclude_id`.
698    /// Returns at most 50 results ordered by creation time (most recent siblings first).
699    pub fn find_memory_ids_by_tag(
700        &self,
701        tag: &str,
702        namespace: Option<&str>,
703        exclude_id: &str,
704    ) -> Result<Vec<String>, CodememError> {
705        let conn = self.conn()?;
706
707        // Use json_each() for exact tag matching instead of LIKE (safe against %, _, " in tags).
708        let (sql, params_vec): (String, Vec<Box<dyn rusqlite::types::ToSql>>) =
709            if let Some(ns) = namespace {
710                (
711                    "SELECT DISTINCT m.id FROM memories m, json_each(m.tags) t \
712                 WHERE t.value = ?1 AND m.namespace IS ?2 AND m.id != ?3 \
713                 ORDER BY m.created_at DESC LIMIT 50"
714                        .to_string(),
715                    vec![
716                        Box::new(tag.to_string()) as Box<dyn rusqlite::types::ToSql>,
717                        Box::new(ns.to_string()),
718                        Box::new(exclude_id.to_string()),
719                    ],
720                )
721            } else {
722                (
723                    "SELECT DISTINCT m.id FROM memories m, json_each(m.tags) t \
724                 WHERE t.value = ?1 AND m.namespace IS NULL AND m.id != ?2 \
725                 ORDER BY m.created_at DESC LIMIT 50"
726                        .to_string(),
727                    vec![
728                        Box::new(tag.to_string()) as Box<dyn rusqlite::types::ToSql>,
729                        Box::new(exclude_id.to_string()),
730                    ],
731                )
732            };
733
734        let refs: Vec<&dyn rusqlite::types::ToSql> =
735            params_vec.iter().map(|p| p.as_ref()).collect();
736
737        let mut stmt = conn
738            .prepare(&sql)
739            .map_err(|e| CodememError::Storage(e.to_string()))?;
740
741        let ids = stmt
742            .query_map(refs.as_slice(), |row| row.get(0))
743            .map_err(|e| CodememError::Storage(e.to_string()))?
744            .collect::<Result<Vec<String>, _>>()
745            .map_err(|e| CodememError::Storage(e.to_string()))?;
746
747        Ok(ids)
748    }
749
750    // ── Graph Cleanup ───────────────────────────────────────────────
751
752    /// Delete all graph nodes, their edges, and their embeddings where the
753    /// node ID starts with the given prefix. Returns count of nodes deleted.
754    /// Wrapped in a transaction so all three DELETEs are atomic.
755    pub fn delete_graph_nodes_by_prefix(&self, prefix: &str) -> Result<usize, CodememError> {
756        let conn = self.conn()?;
757        let like_pattern = format!("{prefix}%");
758
759        let tx = conn
760            .unchecked_transaction()
761            .map_err(|e| CodememError::Storage(e.to_string()))?;
762
763        // Delete edges where src or dst matches prefix
764        tx.execute(
765            "DELETE FROM graph_edges WHERE src LIKE ?1 OR dst LIKE ?1",
766            params![like_pattern],
767        )
768        .map_err(|e| CodememError::Storage(e.to_string()))?;
769
770        // Delete embeddings for matching nodes
771        tx.execute(
772            "DELETE FROM memory_embeddings WHERE memory_id LIKE ?1",
773            params![like_pattern],
774        )
775        .map_err(|e| CodememError::Storage(e.to_string()))?;
776
777        // Delete the nodes themselves
778        let rows = tx
779            .execute(
780                "DELETE FROM graph_nodes WHERE id LIKE ?1",
781                params![like_pattern],
782            )
783            .map_err(|e| CodememError::Storage(e.to_string()))?;
784
785        tx.commit()
786            .map_err(|e| CodememError::Storage(e.to_string()))?;
787
788        Ok(rows)
789    }
790}
791
792#[cfg(test)]
793#[path = "tests/queries_tests.rs"]
794mod tests;