1use crate::error::MemoryError;
27use crate::types::{GraphEdge, GraphEdgeType};
28use chrono::Utc;
29use rusqlite::{params, Connection};
30use std::collections::HashSet;
31
32#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
34pub struct StoredGraphEdge {
35 pub id: String,
37 pub source: String,
39 pub target: String,
41 pub edge_type: String,
43 #[serde(skip)]
45 pub edge_type_parsed: Option<GraphEdgeType>,
46 pub weight: f64,
48 pub metadata: Option<String>,
50 pub content_digest: String,
52 pub recorded_at: String,
54 pub is_invalidated: bool,
56 pub invalidated_at: Option<String>,
58 pub invalidation_reason: Option<String>,
60}
61
62#[derive(Debug, Clone)]
64pub struct AddGraphEdgeParams {
65 pub source: String,
67 pub target: String,
69 pub edge_type: GraphEdgeType,
71 pub weight: f64,
73 pub metadata: Option<serde_json::Value>,
75}
76
77pub(crate) fn insert_graph_edge(
79 conn: &Connection,
80 params: &AddGraphEdgeParams,
81) -> Result<StoredGraphEdge, MemoryError> {
82 let edge_type_json = serde_json::to_string(¶ms.edge_type)
83 .map_err(|e| MemoryError::Other(format!("failed to serialize edge_type: {e}")))?;
84
85 let metadata_json = match ¶ms.metadata {
86 Some(v) => Some(serde_json::to_string(v)
87 .map_err(|e| MemoryError::Other(format!("failed to serialize metadata: {e}")))?),
88 None => None,
89 };
90
91 let content_digest = compute_edge_digest(
92 ¶ms.source,
93 ¶ms.target,
94 &edge_type_json,
95 params.weight,
96 &metadata_json,
97 );
98
99 let existing: Option<(String, String)> = conn
101 .query_row(
102 "SELECT id, recorded_at FROM graph_edges WHERE content_digest = ?1 AND is_invalidated = 0",
103 params![&content_digest],
104 |row| Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)),
105 )
106 .ok();
107
108 if let Some((id, recorded_at)) = existing {
109 return Ok(StoredGraphEdge {
110 id,
111 source: params.source.clone(),
112 target: params.target.clone(),
113 edge_type: edge_type_json,
114 edge_type_parsed: Some(params.edge_type.clone()),
115 weight: params.weight,
116 metadata: metadata_json,
117 content_digest,
118 recorded_at,
119 is_invalidated: false,
120 invalidated_at: None,
121 invalidation_reason: None,
122 });
123 }
124
125 let id = uuid::Uuid::new_v4().to_string();
126 let recorded_at = Utc::now().format("%Y-%m-%d %H:%M:%S%.6f").to_string();
127
128 conn.execute(
129 "INSERT INTO graph_edges (id, source, target, edge_type, weight, metadata, content_digest, recorded_at)
130 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
131 params![
132 &id,
133 ¶ms.source,
134 ¶ms.target,
135 &edge_type_json,
136 params.weight,
137 metadata_json.as_deref(),
138 &content_digest,
139 &recorded_at,
140 ],
141 )
142 .map_err(|e| MemoryError::Database(e))?;
143
144 Ok(StoredGraphEdge {
145 id,
146 source: params.source.clone(),
147 target: params.target.clone(),
148 edge_type: edge_type_json,
149 edge_type_parsed: Some(params.edge_type.clone()),
150 weight: params.weight,
151 metadata: metadata_json,
152 content_digest,
153 recorded_at,
154 is_invalidated: false,
155 invalidated_at: None,
156 invalidation_reason: None,
157 })
158}
159
160pub(crate) fn list_graph_edges_for_node(
163 conn: &Connection,
164 node_id: &str,
165) -> Result<Vec<StoredGraphEdge>, MemoryError> {
166 let mut stmt = conn.prepare(
167 "SELECT id, source, target, edge_type, weight, metadata, content_digest, recorded_at,
168 is_invalidated, invalidated_at, invalidation_reason
169 FROM graph_edges
170 WHERE (source = ?1 OR target = ?1) AND is_invalidated = 0
171 ORDER BY recorded_at ASC",
172 )?;
173 let rows = stmt
174 .query_map(params![node_id], |row| {
175 Ok(StoredGraphEdge {
176 id: row.get(0)?,
177 source: row.get(1)?,
178 target: row.get(2)?,
179 edge_type: row.get(3)?,
180 edge_type_parsed: None,
181 weight: row.get(4)?,
182 metadata: row.get(5)?,
183 content_digest: row.get(6)?,
184 recorded_at: row.get(7)?,
185 is_invalidated: row.get::<_, i64>(8)? != 0,
186 invalidated_at: row.get(9)?,
187 invalidation_reason: row.get(10)?,
188 })
189 })?
190 .collect::<Result<Vec<_>, _>>()?;
191 Ok(rows)
192}
193
194pub(crate) fn list_all_graph_edges(conn: &Connection) -> Result<Vec<StoredGraphEdge>, MemoryError> {
196 let mut stmt = conn.prepare(
197 "SELECT id, source, target, edge_type, weight, metadata, content_digest, recorded_at,
198 is_invalidated, invalidated_at, invalidation_reason
199 FROM graph_edges
200 WHERE is_invalidated = 0
201 ORDER BY recorded_at ASC",
202 )?;
203 let rows = stmt
204 .query_map([], |row| {
205 Ok(StoredGraphEdge {
206 id: row.get(0)?,
207 source: row.get(1)?,
208 target: row.get(2)?,
209 edge_type: row.get(3)?,
210 edge_type_parsed: None,
211 weight: row.get(4)?,
212 metadata: row.get(5)?,
213 content_digest: row.get(6)?,
214 recorded_at: row.get(7)?,
215 is_invalidated: row.get::<_, i64>(8)? != 0,
216 invalidated_at: row.get(9)?,
217 invalidation_reason: row.get(10)?,
218 })
219 })?
220 .collect::<Result<Vec<_>, _>>()?;
221 Ok(rows)
222}
223
224pub(crate) fn list_graph_edges_for_neighborhood(
234 conn: &Connection,
235 seed_ids: &[String],
236 max_hops: usize,
237 max_nodes: usize,
238) -> Result<Vec<StoredGraphEdge>, MemoryError> {
239 if seed_ids.is_empty() || max_hops == 0 {
240 return Ok(Vec::new());
241 }
242
243 let mut visited: HashSet<String> = seed_ids.iter().cloned().collect();
244 let mut all_edges: Vec<StoredGraphEdge> = Vec::new();
245 let mut frontier: Vec<String> = seed_ids.iter().cloned().collect();
246
247 for _hop in 0..max_hops {
248 if frontier.is_empty() || visited.len() >= max_nodes {
249 break;
250 }
251
252 let mut next_frontier: Vec<String> = Vec::new();
253
254 for node_id in &frontier {
255 let edges = list_graph_edges_for_node(conn, node_id)?;
256 for edge in edges {
257 let other = if edge.source == *node_id {
259 &edge.target
260 } else {
261 &edge.source
262 };
263
264 if !visited.contains(other) {
265 visited.insert(other.clone());
266 next_frontier.push(other.clone());
267 }
268
269 if !all_edges.iter().any(|e: &StoredGraphEdge| e.id == edge.id) {
271 all_edges.push(edge);
272 }
273 }
274 }
275
276 frontier = next_frontier;
277 if visited.len() >= max_nodes {
278 break;
279 }
280 }
281
282 all_edges.sort_by(|a, b| a.recorded_at.cmp(&b.recorded_at));
284 Ok(all_edges)
285}
286
287pub(crate) fn invalidate_graph_edge(
289 conn: &Connection,
290 edge_id: &str,
291 reason: &str,
292) -> Result<(), MemoryError> {
293 let invalidated_at = Utc::now().format("%Y-%m-%d %H:%M:%S%.6f").to_string();
294 let count = conn.execute(
295 "UPDATE graph_edges SET is_invalidated = 1, invalidated_at = ?1, invalidation_reason = ?2
296 WHERE id = ?3 AND is_invalidated = 0",
297 params![&invalidated_at, reason, edge_id],
298 )?;
299 if count == 0 {
300 return Err(MemoryError::Other(format!(
301 "graph edge {} not found or already invalidated",
302 edge_id
303 )));
304 }
305 Ok(())
306}
307
308#[allow(dead_code)] pub(crate) fn stored_edges_for_node(
312 conn: &Connection,
313 node_id: &str,
314) -> Result<Vec<GraphEdge>, MemoryError> {
315 let rows = list_graph_edges_for_node(conn, node_id)?;
316 let mut edges = Vec::new();
317 for row in rows {
318 let edge_type: GraphEdgeType = serde_json::from_str(&row.edge_type)
319 .map_err(|e| MemoryError::Other(format!("failed to deserialize edge_type: {e}")))?;
320 let metadata: Option<serde_json::Value> = match &row.metadata {
321 Some(s) => Some(serde_json::from_str(s).map_err(|e| {
322 MemoryError::Other(format!("failed to deserialize metadata: {e}"))
323 })?),
324 None => None,
325 };
326 edges.push(GraphEdge {
327 source: row.source,
328 target: row.target,
329 edge_type,
330 weight: row.weight,
331 metadata,
332 });
333 }
334 Ok(edges)
335}
336
337pub(crate) fn stored_outgoing_edges(
340 conn: &Connection,
341 node_id: &str,
342) -> Result<Vec<GraphEdge>, MemoryError> {
343 let mut stmt = conn.prepare(
344 "SELECT id, source, target, edge_type, weight, metadata, content_digest, recorded_at,
345 is_invalidated, invalidated_at, invalidation_reason
346 FROM graph_edges
347 WHERE source = ?1 AND is_invalidated = 0
348 ORDER BY recorded_at ASC",
349 )?;
350 let rows = stmt
351 .query_map(params![node_id], |row| {
352 Ok(StoredGraphEdge {
353 id: row.get(0)?,
354 source: row.get(1)?,
355 target: row.get(2)?,
356 edge_type: row.get(3)?,
357 edge_type_parsed: None,
358 weight: row.get(4)?,
359 metadata: row.get(5)?,
360 content_digest: row.get(6)?,
361 recorded_at: row.get(7)?,
362 is_invalidated: row.get::<_, i64>(8)? != 0,
363 invalidated_at: row.get(9)?,
364 invalidation_reason: row.get(10)?,
365 })
366 })?
367 .collect::<Result<Vec<_>, _>>()?;
368 rows_to_graph_edges(rows)
369}
370
371pub(crate) fn stored_incoming_edges(
374 conn: &Connection,
375 node_id: &str,
376) -> Result<Vec<GraphEdge>, MemoryError> {
377 let mut stmt = conn.prepare(
378 "SELECT id, source, target, edge_type, weight, metadata, content_digest, recorded_at,
379 is_invalidated, invalidated_at, invalidation_reason
380 FROM graph_edges
381 WHERE target = ?1 AND is_invalidated = 0
382 ORDER BY recorded_at ASC",
383 )?;
384 let rows = stmt
385 .query_map(params![node_id], |row| {
386 Ok(StoredGraphEdge {
387 id: row.get(0)?,
388 source: row.get(1)?,
389 target: row.get(2)?,
390 edge_type: row.get(3)?,
391 edge_type_parsed: None,
392 weight: row.get(4)?,
393 metadata: row.get(5)?,
394 content_digest: row.get(6)?,
395 recorded_at: row.get(7)?,
396 is_invalidated: row.get::<_, i64>(8)? != 0,
397 invalidated_at: row.get(9)?,
398 invalidation_reason: row.get(10)?,
399 })
400 })?
401 .collect::<Result<Vec<_>, _>>()?;
402 rows_to_graph_edges(rows)
403}
404
405fn rows_to_graph_edges(rows: Vec<StoredGraphEdge>) -> Result<Vec<GraphEdge>, MemoryError> {
407 let mut edges = Vec::new();
408 for row in rows {
409 let edge_type: GraphEdgeType = serde_json::from_str(&row.edge_type)
410 .map_err(|e| MemoryError::Other(format!("failed to deserialize edge_type: {e}")))?;
411 let metadata: Option<serde_json::Value> = match &row.metadata {
412 Some(s) => Some(serde_json::from_str(s).map_err(|e| {
413 MemoryError::Other(format!("failed to deserialize metadata: {e}"))
414 })?),
415 None => None,
416 };
417 edges.push(GraphEdge {
418 source: row.source,
419 target: row.target,
420 edge_type,
421 weight: row.weight,
422 metadata,
423 });
424 }
425 Ok(edges)
426}
427pub(crate) fn count_graph_edges(conn: &Connection) -> Result<usize, MemoryError> {
429 let count: i64 = conn
430 .query_row(
431 "SELECT COUNT(*) FROM graph_edges WHERE is_invalidated = 0",
432 [],
433 |row| row.get(0),
434 )
435 .map_err(|e| MemoryError::Database(e))?;
436 Ok(count as usize)
437}
438
439fn compute_edge_digest(
440 source: &str,
441 target: &str,
442 edge_type_json: &str,
443 weight: f64,
444 metadata_json: &Option<String>,
445) -> String {
446 let mut builder = stack_ids::DigestBuilder::new();
447 builder.update(source.as_bytes());
448 builder.update(target.as_bytes());
449 builder.update(edge_type_json.as_bytes());
450 builder.update(&weight.to_le_bytes());
451 if let Some(meta) = metadata_json {
452 builder.update(meta.as_bytes());
453 }
454 builder.finalize().0
455}