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, valid_until, last_verified";
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        valid_until: row.get("valid_until"),
73        last_verified: row.get("last_verified"),
74    }
75}
76
77/// Map an `edges` SELECT row into a [`GraphEdge`].
78fn row_to_edge(row: &PgRow) -> GraphEdge {
79    GraphEdge {
80        id: row.get("id"),
81        source: row.get("source"),
82        target: row.get("target"),
83        relation: row.get("relation"),
84        weight: row.get("weight"),
85        ts: row.get("ts"),
86    }
87}
88
89/// Build escaped, wrapped ILIKE patterns for each whitespace-separated term in
90/// `query` — e.g. `"rust db"` → `["%rust%", "%db%"]`, `"100%"` → `["%100\\%%"]`.
91/// Pure (no connection); identical to `pg.rs::search_patterns` — duplicated here
92/// because that helper is private and this module must not edit `pg.rs`.
93fn search_patterns(query: &str) -> Vec<String> {
94    query
95        .split_whitespace()
96        .map(|t| format!("%{}%", escape_like(t)))
97        .collect()
98}
99
100/// Validate a PostgreSQL table-name prefix: empty (default, backward-compatible)
101/// or a non-empty run of ASCII alphanumeric / underscore characters. Identical
102/// to `pg.rs::is_identifier_safe` — duplicated here because the `pg` module is
103/// gated behind `graph-pg` (not implied by `graph-pg-sqlx`), and pulling in the
104/// synchronous `postgres`+`clap` deps for one pure function would bloat klr's
105/// async-only dependency tree.
106fn is_identifier_safe(prefix: &str) -> bool {
107    let mut chars = prefix.chars();
108    match chars.next() {
109        None => true,
110        Some(first) if first == '_' || first.is_ascii_alphabetic() => {
111            chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
112        }
113        Some(_) => false,
114    }
115}
116
117/// Async PostgreSQL graph backend over a shared `sqlx::PgPool`.
118///
119/// Every constructor applies the schema and runs pending migrations on connect,
120/// matching `PgGraph`. An optional `table_prefix` namespaces every backing table
121/// and index so several graphs can coexist in one database (identical semantics
122/// to `PgGraph`'s `*_with_prefix` constructors). The prefix is validated by
123/// `is_identifier_safe` before any SQL is emitted, keeping the identifier
124/// interpolation injection-safe.
125///
126/// # Transaction exposure
127///
128/// [`pool`](Self::pool) plus [`append_edges_in_tx`](Self::append_edges_in_tx) and
129/// [`remove_edges_for_node_in_tx`](Self::remove_edges_for_node_in_tx) let a caller
130/// coordinate a multi-table transaction (e.g. klr's prune: delete chunks +
131/// vectors + edges atomically). Begin a tx on `pool()`, run cross-table DML,
132/// call the `*_in_tx` helpers with the same `&mut PgConnection`, then commit or
133/// rollback — mirroring `PgVectorIndex::remove_in_tx`.
134pub struct SqlxPgGraph {
135    pool: PgPool,
136    table_prefix: String,
137}
138
139impl SqlxPgGraph {
140    /// Adopt a caller-owned `PgPool`, apply schema + migrations, and return a
141    /// ready backend. Uses the default (empty) table prefix.
142    pub async fn from_pool(pool: PgPool) -> Result<Self> {
143        Self::from_pool_with_prefix(pool, "").await
144    }
145
146    /// Like [`from_pool`](Self::from_pool), but every table/index is namespaced
147    /// under `prefix`. The prefix is validated by `is_identifier_safe` before
148    /// any SQL runs; an unsafe value returns [`KernelError::Store`].
149    pub async fn from_pool_with_prefix(pool: PgPool, prefix: &str) -> Result<Self> {
150        if !is_identifier_safe(prefix) {
151            return Err(KernelError::Store(format!(
152                "invalid table prefix {prefix:?}: only ASCII letters, digits, and underscore are allowed (and the first character must not be a digit)"
153            )));
154        }
155        let graph = Self {
156            pool,
157            table_prefix: prefix.to_string(),
158        };
159        graph.init_schema().await?;
160        graph.migrate().await?;
161        Ok(graph)
162    }
163
164    /// Connect to `url` (libpq connstring / `postgresql://…`) with an 8-connection
165    /// pool, apply schema + migrations, and return a ready backend. Uses the
166    /// default (empty) table prefix.
167    pub async fn connect(url: &str) -> Result<Self> {
168        Self::connect_with_prefix(url, "").await
169    }
170
171    /// Like [`connect`](Self::connect), but every table/index is namespaced
172    /// under `prefix`.
173    pub async fn connect_with_prefix(url: &str, prefix: &str) -> Result<Self> {
174        let pool = PgPoolOptions::new()
175            .max_connections(8)
176            .connect(url)
177            .await
178            .map_err(pg_err)?;
179        Self::from_pool_with_prefix(pool, prefix).await
180    }
181
182    /// The underlying connection pool.
183    ///
184    /// Callers that need cross-table transactional consistency — e.g. pruning a
185    /// law's chunks plus their graph edges atomically — can `pool().begin()` and
186    /// run their own DML, then call the [`append_edges_in_tx`](Self::append_edges_in_tx)
187    /// / [`remove_edges_for_node_in_tx`](Self::remove_edges_for_node_in_tx)
188    /// helpers on the same `&mut PgConnection`.
189    pub fn pool(&self) -> &PgPool {
190        &self.pool
191    }
192
193    // ── table-name helpers (identical to PgGraph) ──────────────────────
194
195    fn nodes_tbl(&self) -> String {
196        format!("{}nodes", self.table_prefix)
197    }
198
199    fn edges_tbl(&self) -> String {
200        format!("{}edges", self.table_prefix)
201    }
202
203    fn meta_tbl(&self) -> String {
204        format!("{}_meta", self.table_prefix)
205    }
206
207    // ── schema init + migrate ──────────────────────────────────────────
208
209    /// Create the graph schema if absent. Idempotent — DDL ported verbatim from
210    /// `pg.rs::init_schema` (same tables, same prefixed index names). Each DDL
211    /// statement is a separate round trip because sqlx's extended-protocol
212    /// `query()` accepts one statement per call; init runs once per connect.
213    async fn init_schema(&self) -> Result<()> {
214        let p = &self.table_prefix;
215        let nodes = self.nodes_tbl();
216        let edges = self.edges_tbl();
217        let meta = self.meta_tbl();
218        let idx = |name: &str| format!("{p}{name}");
219
220        let stmts: Vec<String> = vec![
221            format!(
222                "CREATE TABLE IF NOT EXISTS {nodes} (
223                    id           TEXT PRIMARY KEY,
224                    node_type    TEXT NOT NULL,
225                    title        TEXT NOT NULL,
226                    tags         TEXT NOT NULL DEFAULT '',
227                    projects     TEXT NOT NULL DEFAULT '',
228                    agents       TEXT NOT NULL DEFAULT '',
229                    created      TEXT NOT NULL,
230                    updated      TEXT NOT NULL,
231                    body         TEXT NOT NULL DEFAULT '',
232                    importance   DOUBLE PRECISION NOT NULL DEFAULT 0.5,
233                    access_count BIGINT NOT NULL DEFAULT 0,
234                    accessed_at  TEXT NOT NULL DEFAULT '',
235                    valid_until  TEXT NOT NULL DEFAULT '',
236                    last_verified TEXT NOT NULL DEFAULT ''
237                )"
238            ),
239            format!(
240                "CREATE TABLE IF NOT EXISTS {edges} (
241                    id       TEXT PRIMARY KEY,
242                    source   TEXT NOT NULL,
243                    target   TEXT NOT NULL,
244                    relation TEXT NOT NULL DEFAULT 'related',
245                    weight   DOUBLE PRECISION NOT NULL DEFAULT 1.0,
246                    ts       TEXT NOT NULL
247                )"
248            ),
249            format!(
250                "CREATE INDEX IF NOT EXISTS {} ON {edges}(source)",
251                idx("idx_edges_source")
252            ),
253            format!(
254                "CREATE INDEX IF NOT EXISTS {} ON {edges}(target)",
255                idx("idx_edges_target")
256            ),
257            format!(
258                "CREATE UNIQUE INDEX IF NOT EXISTS {} ON {edges}(source, target, relation)",
259                idx("idx_edges_src_tgt_rel")
260            ),
261            format!(
262                "CREATE INDEX IF NOT EXISTS {} ON {edges}(source, relation)",
263                idx("idx_edges_src_rel")
264            ),
265            format!(
266                "CREATE INDEX IF NOT EXISTS {} ON {edges}(target, relation)",
267                idx("idx_edges_tgt_rel")
268            ),
269            format!(
270                "CREATE INDEX IF NOT EXISTS {} ON {nodes}(node_type)",
271                idx("idx_nodes_type")
272            ),
273            format!(
274                "CREATE INDEX IF NOT EXISTS {} ON {nodes}(updated DESC)",
275                idx("idx_nodes_updated")
276            ),
277            format!(
278                "CREATE INDEX IF NOT EXISTS {} ON {nodes}(importance DESC)",
279                idx("idx_nodes_importance")
280            ),
281            format!(
282                "CREATE INDEX IF NOT EXISTS {} ON {nodes}(accessed_at DESC)",
283                idx("idx_nodes_accessed")
284            ),
285            format!(
286                "CREATE INDEX IF NOT EXISTS {} ON {nodes}(created)",
287                idx("idx_nodes_created")
288            ),
289            format!(
290                "CREATE TABLE IF NOT EXISTS {meta} (key TEXT PRIMARY KEY, value TEXT NOT NULL)"
291            ),
292            format!(
293                "ALTER TABLE {nodes} ADD COLUMN IF NOT EXISTS valid_until TEXT NOT NULL DEFAULT ''"
294            ),
295            format!(
296                "ALTER TABLE {nodes} ADD COLUMN IF NOT EXISTS last_verified TEXT NOT NULL DEFAULT ''"
297            ),
298            format!(
299                "INSERT INTO {meta} (key, value) VALUES ('graph_schema_version', '{}')
300                 ON CONFLICT (key) DO NOTHING",
301                GRAPH_SCHEMA_VERSION
302            ),
303        ];
304
305        for ddl in &stmts {
306            sqlx::query(ddl).execute(&self.pool).await.map_err(pg_err)?;
307        }
308        Ok(())
309    }
310
311    /// Recorded graph schema version from `{prefix}_meta`, or `0` if unset.
312    pub async fn current_version(&self) -> Result<u32> {
313        let meta = self.meta_tbl();
314        let row = sqlx::query(&format!(
315            "SELECT value FROM {meta} WHERE key = 'graph_schema_version'"
316        ))
317        .fetch_optional(&self.pool)
318        .await
319        .map_err(pg_err)?;
320        match row {
321            Some(r) => {
322                let s: String = r.get("value");
323                Ok(s.parse().unwrap_or(0))
324            }
325            None => Ok(0),
326        }
327    }
328
329    /// Apply pending migrations up to [`GRAPH_SCHEMA_VERSION`]. No-op when current.
330    ///
331    /// Runs in a single transaction with rollback on failure — matching
332    /// `pg.rs::migrate`. Returns the new version.
333    pub async fn migrate(&self) -> Result<u32> {
334        let current = self.current_version().await?;
335        if current >= GRAPH_SCHEMA_VERSION {
336            return Ok(current);
337        }
338        let p = &self.table_prefix;
339        let nodes = self.nodes_tbl();
340        let edges = self.edges_tbl();
341        let meta = self.meta_tbl();
342        let idx = |name: &str| format!("{p}{name}");
343
344        let mut tx = self.pool.begin().await.map_err(pg_err)?;
345        let mut v = current;
346        // v1 -> v2: index nodes by creation timestamp.
347        if v < 2 {
348            sqlx::query(&format!(
349                "CREATE INDEX IF NOT EXISTS {} ON {nodes}(created)",
350                idx("idx_nodes_created")
351            ))
352            .execute(&mut *tx)
353            .await
354            .map_err(pg_err)?;
355            v = 2;
356        }
357        // v2 -> v3: composite indexes for relation-filtered directed edge lookups.
358        if v < 3 {
359            sqlx::query(&format!(
360                "CREATE INDEX IF NOT EXISTS {} ON {edges}(source, relation)",
361                idx("idx_edges_src_rel")
362            ))
363            .execute(&mut *tx)
364            .await
365            .map_err(pg_err)?;
366            sqlx::query(&format!(
367                "CREATE INDEX IF NOT EXISTS {} ON {edges}(target, relation)",
368                idx("idx_edges_tgt_rel")
369            ))
370            .execute(&mut *tx)
371            .await
372            .map_err(pg_err)?;
373            v = 3;
374        }
375        // v3 -> v4: temporal validity columns.
376        if v < 4 {
377            sqlx::query(&format!(
378                "ALTER TABLE {nodes} ADD COLUMN IF NOT EXISTS valid_until TEXT NOT NULL DEFAULT ''"
379            ))
380            .execute(&mut *tx)
381            .await
382            .map_err(pg_err)?;
383            sqlx::query(&format!(
384                "ALTER TABLE {nodes} ADD COLUMN IF NOT EXISTS last_verified TEXT NOT NULL DEFAULT ''"
385            ))
386            .execute(&mut *tx)
387            .await
388            .map_err(pg_err)?;
389            v = 4;
390        }
391        sqlx::query(&format!(
392            "UPDATE {meta} SET value = $1 WHERE key = 'graph_schema_version'"
393        ))
394        .bind(v.to_string())
395        .execute(&mut *tx)
396        .await
397        .map_err(pg_err)?;
398        tx.commit().await.map_err(pg_err)?;
399        Ok(v)
400    }
401
402    // ── basic CRUD ─────────────────────────────────────────────────────
403
404    /// Insert or update a node (upsert by `id`). Mirrors `PgGraph::upsert_node`.
405    pub async fn upsert_node(&self, node: &GraphNode) -> Result<()> {
406        let tags = join_csv(&node.tags);
407        let projects = join_csv(&node.projects);
408        let agents = join_csv(&node.agents);
409        let nodes = self.nodes_tbl();
410        sqlx::query(&format!(
411            "INSERT INTO {nodes} (id, node_type, title, tags, projects, agents, created, updated, body, importance, access_count, accessed_at, valid_until, last_verified)
412             VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)
413             ON CONFLICT (id) DO UPDATE SET
414               node_type=EXCLUDED.node_type, title=EXCLUDED.title, tags=EXCLUDED.tags,
415               projects=EXCLUDED.projects, agents=EXCLUDED.agents, created=EXCLUDED.created,
416               updated=EXCLUDED.updated, body=EXCLUDED.body, importance=EXCLUDED.importance,
417               access_count=EXCLUDED.access_count, accessed_at=EXCLUDED.accessed_at,
418               valid_until=EXCLUDED.valid_until, last_verified=EXCLUDED.last_verified"
419        ))
420        .bind(node.id.as_str())
421        .bind(node.node_type.as_str())
422        .bind(node.title.as_str())
423        .bind(tags.as_str())
424        .bind(projects.as_str())
425        .bind(agents.as_str())
426        .bind(node.created.as_str())
427        .bind(node.updated.as_str())
428        .bind(node.body.as_str())
429        .bind(node.importance)
430        .bind(node.access_count)
431        .bind(node.accessed_at.as_str())
432        .bind(node.valid_until.as_str())
433        .bind(node.last_verified.as_str())
434        .execute(&self.pool)
435        .await
436        .map_err(pg_err)?;
437        Ok(())
438    }
439
440    /// Read a node by `id`, or `None` if absent.
441    pub async fn read_node(&self, id: &str) -> Result<Option<GraphNode>> {
442        let nodes = self.nodes_tbl();
443        let row = sqlx::query(&format!("SELECT {NODE_COLUMNS} FROM {nodes} WHERE id = $1"))
444            .bind(id)
445            .fetch_optional(&self.pool)
446            .await
447            .map_err(pg_err)?;
448        Ok(row.as_ref().map(row_to_node))
449    }
450
451    /// Delete a node by `id`. Returns `true` if a row was removed.
452    pub async fn delete_node(&self, id: &str) -> Result<bool> {
453        let nodes = self.nodes_tbl();
454        let res = sqlx::query(&format!("DELETE FROM {nodes} WHERE id = $1"))
455            .bind(id)
456            .execute(&self.pool)
457            .await
458            .map_err(pg_err)?;
459        Ok(res.rows_affected() > 0)
460    }
461
462    /// Append a single edge (`ON CONFLICT DO NOTHING`). Mirrors `PgGraph::append_edge`.
463    pub async fn append_edge(&self, edge: &GraphEdge) -> Result<()> {
464        let edges = self.edges_tbl();
465        sqlx::query(&format!(
466            "INSERT INTO {edges} (id, source, target, relation, weight, ts)
467             VALUES ($1,$2,$3,$4,$5,$6) ON CONFLICT DO NOTHING"
468        ))
469        .bind(edge.id.as_str())
470        .bind(edge.source.as_str())
471        .bind(edge.target.as_str())
472        .bind(edge.relation.as_str())
473        .bind(edge.weight)
474        .bind(edge.ts.as_str())
475        .execute(&self.pool)
476        .await
477        .map_err(pg_err)?;
478        Ok(())
479    }
480
481    /// All edges touching `node_id` (either endpoint). Mirrors `PgGraph::edges_for_node`.
482    pub async fn edges_for_node(&self, node_id: &str) -> Result<Vec<GraphEdge>> {
483        let edges = self.edges_tbl();
484        let rows = sqlx::query(&format!(
485            "SELECT id, source, target, relation, weight, ts FROM {edges} \
486             WHERE source = $1 OR target = $1"
487        ))
488        .bind(node_id)
489        .fetch_all(&self.pool)
490        .await
491        .map_err(pg_err)?;
492        Ok(rows.iter().map(row_to_edge).collect())
493    }
494
495    /// Delete a single edge by `id`. Returns `true` if a row was removed.
496    pub async fn delete_edge(&self, id: &str) -> Result<bool> {
497        let edges = self.edges_tbl();
498        let res = sqlx::query(&format!("DELETE FROM {edges} WHERE id = $1"))
499            .bind(id)
500            .execute(&self.pool)
501            .await
502            .map_err(pg_err)?;
503        Ok(res.rows_affected() > 0)
504    }
505
506    // ── klr core: batch edges + directed traversal ─────────────────────
507
508    /// Batch-insert edges in chunks of 5000, each chunk its own transaction
509    /// (`ON CONFLICT DO NOTHING` per row). Mirrors `PgGraph::append_edges`.
510    /// sqlx caches the prepared statement per connection, so the repeated
511    /// `INSERT` reuses one parse/plan per chunk.
512    pub async fn append_edges(&self, edges: &[GraphEdge]) -> Result<()> {
513        if edges.is_empty() {
514            return Ok(());
515        }
516        let edges_tbl = self.edges_tbl();
517        let sql = format!(
518            "INSERT INTO {edges_tbl} (id, source, target, relation, weight, ts)
519             VALUES ($1,$2,$3,$4,$5,$6) ON CONFLICT DO NOTHING"
520        );
521        const CHUNK: usize = 5000;
522        for chunk in edges.chunks(CHUNK) {
523            let mut tx = self.pool.begin().await.map_err(pg_err)?;
524            for e in chunk {
525                sqlx::query(&sql)
526                    .bind(e.id.as_str())
527                    .bind(e.source.as_str())
528                    .bind(e.target.as_str())
529                    .bind(e.relation.as_str())
530                    .bind(e.weight)
531                    .bind(e.ts.as_str())
532                    .execute(&mut *tx)
533                    .await
534                    .map_err(pg_err)?;
535            }
536            tx.commit().await.map_err(pg_err)?;
537        }
538        Ok(())
539    }
540
541    /// Append edges **within a caller-provided transaction** (commit is the
542    /// caller's responsibility — this method does not commit). Enables atomic
543    /// cross-table writes: begin a tx on [`pool`](Self::pool), write chunks /
544    /// vectors / other relations, call this with the same `&mut PgConnection`,
545    /// then `commit()`. A failure anywhere rolls back the whole set.
546    pub async fn append_edges_in_tx(
547        &self,
548        tx: &mut sqlx::PgConnection,
549        edges: &[GraphEdge],
550    ) -> Result<()> {
551        if edges.is_empty() {
552            return Ok(());
553        }
554        let edges_tbl = self.edges_tbl();
555        let sql = format!(
556            "INSERT INTO {edges_tbl} (id, source, target, relation, weight, ts)
557             VALUES ($1,$2,$3,$4,$5,$6) ON CONFLICT DO NOTHING"
558        );
559        for e in edges {
560            sqlx::query(&sql)
561                .bind(e.id.as_str())
562                .bind(e.source.as_str())
563                .bind(e.target.as_str())
564                .bind(e.relation.as_str())
565                .bind(e.weight)
566                .bind(e.ts.as_str())
567                .execute(&mut *tx)
568                .await
569                .map_err(pg_err)?;
570        }
571        Ok(())
572    }
573
574    /// Directed, optionally relation-filtered edge lookup. Mirrors
575    /// `PgGraph::edges_for_node_dir`. `dir` selects out / in / both; `relation`
576    /// further restricts to a single relationship type when `Some`.
577    pub async fn edges_for_node_dir(
578        &self,
579        node_id: &str,
580        dir: EdgeDirection,
581        relation: Option<&str>,
582    ) -> Result<Vec<GraphEdge>> {
583        let edges = self.edges_tbl();
584        let dir_clause = match dir {
585            EdgeDirection::Out => "source = $1",
586            EdgeDirection::In => "target = $1",
587            EdgeDirection::Both => "(source = $1 OR target = $1)",
588        };
589        let rows = if let Some(r) = relation {
590            let sql = format!(
591                "SELECT id, source, target, relation, weight, ts FROM {edges} \
592                 WHERE {dir_clause} AND relation = $2 ORDER BY weight DESC"
593            );
594            sqlx::query(&sql)
595                .bind(node_id)
596                .bind(r)
597                .fetch_all(&self.pool)
598                .await
599                .map_err(pg_err)?
600        } else {
601            let sql = format!(
602                "SELECT id, source, target, relation, weight, ts FROM {edges} \
603                 WHERE {dir_clause} ORDER BY weight DESC"
604            );
605            sqlx::query(&sql)
606                .bind(node_id)
607                .fetch_all(&self.pool)
608                .await
609                .map_err(pg_err)?
610        };
611        Ok(rows.iter().map(row_to_edge).collect())
612    }
613
614    /// Aggregate neighbor weights from seed nodes, optionally direction- and
615    /// relation-filtered. Mirrors `PgGraph::neighbors_weighted`: walks the
616    /// source/target halves, uses `= ANY($1)` for the seed set, `GROUP BY` the
617    /// opposite endpoint, and `SUM(weight)`. Seeds themselves are excluded.
618    /// Capped at 100 seeds.
619    pub async fn neighbors_weighted(
620        &self,
621        seed_ids: &[String],
622        dir: EdgeDirection,
623        relation: Option<&str>,
624    ) -> Result<Vec<(String, f64)>> {
625        if seed_ids.is_empty() {
626            return Ok(vec![]);
627        }
628        let edges = self.edges_tbl();
629        const MAX_SEEDS: usize = 100;
630        let seed_ids = if seed_ids.len() > MAX_SEEDS {
631            &seed_ids[..MAX_SEEDS]
632        } else {
633            seed_ids
634        };
635        let seed_arr: Vec<String> = seed_ids.to_vec();
636        let seed_set: HashSet<&str> = seed_ids.iter().map(String::as_str).collect();
637
638        let halves: &[&str] = match dir {
639            EdgeDirection::Out => &["source"],
640            EdgeDirection::In => &["target"],
641            EdgeDirection::Both => &["source", "target"],
642        };
643        let mut weights: HashMap<String, f64> = HashMap::new();
644
645        for &follow in halves {
646            // `follow` is the column the seed matches; the neighbor is the
647            // opposite endpoint.
648            let select_col = if follow == "source" {
649                "target"
650            } else {
651                "source"
652            };
653            let rel_clause = relation.map(|_| " AND relation = $2").unwrap_or("");
654            let sql = format!(
655                "SELECT {select_col} AS nb, SUM(weight) AS w FROM {edges} \
656                 WHERE {follow} = ANY($1){rel_clause} GROUP BY {select_col}"
657            );
658            let rows = if let Some(r) = relation {
659                sqlx::query(&sql)
660                    .bind(&seed_arr)
661                    .bind(r)
662                    .fetch_all(&self.pool)
663                    .await
664                    .map_err(pg_err)?
665            } else {
666                sqlx::query(&sql)
667                    .bind(&seed_arr)
668                    .fetch_all(&self.pool)
669                    .await
670                    .map_err(pg_err)?
671            };
672            for row in &rows {
673                let nb: String = row.get("nb");
674                let w: f64 = row.get("w");
675                if !seed_set.contains(nb.as_str()) {
676                    *weights.entry(nb).or_default() += w;
677                }
678            }
679        }
680
681        let mut result: Vec<(String, f64)> = weights.into_iter().collect();
682        result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
683        Ok(result)
684    }
685
686    /// Remove every edge touching `node_id` (either endpoint).
687    /// Mirrors `PgGraph::remove_edges_for_node`.
688    pub async fn remove_edges_for_node(&self, node_id: &str) -> Result<()> {
689        let edges = self.edges_tbl();
690        sqlx::query(&format!(
691            "DELETE FROM {edges} WHERE source = $1 OR target = $1"
692        ))
693        .bind(node_id)
694        .execute(&self.pool)
695        .await
696        .map_err(pg_err)?;
697        Ok(())
698    }
699
700    /// Remove edges for `node_id` **within a caller-provided transaction** (no
701    /// commit here). Same atomic-prune pattern as [`Self::append_edges_in_tx`].
702    pub async fn remove_edges_for_node_in_tx(
703        &self,
704        tx: &mut sqlx::PgConnection,
705        node_id: &str,
706    ) -> Result<()> {
707        let edges = self.edges_tbl();
708        sqlx::query(&format!(
709            "DELETE FROM {edges} WHERE source = $1 OR target = $1"
710        ))
711        .bind(node_id)
712        .execute(tx)
713        .await
714        .map_err(pg_err)?;
715        Ok(())
716    }
717
718    // ── search / traversal ─────────────────────────────────────────────
719
720    /// ILIKE substring search over `(title, body, tags)` for each
721    /// whitespace-separated term (AND of terms). Mirrors `PgGraph::search_nodes`
722    /// — extension-free vanilla PostgreSQL ILIKE with escaped, wrapped patterns.
723    pub async fn search_nodes(&self, query: &str, limit: usize) -> Result<Vec<GraphNode>> {
724        let terms = search_patterns(query);
725        if terms.is_empty() {
726            return Ok(vec![]);
727        }
728        let nodes = self.nodes_tbl();
729        let mut qb = QueryBuilder::<Postgres>::new("SELECT ");
730        qb.push(NODE_COLUMNS);
731        qb.push(" FROM ");
732        qb.push(nodes.as_str());
733        qb.push(" WHERE ");
734        for (i, term) in terms.iter().enumerate() {
735            if i > 0 {
736                qb.push(" AND ");
737            }
738            qb.push("(title || ' ' || body || ' ' || tags) ILIKE ");
739            qb.push_bind(term.clone());
740            qb.push(" ESCAPE '\\'");
741        }
742        qb.push(" ORDER BY importance DESC, updated DESC LIMIT ");
743        qb.push_bind(limit as i64);
744        let rows = qb.build().fetch_all(&self.pool).await.map_err(pg_err)?;
745        Ok(rows.iter().map(row_to_node).collect())
746    }
747
748    /// Dynamic filter by tag / node_type / project, ranked by `updated` DESC.
749    /// Mirrors `PgGraph::query_nodes`.
750    pub async fn query_nodes(
751        &self,
752        tag: Option<&str>,
753        node_type: Option<&str>,
754        project: Option<&str>,
755        limit: usize,
756    ) -> Result<Vec<GraphNode>> {
757        let limit = limit.min(200) as i64;
758        let nodes = self.nodes_tbl();
759        let mut qb = QueryBuilder::<Postgres>::new("SELECT ");
760        qb.push(NODE_COLUMNS);
761        qb.push(" FROM ");
762        qb.push(nodes.as_str());
763        let mut filtered = false;
764        if let Some(t) = tag {
765            qb.push(" WHERE (',' || tags || ',') ILIKE ('%,' || ");
766            qb.push_bind(escape_like(t));
767            qb.push(" || ',%') ESCAPE '\\'");
768            filtered = true;
769        }
770        if let Some(nt) = node_type {
771            qb.push(if filtered {
772                " AND node_type = "
773            } else {
774                " WHERE node_type = "
775            });
776            qb.push_bind(nt);
777            filtered = true;
778        }
779        if let Some(p) = project {
780            qb.push(if filtered { " AND " } else { " WHERE " });
781            qb.push("(',' || projects || ',') ILIKE ('%,' || ");
782            qb.push_bind(escape_like(p));
783            qb.push(" || ',%') ESCAPE '\\'");
784        }
785        qb.push(" ORDER BY updated DESC LIMIT ");
786        qb.push_bind(limit);
787        let rows = qb.build().fetch_all(&self.pool).await.map_err(pg_err)?;
788        Ok(rows.iter().map(row_to_node).collect())
789    }
790
791    /// Composite recall — rank nodes by recency, importance, access, FTS, and a
792    /// PageRank centrality boost over the induced candidate subgraph. Mirrors
793    /// `PgGraph::smart_recall` (`pg.rs:573`): identical composite weights and
794    /// graph-boost math. Only the driver differs — `ANY($1)` array binds replace
795    /// the dynamic placeholder lists, so the PageRank pass loads the subgraph
796    /// edges in one round-trip instead of one-per-candidate.
797    pub async fn smart_recall(
798        &self,
799        project: Option<&str>,
800        hint: Option<&str>,
801        limit: usize,
802    ) -> Result<Vec<ScoredNode>> {
803        let now_secs = SystemTime::now()
804            .duration_since(UNIX_EPOCH)
805            .unwrap_or_default()
806            .as_secs();
807
808        // FTS match set (ILIKE), used as a binary boost signal.
809        let fts_ids: HashSet<String> = match hint {
810            Some(h) if !h.is_empty() => self
811                .search_nodes(h, limit * 4)
812                .await?
813                .into_iter()
814                .map(|n| n.id)
815                .collect(),
816            _ => HashSet::new(),
817        };
818
819        // Candidate fetch (broad set), excluding stale nodes.
820        let candidate_limit = (limit * 4).max(40) as i64;
821        let nodes = self.nodes_tbl();
822        let mut qb = QueryBuilder::<Postgres>::new("SELECT ");
823        qb.push(NODE_COLUMNS);
824        qb.push(" FROM ");
825        qb.push(nodes.as_str());
826        qb.push(" WHERE (',' || tags || ',') NOT ILIKE '%,stale,%'");
827        if let Some(p) = project {
828            qb.push(" AND (',' || projects || ',') ILIKE ('%,' || ");
829            qb.push_bind(escape_like(p));
830            qb.push(" || ',%') ESCAPE '\\'");
831        }
832        qb.push(" ORDER BY importance DESC, updated DESC LIMIT ");
833        qb.push_bind(candidate_limit);
834        let rows = qb.build().fetch_all(&self.pool).await.map_err(pg_err)?;
835        let candidates: Vec<GraphNode> = rows.iter().map(row_to_node).collect();
836
837        // Composite scoring — identical weights/recency as the SQLite/PgGraph backends.
838        let mut scored: Vec<ScoredNode> = candidates
839            .into_iter()
840            .map(|node| {
841                let recency = compute_recency(&node.updated, now_secs);
842                let importance = node.importance;
843                let access_freq = (node.access_count.max(0) as f64 / 20.0).min(1.0);
844                let fts_match = if fts_ids.contains(&node.id) { 1.0 } else { 0.0 };
845                let score = W_RECENCY * recency
846                    + W_IMPORTANCE * importance
847                    + W_ACCESS * access_freq
848                    + W_FTS * fts_match;
849                ScoredNode { node, score }
850            })
851            .collect();
852        scored.sort_by(|a, b| {
853            b.score
854                .partial_cmp(&a.score)
855                .unwrap_or(std::cmp::Ordering::Equal)
856        });
857        scored.truncate(limit);
858
859        // Graph-boost pass: PageRank centrality over the induced subgraph of the
860        // top candidates. Shares the pagerank math with PgGraph/SQLite (zero drift)
861        // — only the edge-load SQL differs. One round-trip via `ANY($1)`.
862        if scored.len() > 1 {
863            const MAX_GRAPH_BOOST_PARTICIPANTS: usize = 100;
864            let candidate_ids: Vec<String> = scored
865                .iter()
866                .take(MAX_GRAPH_BOOST_PARTICIPANTS)
867                .map(|sn| sn.node.id.clone())
868                .collect();
869            let edges = self.edges_tbl();
870            let sql = format!(
871                "SELECT id, source, target, relation, weight, ts FROM {edges} \
872                 WHERE source = ANY($1) AND target = ANY($1)"
873            );
874            let sub_edges: Vec<GraphEdge> = sqlx::query(&sql)
875                .bind(&candidate_ids)
876                .fetch_all(&self.pool)
877                .await
878                .map(|rows| rows.iter().map(row_to_edge).collect())
879                .unwrap_or_default();
880            let csr = CsrGraph::from_edges(&candidate_ids, &sub_edges);
881            let pr = pagerank_default(&csr);
882            let max_pr = pr.iter().copied().fold(0.0_f64, f64::max).max(1e-12);
883            let pr_map: HashMap<String, f64> = candidate_ids
884                .iter()
885                .zip(pr.iter())
886                .map(|(id, &s)| (id.clone(), s / max_pr))
887                .collect();
888            for sn in &mut scored {
889                let boost = pr_map.get(&sn.node.id).copied().unwrap_or(0.0);
890                sn.score += W_GRAPH * boost;
891            }
892            scored.sort_by(|a, b| {
893                b.score
894                    .partial_cmp(&a.score)
895                    .unwrap_or(std::cmp::Ordering::Equal)
896            });
897        }
898
899        // Touch retrieved nodes in one statement (access_count++,
900        // accessed_at = now) rather than N round-trips.
901        if !scored.is_empty() {
902            let now = now_iso();
903            let ids: Vec<String> = scored.iter().map(|sn| sn.node.id.clone()).collect();
904            let sql = format!(
905                "UPDATE {nodes} SET access_count = access_count + 1, accessed_at = $1 WHERE id = ANY($2)"
906            );
907            let _ = sqlx::query(&sql)
908                .bind(now)
909                .bind(&ids)
910                .execute(&self.pool)
911                .await;
912        }
913
914        Ok(scored)
915    }
916
917    /// Bidirectional BFS up to `depth` hops from `start_id` via a recursive CTE.
918    /// Mirrors `PgGraph::related_nodes`. Returns distinct neighbor IDs (the start
919    /// node is excluded), capped at 500.
920    pub async fn related_nodes(&self, start_id: &str, depth: usize) -> Result<Vec<String>> {
921        let edges = self.edges_tbl();
922        // PostgreSQL requires a single recursive term: the bidirectional seed
923        // is folded into a subquery, then one recursive step follows edges in
924        // either direction (CASE picks the opposite endpoint).
925        let sql = format!(
926            "WITH RECURSIVE bfs(node_id, lvl) AS (
927                SELECT nb.node_id, 1 FROM (
928                    SELECT target AS node_id FROM {edges} WHERE source = $1
929                    UNION
930                    SELECT source AS node_id FROM {edges} WHERE target = $1
931                ) nb
932                UNION
933                SELECT CASE WHEN e.source = bfs.node_id THEN e.target ELSE e.source END,
934                       bfs.lvl + 1
935                FROM bfs
936                JOIN {edges} e ON e.source = bfs.node_id OR e.target = bfs.node_id
937                WHERE bfs.lvl < $2
938            )
939            SELECT DISTINCT node_id FROM bfs WHERE node_id <> $1 LIMIT 500"
940        );
941        let rows = sqlx::query(&sql)
942            .bind(start_id)
943            .bind(depth as i32)
944            .fetch_all(&self.pool)
945            .await
946            .map_err(pg_err)?;
947        Ok(rows
948            .iter()
949            .map(|r| {
950                let id: String = r.get("node_id");
951                id
952            })
953            .collect())
954    }
955}
956
957#[cfg(test)]
958mod tests {
959    use super::*;
960
961    /// `LLMKERNEL_PG_URL` unset → self-skip (pgvector / pg.rs pattern).
962    fn pg_url() -> Option<String> {
963        std::env::var("LLMKERNEL_PG_URL").ok()
964    }
965
966    fn sample_node(id: &str) -> GraphNode {
967        GraphNode {
968            id: id.to_string(),
969            node_type: "concept".to_string(),
970            title: format!("Node {id}"),
971            body: "sqlx pg test body".to_string(),
972            tags: vec!["sqlx".to_string()],
973            projects: vec![],
974            agents: vec![],
975            created: "2026-01-01T00:00:00Z".to_string(),
976            updated: "2026-01-01T00:00:00Z".to_string(),
977            importance: 0.5,
978            access_count: 0,
979            accessed_at: String::new(),
980            ..Default::default()
981        }
982    }
983
984    fn sample_edge(id: &str, src: &str, tgt: &str, rel: &str, w: f64) -> GraphEdge {
985        GraphEdge {
986            id: id.to_string(),
987            source: src.to_string(),
988            target: tgt.to_string(),
989            relation: rel.to_string(),
990            weight: w,
991            ts: "2026-01-01T00:00:00Z".to_string(),
992        }
993    }
994
995    /// Drop the prefixed tables so parallel tests with distinct prefixes don't
996    /// accumulate state.
997    async fn cleanup(pool: &PgPool, prefix: &str) {
998        let _ = sqlx::query(&format!("DROP TABLE IF EXISTS {}nodes", prefix))
999            .execute(pool)
1000            .await;
1001        let _ = sqlx::query(&format!("DROP TABLE IF EXISTS {}edges", prefix))
1002            .execute(pool)
1003            .await;
1004        let _ = sqlx::query(&format!("DROP TABLE IF EXISTS {}_meta", prefix))
1005            .execute(pool)
1006            .await;
1007    }
1008
1009    // ── offline (no server) ────────────────────────────────────────────
1010
1011    /// The ILIKE pattern transform is identical to pg.rs (escapes LIKE
1012    /// wildcards, wraps each term).
1013    #[test]
1014    fn search_patterns_escapes_and_wraps() {
1015        assert!(search_patterns("").is_empty());
1016        assert_eq!(search_patterns("rust"), vec!["%rust%".to_string()]);
1017        assert_eq!(search_patterns("100%"), vec!["%100\\%%".to_string()]);
1018        assert_eq!(search_patterns("a_b"), vec!["%a\\_b%".to_string()]);
1019        assert_eq!(
1020            search_patterns("rust db"),
1021            vec!["%rust%".to_string(), "%db%".to_string()]
1022        );
1023    }
1024
1025    // ── live (LLMKERNEL_PG_URL gated, throwaway prefixed tables) ───────
1026
1027    /// Basic CRUD round-trip: upsert/read/update/delete nodes + edges.
1028    #[tokio::test]
1029    async fn basic_crud_roundtrip() {
1030        let Some(url) = pg_url() else {
1031            eprintln!("skip: LLMKERNEL_PG_URL unset");
1032            return;
1033        };
1034        let prefix = "lk_sg1_";
1035        let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1036            .await
1037            .expect("connect");
1038
1039        assert!(g.read_node("n1").await.unwrap().is_none());
1040        g.upsert_node(&sample_node("n1")).await.unwrap();
1041        let loaded = g.read_node("n1").await.unwrap().unwrap();
1042        assert_eq!(loaded.title, "Node n1");
1043        assert_eq!(loaded.tags, vec!["sqlx".to_string()]);
1044
1045        // upsert = update on conflict
1046        let mut updated = sample_node("n1");
1047        updated.title = "Updated".into();
1048        g.upsert_node(&updated).await.unwrap();
1049        assert_eq!(g.read_node("n1").await.unwrap().unwrap().title, "Updated");
1050
1051        // single edge append + edges_for_node + delete_edge
1052        g.upsert_node(&sample_node("n2")).await.unwrap();
1053        g.append_edge(&sample_edge("e1", "n1", "n2", "related", 1.0))
1054            .await
1055            .unwrap();
1056        assert_eq!(g.edges_for_node("n1").await.unwrap().len(), 1);
1057        assert!(g.delete_edge("e1").await.unwrap());
1058        assert!(!g.delete_edge("e1").await.unwrap());
1059        assert_eq!(g.edges_for_node("n1").await.unwrap().len(), 0);
1060
1061        // delete node
1062        assert!(g.delete_node("n1").await.unwrap());
1063        assert!(!g.delete_node("n1").await.unwrap());
1064        assert!(g.read_node("n1").await.unwrap().is_none());
1065
1066        cleanup(g.pool(), prefix).await;
1067    }
1068
1069    /// Batch edges with `ON CONFLICT DO NOTHING` dedup: a fresh-id edge sharing
1070    /// a duplicate (source, target, relation) triple is silently ignored.
1071    #[tokio::test]
1072    async fn batch_edges_dedup() {
1073        let Some(url) = pg_url() else {
1074            eprintln!("skip: LLMKERNEL_PG_URL unset");
1075            return;
1076        };
1077        let prefix = "lk_sg2_";
1078        let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1079            .await
1080            .expect("connect");
1081
1082        for n in &["a", "b", "c"] {
1083            g.upsert_node(&sample_node(n)).await.unwrap();
1084        }
1085
1086        let edges = vec![
1087            sample_edge("e1", "a", "b", "cites", 1.0),
1088            sample_edge("e2", "a", "c", "cites", 0.8),
1089            // duplicate (src,tgt,rel) of e1 — ON CONFLICT DO NOTHING
1090            sample_edge("e1dup", "a", "b", "cites", 2.0),
1091        ];
1092        g.append_edges(&edges).await.unwrap();
1093
1094        // Only e1 and e2 survived (e1dup deduped by the unique index).
1095        let out = g
1096            .edges_for_node_dir("a", EdgeDirection::Out, Some("cites"))
1097            .await
1098            .unwrap();
1099        assert_eq!(out.len(), 2, "duplicate (src,tgt,rel) edge ignored");
1100
1101        // empty slice is a no-op
1102        g.append_edges(&[]).await.unwrap();
1103
1104        cleanup(g.pool(), prefix).await;
1105    }
1106
1107    /// Directed, relation-filtered edge lookup + weighted neighbor aggregation.
1108    #[tokio::test]
1109    async fn edges_dir_and_neighbors() {
1110        let Some(url) = pg_url() else {
1111            eprintln!("skip: LLMKERNEL_PG_URL unset");
1112            return;
1113        };
1114        let prefix = "lk_sg3_";
1115        let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1116            .await
1117            .expect("connect");
1118
1119        for n in &["seed", "t1", "t2", "t3"] {
1120            g.upsert_node(&sample_node(n)).await.unwrap();
1121        }
1122
1123        // seed cites t1/t2/t3 (out), t1 cites seed (in), seed related t1 (other rel)
1124        g.append_edges(&[
1125            sample_edge("e1", "seed", "t1", "cites", 1.0),
1126            sample_edge("e2", "seed", "t2", "cites", 0.5),
1127            sample_edge("e3", "seed", "t3", "cites", 0.3),
1128        ])
1129        .await
1130        .unwrap();
1131        g.append_edge(&sample_edge("e4", "t1", "seed", "cites", 2.0))
1132            .await
1133            .unwrap();
1134        g.append_edge(&sample_edge("e5", "seed", "t1", "related", 1.0))
1135            .await
1136            .unwrap();
1137
1138        // Out + "cites": 3 (t1, t2, t3)
1139        let out_cites = g
1140            .edges_for_node_dir("seed", EdgeDirection::Out, Some("cites"))
1141            .await
1142            .unwrap();
1143        assert_eq!(out_cites.len(), 3);
1144
1145        // In + "cites": 1 (from t1)
1146        let in_cites = g
1147            .edges_for_node_dir("seed", EdgeDirection::In, Some("cites"))
1148            .await
1149            .unwrap();
1150        assert_eq!(in_cites.len(), 1);
1151        assert_eq!(in_cites[0].source, "t1");
1152
1153        // neighbors_weighted Out "cites" → t1(1.0), t2(0.5), t3(0.3), desc by weight
1154        let neighbors = g
1155            .neighbors_weighted(&["seed".to_string()], EdgeDirection::Out, Some("cites"))
1156            .await
1157            .unwrap();
1158        assert_eq!(neighbors.len(), 3);
1159        assert_eq!(neighbors[0].0, "t1");
1160        assert!((neighbors[0].1 - 1.0).abs() < 1e-9);
1161        assert_eq!(neighbors[2].0, "t3");
1162
1163        cleanup(g.pool(), prefix).await;
1164    }
1165
1166    /// `remove_edges_for_node` drops only the edges touching the target node.
1167    #[tokio::test]
1168    async fn remove_edges_for_node_drops_only_touching() {
1169        let Some(url) = pg_url() else {
1170            eprintln!("skip: LLMKERNEL_PG_URL unset");
1171            return;
1172        };
1173        let prefix = "lk_sg4_";
1174        let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1175            .await
1176            .expect("connect");
1177
1178        for n in &["x", "y", "z"] {
1179            g.upsert_node(&sample_node(n)).await.unwrap();
1180        }
1181        g.append_edges(&[
1182            sample_edge("e1", "x", "y", "cites", 1.0),
1183            sample_edge("e2", "z", "x", "cites", 1.0),
1184            sample_edge("e3", "y", "z", "cites", 1.0),
1185        ])
1186        .await
1187        .unwrap();
1188        assert_eq!(g.edges_for_node("x").await.unwrap().len(), 2);
1189
1190        g.remove_edges_for_node("x").await.unwrap();
1191        assert_eq!(g.edges_for_node("x").await.unwrap().len(), 0);
1192        // y-z edge survives (neither endpoint is x)
1193        assert_eq!(g.edges_for_node("y").await.unwrap().len(), 1);
1194
1195        cleanup(g.pool(), prefix).await;
1196    }
1197
1198    /// Transaction atomicity: `*_in_tx` commit and rollback paths.
1199    /// This is the core guarantee klr relies on for atomic prune.
1200    #[tokio::test]
1201    async fn tx_atomicity_commit_and_rollback() {
1202        let Some(url) = pg_url() else {
1203            eprintln!("skip: LLMKERNEL_PG_URL unset");
1204            return;
1205        };
1206        let prefix = "lk_sg5_";
1207        let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1208            .await
1209            .expect("connect");
1210
1211        for n in &["p", "q"] {
1212            g.upsert_node(&sample_node(n)).await.unwrap();
1213        }
1214
1215        // Commit path: append_edges_in_tx → commit → edge visible
1216        let mut tx = g.pool().begin().await.expect("begin tx");
1217        g.append_edges_in_tx(&mut tx, &[sample_edge("e1", "p", "q", "cites", 1.0)])
1218            .await
1219            .unwrap();
1220        tx.commit().await.expect("commit");
1221        assert_eq!(g.edges_for_node("p").await.unwrap().len(), 1);
1222
1223        // Rollback path: append_edges_in_tx → rollback → edge NOT visible
1224        let mut tx2 = g.pool().begin().await.expect("begin tx2");
1225        g.append_edges_in_tx(&mut tx2, &[sample_edge("e2", "p", "q", "cites", 1.0)])
1226            .await
1227            .unwrap();
1228        tx2.rollback().await.expect("rollback");
1229        assert_eq!(
1230            g.edges_for_node("p").await.unwrap().len(),
1231            1,
1232            "rolled-back edge not visible"
1233        );
1234
1235        // remove_edges_for_node_in_tx rollback → edge survives
1236        let mut tx3 = g.pool().begin().await.expect("begin tx3");
1237        g.remove_edges_for_node_in_tx(&mut tx3, "p").await.unwrap();
1238        tx3.rollback().await.expect("rollback tx3");
1239        assert_eq!(g.edges_for_node("p").await.unwrap().len(), 1);
1240
1241        // remove_edges_for_node_in_tx commit → edge removed
1242        let mut tx4 = g.pool().begin().await.expect("begin tx4");
1243        g.remove_edges_for_node_in_tx(&mut tx4, "p").await.unwrap();
1244        tx4.commit().await.expect("commit tx4");
1245        assert_eq!(g.edges_for_node("p").await.unwrap().len(), 0);
1246
1247        cleanup(g.pool(), prefix).await;
1248    }
1249
1250    /// ILIKE search + recursive BFS traversal + version helpers.
1251    #[tokio::test]
1252    async fn search_related_and_version() {
1253        let Some(url) = pg_url() else {
1254            eprintln!("skip: LLMKERNEL_PG_URL unset");
1255            return;
1256        };
1257        let prefix = "lk_sg6_";
1258        let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1259            .await
1260            .expect("connect");
1261
1262        let mut rust = sample_node("rust");
1263        rust.title = "Rust ownership".into();
1264        rust.body = "borrow checker".into();
1265        g.upsert_node(&rust).await.unwrap();
1266        g.upsert_node(&sample_node("py")).await.unwrap();
1267
1268        // search_nodes
1269        let hits = g.search_nodes("rust", 10).await.unwrap();
1270        assert_eq!(hits.len(), 1);
1271        assert_eq!(hits[0].id, "rust");
1272        assert!(g.search_nodes("", 10).await.unwrap().is_empty());
1273
1274        // related_nodes via BFS
1275        g.append_edge(&sample_edge("e1", "rust", "py", "related", 1.0))
1276            .await
1277            .unwrap();
1278        let related = g.related_nodes("rust", 2).await.unwrap();
1279        assert!(related.contains(&"py".to_string()));
1280
1281        // version helpers (schema applied on connect)
1282        assert_eq!(g.current_version().await.unwrap(), GRAPH_SCHEMA_VERSION);
1283        assert_eq!(g.migrate().await.unwrap(), GRAPH_SCHEMA_VERSION);
1284
1285        cleanup(g.pool(), prefix).await;
1286    }
1287
1288    /// An invalid table prefix is rejected at construction (before any SQL).
1289    #[tokio::test]
1290    async fn invalid_prefix_rejected() {
1291        let Some(url) = pg_url() else {
1292            eprintln!("skip: LLMKERNEL_PG_URL unset");
1293            return;
1294        };
1295        assert!(
1296            SqlxPgGraph::connect_with_prefix(&url, "lk; drop")
1297                .await
1298                .is_err()
1299        );
1300        assert!(SqlxPgGraph::connect_with_prefix(&url, "1lk").await.is_err());
1301    }
1302
1303    /// `smart_recall` ranks by the composite score (recency + importance + FTS
1304    /// + PageRank boost); `query_nodes` filters by tag/project. An FTS-matched
1305    /// high-importance node outranks a low-importance match.
1306    #[tokio::test]
1307    async fn smart_recall_and_query_nodes() {
1308        let Some(url) = pg_url() else {
1309            eprintln!("skip: LLMKERNEL_PG_URL unset");
1310            return;
1311        };
1312        let prefix = "lk_sgrecall_";
1313        let g = SqlxPgGraph::connect_with_prefix(&url, prefix)
1314            .await
1315            .expect("connect");
1316
1317        // a: high importance, matches "rust"
1318        let mut a = sample_node("a");
1319        a.importance = 0.9;
1320        a.title = "rust memory".into();
1321        a.tags = vec!["lang".into()];
1322        a.projects = vec!["proj".into()];
1323        // b: low importance, matches "rust"
1324        let mut b = sample_node("b");
1325        b.importance = 0.3;
1326        b.title = "rust kernel".into();
1327        b.tags = vec!["lang".into()];
1328        b.projects = vec!["proj".into()];
1329        // c: high importance, does NOT match "rust", different project
1330        let mut c = sample_node("c");
1331        c.importance = 0.8;
1332        c.title = "python runtime".into();
1333        c.tags = vec!["lang".into()];
1334
1335        g.upsert_node(&a).await.unwrap();
1336        g.upsert_node(&b).await.unwrap();
1337        g.upsert_node(&c).await.unwrap();
1338        // citation edges → induce a subgraph for the PageRank boost pass
1339        g.append_edge(&sample_edge("e1", "a", "b", "cites", 1.0))
1340            .await
1341            .unwrap();
1342        g.append_edge(&sample_edge("e2", "b", "c", "cites", 0.7))
1343            .await
1344            .unwrap();
1345
1346        // smart_recall scoped to project "proj" with a "rust" hint: a (0.9 +
1347        // FTS) outranks b (0.3 + FTS); c is excluded by the project filter.
1348        let recalled = g.smart_recall(Some("proj"), Some("rust"), 5).await.unwrap();
1349        let ids: Vec<&str> = recalled.iter().map(|s| s.node.id.as_str()).collect();
1350        assert!(ids.contains(&"a"));
1351        assert!(ids.contains(&"b"));
1352        assert!(!ids.contains(&"c"), "c excluded by project scope");
1353        let pos_a = ids.iter().position(|&x| x == "a").unwrap();
1354        let pos_b = ids.iter().position(|&x| x == "b").unwrap();
1355        assert!(pos_a < pos_b, "high-importance rust match ranks first");
1356        assert!(recalled[0].score > 0.0);
1357
1358        // query_nodes by tag returns all lang nodes
1359        let tagged = g.query_nodes(Some("lang"), None, None, 10).await.unwrap();
1360        assert_eq!(tagged.len(), 3);
1361        // project filter narrows to a, b
1362        let proj = g.query_nodes(None, None, Some("proj"), 10).await.unwrap();
1363        let pids: Vec<&str> = proj.iter().map(|n| n.id.as_str()).collect();
1364        assert!(pids.contains(&"a") && pids.contains(&"b"));
1365        assert!(!pids.contains(&"c"));
1366
1367        cleanup(g.pool(), prefix).await;
1368    }
1369}