Skip to main content

llm_kernel/graph/
store.rs

1//! Node and edge CRUD operations.
2
3use rusqlite::{Connection, params};
4
5use crate::error::{KernelError, Result};
6
7use super::types::{GraphEdge, GraphNode, NODE_COLUMNS, join_csv, row_to_node};
8
9// ── Node CRUD ─────────────────────────────────────────
10
11/// Insert or replace a node.
12pub fn upsert_node(conn: &Connection, node: &GraphNode) -> Result<()> {
13    conn.execute(
14        "INSERT OR REPLACE INTO nodes
15         (id, type, title, tags, projects, agents, created, updated, body, importance, access_count, accessed_at)
16         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)",
17        params![
18            node.id,
19            node.node_type,
20            node.title,
21            join_csv(&node.tags),
22            join_csv(&node.projects),
23            join_csv(&node.agents),
24            node.created,
25            node.updated,
26            node.body,
27            node.importance,
28            node.access_count,
29            node.accessed_at,
30        ],
31    )
32    .map_err(|e| KernelError::Store(e.to_string()))?;
33    Ok(())
34}
35
36/// Read a single node by ID. Returns `None` if not found.
37pub fn read_node(conn: &Connection, id: &str) -> Result<Option<GraphNode>> {
38    let sql = format!("SELECT {NODE_COLUMNS} FROM nodes WHERE id = ?1");
39    match conn.query_row(&sql, params![id], row_to_node) {
40        Ok(node) => Ok(Some(node)),
41        Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
42        Err(e) => Err(KernelError::Store(e.to_string())),
43    }
44}
45
46/// Batch-read multiple nodes by ID.
47pub fn read_nodes(conn: &Connection, ids: &[&str]) -> Result<Vec<GraphNode>> {
48    if ids.is_empty() {
49        return Ok(vec![]);
50    }
51    let ph = ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
52    let sql = format!("SELECT {NODE_COLUMNS} FROM nodes WHERE id IN ({ph})");
53    let mut stmt = conn
54        .prepare(&sql)
55        .map_err(|e| KernelError::Store(e.to_string()))?;
56    let nodes: Vec<GraphNode> = stmt
57        .query_map(rusqlite::params_from_iter(ids.iter()), row_to_node)
58        .map_err(|e| KernelError::Store(e.to_string()))?
59        .filter_map(|r| r.ok())
60        .collect();
61    Ok(nodes)
62}
63
64/// Delete a node by ID. Returns whether a row was deleted.
65pub fn delete_node(conn: &Connection, id: &str) -> Result<bool> {
66    let changed = conn
67        .execute("DELETE FROM nodes WHERE id = ?1", params![id])
68        .map_err(|e| KernelError::Store(e.to_string()))?;
69    Ok(changed > 0)
70}
71
72/// List all node IDs.
73pub fn list_node_ids(conn: &Connection) -> Result<Vec<String>> {
74    let mut stmt = conn
75        .prepare("SELECT id FROM nodes")
76        .map_err(|e| KernelError::Store(e.to_string()))?;
77    let ids: Vec<String> = stmt
78        .query_map([], |row| row.get(0))
79        .map_err(|e| KernelError::Store(e.to_string()))?
80        .filter_map(|r| r.ok())
81        .collect();
82    Ok(ids)
83}
84
85/// Read nodes with optional limit, ordered by updated DESC.
86pub fn read_nodes_limited(conn: &Connection, limit: usize) -> Result<Vec<GraphNode>> {
87    let sql = format!("SELECT {NODE_COLUMNS} FROM nodes ORDER BY updated DESC LIMIT ?");
88    let mut stmt = conn
89        .prepare(&sql)
90        .map_err(|e| KernelError::Store(e.to_string()))?;
91    let nodes: Vec<GraphNode> = stmt
92        .query_map(params![limit as i64], row_to_node)
93        .map_err(|e| KernelError::Store(e.to_string()))?
94        .filter_map(|r| r.ok())
95        .collect();
96    Ok(nodes)
97}
98
99// ── Edge CRUD ─────────────────────────────────────────
100
101/// Append an edge (INSERT OR IGNORE — duplicates by ID are skipped).
102pub fn append_edge(conn: &Connection, edge: &GraphEdge) -> Result<()> {
103    conn.execute(
104        "INSERT OR IGNORE INTO edges (id, source, target, relation, weight, ts)
105         VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
106        params![
107            edge.id,
108            edge.source,
109            edge.target,
110            edge.relation,
111            edge.weight,
112            edge.ts
113        ],
114    )
115    .map_err(|e| KernelError::Store(e.to_string()))?;
116    Ok(())
117}
118
119/// Read edges, capped at `limit`.
120pub fn read_edges(conn: &Connection, limit: usize) -> Result<Vec<GraphEdge>> {
121    let mut stmt = conn
122        .prepare("SELECT id, source, target, relation, weight, ts FROM edges LIMIT ?1")
123        .map_err(|e| KernelError::Store(e.to_string()))?;
124    let edges: Vec<GraphEdge> = stmt
125        .query_map(params![limit as i64], |row| {
126            Ok(GraphEdge {
127                id: row.get(0)?,
128                source: row.get(1)?,
129                target: row.get(2)?,
130                relation: row.get(3)?,
131                weight: row.get(4)?,
132                ts: row.get(5)?,
133            })
134        })
135        .map_err(|e| KernelError::Store(e.to_string()))?
136        .filter_map(|r| r.ok())
137        .collect();
138    Ok(edges)
139}
140
141/// Read edges whose source AND target are both in `ids` — the induced subgraph
142/// over a candidate node set.
143///
144/// Used to build the candidate subgraph for PageRank boosting in
145/// [`smart_recall`](super::recall::smart_recall). `ids.len()` must stay under
146/// SQLite's bind-variable limit (999 by default); `smart_recall` caps at 100.
147pub fn edges_among(conn: &Connection, ids: &[&str]) -> Result<Vec<GraphEdge>> {
148    if ids.is_empty() {
149        return Ok(Vec::new());
150    }
151    let ph = ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
152    let sql = format!(
153        "SELECT id, source, target, relation, weight, ts FROM edges \
154         WHERE source IN ({ph}) AND target IN ({ph})"
155    );
156    let mut stmt = conn
157        .prepare(&sql)
158        .map_err(|e| KernelError::Store(e.to_string()))?;
159    let edges: Vec<GraphEdge> = stmt
160        .query_map(
161            rusqlite::params_from_iter(ids.iter().chain(ids.iter()).copied()),
162            |row| {
163                Ok(GraphEdge {
164                    id: row.get(0)?,
165                    source: row.get(1)?,
166                    target: row.get(2)?,
167                    relation: row.get(3)?,
168                    weight: row.get(4)?,
169                    ts: row.get(5)?,
170                })
171            },
172        )
173        .map_err(|e| KernelError::Store(e.to_string()))?
174        .filter_map(|r| r.ok())
175        .collect();
176    Ok(edges)
177}
178
179/// Delete an edge by ID.
180pub fn delete_edge(conn: &Connection, id: &str) -> Result<bool> {
181    let changed = conn
182        .execute("DELETE FROM edges WHERE id = ?1", params![id])
183        .map_err(|e| KernelError::Store(e.to_string()))?;
184    Ok(changed > 0)
185}
186
187/// Delete all edges connected to a node (source or target).
188pub fn remove_edges_for_node(conn: &Connection, node_id: &str) -> Result<()> {
189    conn.execute(
190        "DELETE FROM edges WHERE source = ?1 OR target = ?1",
191        params![node_id],
192    )
193    .map_err(|e| KernelError::Store(e.to_string()))?;
194    Ok(())
195}
196
197/// Read edges where the given node is source or target.
198pub fn edges_for_node(conn: &Connection, node_id: &str) -> Result<Vec<GraphEdge>> {
199    let mut stmt = conn
200        .prepare(
201            "SELECT id, source, target, relation, weight, ts FROM edges WHERE source = ?1 OR target = ?1",
202        )
203        .map_err(|e| KernelError::Store(e.to_string()))?;
204    let edges: Vec<GraphEdge> = stmt
205        .query_map(params![node_id], |row| {
206            Ok(GraphEdge {
207                id: row.get(0)?,
208                source: row.get(1)?,
209                target: row.get(2)?,
210                relation: row.get(3)?,
211                weight: row.get(4)?,
212                ts: row.get(5)?,
213            })
214        })
215        .map_err(|e| KernelError::Store(e.to_string()))?
216        .filter_map(|r| r.ok())
217        .collect();
218    Ok(edges)
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use crate::graph::schema::init_graph_schema;
225    use rusqlite::Connection;
226
227    fn mem_db() -> Connection {
228        let conn = Connection::open_in_memory().unwrap();
229        init_graph_schema(&conn).unwrap();
230        conn
231    }
232
233    fn test_node(id: &str) -> GraphNode {
234        GraphNode {
235            id: id.to_string(),
236            node_type: "concept".to_string(),
237            title: format!("Node {id}"),
238            body: "test body".to_string(),
239            tags: vec!["test".to_string()],
240            projects: vec![],
241            agents: vec![],
242            created: "2026-01-01T00:00:00Z".to_string(),
243            updated: "2026-01-01T00:00:00Z".to_string(),
244            importance: 0.7,
245            access_count: 0,
246            accessed_at: String::new(),
247        }
248    }
249
250    #[test]
251    fn upsert_and_read_node() {
252        let conn = mem_db();
253        let node = test_node("n1");
254        upsert_node(&conn, &node).unwrap();
255        let loaded = read_node(&conn, "n1").unwrap().unwrap();
256        assert_eq!(loaded.id, "n1");
257        assert_eq!(loaded.title, "Node n1");
258        assert_eq!(loaded.tags, vec!["test"]);
259    }
260
261    #[test]
262    fn read_missing_node_returns_none() {
263        let conn = mem_db();
264        assert!(read_node(&conn, "nope").unwrap().is_none());
265    }
266
267    #[test]
268    fn delete_node_returns_true_when_exists() {
269        let conn = mem_db();
270        upsert_node(&conn, &test_node("n1")).unwrap();
271        assert!(delete_node(&conn, "n1").unwrap());
272        assert!(!delete_node(&conn, "n1").unwrap());
273    }
274
275    #[test]
276    fn list_node_ids_returns_all() {
277        let conn = mem_db();
278        upsert_node(&conn, &test_node("a")).unwrap();
279        upsert_node(&conn, &test_node("b")).unwrap();
280        let ids = list_node_ids(&conn).unwrap();
281        assert_eq!(ids.len(), 2);
282    }
283
284    #[test]
285    fn append_and_read_edges() {
286        let conn = mem_db();
287        let edge = GraphEdge {
288            id: "e1".to_string(),
289            source: "a".to_string(),
290            target: "b".to_string(),
291            relation: "related".to_string(),
292            weight: 1.0,
293            ts: "2026-01-01T00:00:00Z".to_string(),
294        };
295        append_edge(&conn, &edge).unwrap();
296        let edges = read_edges(&conn, 10).unwrap();
297        assert_eq!(edges.len(), 1);
298        assert_eq!(edges[0].source, "a");
299    }
300
301    #[test]
302    fn edges_for_node_returns_both_directions() {
303        let conn = mem_db();
304        append_edge(
305            &conn,
306            &GraphEdge {
307                id: "e1".into(),
308                source: "a".into(),
309                target: "b".into(),
310                relation: "related".into(),
311                weight: 1.0,
312                ts: "2026-01-01T00:00:00Z".into(),
313            },
314        )
315        .unwrap();
316        append_edge(
317            &conn,
318            &GraphEdge {
319                id: "e2".into(),
320                source: "c".into(),
321                target: "a".into(),
322                relation: "related".into(),
323                weight: 1.0,
324                ts: "2026-01-01T00:00:00Z".into(),
325            },
326        )
327        .unwrap();
328        let edges = edges_for_node(&conn, "a").unwrap();
329        assert_eq!(edges.len(), 2);
330    }
331
332    #[test]
333    fn test_remove_edges_for_node() {
334        let conn = mem_db();
335        append_edge(
336            &conn,
337            &GraphEdge {
338                id: "e1".into(),
339                source: "a".into(),
340                target: "b".into(),
341                relation: "related".into(),
342                weight: 1.0,
343                ts: "2026-01-01T00:00:00Z".into(),
344            },
345        )
346        .unwrap();
347        remove_edges_for_node(&conn, "a").unwrap();
348        assert!(read_edges(&conn, 10).unwrap().is_empty());
349    }
350}