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