1use crate::error::{BrainError, Result};
7use crate::fts::escape_fts5_query;
8use crate::types::{Edge, Node, NodeType};
9use rusqlite::{params, Connection, OptionalExtension};
10use std::collections::{HashMap, HashSet};
11use std::path::Path;
12
13#[cfg(feature = "ast")]
14use crate::ast::SymbolAnchor;
15
16pub struct Database {
21 pub(crate) conn: Connection,
22}
23
24impl Database {
25 pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
27 let conn = Connection::open(path.as_ref())?;
28 let db = Self { conn };
29 db.configure_connection()?;
30 super::migrations::migrate(&db.conn)?;
31 Ok(db)
32 }
33
34 pub fn open_in_memory() -> Result<Self> {
36 let conn = Connection::open_in_memory()?;
37 let db = Self { conn };
38 db.configure_connection()?;
39 super::migrations::migrate(&db.conn)?;
40 Ok(db)
41 }
42
43 fn configure_connection(&self) -> Result<()> {
44 self.conn.pragma_update(None, "foreign_keys", true)?;
45 self.conn.pragma_update(None, "busy_timeout", 5000i32)?;
46 let _ = self.conn.pragma_update(None, "journal_mode", "WAL");
48 let _ = self.conn.pragma_update(None, "synchronous", "NORMAL");
49 Ok(())
50 }
51
52 pub fn with_transaction<T, F>(&self, f: F) -> Result<T>
54 where
55 F: FnOnce(&Connection) -> Result<T>,
56 {
57 let tx = self.conn.unchecked_transaction()?;
58 let result = f(&tx)?;
59 tx.commit()?;
60 Ok(result)
61 }
62
63 pub fn insert_node(&self, node: &Node) -> Result<()> {
66 self.insert_node_on(&self.conn, node)
67 }
68
69 pub(crate) fn insert_node_on(&self, conn: &Connection, node: &Node) -> Result<()> {
70 let existing_created: Option<i64> = conn
72 .query_row(
73 "SELECT created_at FROM nodes WHERE id = ?1",
74 [&node.id],
75 |row| row.get(0),
76 )
77 .optional()?;
78
79 let created_at = existing_created.unwrap_or(node.created_at);
80
81 let scope = if node.scope.trim().is_empty() {
82 crate::scopes::MAIN_SCOPE
83 } else {
84 node.scope.as_str()
85 };
86 conn.execute(
87 "INSERT INTO nodes (id, node_type, title, file_path, symbol_hash, summary, content_hash, scope, created_at, updated_at)
88 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
89 ON CONFLICT(id) DO UPDATE SET
90 node_type = excluded.node_type,
91 title = excluded.title,
92 file_path = excluded.file_path,
93 symbol_hash = excluded.symbol_hash,
94 summary = excluded.summary,
95 content_hash = excluded.content_hash,
96 scope = excluded.scope,
97 updated_at = excluded.updated_at",
98 params![
99 node.id,
100 node.node_type.as_str(),
101 node.title,
102 node.file_path,
103 node.symbol_hash.map(|h| h as i64),
104 node.summary,
105 node.content_hash,
106 scope,
107 created_at,
108 node.updated_at,
109 ],
110 )?;
111 Ok(())
112 }
113
114 pub fn set_node_scope(&self, node_id: &str, scope: &str) -> Result<()> {
116 let scope = if scope.trim().is_empty() {
117 crate::scopes::MAIN_SCOPE
118 } else {
119 scope
120 };
121 self.conn.execute(
122 "UPDATE nodes SET scope = ?1 WHERE id = ?2",
123 params![scope, node_id],
124 )?;
125 let title: Option<String> = self
127 .conn
128 .query_row(
129 "SELECT title FROM nodes WHERE id = ?1",
130 [node_id],
131 |row| row.get(0),
132 )
133 .optional()?;
134 let content = self.get_fts_content(node_id)?.unwrap_or_default();
135 let mut tags = self.get_tags_for_node(node_id)?;
136 tags.retain(|t| !t.starts_with("scope:"));
137 tags.push(format!("scope:{scope}"));
138 let tags_str = tags.join(" ");
139 if let Some(title) = title {
140 self.index_fts(node_id, &title, &content, &tags_str)?;
141 }
142 self.replace_node_tags(node_id, &tags)?;
144 Ok(())
145 }
146
147 fn get_tags_for_node(&self, node_id: &str) -> Result<Vec<String>> {
148 let mut stmt = self
149 .conn
150 .prepare("SELECT tag FROM node_tags WHERE node_id = ?1")?;
151 let rows = stmt.query_map([node_id], |row| row.get(0))?;
152 let mut v = Vec::new();
153 for r in rows {
154 v.push(r?);
155 }
156 Ok(v)
157 }
158
159 pub fn get_node_scope(&self, node_id: &str) -> Result<Option<String>> {
161 let v: Option<String> = self
162 .conn
163 .query_row(
164 "SELECT scope FROM nodes WHERE id = ?1",
165 [node_id],
166 |row| row.get(0),
167 )
168 .optional()?;
169 Ok(v)
170 }
171
172 pub fn get_content_hash(&self, node_id: &str) -> Result<Option<String>> {
174 let v: Option<String> = self
175 .conn
176 .query_row(
177 "SELECT content_hash FROM nodes WHERE id = ?1",
178 [node_id],
179 |row| row.get(0),
180 )
181 .optional()?;
182 Ok(v)
183 }
184
185 pub fn insert_edge(&self, edge: &Edge) -> Result<()> {
187 self.insert_edge_on(&self.conn, edge)
188 }
189
190 pub(crate) fn insert_edge_on(&self, conn: &Connection, edge: &Edge) -> Result<()> {
191 conn.execute(
192 "INSERT INTO edges (source_id, target_id, relation_type, weight, decay_rate, created_at)
193 VALUES (?1, ?2, ?3, ?4, ?5, ?6)
194 ON CONFLICT(source_id, target_id, relation_type) DO UPDATE SET
195 weight = excluded.weight,
196 decay_rate = excluded.decay_rate",
197 params![
198 edge.source_id,
199 edge.target_id,
200 edge.relation_type,
201 edge.weight,
202 edge.decay_rate,
203 edge.created_at,
204 ],
205 )?;
206 Ok(())
207 }
208
209 pub fn insert_pending_link(
211 &self,
212 source_id: &str,
213 raw_target: &str,
214 relation_type: &str,
215 created_at: i64,
216 ) -> Result<()> {
217 self.insert_pending_link_on(&self.conn, source_id, raw_target, relation_type, created_at)
218 }
219
220 pub(crate) fn insert_pending_link_on(
221 &self,
222 conn: &Connection,
223 source_id: &str,
224 raw_target: &str,
225 relation_type: &str,
226 created_at: i64,
227 ) -> Result<()> {
228 conn.execute(
229 "INSERT INTO pending_links (source_id, raw_target, relation_type, created_at)
230 VALUES (?1, ?2, ?3, ?4)
231 ON CONFLICT(source_id, raw_target, relation_type) DO NOTHING",
232 params![source_id, raw_target, relation_type, created_at],
233 )?;
234 Ok(())
235 }
236
237 pub fn clear_pending_links_for(&self, source_id: &str) -> Result<()> {
239 self.conn
240 .execute("DELETE FROM pending_links WHERE source_id = ?1", [source_id])?;
241 Ok(())
242 }
243
244 pub(crate) fn clear_pending_links_for_on(
245 &self,
246 conn: &Connection,
247 source_id: &str,
248 ) -> Result<()> {
249 conn.execute("DELETE FROM pending_links WHERE source_id = ?1", [source_id])?;
250 Ok(())
251 }
252
253 pub(crate) fn clear_edges_from_on(
255 &self,
256 conn: &Connection,
257 source_id: &str,
258 relation_type: &str,
259 ) -> Result<()> {
260 conn.execute(
261 "DELETE FROM edges WHERE source_id = ?1 AND relation_type = ?2",
262 params![source_id, relation_type],
263 )?;
264 Ok(())
265 }
266
267 pub fn clear_all_auto_edges(&self) -> Result<()> {
269 self.conn
270 .execute("DELETE FROM edges WHERE relation_type LIKE 'auto_%'", [])?;
271 Ok(())
272 }
273
274 pub fn clear_auto_edges_involving(&self, node_id: &str) -> Result<()> {
276 self.conn.execute(
277 "DELETE FROM edges WHERE relation_type LIKE 'auto_%' AND (source_id = ?1 OR target_id = ?1)",
278 [node_id],
279 )?;
280 Ok(())
281 }
282
283 #[cfg(feature = "ast")]
285 pub fn insert_symbol_anchor(&self, anchor: &SymbolAnchor) -> Result<()> {
286 self.conn.execute(
287 "INSERT INTO symbol_anchors (symbol_hash, crate_name, module_path, symbol_name, file_path, start_line, end_line, doc_comment)
288 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
289 ON CONFLICT(symbol_hash) DO UPDATE SET
290 crate_name = excluded.crate_name,
291 module_path = excluded.module_path,
292 symbol_name = excluded.symbol_name,
293 file_path = excluded.file_path,
294 start_line = excluded.start_line,
295 end_line = excluded.end_line,
296 doc_comment = excluded.doc_comment",
297 params![
298 anchor.symbol_hash as i64,
299 anchor.crate_name,
300 anchor.module_path,
301 anchor.symbol_name,
302 anchor.file_path,
303 anchor.start_line,
304 anchor.end_line,
305 anchor.doc_comment,
306 ],
307 )?;
308 Ok(())
309 }
310
311 pub fn replace_node_tags(&self, node_id: &str, tags: &[String]) -> Result<()> {
313 self.replace_node_tags_on(&self.conn, node_id, tags)
314 }
315
316 pub(crate) fn replace_node_tags_on(
317 &self,
318 conn: &Connection,
319 node_id: &str,
320 tags: &[String],
321 ) -> Result<()> {
322 conn.execute("DELETE FROM node_tags WHERE node_id = ?1", [node_id])?;
323 for tag in tags {
324 conn.execute(
325 "INSERT OR IGNORE INTO node_tags (node_id, tag) VALUES (?1, ?2)",
326 params![node_id, tag],
327 )?;
328 }
329 Ok(())
330 }
331
332 pub fn replace_node_aliases(&self, node_id: &str, aliases: &[String]) -> Result<()> {
334 self.replace_node_aliases_on(&self.conn, node_id, aliases)
335 }
336
337 pub(crate) fn replace_node_aliases_on(
338 &self,
339 conn: &Connection,
340 node_id: &str,
341 aliases: &[String],
342 ) -> Result<()> {
343 conn.execute("DELETE FROM node_aliases WHERE node_id = ?1", [node_id])?;
344 for alias in aliases {
345 let key = alias.trim().to_lowercase();
346 if key.is_empty() {
347 continue;
348 }
349 conn.execute(
350 "INSERT INTO node_aliases (alias, node_id) VALUES (?1, ?2)
351 ON CONFLICT(alias) DO UPDATE SET node_id = excluded.node_id",
352 params![key, node_id],
353 )?;
354 }
355 Ok(())
356 }
357
358 pub fn index_fts(&self, node_id: &str, title: &str, content: &str, tags: &str) -> Result<()> {
360 self.index_fts_on(&self.conn, node_id, title, content, tags)
361 }
362
363 pub fn get_fts_content(&self, node_id: &str) -> Result<Option<String>> {
365 let mut stmt = self
366 .conn
367 .prepare("SELECT content FROM node_fts WHERE node_id = ?1 LIMIT 1")?;
368 let content: Option<String> = stmt
369 .query_row(params![node_id], |row| row.get(0))
370 .optional()?;
371 Ok(content)
372 }
373
374 pub(crate) fn index_fts_on(
375 &self,
376 conn: &Connection,
377 node_id: &str,
378 title: &str,
379 content: &str,
380 tags: &str,
381 ) -> Result<()> {
382 conn.execute("DELETE FROM node_fts WHERE node_id = ?1", [node_id])?;
383 conn.execute(
384 "INSERT INTO node_fts (node_id, title, content, tags) VALUES (?1, ?2, ?3, ?4)",
385 params![node_id, title, content, tags],
386 )?;
387 Ok(())
388 }
389
390 pub fn search_fts(&self, query: &str) -> Result<Vec<Node>> {
392 let escaped = escape_fts5_query(query)?;
393 let mut stmt = self.conn.prepare(
394 "SELECT n.id, n.node_type, n.title, n.file_path, n.symbol_hash, n.summary,
395 n.content_hash, n.scope, n.created_at, n.updated_at
396 FROM node_fts f
397 JOIN nodes n ON n.id = f.node_id
398 WHERE node_fts MATCH ?1
399 ORDER BY rank
400 LIMIT 50",
401 )?;
402
403 let rows = stmt.query_map(params![escaped], Self::map_node_row)?;
404
405 let mut results = Vec::new();
406 for r in rows {
407 results.push(r.map_err(BrainError::from)?);
408 }
409 Ok(results)
410 }
411
412 pub fn list_node_ids_by_type(&self, node_type: &str) -> Result<Vec<String>> {
414 let mut stmt = self
415 .conn
416 .prepare("SELECT id FROM nodes WHERE node_type = ?1 ORDER BY id")?;
417 let rows = stmt.query_map(params![node_type], |row| row.get(0))?;
418 let mut out = Vec::new();
419 for r in rows {
420 out.push(r?);
421 }
422 Ok(out)
423 }
424
425 pub fn get_node(&self, id: &str) -> Result<Option<Node>> {
427 let mut stmt = self.conn.prepare(
428 "SELECT id, node_type, title, file_path, symbol_hash, summary, content_hash, scope, created_at, updated_at
429 FROM nodes WHERE id = ?1",
430 )?;
431 let node = stmt
432 .query_row(params![id], Self::map_node_row)
433 .optional()?;
434 Ok(node)
435 }
436
437 pub(crate) fn map_node_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<Node> {
438 let type_str: String = row.get(1)?;
439 let node_type = NodeType::parse(&type_str).unwrap_or(NodeType::Concept);
440 let symbol_hash_raw: Option<i64> = row.get(4)?;
441 let scope: String = row
442 .get::<_, Option<String>>(7)?
443 .filter(|s| !s.is_empty())
444 .unwrap_or_else(|| crate::scopes::MAIN_SCOPE.to_string());
445 Ok(Node {
446 id: row.get(0)?,
447 node_type,
448 title: row.get(2)?,
449 file_path: row.get(3)?,
450 symbol_hash: symbol_hash_raw.map(|h| h as u64),
451 summary: row.get(5)?,
452 content_hash: row.get(6)?,
453 scope,
454 created_at: row.get(8)?,
455 updated_at: row.get(9)?,
456 })
457 }
458
459 pub fn count_nodes(&self) -> Result<usize> {
461 let count: usize = self
462 .conn
463 .query_row("SELECT COUNT(*) FROM nodes", [], |row| row.get(0))?;
464 Ok(count)
465 }
466
467 pub fn count_edges(&self) -> Result<usize> {
469 let count: usize = self
470 .conn
471 .query_row("SELECT COUNT(*) FROM edges", [], |row| row.get(0))?;
472 Ok(count)
473 }
474
475 pub fn count_fts_rows(&self) -> Result<usize> {
477 let count: usize = self
478 .conn
479 .query_row("SELECT COUNT(*) FROM node_fts", [], |row| row.get(0))?;
480 Ok(count)
481 }
482
483 pub fn count_pending_links(&self) -> Result<usize> {
485 let count: usize = self
486 .conn
487 .query_row("SELECT COUNT(*) FROM pending_links", [], |row| row.get(0))?;
488 Ok(count)
489 }
490
491 pub fn count_symbol_anchors(&self) -> Result<usize> {
493 let count: usize = self
494 .conn
495 .query_row("SELECT COUNT(*) FROM symbol_anchors", [], |row| row.get(0))?;
496 Ok(count)
497 }
498
499 pub fn get_all_node_ids(&self) -> Result<Vec<String>> {
501 let mut stmt = self
502 .conn
503 .prepare("SELECT id FROM nodes ORDER BY id ASC")?;
504 let rows = stmt.query_map([], |row| row.get(0))?;
505 let mut ids = Vec::new();
506 for r in rows {
507 ids.push(r?);
508 }
509 Ok(ids)
510 }
511
512 pub fn get_all_edges(&self) -> Result<Vec<Edge>> {
514 let mut stmt = self.conn.prepare(
515 "SELECT source_id, target_id, relation_type, weight, decay_rate, created_at FROM edges",
516 )?;
517 let rows = stmt.query_map([], |row| {
518 let weight: f64 = row.get(3)?;
519 let decay: f64 = row.get(4)?;
520 Ok(Edge {
521 source_id: row.get(0)?,
522 target_id: row.get(1)?,
523 relation_type: row.get(2)?,
524 weight: weight as f32,
525 decay_rate: decay as f32,
526 created_at: row.get(5)?,
527 })
528 })?;
529
530 let mut edges = Vec::new();
531 for r in rows {
532 edges.push(r?);
533 }
534 Ok(edges)
535 }
536
537 pub fn get_csr_edges(&self) -> Result<Vec<(String, String, f32)>> {
539 Ok(self
540 .get_all_edges()?
541 .into_iter()
542 .map(|e| (e.source_id, e.target_id, e.weight))
543 .collect())
544 }
545
546 pub fn get_all_nodes(&self) -> Result<Vec<Node>> {
548 let mut stmt = self.conn.prepare(
549 "SELECT id, node_type, title, file_path, symbol_hash, summary, content_hash, scope, created_at, updated_at
550 FROM nodes ORDER BY id ASC",
551 )?;
552 let rows = stmt.query_map([], Self::map_node_row)?;
553 let mut nodes = Vec::new();
554 for r in rows {
555 nodes.push(r?);
556 }
557 Ok(nodes)
558 }
559
560 #[allow(clippy::type_complexity)]
562 pub fn link_resolution_maps(
563 &self,
564 ) -> Result<(
565 HashSet<String>,
566 HashMap<String, String>,
567 HashMap<String, String>,
568 )> {
569 let mut ids = HashSet::new();
570 let mut titles: HashMap<String, String> = HashMap::new();
571 {
572 let mut stmt = self.conn.prepare("SELECT id, title FROM nodes")?;
573 let rows = stmt.query_map([], |row| {
574 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
575 })?;
576 for r in rows {
577 let (id, title) = r?;
578 titles
579 .entry(title.to_lowercase())
580 .or_insert_with(|| id.clone());
581 ids.insert(id);
582 }
583 }
584
585 let mut aliases: HashMap<String, String> = HashMap::new();
586 {
587 let mut stmt = self.conn.prepare("SELECT alias, node_id FROM node_aliases")?;
588 let rows = stmt.query_map([], |row| {
589 Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
590 })?;
591 for r in rows {
592 let (alias, node_id) = r?;
593 aliases.insert(alias, node_id);
594 }
595 }
596
597 Ok((ids, aliases, titles))
598 }
599
600 pub fn resolve_pending_links(&self) -> Result<(usize, usize)> {
602 let (ids, aliases, titles) = self.link_resolution_maps()?;
603 let pending: Vec<(String, String, String, i64)> = {
604 let mut stmt = self.conn.prepare(
605 "SELECT source_id, raw_target, relation_type, created_at FROM pending_links",
606 )?;
607 let rows = stmt.query_map([], |row| {
608 Ok((
609 row.get::<_, String>(0)?,
610 row.get::<_, String>(1)?,
611 row.get::<_, String>(2)?,
612 row.get::<_, i64>(3)?,
613 ))
614 })?;
615 let mut v = Vec::new();
616 for r in rows {
617 v.push(r?);
618 }
619 v
620 };
621
622 let symbol_ids: std::collections::HashSet<String> = ids
624 .iter()
625 .filter(|id| id.starts_with("symbol/"))
626 .cloned()
627 .collect();
628
629 let mut resolved = 0usize;
630 for (source_id, raw_target, relation_type, created_at) in &pending {
631 let target_id = if let Some(sym_path) = raw_target.strip_prefix("symbol:") {
632 crate::symbols::parse_symbol_path(sym_path)
633 .and_then(|s| crate::symbols::resolve_symbol_ref(&s, &symbol_ids))
634 .or_else(|| {
635 crate::symbols::parse_symbol_path(sym_path).and_then(|s| {
637 crate::id::resolve_link_target(
638 &s.symbol_name,
639 &ids,
640 &aliases,
641 &titles,
642 )
643 })
644 })
645 } else {
646 crate::id::resolve_link_target(raw_target, &ids, &aliases, &titles)
647 };
648
649 if let Some(target_id) = target_id {
650 let edge = Edge {
651 source_id: source_id.clone(),
652 target_id,
653 relation_type: relation_type.clone(),
654 weight: 1.0,
655 decay_rate: 0.0,
656 created_at: *created_at,
657 };
658 self.insert_edge(&edge)?;
659 self.conn.execute(
660 "DELETE FROM pending_links WHERE source_id = ?1 AND raw_target = ?2 AND relation_type = ?3",
661 params![source_id, raw_target, relation_type],
662 )?;
663 resolved += 1;
664 }
665 }
666
667 let still = self.count_pending_links()?;
668 Ok((resolved, still))
669 }
670}
671
672#[cfg(test)]
673mod tests {
674 use super::*;
675 use crate::types::NodeType;
676
677 fn sample_node(id: &str) -> Node {
678 Node {
679 id: id.to_string(),
680 node_type: NodeType::Concept,
681 title: id.to_string(),
682 file_path: Some(format!("{id}.md")),
683 symbol_hash: None,
684 summary: Some("summary".into()),
685 content_hash: Some("abc".into()),
686 scope: crate::scopes::MAIN_SCOPE.to_string(),
687 created_at: 100,
688 updated_at: 100,
689 }
690 }
691
692 #[test]
693 fn fts_idempotent_reindex() {
694 let db = Database::open_in_memory().unwrap();
695 let n = sample_node("docs/raft");
696 db.insert_node(&n).unwrap();
697 db.index_fts(&n.id, &n.title, "raft consensus protocol", "raft").unwrap();
698 db.index_fts(&n.id, &n.title, "raft consensus protocol v2", "raft").unwrap();
699 assert_eq!(db.count_fts_rows().unwrap(), 1);
700 let hits = db.search_fts("raft").unwrap();
701 assert_eq!(hits.len(), 1);
702 }
703
704 #[test]
705 fn edge_requires_endpoints() {
706 let db = Database::open_in_memory().unwrap();
707 let edge = Edge {
708 source_id: "a".into(),
709 target_id: "b".into(),
710 relation_type: "relates_to".into(),
711 weight: 1.0,
712 decay_rate: 0.0,
713 created_at: 1,
714 };
715 assert!(db.insert_edge(&edge).is_err());
716 }
717
718 #[test]
719 fn full_edge_roundtrip() {
720 let db = Database::open_in_memory().unwrap();
721 db.insert_node(&sample_node("a")).unwrap();
722 db.insert_node(&sample_node("b")).unwrap();
723 let edge = Edge {
724 source_id: "a".into(),
725 target_id: "b".into(),
726 relation_type: "implements".into(),
727 weight: 0.75,
728 decay_rate: 0.1,
729 created_at: 42,
730 };
731 db.insert_edge(&edge).unwrap();
732 let edges = db.get_all_edges().unwrap();
733 assert_eq!(edges.len(), 1);
734 assert_eq!(edges[0].relation_type, "implements");
735 assert!((edges[0].weight - 0.75).abs() < 1e-6);
736 assert!((edges[0].decay_rate - 0.1).abs() < 1e-6);
737 assert_eq!(edges[0].created_at, 42);
738 }
739
740 #[test]
741 fn preserves_created_at_on_upsert() {
742 let db = Database::open_in_memory().unwrap();
743 let mut n = sample_node("x");
744 n.created_at = 10;
745 n.updated_at = 10;
746 db.insert_node(&n).unwrap();
747 n.created_at = 999;
748 n.updated_at = 20;
749 n.title = "updated".into();
750 db.insert_node(&n).unwrap();
751 let got = db.get_node("x").unwrap().unwrap();
752 assert_eq!(got.created_at, 10);
753 assert_eq!(got.updated_at, 20);
754 assert_eq!(got.title, "updated");
755 }
756}