Skip to main content

llm_kernel/graph/
sqlx_pg.rs

1//! Async PostgreSQL graph backend over `sqlx::PgPool` (feature `graph-pg-sqlx`).
2//!
3//! `SqlxPgGraph` mirrors the SQL/DDL/semantics of the synchronous `PgGraph`
4//! (`graph-pg`), but drives every operation asynchronously over a shared
5//! `sqlx::PgPool` instead of one mutex-guarded `postgres::Client`. The SQL is
6//! ported verbatim from `pg.rs` — only the driver changes (sqlx prepared-statement
7//! cache + connection pool instead of blocking `postgres::Client`).
8//!
9//! # Why a second PostgreSQL backend?
10//!
11//! Consumers that already own a `sqlx::PgPool` (notably the `klr`
12//! citation graph) cannot share a pool or transaction with `PgGraph`, which is
13//! built on a synchronous `postgres::Client`. `SqlxPgGraph` lets such consumers
14//! adopt the llm-kernel graph on their own pool and — critically — **expose the
15//! transaction** so a multi-table prune (chunks + vectors + edges in one tx)
16//! stays atomic. See [`SqlxPgGraph::pool`] and the `*_in_tx` methods.
17//!
18//! This backend is intentionally **non-breaking**: it is a new module that does
19//! not touch `GraphBackend`, `PgGraph`, or any other file. The SQL/schema is
20//! identical to `PgGraph`, so the two backends can share one database (use a
21//! `table_prefix` to namespace tables when both run side by side).
22//!
23//! # Scope
24//!
25//! Inherent async methods cover klr's needs — batch edge writes, directed /
26//! relation-filtered lookups, weighted neighbor aggregation, search, traversal,
27//! and basic node/edge CRUD — plus `query_nodes` and `smart_recall` (composite
28//! recall with a PageRank centrality boost, ported from `PgGraph`).
29
30use 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
45/// Standard node SELECT columns (positional order — keep in sync with [`row_to_node`]).
46const NODE_COLUMNS: &str = "id, node_type, title, tags, projects, agents, \
47     created, updated, body, importance, access_count, accessed_at";
48
49/// Map a `sqlx::Error` into a [`KernelError::Store`].
50fn pg_err(e: sqlx::Error) -> KernelError {
51    KernelError::Store(format!("sqlx: {e:?}"))
52}
53
54/// Map a `nodes` SELECT row into a [`GraphNode`]. Column names match [`NODE_COLUMNS`].
55fn 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
75/// Map an `edges` SELECT row into a [`GraphEdge`].
76fn 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
87/// Build escaped, wrapped ILIKE patterns for each whitespace-separated term in
88/// `query` — e.g. `"rust db"` → `["%rust%", "%db%"]`, `"100%"` → `["%100\\%%"]`.
89/// Pure (no connection); identical to `pg.rs::search_patterns` — duplicated here
90/// because that helper is private and this module must not edit `pg.rs`.
91fn search_patterns(query: &str) -> Vec<String> {
92    query
93        .split_whitespace()
94        .map(|t| format!("%{}%", escape_like(t)))
95        .collect()
96}
97
98/// Validate a PostgreSQL table-name prefix: empty (default, backward-compatible)
99/// or a non-empty run of ASCII alphanumeric / underscore characters. Identical
100/// to `pg.rs::is_identifier_safe` — duplicated here because the `pg` module is
101/// gated behind `graph-pg` (not implied by `graph-pg-sqlx`), and pulling in the
102/// synchronous `postgres`+`clap` deps for one pure function would bloat klr's
103/// async-only dependency tree.
104fn 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
115/// Async PostgreSQL graph backend over a shared `sqlx::PgPool`.
116///
117/// Every constructor applies the schema and runs pending migrations on connect,
118/// matching `PgGraph`. An optional `table_prefix` namespaces every backing table
119/// and index so several graphs can coexist in one database (identical semantics
120/// to `PgGraph`'s `*_with_prefix` constructors). The prefix is validated by
121/// `is_identifier_safe` before any SQL is emitted, keeping the identifier
122/// interpolation injection-safe.
123///
124/// # Transaction exposure
125///
126/// [`pool`](Self::pool) plus [`append_edges_in_tx`](Self::append_edges_in_tx) and
127/// [`remove_edges_for_node_in_tx`](Self::remove_edges_for_node_in_tx) let a caller
128/// coordinate a multi-table transaction (e.g. klr's prune: delete chunks +
129/// vectors + edges atomically). Begin a tx on `pool()`, run cross-table DML,
130/// call the `*_in_tx` helpers with the same `&mut PgConnection`, then commit or
131/// rollback — mirroring `PgVectorIndex::remove_in_tx`.
132pub struct SqlxPgGraph {
133    pool: PgPool,
134    table_prefix: String,
135}
136
137impl SqlxPgGraph {
138    /// Adopt a caller-owned `PgPool`, apply schema + migrations, and return a
139    /// ready backend. Uses the default (empty) table prefix.
140    pub async fn from_pool(pool: PgPool) -> Result<Self> {
141        Self::from_pool_with_prefix(pool, "").await
142    }
143
144    /// Like [`from_pool`](Self::from_pool), but every table/index is namespaced
145    /// under `prefix`. The prefix is validated by `is_identifier_safe` before
146    /// any SQL runs; an unsafe value returns [`KernelError::Store`].
147    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    /// Connect to `url` (libpq connstring / `postgresql://…`) with an 8-connection
163    /// pool, apply schema + migrations, and return a ready backend. Uses the
164    /// default (empty) table prefix.
165    pub async fn connect(url: &str) -> Result<Self> {
166        Self::connect_with_prefix(url, "").await
167    }
168
169    /// Like [`connect`](Self::connect), but every table/index is namespaced
170    /// under `prefix`.
171    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    /// The underlying connection pool.
181    ///
182    /// Callers that need cross-table transactional consistency — e.g. pruning a
183    /// law's chunks plus their graph edges atomically — can `pool().begin()` and
184    /// run their own DML, then call the [`append_edges_in_tx`](Self::append_edges_in_tx)
185    /// / [`remove_edges_for_node_in_tx`](Self::remove_edges_for_node_in_tx)
186    /// helpers on the same `&mut PgConnection`.
187    pub fn pool(&self) -> &PgPool {
188        &self.pool
189    }
190
191    // ── table-name helpers (identical to PgGraph) ──────────────────────
192
193    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    // ── schema init + migrate ──────────────────────────────────────────
206
207    /// Create the graph schema if absent. Idempotent — DDL ported verbatim from
208    /// `pg.rs::init_schema` (same tables, same prefixed index names). Each DDL
209    /// statement is a separate round trip because sqlx's extended-protocol
210    /// `query()` accepts one statement per call; init runs once per connect.
211    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    /// Recorded graph schema version from `{prefix}_meta`, or `0` if unset.
302    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    /// Apply pending migrations up to [`GRAPH_SCHEMA_VERSION`]. No-op when current.
320    ///
321    /// Runs in a single transaction with rollback on failure — matching
322    /// `pg.rs::migrate`. Returns the new version.
323    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        // v1 -> v2: index nodes by creation timestamp.
337        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        // v2 -> v3: composite indexes for relation-filtered directed edge lookups.
348        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    // ── basic CRUD ─────────────────────────────────────────────────────
377
378    /// Insert or update a node (upsert by `id`). Mirrors `PgGraph::upsert_node`.
379    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    /// Read a node by `id`, or `None` if absent.
412    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    /// Delete a node by `id`. Returns `true` if a row was removed.
423    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    /// Append a single edge (`ON CONFLICT DO NOTHING`). Mirrors `PgGraph::append_edge`.
434    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    /// All edges touching `node_id` (either endpoint). Mirrors `PgGraph::edges_for_node`.
453    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    /// Delete a single edge by `id`. Returns `true` if a row was removed.
467    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    // ── klr core: batch edges + directed traversal ─────────────────────
478
479    /// Batch-insert edges in chunks of 5000, each chunk its own transaction
480    /// (`ON CONFLICT DO NOTHING` per row). Mirrors `PgGraph::append_edges`.
481    /// sqlx caches the prepared statement per connection, so the repeated
482    /// `INSERT` reuses one parse/plan per chunk.
483    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    /// Append edges **within a caller-provided transaction** (commit is the
513    /// caller's responsibility — this method does not commit). Enables atomic
514    /// cross-table writes: begin a tx on [`pool`](Self::pool), write chunks /
515    /// vectors / other relations, call this with the same `&mut PgConnection`,
516    /// then `commit()`. A failure anywhere rolls back the whole set.
517    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    /// Directed, optionally relation-filtered edge lookup. Mirrors
546    /// `PgGraph::edges_for_node_dir`. `dir` selects out / in / both; `relation`
547    /// further restricts to a single relationship type when `Some`.
548    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    /// Aggregate neighbor weights from seed nodes, optionally direction- and
586    /// relation-filtered. Mirrors `PgGraph::neighbors_weighted`: walks the
587    /// source/target halves, uses `= ANY($1)` for the seed set, `GROUP BY` the
588    /// opposite endpoint, and `SUM(weight)`. Seeds themselves are excluded.
589    /// Capped at 100 seeds.
590    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            // `follow` is the column the seed matches; the neighbor is the
618            // opposite endpoint.
619            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    /// Remove every edge touching `node_id` (either endpoint).
658    /// Mirrors `PgGraph::remove_edges_for_node`.
659    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    /// Remove edges for `node_id` **within a caller-provided transaction** (no
672    /// commit here). Same atomic-prune pattern as [`Self::append_edges_in_tx`].
673    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    // ── search / traversal ─────────────────────────────────────────────
690
691    /// ILIKE substring search over `(title, body, tags)` for each
692    /// whitespace-separated term (AND of terms). Mirrors `PgGraph::search_nodes`
693    /// — extension-free vanilla PostgreSQL ILIKE with escaped, wrapped patterns.
694    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    /// Dynamic filter by tag / node_type / project, ranked by `updated` DESC.
720    /// Mirrors `PgGraph::query_nodes`.
721    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    /// Composite recall — rank nodes by recency, importance, access, FTS, and a
763    /// PageRank centrality boost over the induced candidate subgraph. Mirrors
764    /// `PgGraph::smart_recall` (`pg.rs:573`): identical composite weights and
765    /// graph-boost math. Only the driver differs — `ANY($1)` array binds replace
766    /// the dynamic placeholder lists, so the PageRank pass loads the subgraph
767    /// edges in one round-trip instead of one-per-candidate.
768    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        // FTS match set (ILIKE), used as a binary boost signal.
780        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        // Candidate fetch (broad set), excluding stale nodes.
791        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        // Composite scoring — identical weights/recency as the SQLite/PgGraph backends.
809        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        // Graph-boost pass: PageRank centrality over the induced subgraph of the
831        // top candidates. Shares the pagerank math with PgGraph/SQLite (zero drift)
832        // — only the edge-load SQL differs. One round-trip via `ANY($1)`.
833        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        // Touch retrieved nodes in one statement (access_count++,
871        // accessed_at = now) rather than N round-trips.
872        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    /// Bidirectional BFS up to `depth` hops from `start_id` via a recursive CTE.
889    /// Mirrors `PgGraph::related_nodes`. Returns distinct neighbor IDs (the start
890    /// node is excluded), capped at 500.
891    pub async fn related_nodes(&self, start_id: &str, depth: usize) -> Result<Vec<String>> {
892        let edges = self.edges_tbl();
893        // PostgreSQL requires a single recursive term: the bidirectional seed
894        // is folded into a subquery, then one recursive step follows edges in
895        // either direction (CASE picks the opposite endpoint).
896        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    /// `LLMKERNEL_PG_URL` unset → self-skip (pgvector / pg.rs pattern).
933    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    /// Drop the prefixed tables so parallel tests with distinct prefixes don't
966    /// accumulate state.
967    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    // ── offline (no server) ────────────────────────────────────────────
980
981    /// The ILIKE pattern transform is identical to pg.rs (escapes LIKE
982    /// wildcards, wraps each term).
983    #[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    // ── live (LLMKERNEL_PG_URL gated, throwaway prefixed tables) ───────
996
997    /// Basic CRUD round-trip: upsert/read/update/delete nodes + edges.
998    #[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        // upsert = update on conflict
1016        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        // single edge append + edges_for_node + delete_edge
1022        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        // delete node
1032        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    /// Batch edges with `ON CONFLICT DO NOTHING` dedup: a fresh-id edge sharing
1040    /// a duplicate (source, target, relation) triple is silently ignored.
1041    #[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            // duplicate (src,tgt,rel) of e1 — ON CONFLICT DO NOTHING
1060            sample_edge("e1dup", "a", "b", "cites", 2.0),
1061        ];
1062        g.append_edges(&edges).await.unwrap();
1063
1064        // Only e1 and e2 survived (e1dup deduped by the unique index).
1065        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        // empty slice is a no-op
1072        g.append_edges(&[]).await.unwrap();
1073
1074        cleanup(g.pool(), prefix).await;
1075    }
1076
1077    /// Directed, relation-filtered edge lookup + weighted neighbor aggregation.
1078    #[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        // seed cites t1/t2/t3 (out), t1 cites seed (in), seed related t1 (other rel)
1094        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        // Out + "cites": 3 (t1, t2, t3)
1109        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        // In + "cites": 1 (from t1)
1116        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        // neighbors_weighted Out "cites" → t1(1.0), t2(0.5), t3(0.3), desc by weight
1124        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    /// `remove_edges_for_node` drops only the edges touching the target node.
1137    #[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        // y-z edge survives (neither endpoint is x)
1163        assert_eq!(g.edges_for_node("y").await.unwrap().len(), 1);
1164
1165        cleanup(g.pool(), prefix).await;
1166    }
1167
1168    /// Transaction atomicity: `*_in_tx` commit and rollback paths.
1169    /// This is the core guarantee klr relies on for atomic prune.
1170    #[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        // Commit path: append_edges_in_tx → commit → edge visible
1186        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        // Rollback path: append_edges_in_tx → rollback → edge NOT visible
1194        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        // remove_edges_for_node_in_tx rollback → edge survives
1206        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        // remove_edges_for_node_in_tx commit → edge removed
1212        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    /// ILIKE search + recursive BFS traversal + version helpers.
1221    #[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        // search_nodes
1239        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        // related_nodes via BFS
1245        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        // version helpers (schema applied on connect)
1252        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    /// An invalid table prefix is rejected at construction (before any SQL).
1259    #[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    /// `smart_recall` ranks by the composite score (recency + importance + FTS
1274    /// + PageRank boost); `query_nodes` filters by tag/project. An FTS-matched
1275    /// high-importance node outranks a low-importance match.
1276    #[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        // a: high importance, matches "rust"
1288        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        // b: low importance, matches "rust"
1294        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        // c: high importance, does NOT match "rust", different project
1300        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        // citation edges → induce a subgraph for the PageRank boost pass
1309        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        // smart_recall scoped to project "proj" with a "rust" hint: a (0.9 +
1317        // FTS) outranks b (0.3 + FTS); c is excluded by the project filter.
1318        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        // query_nodes by tag returns all lang nodes
1329        let tagged = g.query_nodes(Some("lang"), None, None, 10).await.unwrap();
1330        assert_eq!(tagged.len(), 3);
1331        // project filter narrows to a, b
1332        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}