1use std::collections::{HashMap, HashSet};
31use std::time::{SystemTime, UNIX_EPOCH};
32
33use sqlx::postgres::{PgPoolOptions, PgRow};
34use sqlx::{PgPool, Postgres, QueryBuilder, Row};
35
36use super::algo::{CsrGraph, pagerank_default};
37use super::lifecycle::now_iso;
38use super::recall::{W_ACCESS, W_FTS, W_GRAPH, W_IMPORTANCE, W_RECENCY, compute_recency};
39use super::schema::GRAPH_SCHEMA_VERSION;
40use super::types::{
41 EdgeDirection, GraphEdge, GraphNode, ScoredNode, escape_like, join_csv, split_csv,
42};
43use crate::error::{KernelError, Result};
44
45const NODE_COLUMNS: &str = "id, node_type, title, tags, projects, agents, \
47 created, updated, body, importance, access_count, accessed_at";
48
49fn pg_err(e: sqlx::Error) -> KernelError {
51 KernelError::Store(format!("sqlx: {e:?}"))
52}
53
54fn row_to_node(row: &PgRow) -> GraphNode {
56 let tags: String = row.get("tags");
57 let projects: String = row.get("projects");
58 let agents: String = row.get("agents");
59 GraphNode {
60 id: row.get("id"),
61 node_type: row.get("node_type"),
62 title: row.get("title"),
63 tags: split_csv(&tags),
64 projects: split_csv(&projects),
65 agents: split_csv(&agents),
66 created: row.get("created"),
67 updated: row.get("updated"),
68 body: row.get("body"),
69 importance: row.get("importance"),
70 access_count: row.get("access_count"),
71 accessed_at: row.get("accessed_at"),
72 }
73}
74
75fn row_to_edge(row: &PgRow) -> GraphEdge {
77 GraphEdge {
78 id: row.get("id"),
79 source: row.get("source"),
80 target: row.get("target"),
81 relation: row.get("relation"),
82 weight: row.get("weight"),
83 ts: row.get("ts"),
84 }
85}
86
87fn search_patterns(query: &str) -> Vec<String> {
92 query
93 .split_whitespace()
94 .map(|t| format!("%{}%", escape_like(t)))
95 .collect()
96}
97
98fn is_identifier_safe(prefix: &str) -> bool {
105 let mut chars = prefix.chars();
106 match chars.next() {
107 None => true,
108 Some(first) if first == '_' || first.is_ascii_alphabetic() => {
109 chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
110 }
111 Some(_) => false,
112 }
113}
114
115pub struct SqlxPgGraph {
133 pool: PgPool,
134 table_prefix: String,
135}
136
137impl SqlxPgGraph {
138 pub async fn from_pool(pool: PgPool) -> Result<Self> {
141 Self::from_pool_with_prefix(pool, "").await
142 }
143
144 pub async fn from_pool_with_prefix(pool: PgPool, prefix: &str) -> Result<Self> {
148 if !is_identifier_safe(prefix) {
149 return Err(KernelError::Store(format!(
150 "invalid table prefix {prefix:?}: only ASCII letters, digits, and underscore are allowed (and the first character must not be a digit)"
151 )));
152 }
153 let graph = Self {
154 pool,
155 table_prefix: prefix.to_string(),
156 };
157 graph.init_schema().await?;
158 graph.migrate().await?;
159 Ok(graph)
160 }
161
162 pub async fn connect(url: &str) -> Result<Self> {
166 Self::connect_with_prefix(url, "").await
167 }
168
169 pub async fn connect_with_prefix(url: &str, prefix: &str) -> Result<Self> {
172 let pool = PgPoolOptions::new()
173 .max_connections(8)
174 .connect(url)
175 .await
176 .map_err(pg_err)?;
177 Self::from_pool_with_prefix(pool, prefix).await
178 }
179
180 pub fn pool(&self) -> &PgPool {
188 &self.pool
189 }
190
191 fn nodes_tbl(&self) -> String {
194 format!("{}nodes", self.table_prefix)
195 }
196
197 fn edges_tbl(&self) -> String {
198 format!("{}edges", self.table_prefix)
199 }
200
201 fn meta_tbl(&self) -> String {
202 format!("{}_meta", self.table_prefix)
203 }
204
205 async fn init_schema(&self) -> Result<()> {
212 let p = &self.table_prefix;
213 let nodes = self.nodes_tbl();
214 let edges = self.edges_tbl();
215 let meta = self.meta_tbl();
216 let idx = |name: &str| format!("{p}{name}");
217
218 let stmts: Vec<String> = vec![
219 format!(
220 "CREATE TABLE IF NOT EXISTS {nodes} (
221 id TEXT PRIMARY KEY,
222 node_type TEXT NOT NULL,
223 title TEXT NOT NULL,
224 tags TEXT NOT NULL DEFAULT '',
225 projects TEXT NOT NULL DEFAULT '',
226 agents TEXT NOT NULL DEFAULT '',
227 created TEXT NOT NULL,
228 updated TEXT NOT NULL,
229 body TEXT NOT NULL DEFAULT '',
230 importance DOUBLE PRECISION NOT NULL DEFAULT 0.5,
231 access_count BIGINT NOT NULL DEFAULT 0,
232 accessed_at TEXT NOT NULL DEFAULT ''
233 )"
234 ),
235 format!(
236 "CREATE TABLE IF NOT EXISTS {edges} (
237 id TEXT PRIMARY KEY,
238 source TEXT NOT NULL,
239 target TEXT NOT NULL,
240 relation TEXT NOT NULL DEFAULT 'related',
241 weight DOUBLE PRECISION NOT NULL DEFAULT 1.0,
242 ts TEXT NOT NULL
243 )"
244 ),
245 format!(
246 "CREATE INDEX IF NOT EXISTS {} ON {edges}(source)",
247 idx("idx_edges_source")
248 ),
249 format!(
250 "CREATE INDEX IF NOT EXISTS {} ON {edges}(target)",
251 idx("idx_edges_target")
252 ),
253 format!(
254 "CREATE UNIQUE INDEX IF NOT EXISTS {} ON {edges}(source, target, relation)",
255 idx("idx_edges_src_tgt_rel")
256 ),
257 format!(
258 "CREATE INDEX IF NOT EXISTS {} ON {edges}(source, relation)",
259 idx("idx_edges_src_rel")
260 ),
261 format!(
262 "CREATE INDEX IF NOT EXISTS {} ON {edges}(target, relation)",
263 idx("idx_edges_tgt_rel")
264 ),
265 format!(
266 "CREATE INDEX IF NOT EXISTS {} ON {nodes}(node_type)",
267 idx("idx_nodes_type")
268 ),
269 format!(
270 "CREATE INDEX IF NOT EXISTS {} ON {nodes}(updated DESC)",
271 idx("idx_nodes_updated")
272 ),
273 format!(
274 "CREATE INDEX IF NOT EXISTS {} ON {nodes}(importance DESC)",
275 idx("idx_nodes_importance")
276 ),
277 format!(
278 "CREATE INDEX IF NOT EXISTS {} ON {nodes}(accessed_at DESC)",
279 idx("idx_nodes_accessed")
280 ),
281 format!(
282 "CREATE INDEX IF NOT EXISTS {} ON {nodes}(created)",
283 idx("idx_nodes_created")
284 ),
285 format!(
286 "CREATE TABLE IF NOT EXISTS {meta} (key TEXT PRIMARY KEY, value TEXT NOT NULL)"
287 ),
288 format!(
289 "INSERT INTO {meta} (key, value) VALUES ('graph_schema_version', '{}')
290 ON CONFLICT (key) DO NOTHING",
291 GRAPH_SCHEMA_VERSION
292 ),
293 ];
294
295 for ddl in &stmts {
296 sqlx::query(ddl).execute(&self.pool).await.map_err(pg_err)?;
297 }
298 Ok(())
299 }
300
301 pub async fn current_version(&self) -> Result<u32> {
303 let meta = self.meta_tbl();
304 let row = sqlx::query(&format!(
305 "SELECT value FROM {meta} WHERE key = 'graph_schema_version'"
306 ))
307 .fetch_optional(&self.pool)
308 .await
309 .map_err(pg_err)?;
310 match row {
311 Some(r) => {
312 let s: String = r.get("value");
313 Ok(s.parse().unwrap_or(0))
314 }
315 None => Ok(0),
316 }
317 }
318
319 pub async fn migrate(&self) -> Result<u32> {
324 let current = self.current_version().await?;
325 if current >= GRAPH_SCHEMA_VERSION {
326 return Ok(current);
327 }
328 let p = &self.table_prefix;
329 let nodes = self.nodes_tbl();
330 let edges = self.edges_tbl();
331 let meta = self.meta_tbl();
332 let idx = |name: &str| format!("{p}{name}");
333
334 let mut tx = self.pool.begin().await.map_err(pg_err)?;
335 let mut v = current;
336 if v < 2 {
338 sqlx::query(&format!(
339 "CREATE INDEX IF NOT EXISTS {} ON {nodes}(created)",
340 idx("idx_nodes_created")
341 ))
342 .execute(&mut *tx)
343 .await
344 .map_err(pg_err)?;
345 v = 2;
346 }
347 if v < 3 {
349 sqlx::query(&format!(
350 "CREATE INDEX IF NOT EXISTS {} ON {edges}(source, relation)",
351 idx("idx_edges_src_rel")
352 ))
353 .execute(&mut *tx)
354 .await
355 .map_err(pg_err)?;
356 sqlx::query(&format!(
357 "CREATE INDEX IF NOT EXISTS {} ON {edges}(target, relation)",
358 idx("idx_edges_tgt_rel")
359 ))
360 .execute(&mut *tx)
361 .await
362 .map_err(pg_err)?;
363 v = 3;
364 }
365 sqlx::query(&format!(
366 "UPDATE {meta} SET value = $1 WHERE key = 'graph_schema_version'"
367 ))
368 .bind(v.to_string())
369 .execute(&mut *tx)
370 .await
371 .map_err(pg_err)?;
372 tx.commit().await.map_err(pg_err)?;
373 Ok(v)
374 }
375
376 pub async fn upsert_node(&self, node: &GraphNode) -> Result<()> {
380 let tags = join_csv(&node.tags);
381 let projects = join_csv(&node.projects);
382 let agents = join_csv(&node.agents);
383 let nodes = self.nodes_tbl();
384 sqlx::query(&format!(
385 "INSERT INTO {nodes} (id, node_type, title, tags, projects, agents, created, updated, body, importance, access_count, accessed_at)
386 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)
387 ON CONFLICT (id) DO UPDATE SET
388 node_type=EXCLUDED.node_type, title=EXCLUDED.title, tags=EXCLUDED.tags,
389 projects=EXCLUDED.projects, agents=EXCLUDED.agents, created=EXCLUDED.created,
390 updated=EXCLUDED.updated, body=EXCLUDED.body, importance=EXCLUDED.importance,
391 access_count=EXCLUDED.access_count, accessed_at=EXCLUDED.accessed_at"
392 ))
393 .bind(node.id.as_str())
394 .bind(node.node_type.as_str())
395 .bind(node.title.as_str())
396 .bind(tags.as_str())
397 .bind(projects.as_str())
398 .bind(agents.as_str())
399 .bind(node.created.as_str())
400 .bind(node.updated.as_str())
401 .bind(node.body.as_str())
402 .bind(node.importance)
403 .bind(node.access_count)
404 .bind(node.accessed_at.as_str())
405 .execute(&self.pool)
406 .await
407 .map_err(pg_err)?;
408 Ok(())
409 }
410
411 pub async fn read_node(&self, id: &str) -> Result<Option<GraphNode>> {
413 let nodes = self.nodes_tbl();
414 let row = sqlx::query(&format!("SELECT {NODE_COLUMNS} FROM {nodes} WHERE id = $1"))
415 .bind(id)
416 .fetch_optional(&self.pool)
417 .await
418 .map_err(pg_err)?;
419 Ok(row.as_ref().map(row_to_node))
420 }
421
422 pub async fn delete_node(&self, id: &str) -> Result<bool> {
424 let nodes = self.nodes_tbl();
425 let res = sqlx::query(&format!("DELETE FROM {nodes} WHERE id = $1"))
426 .bind(id)
427 .execute(&self.pool)
428 .await
429 .map_err(pg_err)?;
430 Ok(res.rows_affected() > 0)
431 }
432
433 pub async fn append_edge(&self, edge: &GraphEdge) -> Result<()> {
435 let edges = self.edges_tbl();
436 sqlx::query(&format!(
437 "INSERT INTO {edges} (id, source, target, relation, weight, ts)
438 VALUES ($1,$2,$3,$4,$5,$6) ON CONFLICT DO NOTHING"
439 ))
440 .bind(edge.id.as_str())
441 .bind(edge.source.as_str())
442 .bind(edge.target.as_str())
443 .bind(edge.relation.as_str())
444 .bind(edge.weight)
445 .bind(edge.ts.as_str())
446 .execute(&self.pool)
447 .await
448 .map_err(pg_err)?;
449 Ok(())
450 }
451
452 pub async fn edges_for_node(&self, node_id: &str) -> Result<Vec<GraphEdge>> {
454 let edges = self.edges_tbl();
455 let rows = sqlx::query(&format!(
456 "SELECT id, source, target, relation, weight, ts FROM {edges} \
457 WHERE source = $1 OR target = $1"
458 ))
459 .bind(node_id)
460 .fetch_all(&self.pool)
461 .await
462 .map_err(pg_err)?;
463 Ok(rows.iter().map(row_to_edge).collect())
464 }
465
466 pub async fn delete_edge(&self, id: &str) -> Result<bool> {
468 let edges = self.edges_tbl();
469 let res = sqlx::query(&format!("DELETE FROM {edges} WHERE id = $1"))
470 .bind(id)
471 .execute(&self.pool)
472 .await
473 .map_err(pg_err)?;
474 Ok(res.rows_affected() > 0)
475 }
476
477 pub async fn append_edges(&self, edges: &[GraphEdge]) -> Result<()> {
484 if edges.is_empty() {
485 return Ok(());
486 }
487 let edges_tbl = self.edges_tbl();
488 let sql = format!(
489 "INSERT INTO {edges_tbl} (id, source, target, relation, weight, ts)
490 VALUES ($1,$2,$3,$4,$5,$6) ON CONFLICT DO NOTHING"
491 );
492 const CHUNK: usize = 5000;
493 for chunk in edges.chunks(CHUNK) {
494 let mut tx = self.pool.begin().await.map_err(pg_err)?;
495 for e in chunk {
496 sqlx::query(&sql)
497 .bind(e.id.as_str())
498 .bind(e.source.as_str())
499 .bind(e.target.as_str())
500 .bind(e.relation.as_str())
501 .bind(e.weight)
502 .bind(e.ts.as_str())
503 .execute(&mut *tx)
504 .await
505 .map_err(pg_err)?;
506 }
507 tx.commit().await.map_err(pg_err)?;
508 }
509 Ok(())
510 }
511
512 pub async fn append_edges_in_tx(
518 &self,
519 tx: &mut sqlx::PgConnection,
520 edges: &[GraphEdge],
521 ) -> Result<()> {
522 if edges.is_empty() {
523 return Ok(());
524 }
525 let edges_tbl = self.edges_tbl();
526 let sql = format!(
527 "INSERT INTO {edges_tbl} (id, source, target, relation, weight, ts)
528 VALUES ($1,$2,$3,$4,$5,$6) ON CONFLICT DO NOTHING"
529 );
530 for e in edges {
531 sqlx::query(&sql)
532 .bind(e.id.as_str())
533 .bind(e.source.as_str())
534 .bind(e.target.as_str())
535 .bind(e.relation.as_str())
536 .bind(e.weight)
537 .bind(e.ts.as_str())
538 .execute(&mut *tx)
539 .await
540 .map_err(pg_err)?;
541 }
542 Ok(())
543 }
544
545 pub async fn edges_for_node_dir(
549 &self,
550 node_id: &str,
551 dir: EdgeDirection,
552 relation: Option<&str>,
553 ) -> Result<Vec<GraphEdge>> {
554 let edges = self.edges_tbl();
555 let dir_clause = match dir {
556 EdgeDirection::Out => "source = $1",
557 EdgeDirection::In => "target = $1",
558 EdgeDirection::Both => "(source = $1 OR target = $1)",
559 };
560 let rows = if let Some(r) = relation {
561 let sql = format!(
562 "SELECT id, source, target, relation, weight, ts FROM {edges} \
563 WHERE {dir_clause} AND relation = $2 ORDER BY weight DESC"
564 );
565 sqlx::query(&sql)
566 .bind(node_id)
567 .bind(r)
568 .fetch_all(&self.pool)
569 .await
570 .map_err(pg_err)?
571 } else {
572 let sql = format!(
573 "SELECT id, source, target, relation, weight, ts FROM {edges} \
574 WHERE {dir_clause} ORDER BY weight DESC"
575 );
576 sqlx::query(&sql)
577 .bind(node_id)
578 .fetch_all(&self.pool)
579 .await
580 .map_err(pg_err)?
581 };
582 Ok(rows.iter().map(row_to_edge).collect())
583 }
584
585 pub async fn neighbors_weighted(
591 &self,
592 seed_ids: &[String],
593 dir: EdgeDirection,
594 relation: Option<&str>,
595 ) -> Result<Vec<(String, f64)>> {
596 if seed_ids.is_empty() {
597 return Ok(vec![]);
598 }
599 let edges = self.edges_tbl();
600 const MAX_SEEDS: usize = 100;
601 let seed_ids = if seed_ids.len() > MAX_SEEDS {
602 &seed_ids[..MAX_SEEDS]
603 } else {
604 seed_ids
605 };
606 let seed_arr: Vec<String> = seed_ids.to_vec();
607 let seed_set: HashSet<&str> = seed_ids.iter().map(String::as_str).collect();
608
609 let halves: &[&str] = match dir {
610 EdgeDirection::Out => &["source"],
611 EdgeDirection::In => &["target"],
612 EdgeDirection::Both => &["source", "target"],
613 };
614 let mut weights: HashMap<String, f64> = HashMap::new();
615
616 for &follow in halves {
617 let select_col = if follow == "source" {
620 "target"
621 } else {
622 "source"
623 };
624 let rel_clause = relation.map(|_| " AND relation = $2").unwrap_or("");
625 let sql = format!(
626 "SELECT {select_col} AS nb, SUM(weight) AS w FROM {edges} \
627 WHERE {follow} = ANY($1){rel_clause} GROUP BY {select_col}"
628 );
629 let rows = if let Some(r) = relation {
630 sqlx::query(&sql)
631 .bind(&seed_arr)
632 .bind(r)
633 .fetch_all(&self.pool)
634 .await
635 .map_err(pg_err)?
636 } else {
637 sqlx::query(&sql)
638 .bind(&seed_arr)
639 .fetch_all(&self.pool)
640 .await
641 .map_err(pg_err)?
642 };
643 for row in &rows {
644 let nb: String = row.get("nb");
645 let w: f64 = row.get("w");
646 if !seed_set.contains(nb.as_str()) {
647 *weights.entry(nb).or_default() += w;
648 }
649 }
650 }
651
652 let mut result: Vec<(String, f64)> = weights.into_iter().collect();
653 result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
654 Ok(result)
655 }
656
657 pub async fn remove_edges_for_node(&self, node_id: &str) -> Result<()> {
660 let edges = self.edges_tbl();
661 sqlx::query(&format!(
662 "DELETE FROM {edges} WHERE source = $1 OR target = $1"
663 ))
664 .bind(node_id)
665 .execute(&self.pool)
666 .await
667 .map_err(pg_err)?;
668 Ok(())
669 }
670
671 pub async fn remove_edges_for_node_in_tx(
674 &self,
675 tx: &mut sqlx::PgConnection,
676 node_id: &str,
677 ) -> Result<()> {
678 let edges = self.edges_tbl();
679 sqlx::query(&format!(
680 "DELETE FROM {edges} WHERE source = $1 OR target = $1"
681 ))
682 .bind(node_id)
683 .execute(tx)
684 .await
685 .map_err(pg_err)?;
686 Ok(())
687 }
688
689 pub async fn search_nodes(&self, query: &str, limit: usize) -> Result<Vec<GraphNode>> {
695 let terms = search_patterns(query);
696 if terms.is_empty() {
697 return Ok(vec![]);
698 }
699 let nodes = self.nodes_tbl();
700 let mut qb = QueryBuilder::<Postgres>::new("SELECT ");
701 qb.push(NODE_COLUMNS);
702 qb.push(" FROM ");
703 qb.push(nodes.as_str());
704 qb.push(" WHERE ");
705 for (i, term) in terms.iter().enumerate() {
706 if i > 0 {
707 qb.push(" AND ");
708 }
709 qb.push("(title || ' ' || body || ' ' || tags) ILIKE ");
710 qb.push_bind(term.clone());
711 qb.push(" ESCAPE '\\'");
712 }
713 qb.push(" ORDER BY importance DESC, updated DESC LIMIT ");
714 qb.push_bind(limit as i64);
715 let rows = qb.build().fetch_all(&self.pool).await.map_err(pg_err)?;
716 Ok(rows.iter().map(row_to_node).collect())
717 }
718
719 pub async fn query_nodes(
722 &self,
723 tag: Option<&str>,
724 node_type: Option<&str>,
725 project: Option<&str>,
726 limit: usize,
727 ) -> Result<Vec<GraphNode>> {
728 let limit = limit.min(200) as i64;
729 let nodes = self.nodes_tbl();
730 let mut qb = QueryBuilder::<Postgres>::new("SELECT ");
731 qb.push(NODE_COLUMNS);
732 qb.push(" FROM ");
733 qb.push(nodes.as_str());
734 let mut filtered = false;
735 if let Some(t) = tag {
736 qb.push(" WHERE (',' || tags || ',') ILIKE ('%,' || ");
737 qb.push_bind(escape_like(t));
738 qb.push(" || ',%') ESCAPE '\\'");
739 filtered = true;
740 }
741 if let Some(nt) = node_type {
742 qb.push(if filtered {
743 " AND node_type = "
744 } else {
745 " WHERE node_type = "
746 });
747 qb.push_bind(nt);
748 filtered = true;
749 }
750 if let Some(p) = project {
751 qb.push(if filtered { " AND " } else { " WHERE " });
752 qb.push("(',' || projects || ',') ILIKE ('%,' || ");
753 qb.push_bind(escape_like(p));
754 qb.push(" || ',%') ESCAPE '\\'");
755 }
756 qb.push(" ORDER BY updated DESC LIMIT ");
757 qb.push_bind(limit);
758 let rows = qb.build().fetch_all(&self.pool).await.map_err(pg_err)?;
759 Ok(rows.iter().map(row_to_node).collect())
760 }
761
762 pub async fn smart_recall(
769 &self,
770 project: Option<&str>,
771 hint: Option<&str>,
772 limit: usize,
773 ) -> Result<Vec<ScoredNode>> {
774 let now_secs = SystemTime::now()
775 .duration_since(UNIX_EPOCH)
776 .unwrap_or_default()
777 .as_secs();
778
779 let fts_ids: HashSet<String> = match hint {
781 Some(h) if !h.is_empty() => self
782 .search_nodes(h, limit * 4)
783 .await?
784 .into_iter()
785 .map(|n| n.id)
786 .collect(),
787 _ => HashSet::new(),
788 };
789
790 let candidate_limit = (limit * 4).max(40) as i64;
792 let nodes = self.nodes_tbl();
793 let mut qb = QueryBuilder::<Postgres>::new("SELECT ");
794 qb.push(NODE_COLUMNS);
795 qb.push(" FROM ");
796 qb.push(nodes.as_str());
797 qb.push(" WHERE (',' || tags || ',') NOT ILIKE '%,stale,%'");
798 if let Some(p) = project {
799 qb.push(" AND (',' || projects || ',') ILIKE ('%,' || ");
800 qb.push_bind(escape_like(p));
801 qb.push(" || ',%') ESCAPE '\\'");
802 }
803 qb.push(" ORDER BY importance DESC, updated DESC LIMIT ");
804 qb.push_bind(candidate_limit);
805 let rows = qb.build().fetch_all(&self.pool).await.map_err(pg_err)?;
806 let candidates: Vec<GraphNode> = rows.iter().map(row_to_node).collect();
807
808 let mut scored: Vec<ScoredNode> = candidates
810 .into_iter()
811 .map(|node| {
812 let recency = compute_recency(&node.updated, now_secs);
813 let importance = node.importance;
814 let access_freq = (node.access_count.max(0) as f64 / 20.0).min(1.0);
815 let fts_match = if fts_ids.contains(&node.id) { 1.0 } else { 0.0 };
816 let score = W_RECENCY * recency
817 + W_IMPORTANCE * importance
818 + W_ACCESS * access_freq
819 + W_FTS * fts_match;
820 ScoredNode { node, score }
821 })
822 .collect();
823 scored.sort_by(|a, b| {
824 b.score
825 .partial_cmp(&a.score)
826 .unwrap_or(std::cmp::Ordering::Equal)
827 });
828 scored.truncate(limit);
829
830 if scored.len() > 1 {
834 const MAX_GRAPH_BOOST_PARTICIPANTS: usize = 100;
835 let candidate_ids: Vec<String> = scored
836 .iter()
837 .take(MAX_GRAPH_BOOST_PARTICIPANTS)
838 .map(|sn| sn.node.id.clone())
839 .collect();
840 let edges = self.edges_tbl();
841 let sql = format!(
842 "SELECT id, source, target, relation, weight, ts FROM {edges} \
843 WHERE source = ANY($1) AND target = ANY($1)"
844 );
845 let sub_edges: Vec<GraphEdge> = sqlx::query(&sql)
846 .bind(&candidate_ids)
847 .fetch_all(&self.pool)
848 .await
849 .map(|rows| rows.iter().map(row_to_edge).collect())
850 .unwrap_or_default();
851 let csr = CsrGraph::from_edges(&candidate_ids, &sub_edges);
852 let pr = pagerank_default(&csr);
853 let max_pr = pr.iter().copied().fold(0.0_f64, f64::max).max(1e-12);
854 let pr_map: HashMap<String, f64> = candidate_ids
855 .iter()
856 .zip(pr.iter())
857 .map(|(id, &s)| (id.clone(), s / max_pr))
858 .collect();
859 for sn in &mut scored {
860 let boost = pr_map.get(&sn.node.id).copied().unwrap_or(0.0);
861 sn.score += W_GRAPH * boost;
862 }
863 scored.sort_by(|a, b| {
864 b.score
865 .partial_cmp(&a.score)
866 .unwrap_or(std::cmp::Ordering::Equal)
867 });
868 }
869
870 if !scored.is_empty() {
873 let now = now_iso();
874 let ids: Vec<String> = scored.iter().map(|sn| sn.node.id.clone()).collect();
875 let sql = format!(
876 "UPDATE {nodes} SET access_count = access_count + 1, accessed_at = $1 WHERE id = ANY($2)"
877 );
878 let _ = sqlx::query(&sql)
879 .bind(now)
880 .bind(&ids)
881 .execute(&self.pool)
882 .await;
883 }
884
885 Ok(scored)
886 }
887
888 pub async fn related_nodes(&self, start_id: &str, depth: usize) -> Result<Vec<String>> {
892 let edges = self.edges_tbl();
893 let sql = format!(
897 "WITH RECURSIVE bfs(node_id, lvl) AS (
898 SELECT nb.node_id, 1 FROM (
899 SELECT target AS node_id FROM {edges} WHERE source = $1
900 UNION
901 SELECT source AS node_id FROM {edges} WHERE target = $1
902 ) nb
903 UNION
904 SELECT CASE WHEN e.source = bfs.node_id THEN e.target ELSE e.source END,
905 bfs.lvl + 1
906 FROM bfs
907 JOIN {edges} e ON e.source = bfs.node_id OR e.target = bfs.node_id
908 WHERE bfs.lvl < $2
909 )
910 SELECT DISTINCT node_id FROM bfs WHERE node_id <> $1 LIMIT 500"
911 );
912 let rows = sqlx::query(&sql)
913 .bind(start_id)
914 .bind(depth as i32)
915 .fetch_all(&self.pool)
916 .await
917 .map_err(pg_err)?;
918 Ok(rows
919 .iter()
920 .map(|r| {
921 let id: String = r.get("node_id");
922 id
923 })
924 .collect())
925 }
926}
927
928#[cfg(test)]
929mod tests {
930 use super::*;
931
932 fn pg_url() -> Option<String> {
934 std::env::var("LLMKERNEL_PG_URL").ok()
935 }
936
937 fn sample_node(id: &str) -> GraphNode {
938 GraphNode {
939 id: id.to_string(),
940 node_type: "concept".to_string(),
941 title: format!("Node {id}"),
942 body: "sqlx pg test body".to_string(),
943 tags: vec!["sqlx".to_string()],
944 projects: vec![],
945 agents: vec![],
946 created: "2026-01-01T00:00:00Z".to_string(),
947 updated: "2026-01-01T00:00:00Z".to_string(),
948 importance: 0.5,
949 access_count: 0,
950 accessed_at: String::new(),
951 }
952 }
953
954 fn sample_edge(id: &str, src: &str, tgt: &str, rel: &str, w: f64) -> GraphEdge {
955 GraphEdge {
956 id: id.to_string(),
957 source: src.to_string(),
958 target: tgt.to_string(),
959 relation: rel.to_string(),
960 weight: w,
961 ts: "2026-01-01T00:00:00Z".to_string(),
962 }
963 }
964
965 async fn cleanup(pool: &PgPool, prefix: &str) {
968 let _ = sqlx::query(&format!("DROP TABLE IF EXISTS {}nodes", prefix))
969 .execute(pool)
970 .await;
971 let _ = sqlx::query(&format!("DROP TABLE IF EXISTS {}edges", prefix))
972 .execute(pool)
973 .await;
974 let _ = sqlx::query(&format!("DROP TABLE IF EXISTS {}_meta", prefix))
975 .execute(pool)
976 .await;
977 }
978
979 #[test]
984 fn search_patterns_escapes_and_wraps() {
985 assert!(search_patterns("").is_empty());
986 assert_eq!(search_patterns("rust"), vec!["%rust%".to_string()]);
987 assert_eq!(search_patterns("100%"), vec!["%100\\%%".to_string()]);
988 assert_eq!(search_patterns("a_b"), vec!["%a\\_b%".to_string()]);
989 assert_eq!(
990 search_patterns("rust db"),
991 vec!["%rust%".to_string(), "%db%".to_string()]
992 );
993 }
994
995 #[tokio::test]
999 async fn basic_crud_roundtrip() {
1000 let Some(url) = pg_url() else {
1001 eprintln!("skip: LLMKERNEL_PG_URL unset");
1002 return;
1003 };
1004 let prefix = "lk_sg1_";
1005 let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1006 .await
1007 .expect("connect");
1008
1009 assert!(g.read_node("n1").await.unwrap().is_none());
1010 g.upsert_node(&sample_node("n1")).await.unwrap();
1011 let loaded = g.read_node("n1").await.unwrap().unwrap();
1012 assert_eq!(loaded.title, "Node n1");
1013 assert_eq!(loaded.tags, vec!["sqlx".to_string()]);
1014
1015 let mut updated = sample_node("n1");
1017 updated.title = "Updated".into();
1018 g.upsert_node(&updated).await.unwrap();
1019 assert_eq!(g.read_node("n1").await.unwrap().unwrap().title, "Updated");
1020
1021 g.upsert_node(&sample_node("n2")).await.unwrap();
1023 g.append_edge(&sample_edge("e1", "n1", "n2", "related", 1.0))
1024 .await
1025 .unwrap();
1026 assert_eq!(g.edges_for_node("n1").await.unwrap().len(), 1);
1027 assert!(g.delete_edge("e1").await.unwrap());
1028 assert!(!g.delete_edge("e1").await.unwrap());
1029 assert_eq!(g.edges_for_node("n1").await.unwrap().len(), 0);
1030
1031 assert!(g.delete_node("n1").await.unwrap());
1033 assert!(!g.delete_node("n1").await.unwrap());
1034 assert!(g.read_node("n1").await.unwrap().is_none());
1035
1036 cleanup(g.pool(), prefix).await;
1037 }
1038
1039 #[tokio::test]
1042 async fn batch_edges_dedup() {
1043 let Some(url) = pg_url() else {
1044 eprintln!("skip: LLMKERNEL_PG_URL unset");
1045 return;
1046 };
1047 let prefix = "lk_sg2_";
1048 let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1049 .await
1050 .expect("connect");
1051
1052 for n in &["a", "b", "c"] {
1053 g.upsert_node(&sample_node(n)).await.unwrap();
1054 }
1055
1056 let edges = vec![
1057 sample_edge("e1", "a", "b", "cites", 1.0),
1058 sample_edge("e2", "a", "c", "cites", 0.8),
1059 sample_edge("e1dup", "a", "b", "cites", 2.0),
1061 ];
1062 g.append_edges(&edges).await.unwrap();
1063
1064 let out = g
1066 .edges_for_node_dir("a", EdgeDirection::Out, Some("cites"))
1067 .await
1068 .unwrap();
1069 assert_eq!(out.len(), 2, "duplicate (src,tgt,rel) edge ignored");
1070
1071 g.append_edges(&[]).await.unwrap();
1073
1074 cleanup(g.pool(), prefix).await;
1075 }
1076
1077 #[tokio::test]
1079 async fn edges_dir_and_neighbors() {
1080 let Some(url) = pg_url() else {
1081 eprintln!("skip: LLMKERNEL_PG_URL unset");
1082 return;
1083 };
1084 let prefix = "lk_sg3_";
1085 let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1086 .await
1087 .expect("connect");
1088
1089 for n in &["seed", "t1", "t2", "t3"] {
1090 g.upsert_node(&sample_node(n)).await.unwrap();
1091 }
1092
1093 g.append_edges(&[
1095 sample_edge("e1", "seed", "t1", "cites", 1.0),
1096 sample_edge("e2", "seed", "t2", "cites", 0.5),
1097 sample_edge("e3", "seed", "t3", "cites", 0.3),
1098 ])
1099 .await
1100 .unwrap();
1101 g.append_edge(&sample_edge("e4", "t1", "seed", "cites", 2.0))
1102 .await
1103 .unwrap();
1104 g.append_edge(&sample_edge("e5", "seed", "t1", "related", 1.0))
1105 .await
1106 .unwrap();
1107
1108 let out_cites = g
1110 .edges_for_node_dir("seed", EdgeDirection::Out, Some("cites"))
1111 .await
1112 .unwrap();
1113 assert_eq!(out_cites.len(), 3);
1114
1115 let in_cites = g
1117 .edges_for_node_dir("seed", EdgeDirection::In, Some("cites"))
1118 .await
1119 .unwrap();
1120 assert_eq!(in_cites.len(), 1);
1121 assert_eq!(in_cites[0].source, "t1");
1122
1123 let neighbors = g
1125 .neighbors_weighted(&["seed".to_string()], EdgeDirection::Out, Some("cites"))
1126 .await
1127 .unwrap();
1128 assert_eq!(neighbors.len(), 3);
1129 assert_eq!(neighbors[0].0, "t1");
1130 assert!((neighbors[0].1 - 1.0).abs() < 1e-9);
1131 assert_eq!(neighbors[2].0, "t3");
1132
1133 cleanup(g.pool(), prefix).await;
1134 }
1135
1136 #[tokio::test]
1138 async fn remove_edges_for_node_drops_only_touching() {
1139 let Some(url) = pg_url() else {
1140 eprintln!("skip: LLMKERNEL_PG_URL unset");
1141 return;
1142 };
1143 let prefix = "lk_sg4_";
1144 let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1145 .await
1146 .expect("connect");
1147
1148 for n in &["x", "y", "z"] {
1149 g.upsert_node(&sample_node(n)).await.unwrap();
1150 }
1151 g.append_edges(&[
1152 sample_edge("e1", "x", "y", "cites", 1.0),
1153 sample_edge("e2", "z", "x", "cites", 1.0),
1154 sample_edge("e3", "y", "z", "cites", 1.0),
1155 ])
1156 .await
1157 .unwrap();
1158 assert_eq!(g.edges_for_node("x").await.unwrap().len(), 2);
1159
1160 g.remove_edges_for_node("x").await.unwrap();
1161 assert_eq!(g.edges_for_node("x").await.unwrap().len(), 0);
1162 assert_eq!(g.edges_for_node("y").await.unwrap().len(), 1);
1164
1165 cleanup(g.pool(), prefix).await;
1166 }
1167
1168 #[tokio::test]
1171 async fn tx_atomicity_commit_and_rollback() {
1172 let Some(url) = pg_url() else {
1173 eprintln!("skip: LLMKERNEL_PG_URL unset");
1174 return;
1175 };
1176 let prefix = "lk_sg5_";
1177 let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1178 .await
1179 .expect("connect");
1180
1181 for n in &["p", "q"] {
1182 g.upsert_node(&sample_node(n)).await.unwrap();
1183 }
1184
1185 let mut tx = g.pool().begin().await.expect("begin tx");
1187 g.append_edges_in_tx(&mut tx, &[sample_edge("e1", "p", "q", "cites", 1.0)])
1188 .await
1189 .unwrap();
1190 tx.commit().await.expect("commit");
1191 assert_eq!(g.edges_for_node("p").await.unwrap().len(), 1);
1192
1193 let mut tx2 = g.pool().begin().await.expect("begin tx2");
1195 g.append_edges_in_tx(&mut tx2, &[sample_edge("e2", "p", "q", "cites", 1.0)])
1196 .await
1197 .unwrap();
1198 tx2.rollback().await.expect("rollback");
1199 assert_eq!(
1200 g.edges_for_node("p").await.unwrap().len(),
1201 1,
1202 "rolled-back edge not visible"
1203 );
1204
1205 let mut tx3 = g.pool().begin().await.expect("begin tx3");
1207 g.remove_edges_for_node_in_tx(&mut tx3, "p").await.unwrap();
1208 tx3.rollback().await.expect("rollback tx3");
1209 assert_eq!(g.edges_for_node("p").await.unwrap().len(), 1);
1210
1211 let mut tx4 = g.pool().begin().await.expect("begin tx4");
1213 g.remove_edges_for_node_in_tx(&mut tx4, "p").await.unwrap();
1214 tx4.commit().await.expect("commit tx4");
1215 assert_eq!(g.edges_for_node("p").await.unwrap().len(), 0);
1216
1217 cleanup(g.pool(), prefix).await;
1218 }
1219
1220 #[tokio::test]
1222 async fn search_related_and_version() {
1223 let Some(url) = pg_url() else {
1224 eprintln!("skip: LLMKERNEL_PG_URL unset");
1225 return;
1226 };
1227 let prefix = "lk_sg6_";
1228 let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1229 .await
1230 .expect("connect");
1231
1232 let mut rust = sample_node("rust");
1233 rust.title = "Rust ownership".into();
1234 rust.body = "borrow checker".into();
1235 g.upsert_node(&rust).await.unwrap();
1236 g.upsert_node(&sample_node("py")).await.unwrap();
1237
1238 let hits = g.search_nodes("rust", 10).await.unwrap();
1240 assert_eq!(hits.len(), 1);
1241 assert_eq!(hits[0].id, "rust");
1242 assert!(g.search_nodes("", 10).await.unwrap().is_empty());
1243
1244 g.append_edge(&sample_edge("e1", "rust", "py", "related", 1.0))
1246 .await
1247 .unwrap();
1248 let related = g.related_nodes("rust", 2).await.unwrap();
1249 assert!(related.contains(&"py".to_string()));
1250
1251 assert_eq!(g.current_version().await.unwrap(), GRAPH_SCHEMA_VERSION);
1253 assert_eq!(g.migrate().await.unwrap(), GRAPH_SCHEMA_VERSION);
1254
1255 cleanup(g.pool(), prefix).await;
1256 }
1257
1258 #[tokio::test]
1260 async fn invalid_prefix_rejected() {
1261 let Some(url) = pg_url() else {
1262 eprintln!("skip: LLMKERNEL_PG_URL unset");
1263 return;
1264 };
1265 assert!(
1266 SqlxPgGraph::connect_with_prefix(&url, "lk; drop")
1267 .await
1268 .is_err()
1269 );
1270 assert!(SqlxPgGraph::connect_with_prefix(&url, "1lk").await.is_err());
1271 }
1272
1273 #[tokio::test]
1277 async fn smart_recall_and_query_nodes() {
1278 let Some(url) = pg_url() else {
1279 eprintln!("skip: LLMKERNEL_PG_URL unset");
1280 return;
1281 };
1282 let prefix = "lk_sgrecall_";
1283 let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1284 .await
1285 .expect("connect");
1286
1287 let mut a = sample_node("a");
1289 a.importance = 0.9;
1290 a.title = "rust memory".into();
1291 a.tags = vec!["lang".into()];
1292 a.projects = vec!["proj".into()];
1293 let mut b = sample_node("b");
1295 b.importance = 0.3;
1296 b.title = "rust kernel".into();
1297 b.tags = vec!["lang".into()];
1298 b.projects = vec!["proj".into()];
1299 let mut c = sample_node("c");
1301 c.importance = 0.8;
1302 c.title = "python runtime".into();
1303 c.tags = vec!["lang".into()];
1304
1305 g.upsert_node(&a).await.unwrap();
1306 g.upsert_node(&b).await.unwrap();
1307 g.upsert_node(&c).await.unwrap();
1308 g.append_edge(&sample_edge("e1", "a", "b", "cites", 1.0))
1310 .await
1311 .unwrap();
1312 g.append_edge(&sample_edge("e2", "b", "c", "cites", 0.7))
1313 .await
1314 .unwrap();
1315
1316 let recalled = g.smart_recall(Some("proj"), Some("rust"), 5).await.unwrap();
1319 let ids: Vec<&str> = recalled.iter().map(|s| s.node.id.as_str()).collect();
1320 assert!(ids.contains(&"a"));
1321 assert!(ids.contains(&"b"));
1322 assert!(!ids.contains(&"c"), "c excluded by project scope");
1323 let pos_a = ids.iter().position(|&x| x == "a").unwrap();
1324 let pos_b = ids.iter().position(|&x| x == "b").unwrap();
1325 assert!(pos_a < pos_b, "high-importance rust match ranks first");
1326 assert!(recalled[0].score > 0.0);
1327
1328 let tagged = g.query_nodes(Some("lang"), None, None, 10).await.unwrap();
1330 assert_eq!(tagged.len(), 3);
1331 let proj = g.query_nodes(None, None, Some("proj"), 10).await.unwrap();
1333 let pids: Vec<&str> = proj.iter().map(|n| n.id.as_str()).collect();
1334 assert!(pids.contains(&"a") && pids.contains(&"b"));
1335 assert!(!pids.contains(&"c"));
1336
1337 cleanup(g.pool(), prefix).await;
1338 }
1339}