1use std::collections::HashSet;
45use std::sync::{Mutex, MutexGuard};
46use std::time::{SystemTime, UNIX_EPOCH};
47
48use postgres::row::Row;
49use postgres::types::ToSql;
50use postgres::{Client, Config, NoTls};
51
52use super::algo::{CsrGraph, pagerank_default};
53use super::lifecycle::now_iso;
54use super::recall::{W_ACCESS, W_FTS, W_GRAPH, W_IMPORTANCE, W_RECENCY, compute_recency};
55use super::schema::GRAPH_SCHEMA_VERSION;
56use super::types::{EdgeDirection, escape_like, join_csv, split_csv};
57use super::{GraphBackend, GraphEdge, GraphNode, ScoredNode};
58use crate::error::{KernelError, Result};
59
60const NODE_COLUMNS: &str = "id, node_type, title, tags, projects, agents, \
62 created, updated, body, importance, access_count, accessed_at, valid_until, last_verified";
63
64fn pg_err(e: postgres::Error) -> KernelError {
66 KernelError::Store(format!("postgres: {e:?}"))
67}
68
69fn row_to_node(row: &Row) -> GraphNode {
71 GraphNode {
72 id: row.get(0),
73 node_type: row.get(1),
74 title: row.get(2),
75 tags: split_csv(&row.get::<_, String>(3)),
76 projects: split_csv(&row.get::<_, String>(4)),
77 agents: split_csv(&row.get::<_, String>(5)),
78 created: row.get(6),
79 updated: row.get(7),
80 body: row.get(8),
81 importance: row.get(9),
82 access_count: row.get(10),
83 accessed_at: row.get(11),
84 valid_until: row.get(12),
85 last_verified: row.get(13),
86 }
87}
88
89fn row_to_edge(row: &Row) -> GraphEdge {
93 GraphEdge {
94 id: row.get(0),
95 source: row.get(1),
96 target: row.get(2),
97 relation: row.get(3),
98 weight: row.get(4),
99 ts: row.get(5),
100 }
101}
102
103fn search_patterns(query: &str) -> Vec<String> {
107 query
108 .split_whitespace()
109 .map(|t| format!("%{}%", escape_like(t)))
110 .collect()
111}
112
113fn is_identifier_safe(prefix: &str) -> bool {
122 let mut chars = prefix.chars();
123 match chars.next() {
124 None => true,
125 Some(first) if first == '_' || first.is_ascii_alphabetic() => {
126 chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
127 }
128 Some(_) => false,
129 }
130}
131
132fn init_schema(client: &mut Client, prefix: &str) -> Result<()> {
135 let nodes = format!("{prefix}nodes");
136 let edges = format!("{prefix}edges");
137 let meta = format!("{prefix}_meta");
138 let idx_edges_source = format!("{prefix}idx_edges_source");
139 let idx_edges_target = format!("{prefix}idx_edges_target");
140 let idx_edges_src_tgt_rel = format!("{prefix}idx_edges_src_tgt_rel");
141 let idx_edges_src_rel = format!("{prefix}idx_edges_src_rel");
142 let idx_edges_tgt_rel = format!("{prefix}idx_edges_tgt_rel");
143 let idx_nodes_type = format!("{prefix}idx_nodes_type");
144 let idx_nodes_updated = format!("{prefix}idx_nodes_updated");
145 let idx_nodes_importance = format!("{prefix}idx_nodes_importance");
146 let idx_nodes_accessed = format!("{prefix}idx_nodes_accessed");
147 let idx_nodes_created = format!("{prefix}idx_nodes_created");
148 client
149 .batch_execute(&format!(
150 "CREATE TABLE IF NOT EXISTS {nodes} (
151 id TEXT PRIMARY KEY,
152 node_type TEXT NOT NULL,
153 title TEXT NOT NULL,
154 tags TEXT NOT NULL DEFAULT '',
155 projects TEXT NOT NULL DEFAULT '',
156 agents TEXT NOT NULL DEFAULT '',
157 created TEXT NOT NULL,
158 updated TEXT NOT NULL,
159 body TEXT NOT NULL DEFAULT '',
160 importance DOUBLE PRECISION NOT NULL DEFAULT 0.5,
161 access_count BIGINT NOT NULL DEFAULT 0,
162 accessed_at TEXT NOT NULL DEFAULT '',
163 valid_until TEXT NOT NULL DEFAULT '',
164 last_verified TEXT NOT NULL DEFAULT ''
165 );
166 CREATE TABLE IF NOT EXISTS {edges} (
167 id TEXT PRIMARY KEY,
168 source TEXT NOT NULL,
169 target TEXT NOT NULL,
170 relation TEXT NOT NULL DEFAULT 'related',
171 weight DOUBLE PRECISION NOT NULL DEFAULT 1.0,
172 ts TEXT NOT NULL
173 );
174 CREATE INDEX IF NOT EXISTS {idx_edges_source} ON {edges}(source);
175 CREATE INDEX IF NOT EXISTS {idx_edges_target} ON {edges}(target);
176 CREATE UNIQUE INDEX IF NOT EXISTS {idx_edges_src_tgt_rel} ON {edges}(source, target, relation);
177 CREATE INDEX IF NOT EXISTS {idx_edges_src_rel} ON {edges}(source, relation);
178 CREATE INDEX IF NOT EXISTS {idx_edges_tgt_rel} ON {edges}(target, relation);
179 CREATE INDEX IF NOT EXISTS {idx_nodes_type} ON {nodes}(node_type);
180 CREATE INDEX IF NOT EXISTS {idx_nodes_updated} ON {nodes}(updated DESC);
181 CREATE INDEX IF NOT EXISTS {idx_nodes_importance} ON {nodes}(importance DESC);
182 CREATE INDEX IF NOT EXISTS {idx_nodes_accessed} ON {nodes}(accessed_at DESC);
183 CREATE INDEX IF NOT EXISTS {idx_nodes_created} ON {nodes}(created);
184 CREATE TABLE IF NOT EXISTS {meta} (key TEXT PRIMARY KEY, value TEXT NOT NULL);
185 ALTER TABLE {nodes} ADD COLUMN IF NOT EXISTS valid_until TEXT NOT NULL DEFAULT '';
186 ALTER TABLE {nodes} ADD COLUMN IF NOT EXISTS last_verified TEXT NOT NULL DEFAULT '';
187 INSERT INTO {meta} (key, value) VALUES ('graph_schema_version', '4')
188 ON CONFLICT (key) DO NOTHING;",
189 ))
190 .map_err(pg_err)?;
191 Ok(())
192}
193
194fn schema_version(client: &mut Client, prefix: &str) -> Result<u32> {
196 let meta = format!("{prefix}_meta");
197 let row = client
198 .query_opt(
199 &format!("SELECT value FROM {meta} WHERE key = 'graph_schema_version'"),
200 &[],
201 )
202 .map_err(pg_err)?;
203 Ok(row
204 .map(|r| r.get::<_, String>(0))
205 .and_then(|s| s.parse().ok())
206 .unwrap_or(0))
207}
208
209fn migrate(client: &mut Client, current: u32, prefix: &str) -> Result<u32> {
214 if current >= GRAPH_SCHEMA_VERSION {
215 return Ok(current);
216 }
217 let nodes = format!("{prefix}nodes");
218 let edges = format!("{prefix}edges");
219 let meta = format!("{prefix}_meta");
220 let idx_nodes_created = format!("{prefix}idx_nodes_created");
221 let idx_edges_src_rel = format!("{prefix}idx_edges_src_rel");
222 let idx_edges_tgt_rel = format!("{prefix}idx_edges_tgt_rel");
223 let mut tx = client.transaction().map_err(pg_err)?;
224 let mut v = current;
225 if v < 2 {
226 tx.batch_execute(&format!(
227 "CREATE INDEX IF NOT EXISTS {idx_nodes_created} ON {nodes}(created);"
228 ))
229 .map_err(pg_err)?;
230 v = 2;
231 }
232 if v < 3 {
234 tx.batch_execute(&format!(
235 "CREATE INDEX IF NOT EXISTS {idx_edges_src_rel} ON {edges}(source, relation);
236 CREATE INDEX IF NOT EXISTS {idx_edges_tgt_rel} ON {edges}(target, relation);",
237 ))
238 .map_err(pg_err)?;
239 v = 3;
240 }
241 if v < 4 {
243 tx.batch_execute(&format!(
244 "ALTER TABLE {nodes} ADD COLUMN IF NOT EXISTS valid_until TEXT NOT NULL DEFAULT '';
245 ALTER TABLE {nodes} ADD COLUMN IF NOT EXISTS last_verified TEXT NOT NULL DEFAULT '';",
246 ))
247 .map_err(pg_err)?;
248 v = 4;
249 }
250 tx.execute(
251 &format!("UPDATE {meta} SET value = $1 WHERE key = 'graph_schema_version'"),
252 &[&v.to_string()],
253 )
254 .map_err(pg_err)?;
255 tx.commit().map_err(pg_err)?;
256 Ok(v)
257}
258
259pub struct PgGraph {
266 client: Mutex<Client>,
267 table_prefix: String,
268}
269
270impl PgGraph {
271 pub fn connect(url: &str) -> Result<Self> {
275 Self::connect_config(&Self::parse_config(url)?)
276 }
277
278 pub fn connect_with_prefix(url: &str, prefix: &str) -> Result<Self> {
282 Self::connect_config_with_prefix(&Self::parse_config(url)?, prefix)
283 }
284
285 pub fn connect_config(config: &Config) -> Result<Self> {
289 let client = config.connect(NoTls).map_err(pg_err)?;
290 Self::from_client(client)
291 }
292
293 pub fn connect_config_with_prefix(config: &Config, prefix: &str) -> Result<Self> {
296 let client = config.connect(NoTls).map_err(pg_err)?;
297 Self::from_client_with_prefix(client, prefix)
298 }
299
300 #[cfg(feature = "graph-pg-tls")]
306 pub fn connect_tls<T>(url: &str, connector: T) -> Result<Self>
307 where
308 T: postgres::tls::MakeTlsConnect<postgres::Socket> + Send + 'static,
309 T::TlsConnect: Send,
310 T::Stream: Send,
311 <T::TlsConnect as postgres::tls::TlsConnect<postgres::Socket>>::Future: Send,
312 {
313 Self::connect_config_tls(&Self::parse_config(url)?, connector)
314 }
315
316 #[cfg(feature = "graph-pg-tls")]
319 pub fn connect_tls_with_prefix<T>(url: &str, prefix: &str, connector: T) -> Result<Self>
320 where
321 T: postgres::tls::MakeTlsConnect<postgres::Socket> + Send + 'static,
322 T::TlsConnect: Send,
323 T::Stream: Send,
324 <T::TlsConnect as postgres::tls::TlsConnect<postgres::Socket>>::Future: Send,
325 {
326 Self::connect_config_tls_with_prefix(&Self::parse_config(url)?, prefix, connector)
327 }
328
329 #[cfg(feature = "graph-pg-tls")]
333 pub fn connect_config_tls<T>(config: &Config, connector: T) -> Result<Self>
334 where
335 T: postgres::tls::MakeTlsConnect<postgres::Socket> + Send + 'static,
336 T::TlsConnect: Send,
337 T::Stream: Send,
338 <T::TlsConnect as postgres::tls::TlsConnect<postgres::Socket>>::Future: Send,
339 {
340 let client = config.connect(connector).map_err(pg_err)?;
341 Self::from_client(client)
342 }
343
344 #[cfg(feature = "graph-pg-tls")]
347 pub fn connect_config_tls_with_prefix<T>(
348 config: &Config,
349 prefix: &str,
350 connector: T,
351 ) -> Result<Self>
352 where
353 T: postgres::tls::MakeTlsConnect<postgres::Socket> + Send + 'static,
354 T::TlsConnect: Send,
355 T::Stream: Send,
356 <T::TlsConnect as postgres::tls::TlsConnect<postgres::Socket>>::Future: Send,
357 {
358 let client = config.connect(connector).map_err(pg_err)?;
359 Self::from_client_with_prefix(client, prefix)
360 }
361
362 #[cfg(feature = "graph-pg-tls")]
370 pub fn connect_native_tls(url: &str) -> Result<Self> {
371 let tls = native_tls::TlsConnector::new()
372 .map_err(|e| KernelError::Store(format!("native-tls connector: {e}")))?;
373 Self::connect_tls(url, postgres_native_tls::MakeTlsConnector::new(tls))
374 }
375
376 #[cfg(feature = "graph-pg-tls")]
379 pub fn connect_native_tls_with_prefix(url: &str, prefix: &str) -> Result<Self> {
380 let tls = native_tls::TlsConnector::new()
381 .map_err(|e| KernelError::Store(format!("native-tls connector: {e}")))?;
382 Self::connect_tls_with_prefix(url, prefix, postgres_native_tls::MakeTlsConnector::new(tls))
383 }
384
385 fn parse_config(url: &str) -> Result<Config> {
388 url.parse()
389 .map_err(|e| KernelError::Store(format!("invalid postgres config: {e}")))
390 }
391
392 pub fn from_client(client: Client) -> Result<Self> {
399 Self::from_client_with_prefix(client, "")
400 }
401
402 pub fn from_client_with_prefix(mut client: Client, prefix: &str) -> Result<Self> {
407 if !is_identifier_safe(prefix) {
408 return Err(KernelError::Store(format!(
409 "invalid table prefix {prefix:?}: only ASCII letters, digits, and underscore are allowed (and the first character must not be a digit)"
410 )));
411 }
412 init_schema(&mut client, prefix)?;
413 let current = schema_version(&mut client, prefix)?;
414 migrate(&mut client, current, prefix)?;
415 Ok(Self {
416 client: Mutex::new(client),
417 table_prefix: prefix.to_string(),
418 })
419 }
420
421 fn lock(&self) -> MutexGuard<'_, Client> {
422 self.client.lock().unwrap_or_else(|e| e.into_inner())
423 }
424
425 fn nodes_tbl(&self) -> String {
427 format!("{}nodes", self.table_prefix)
428 }
429
430 fn edges_tbl(&self) -> String {
432 format!("{}edges", self.table_prefix)
433 }
434
435 fn meta_tbl(&self) -> String {
437 format!("{}_meta", self.table_prefix)
438 }
439
440 pub fn list_nodes(&self, limit: usize) -> Result<Vec<GraphNode>> {
444 let nodes = self.nodes_tbl();
445 let mut c = self.lock();
446 let sql = format!("SELECT {NODE_COLUMNS} FROM {nodes} ORDER BY updated DESC LIMIT {limit}");
447 let rows = c.query(&sql, &[]).map_err(pg_err)?;
448 Ok(rows.iter().map(row_to_node).collect())
449 }
450
451 pub fn list_edges(&self, limit: usize) -> Result<Vec<GraphEdge>> {
453 let edges = self.edges_tbl();
454 let mut c = self.lock();
455 let rows = c
456 .query(
457 &format!("SELECT id, source, target, relation, weight, ts FROM {edges} LIMIT $1"),
458 &[&(limit as i64)],
459 )
460 .map_err(pg_err)?;
461 Ok(rows.iter().map(row_to_edge).collect())
462 }
463}
464
465impl GraphBackend for PgGraph {
466 fn upsert_node(&self, node: &GraphNode) -> Result<()> {
467 let tags = join_csv(&node.tags);
468 let projects = join_csv(&node.projects);
469 let agents = join_csv(&node.agents);
470 let params: [&(dyn ToSql + Sync); 14] = [
471 &node.id,
472 &node.node_type,
473 &node.title,
474 &tags,
475 &projects,
476 &agents,
477 &node.created,
478 &node.updated,
479 &node.body,
480 &node.importance,
481 &node.access_count,
482 &node.accessed_at,
483 &node.valid_until,
484 &node.last_verified,
485 ];
486 let nodes = self.nodes_tbl();
487 let mut c = self.lock();
488 c.execute(
489 &format!(
490 "INSERT INTO {nodes} (id, node_type, title, tags, projects, agents, created, updated, body, importance, access_count, accessed_at, valid_until, last_verified)
491 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)
492 ON CONFLICT (id) DO UPDATE SET
493 node_type=EXCLUDED.node_type, title=EXCLUDED.title, tags=EXCLUDED.tags,
494 projects=EXCLUDED.projects, agents=EXCLUDED.agents, created=EXCLUDED.created,
495 updated=EXCLUDED.updated, body=EXCLUDED.body, importance=EXCLUDED.importance,
496 access_count=EXCLUDED.access_count, accessed_at=EXCLUDED.accessed_at,
497 valid_until=EXCLUDED.valid_until, last_verified=EXCLUDED.last_verified"
498 ),
499 ¶ms,
500 )
501 .map_err(pg_err)?;
502 Ok(())
503 }
504
505 fn read_node(&self, id: &str) -> Result<Option<GraphNode>> {
506 let nodes = self.nodes_tbl();
507 let mut c = self.lock();
508 let sql = format!("SELECT {NODE_COLUMNS} FROM {nodes} WHERE id = $1");
509 let params: [&(dyn ToSql + Sync); 1] = [&id];
510 let row = c.query_opt(&sql, ¶ms).map_err(pg_err)?;
511 Ok(row.as_ref().map(row_to_node))
512 }
513
514 fn delete_node(&self, id: &str) -> Result<bool> {
515 let nodes = self.nodes_tbl();
516 let mut c = self.lock();
517 let n = c
518 .execute(&format!("DELETE FROM {nodes} WHERE id = $1"), &[&id])
519 .map_err(pg_err)?;
520 Ok(n > 0)
521 }
522
523 fn search_nodes(&self, query: &str, limit: usize) -> Result<Vec<GraphNode>> {
524 let terms = search_patterns(query);
525 if terms.is_empty() {
526 return Ok(vec![]);
527 }
528 let mut conds: Vec<String> = Vec::with_capacity(terms.len());
529 let params: Vec<&(dyn ToSql + Sync)> =
530 terms.iter().map(|s| -> &(dyn ToSql + Sync) { s }).collect();
531 for i in 0..terms.len() {
532 conds.push(format!(
533 "(title || ' ' || body || ' ' || tags) ILIKE ${n} ESCAPE '\\'",
534 n = i + 1
535 ));
536 }
537 let where_clause = conds.join(" AND ");
538 let nodes = self.nodes_tbl();
539 let sql = format!(
540 "SELECT {NODE_COLUMNS} FROM {nodes} WHERE {where_clause} ORDER BY importance DESC, updated DESC LIMIT {limit}"
541 );
542 let mut c = self.lock();
543 let rows = c.query(&sql, ¶ms).map_err(pg_err)?;
544 Ok(rows.iter().map(row_to_node).collect())
545 }
546
547 fn query_nodes(
548 &self,
549 tag: Option<&str>,
550 node_type: Option<&str>,
551 project: Option<&str>,
552 limit: usize,
553 ) -> Result<Vec<GraphNode>> {
554 let limit = limit.min(200) as i64;
555 let mut owned: Vec<String> = Vec::new();
556 let mut conds: Vec<String> = Vec::new();
557 if let Some(t) = tag {
558 owned.push(escape_like(t));
559 conds.push(format!(
560 "(',' || tags || ',') ILIKE ('%,' || ${n} || ',%') ESCAPE '\\'",
561 n = owned.len()
562 ));
563 }
564 if let Some(nt) = node_type {
565 owned.push(nt.to_string());
566 conds.push(format!("node_type = ${n}", n = owned.len()));
567 }
568 if let Some(p) = project {
569 owned.push(escape_like(p));
570 conds.push(format!(
571 "(',' || projects || ',') ILIKE ('%,' || ${n} || ',%') ESCAPE '\\'",
572 n = owned.len()
573 ));
574 }
575 let where_clause = if conds.is_empty() {
576 String::new()
577 } else {
578 format!("WHERE {}", conds.join(" AND "))
579 };
580 let params: Vec<&(dyn ToSql + Sync)> =
581 owned.iter().map(|s| -> &(dyn ToSql + Sync) { s }).collect();
582 let nodes = self.nodes_tbl();
583 let sql = format!(
584 "SELECT {NODE_COLUMNS} FROM {nodes} {where_clause} ORDER BY updated DESC LIMIT {limit}"
585 );
586 let mut c = self.lock();
587 let rows = c.query(&sql, ¶ms).map_err(pg_err)?;
588 Ok(rows.iter().map(row_to_node).collect())
589 }
590
591 fn smart_recall(
592 &self,
593 project: Option<&str>,
594 hint: Option<&str>,
595 limit: usize,
596 ) -> Result<Vec<ScoredNode>> {
597 let now_secs = SystemTime::now()
598 .duration_since(UNIX_EPOCH)
599 .unwrap_or_default()
600 .as_secs();
601
602 let fts_ids: HashSet<String> = match hint {
604 Some(h) if !h.is_empty() => self
605 .search_nodes(h, limit * 4)?
606 .into_iter()
607 .map(|n| n.id)
608 .collect(),
609 _ => HashSet::new(),
610 };
611
612 let candidate_limit = (limit * 4).max(40) as i64;
614 let mut owned: Vec<String> = Vec::new();
615 let mut conds: Vec<String> = vec!["(',' || tags || ',') NOT ILIKE '%,stale,%'".to_string()];
616 if let Some(p) = project {
617 owned.push(escape_like(p));
618 conds.push(format!(
619 "(',' || projects || ',') ILIKE ('%,' || ${n} || ',%') ESCAPE '\\'",
620 n = owned.len()
621 ));
622 }
623 let where_clause = conds.join(" AND ");
624 let params: Vec<&(dyn ToSql + Sync)> =
625 owned.iter().map(|s| -> &(dyn ToSql + Sync) { s }).collect();
626 let nodes = self.nodes_tbl();
627 let edges = self.edges_tbl();
628 let sql = format!(
629 "SELECT {NODE_COLUMNS} FROM {nodes} WHERE {where_clause} ORDER BY importance DESC, updated DESC LIMIT {candidate_limit}"
630 );
631 let mut c = self.lock();
632 let rows = c.query(&sql, ¶ms).map_err(pg_err)?;
633 let candidates: Vec<GraphNode> = rows.iter().map(row_to_node).collect();
634
635 let mut scored: Vec<ScoredNode> = candidates
637 .into_iter()
638 .map(|node| {
639 let recency = compute_recency(&node.updated, now_secs);
640 let importance = node.importance;
641 let access_freq = (node.access_count.max(0) as f64 / 20.0).min(1.0);
642 let fts_match = if fts_ids.contains(&node.id) { 1.0 } else { 0.0 };
643 let score = W_RECENCY * recency
644 + W_IMPORTANCE * importance
645 + W_ACCESS * access_freq
646 + W_FTS * fts_match;
647 ScoredNode { node, score }
648 })
649 .collect();
650 scored.sort_by(|a, b| {
651 b.score
652 .partial_cmp(&a.score)
653 .unwrap_or(std::cmp::Ordering::Equal)
654 });
655 scored.truncate(limit);
656
657 if scored.len() > 1 {
661 const MAX_GRAPH_BOOST_PARTICIPANTS: usize = 100;
662 let candidate_ids: Vec<String> = scored
663 .iter()
664 .take(MAX_GRAPH_BOOST_PARTICIPANTS)
665 .map(|sn| sn.node.id.clone())
666 .collect();
667 let n = candidate_ids.len();
668 let l1: String = (1..=n)
669 .map(|i| format!("${i}"))
670 .collect::<Vec<_>>()
671 .join(",");
672 let l2: String = ((n + 1)..=(2 * n))
673 .map(|i| format!("${i}"))
674 .collect::<Vec<_>>()
675 .join(",");
676 let sql = format!(
677 "SELECT id, source, target, relation, weight, ts FROM {edges} WHERE source IN ({l1}) AND target IN ({l2})"
678 );
679 let mut bp: Vec<&(dyn ToSql + Sync)> = Vec::with_capacity(2 * n);
680 for id in &candidate_ids {
681 bp.push(id);
682 }
683 for id in &candidate_ids {
684 bp.push(id);
685 }
686 let sub_edges: Vec<GraphEdge> = match c.query(&sql, &bp) {
687 Ok(rows) => rows
688 .iter()
689 .map(|r| GraphEdge {
690 id: r.get(0),
691 source: r.get(1),
692 target: r.get(2),
693 relation: r.get(3),
694 weight: r.get(4),
695 ts: r.get(5),
696 })
697 .collect(),
698 Err(_) => Vec::new(),
699 };
700 let csr = CsrGraph::from_edges(&candidate_ids, &sub_edges);
701 let pr = pagerank_default(&csr);
702 let max_pr = pr.iter().copied().fold(0.0_f64, f64::max).max(1e-12);
703 let pr_map: std::collections::HashMap<String, f64> = candidate_ids
704 .iter()
705 .zip(pr.iter())
706 .map(|(id, &s)| (id.clone(), s / max_pr))
707 .collect();
708 for sn in &mut scored {
709 let boost = pr_map.get(&sn.node.id).copied().unwrap_or(0.0);
710 sn.score += W_GRAPH * boost;
711 }
712 scored.sort_by(|a, b| {
713 b.score
714 .partial_cmp(&a.score)
715 .unwrap_or(std::cmp::Ordering::Equal)
716 });
717 }
718
719 if !scored.is_empty() {
722 let now = now_iso();
723 let ids: Vec<&str> = scored.iter().map(|sn| sn.node.id.as_str()).collect();
724 let placeholders: String = (0..ids.len())
725 .map(|i| format!("${}", i + 2))
726 .collect::<Vec<_>>()
727 .join(",");
728 let mut params: Vec<&(dyn ToSql + Sync)> = Vec::with_capacity(ids.len() + 1);
729 params.push(&now);
730 for id in &ids {
731 params.push(id);
732 }
733 let sql = format!(
734 "UPDATE {nodes} SET access_count = access_count + 1, accessed_at = $1 WHERE id IN ({placeholders})"
735 );
736 let _ = c.execute(&sql, ¶ms);
737 }
738
739 Ok(scored)
740 }
741
742 fn related_nodes(&self, start_id: &str, depth: usize) -> Result<Vec<String>> {
743 let edges = self.edges_tbl();
744 let mut c = self.lock();
745 let depth_v = depth as i32;
746 let params: [&(dyn ToSql + Sync); 2] = [&start_id, &depth_v];
747 let rows = c
748 .query(
749 &format!(
753 "WITH RECURSIVE bfs(node_id, lvl) AS (
754 SELECT nb.node_id, 1 FROM (
755 SELECT target AS node_id FROM {edges} WHERE source = $1
756 UNION
757 SELECT source AS node_id FROM {edges} WHERE target = $1
758 ) nb
759 UNION
760 SELECT CASE WHEN e.source = bfs.node_id THEN e.target ELSE e.source END,
761 bfs.lvl + 1
762 FROM bfs
763 JOIN {edges} e ON e.source = bfs.node_id OR e.target = bfs.node_id
764 WHERE bfs.lvl < $2
765 )
766 SELECT DISTINCT node_id FROM bfs WHERE node_id <> $1 LIMIT 500"
767 ),
768 ¶ms,
769 )
770 .map_err(pg_err)?;
771 Ok(rows.iter().map(|r| r.get::<_, String>(0)).collect())
772 }
773
774 fn append_edge(&self, edge: &GraphEdge) -> Result<()> {
775 let params: [&(dyn ToSql + Sync); 6] = [
776 &edge.id,
777 &edge.source,
778 &edge.target,
779 &edge.relation,
780 &edge.weight,
781 &edge.ts,
782 ];
783 let edges = self.edges_tbl();
784 let mut c = self.lock();
785 c.execute(
786 &format!(
787 "INSERT INTO {edges} (id, source, target, relation, weight, ts)
788 VALUES ($1,$2,$3,$4,$5,$6) ON CONFLICT DO NOTHING"
789 ),
790 ¶ms,
791 )
792 .map_err(pg_err)?;
793 Ok(())
794 }
795
796 fn append_edges(&self, edges: &[GraphEdge]) -> Result<()> {
797 if edges.is_empty() {
798 return Ok(());
799 }
800 let edges_tbl = self.edges_tbl();
801 const CHUNK: usize = 5000;
802 let mut c = self.lock();
803 for chunk in edges.chunks(CHUNK) {
804 let mut tx = c.transaction().map_err(pg_err)?;
808 {
809 let stmt = tx
810 .prepare(&format!(
811 "INSERT INTO {edges_tbl} (id, source, target, relation, weight, ts)
812 VALUES ($1,$2,$3,$4,$5,$6) ON CONFLICT DO NOTHING"
813 ))
814 .map_err(pg_err)?;
815 for e in chunk {
816 let params: [&(dyn ToSql + Sync); 6] =
817 [&e.id, &e.source, &e.target, &e.relation, &e.weight, &e.ts];
818 tx.execute(&stmt, ¶ms).map_err(pg_err)?;
819 }
820 }
821 tx.commit().map_err(pg_err)?;
822 }
823 Ok(())
824 }
825
826 fn edges_for_node_dir(
827 &self,
828 node_id: &str,
829 dir: EdgeDirection,
830 relation: Option<&str>,
831 ) -> Result<Vec<GraphEdge>> {
832 let edges = self.edges_tbl();
833 let mut c = self.lock();
834 let dir_clause = match dir {
835 EdgeDirection::Out => "source = $1",
836 EdgeDirection::In => "target = $1",
837 EdgeDirection::Both => "(source = $1 OR target = $1)",
838 };
839 let rows = if let Some(r) = relation {
840 let sql = format!(
841 "SELECT id, source, target, relation, weight, ts FROM {edges} \
842 WHERE {dir_clause} AND relation = $2 ORDER BY weight DESC"
843 );
844 c.query(&sql, &[&node_id, &r]).map_err(pg_err)?
845 } else {
846 let sql = format!(
847 "SELECT id, source, target, relation, weight, ts FROM {edges} \
848 WHERE {dir_clause} ORDER BY weight DESC"
849 );
850 c.query(&sql, &[&node_id]).map_err(pg_err)?
851 };
852 Ok(rows.iter().map(row_to_edge).collect())
853 }
854
855 fn neighbors_weighted(
856 &self,
857 seed_ids: &[String],
858 dir: EdgeDirection,
859 relation: Option<&str>,
860 ) -> Result<Vec<(String, f64)>> {
861 if seed_ids.is_empty() {
862 return Ok(vec![]);
863 }
864 let edges = self.edges_tbl();
865 const MAX_SEEDS: usize = 100;
866 let seed_ids = if seed_ids.len() > MAX_SEEDS {
867 &seed_ids[..MAX_SEEDS]
868 } else {
869 seed_ids
870 };
871 let seed_arr: Vec<String> = seed_ids.to_vec();
872 let seed_set: HashSet<&str> = seed_ids.iter().map(String::as_str).collect();
873 let mut c = self.lock();
874 let mut weights: std::collections::HashMap<String, f64> = std::collections::HashMap::new();
875
876 let halves: &[&str] = match dir {
878 EdgeDirection::Out => &["source"],
879 EdgeDirection::In => &["target"],
880 EdgeDirection::Both => &["source", "target"],
881 };
882 for &follow in halves {
883 let select_col = if follow == "source" {
886 "target"
887 } else {
888 "source"
889 };
890 let rel_clause = relation.map(|_| " AND relation = $2").unwrap_or("");
891 let sql = format!(
892 "SELECT {select_col} AS nb, SUM(weight) AS w FROM {edges} \
893 WHERE {follow} = ANY($1){rel_clause} GROUP BY {select_col}"
894 );
895 let rows = if let Some(r) = relation {
896 c.query(&sql, &[&seed_arr, &r]).map_err(pg_err)?
897 } else {
898 c.query(&sql, &[&seed_arr]).map_err(pg_err)?
899 };
900 for row in &rows {
901 let nb: String = row.get(0);
902 let w: f64 = row.get(1);
903 if !seed_set.contains(nb.as_str()) {
904 *weights.entry(nb).or_default() += w;
905 }
906 }
907 }
908
909 let mut result: Vec<(String, f64)> = weights.into_iter().collect();
910 result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
911 Ok(result)
912 }
913
914 fn edges_for_node(&self, node_id: &str) -> Result<Vec<GraphEdge>> {
915 let edges = self.edges_tbl();
916 let mut c = self.lock();
917 let params: [&(dyn ToSql + Sync); 1] = [&node_id];
918 let rows = c
919 .query(
920 &format!(
921 "SELECT id, source, target, relation, weight, ts FROM {edges} WHERE source = $1 OR target = $1"
922 ),
923 ¶ms,
924 )
925 .map_err(pg_err)?;
926 Ok(rows.iter().map(row_to_edge).collect())
927 }
928
929 fn delete_edge(&self, id: &str) -> Result<bool> {
930 let edges = self.edges_tbl();
931 let mut c = self.lock();
932 let params: [&(dyn ToSql + Sync); 1] = [&id];
933 let n = c
934 .execute(&format!("DELETE FROM {edges} WHERE id = $1"), ¶ms)
935 .map_err(pg_err)?;
936 Ok(n > 0)
937 }
938
939 fn remove_edges_for_node(&self, node_id: &str) -> Result<()> {
940 let edges = self.edges_tbl();
941 let mut c = self.lock();
942 let params: [&(dyn ToSql + Sync); 1] = [&node_id];
943 c.execute(
944 &format!("DELETE FROM {edges} WHERE source = $1 OR target = $1"),
945 ¶ms,
946 )
947 .map_err(pg_err)?;
948 Ok(())
949 }
950
951 fn current_version(&self) -> Result<u32> {
952 let meta = self.meta_tbl();
953 let mut c = self.lock();
954 let row = c
955 .query_opt(
956 &format!("SELECT value FROM {meta} WHERE key = 'graph_schema_version'"),
957 &[],
958 )
959 .map_err(pg_err)?;
960 Ok(row
961 .map(|r| r.get::<_, String>(0))
962 .and_then(|s| s.parse().ok())
963 .unwrap_or(0))
964 }
965
966 fn migrate(&self) -> Result<u32> {
967 let mut c = self.lock();
968 let current = schema_version(&mut c, &self.table_prefix)?;
969 migrate(&mut c, current, &self.table_prefix)
970 }
971}
972
973#[cfg(test)]
974mod tests {
975 use super::*;
976 use crate::graph::{GraphBackend, SqliteGraph};
977 use postgres::{Config, NoTls};
978
979 const TEST_DB: &str = "llm_kernel_pg_test";
980
981 fn sample_node(id: &str) -> GraphNode {
982 GraphNode {
983 id: id.to_string(),
984 node_type: "concept".to_string(),
985 title: format!("Node {id}"),
986 body: "pg backend test body".to_string(),
987 tags: vec!["backend".to_string()],
988 projects: vec![],
989 agents: vec![],
990 created: "2026-01-01T00:00:00Z".to_string(),
991 updated: "2026-01-01T00:00:00Z".to_string(),
992 importance: 0.5,
993 access_count: 0,
994 accessed_at: String::new(),
995 ..Default::default()
996 }
997 }
998
999 fn with_test_db<F: FnOnce(&PgGraph)>(body: F) {
1001 let base = std::env::var("LLMKERNEL_PG_URL").expect("LLMKERNEL_PG_URL set");
1002 let admin_cfg: Config = base
1003 .parse()
1004 .expect("LLMKERNEL_PG_URL is a valid libpq connstring");
1005 {
1006 let mut admin = admin_cfg.connect(NoTls).expect("admin connect");
1007 let _ = admin.batch_execute(&format!("DROP DATABASE IF EXISTS {TEST_DB}"));
1008 admin
1009 .batch_execute(&format!("CREATE DATABASE {TEST_DB}"))
1010 .expect("create test db");
1011 }
1012 let mut test_cfg = admin_cfg.clone();
1013 test_cfg.dbname(TEST_DB);
1014 let graph = PgGraph::connect_config(&test_cfg).expect("connect to test db");
1015 body(&graph);
1016 drop(graph);
1017 let mut admin = admin_cfg.connect(NoTls).expect("admin reconnect");
1018 let _ = admin.batch_execute(&format!("DROP DATABASE IF EXISTS {TEST_DB}"));
1019 }
1020
1021 #[test]
1024 fn search_patterns_escapes_and_wraps() {
1025 assert!(search_patterns("").is_empty());
1026 assert_eq!(search_patterns("rust"), vec!["%rust%".to_string()]);
1027 assert_eq!(search_patterns("100%"), vec!["%100\\%%".to_string()]);
1029 assert_eq!(search_patterns("a_b"), vec!["%a\\_b%".to_string()]);
1030 assert_eq!(
1032 search_patterns("rust db"),
1033 vec!["%rust%".to_string(), "%db%".to_string()]
1034 );
1035 }
1036
1037 #[test]
1042 fn is_identifier_safe_validation() {
1043 assert!(is_identifier_safe(""));
1045 assert!(is_identifier_safe("lk_"));
1046 assert!(is_identifier_safe("graph1"));
1047 assert!(is_identifier_safe("_x"));
1048 assert!(is_identifier_safe("ABC_123"));
1049 assert!(!is_identifier_safe("1lk"));
1051 assert!(!is_identifier_safe("lk nodes"));
1053 assert!(!is_identifier_safe("lk;"));
1054 assert!(!is_identifier_safe("lk' OR 1=1--"));
1055 assert!(!is_identifier_safe("lk-bad"));
1056 assert!(!is_identifier_safe("lk.bad"));
1057 assert!(!is_identifier_safe("데이터"));
1058 }
1059
1060 #[cfg(feature = "graph-pg-tls")]
1066 #[test]
1067 fn connect_native_tls_returns_result_not_panic() {
1068 let base = match std::env::var("LLMKERNEL_PG_URL") {
1069 Ok(u) => u,
1070 Err(_) => {
1071 eprintln!("skipped: LLMKERNEL_PG_URL unset (no live PostgreSQL)");
1072 return;
1073 }
1074 };
1075 match PgGraph::connect_native_tls(&base) {
1076 Ok(g) => {
1077 assert_eq!(g.current_version().unwrap(), GRAPH_SCHEMA_VERSION);
1078 }
1079 Err(e) => {
1080 eprintln!("connect_native_tls returned Err as expected on a non-TLS server: {e}");
1081 }
1082 }
1083 }
1084
1085 #[test]
1088 fn live_pg_graph_backend_conformance() {
1089 if std::env::var("LLMKERNEL_PG_URL").is_err() {
1090 eprintln!("skipped: LLMKERNEL_PG_URL unset (no live PostgreSQL)");
1091 return;
1092 }
1093 with_test_db(|g| {
1094 assert_eq!(g.current_version().unwrap(), GRAPH_SCHEMA_VERSION);
1095 assert_eq!(g.migrate().unwrap(), GRAPH_SCHEMA_VERSION);
1096
1097 let dyn_g: &dyn GraphBackend = g;
1098 assert!(dyn_g.read_node("n1").unwrap().is_none());
1099
1100 g.upsert_node(&sample_node("rust")).unwrap();
1101 let loaded = g.read_node("rust").unwrap().unwrap();
1102 assert_eq!(loaded.title, "Node rust");
1103 assert_eq!(loaded.tags, vec!["backend".to_string()]);
1104
1105 let mut updated = sample_node("rust");
1106 updated.title = "Rust ownership".into();
1107 updated.body = "borrow checker rules".into();
1108 g.upsert_node(&updated).unwrap();
1109 assert_eq!(
1110 g.read_node("rust").unwrap().unwrap().title,
1111 "Rust ownership"
1112 );
1113
1114 assert!(g.delete_node("rust").unwrap());
1115 assert!(!g.delete_node("rust").unwrap());
1116 assert!(g.read_node("rust").unwrap().is_none());
1117
1118 let mut n = sample_node("rust");
1119 n.title = "Rust ownership model".into();
1120 n.body = "borrow checker rules".into();
1121 n.tags = vec!["rust".into(), "memory".into()];
1122 g.upsert_node(&n).unwrap();
1123 let mut other = sample_node("py");
1124 other.title = "Python GIL".into();
1125 g.upsert_node(&other).unwrap();
1126
1127 let hits = g.search_nodes("rust", 10).unwrap();
1128 assert_eq!(hits.len(), 1);
1129 assert_eq!(hits[0].id, "rust");
1130
1131 let tagged = g.query_nodes(Some("rust"), None, None, 10).unwrap();
1132 assert_eq!(tagged.len(), 1);
1133 assert_eq!(tagged[0].id, "rust");
1134
1135 g.append_edge(&GraphEdge {
1136 id: "e1".into(),
1137 source: "rust".into(),
1138 target: "py".into(),
1139 relation: "related".into(),
1140 weight: 1.0,
1141 ts: "2026-01-01T00:00:00Z".into(),
1142 })
1143 .unwrap();
1144 assert_eq!(g.edges_for_node("rust").unwrap().len(), 1);
1145 assert!(
1146 g.related_nodes("rust", 2)
1147 .unwrap()
1148 .contains(&"py".to_string())
1149 );
1150
1151 g.append_edge(&GraphEdge {
1156 id: "e1dup".into(),
1157 source: "rust".into(),
1158 target: "py".into(),
1159 relation: "related".into(),
1160 weight: 2.0,
1161 ts: "2026-01-02T00:00:00Z".into(),
1162 })
1163 .unwrap();
1164 assert_eq!(
1165 g.edges_for_node("rust").unwrap().len(),
1166 1,
1167 "duplicate (src,tgt,rel) edge ignored"
1168 );
1169
1170 g.append_edge(&GraphEdge {
1174 id: "eloop".into(),
1175 source: "rust".into(),
1176 target: "rust".into(),
1177 relation: "self".into(),
1178 weight: 1.0,
1179 ts: "2026-01-01T00:00:00Z".into(),
1180 })
1181 .unwrap();
1182 let related = g.related_nodes("rust", 2).unwrap();
1183 assert!(
1184 !related.contains(&"rust".to_string()),
1185 "start node excluded even with a self-loop"
1186 );
1187
1188 let recalled = g.smart_recall(None, Some("ownership"), 5).unwrap();
1189 assert!(recalled.iter().any(|s| s.node.id == "rust"));
1190 let after = g.read_node("rust").unwrap().unwrap();
1191 assert!(after.access_count >= 1, "access_count incremented");
1192 });
1193 }
1194
1195 #[test]
1199 fn live_prefix_isolation() {
1200 if std::env::var("LLMKERNEL_PG_URL").is_err() {
1201 eprintln!("skipped: LLMKERNEL_PG_URL unset (no live PostgreSQL)");
1202 return;
1203 }
1204 with_test_db(|g_default| {
1205 g_default.upsert_node(&sample_node("default_only")).unwrap();
1207 assert_eq!(
1208 g_default.read_node("default_only").unwrap().unwrap().title,
1209 "Node default_only"
1210 );
1211
1212 let base = std::env::var("LLMKERNEL_PG_URL").unwrap();
1214 let admin_cfg: Config = base.parse().expect("valid connstring");
1215 let mut cfg = admin_cfg.clone();
1216 cfg.dbname(TEST_DB);
1217 let g_prefixed =
1218 PgGraph::connect_config_with_prefix(&cfg, "lk_").expect("prefixed connect");
1219
1220 assert!(g_prefixed.read_node("default_only").unwrap().is_none());
1222 g_prefixed.upsert_node(&sample_node("lk_only")).unwrap();
1223 assert!(g_default.read_node("lk_only").unwrap().is_none());
1224 assert_eq!(
1225 g_prefixed.read_node("lk_only").unwrap().unwrap().title,
1226 "Node lk_only"
1227 );
1228
1229 g_prefixed.upsert_node(&sample_node("lk_peer")).unwrap();
1231 g_prefixed
1232 .append_edge(&GraphEdge {
1233 id: "lk_e1".into(),
1234 source: "lk_only".into(),
1235 target: "lk_peer".into(),
1236 relation: "related".into(),
1237 weight: 1.0,
1238 ts: "2026-01-01T00:00:00Z".into(),
1239 })
1240 .unwrap();
1241 assert_eq!(g_prefixed.edges_for_node("lk_only").unwrap().len(), 1);
1242 assert_eq!(g_default.edges_for_node("lk_only").unwrap().len(), 0);
1243
1244 assert!(g_prefixed.current_version().unwrap() >= 2);
1246
1247 assert!(PgGraph::connect_config_with_prefix(&cfg, "lk; drop").is_err());
1249 assert!(PgGraph::connect_config_with_prefix(&cfg, "1lk").is_err());
1250 });
1251 }
1252
1253 #[test]
1256 fn live_migrate_sqlite_to_postgres_round_trip() {
1257 let base = match std::env::var("LLMKERNEL_PG_URL") {
1258 Ok(u) => u,
1259 Err(_) => {
1260 eprintln!("skipped: LLMKERNEL_PG_URL unset (no live PostgreSQL)");
1261 return;
1262 }
1263 };
1264
1265 let src = SqliteGraph::open_in_memory().expect("sqlite source");
1266 let mut a = sample_node("a");
1267 a.body = "migrate test body".into();
1268 a.tags = vec!["migrate".to_string()];
1269 a.projects = vec!["demo".to_string()];
1270 a.importance = 0.6;
1271 src.upsert_node(&a).unwrap();
1272 src.upsert_node(&sample_node("b")).unwrap();
1273 src.append_edge(&GraphEdge {
1274 id: "e1".into(),
1275 source: "a".into(),
1276 target: "b".into(),
1277 relation: "related".into(),
1278 weight: 1.0,
1279 ts: "2026-01-01T00:00:00Z".into(),
1280 })
1281 .unwrap();
1282 let nodes = src.query_nodes(None, None, None, 200).unwrap();
1283 assert_eq!(nodes.len(), 2);
1284
1285 let admin_cfg: Config = base.parse().expect("valid connstring");
1286 {
1287 let mut admin = admin_cfg.connect(NoTls).expect("admin connect");
1288 let _ = admin.batch_execute("DROP DATABASE IF EXISTS llm_kernel_migrate_test");
1289 admin
1290 .batch_execute("CREATE DATABASE llm_kernel_migrate_test")
1291 .expect("create test db");
1292 }
1293 let mut test_cfg = admin_cfg.clone();
1294 test_cfg.dbname("llm_kernel_migrate_test");
1295 let pg = PgGraph::connect_config(&test_cfg).expect("connect target");
1296
1297 for n in &nodes {
1298 pg.upsert_node(n).unwrap();
1299 }
1300 pg.append_edge(&GraphEdge {
1301 id: "e1".into(),
1302 source: "a".into(),
1303 target: "b".into(),
1304 relation: "related".into(),
1305 weight: 1.0,
1306 ts: "2026-01-01T00:00:00Z".into(),
1307 })
1308 .unwrap();
1309
1310 assert_eq!(pg.list_nodes(100).unwrap().len(), 2);
1311 assert_eq!(pg.list_edges(100).unwrap().len(), 1);
1312 let loaded = pg.read_node("a").unwrap().unwrap();
1313 assert_eq!(loaded.title, "Node a");
1314 assert_eq!(loaded.tags, vec!["migrate".to_string()]);
1315 assert!((loaded.importance - 0.6).abs() < 1e-9);
1316
1317 drop(pg);
1318 let mut admin = admin_cfg.connect(NoTls).expect("admin reconnect");
1319 let _ = admin.batch_execute("DROP DATABASE IF EXISTS llm_kernel_migrate_test");
1320 }
1321}