use crate::error::MemoryError;
use crate::types::{GraphEdge, GraphEdgeType};
use chrono::Utc;
use rusqlite::{params, Connection};
use std::collections::HashSet;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct StoredGraphEdge {
pub id: String,
pub source: String,
pub target: String,
pub edge_type: String,
#[serde(skip)]
pub edge_type_parsed: Option<GraphEdgeType>,
pub weight: f64,
pub metadata: Option<String>,
pub content_digest: String,
pub recorded_at: String,
pub is_invalidated: bool,
pub invalidated_at: Option<String>,
pub invalidation_reason: Option<String>,
}
#[derive(Debug, Clone)]
pub struct AddGraphEdgeParams {
pub source: String,
pub target: String,
pub edge_type: GraphEdgeType,
pub weight: f64,
pub metadata: Option<serde_json::Value>,
}
pub(crate) fn insert_graph_edge(
conn: &Connection,
params: &AddGraphEdgeParams,
) -> Result<StoredGraphEdge, MemoryError> {
let edge_type_json = serde_json::to_string(¶ms.edge_type)
.map_err(|e| MemoryError::Other(format!("failed to serialize edge_type: {e}")))?;
let metadata_json = match ¶ms.metadata {
Some(v) => Some(serde_json::to_string(v)
.map_err(|e| MemoryError::Other(format!("failed to serialize metadata: {e}")))?),
None => None,
};
let content_digest = compute_edge_digest(
¶ms.source,
¶ms.target,
&edge_type_json,
params.weight,
&metadata_json,
);
let existing: Option<(String, String)> = conn
.query_row(
"SELECT id, recorded_at FROM graph_edges WHERE content_digest = ?1 AND is_invalidated = 0",
params![&content_digest],
|row| Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)),
)
.ok();
if let Some((id, recorded_at)) = existing {
return Ok(StoredGraphEdge {
id,
source: params.source.clone(),
target: params.target.clone(),
edge_type: edge_type_json,
edge_type_parsed: Some(params.edge_type.clone()),
weight: params.weight,
metadata: metadata_json,
content_digest,
recorded_at,
is_invalidated: false,
invalidated_at: None,
invalidation_reason: None,
});
}
let id = uuid::Uuid::new_v4().to_string();
let recorded_at = Utc::now().format("%Y-%m-%d %H:%M:%S%.6f").to_string();
conn.execute(
"INSERT INTO graph_edges (id, source, target, edge_type, weight, metadata, content_digest, recorded_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
params![
&id,
¶ms.source,
¶ms.target,
&edge_type_json,
params.weight,
metadata_json.as_deref(),
&content_digest,
&recorded_at,
],
)
.map_err(|e| MemoryError::Database(e))?;
Ok(StoredGraphEdge {
id,
source: params.source.clone(),
target: params.target.clone(),
edge_type: edge_type_json,
edge_type_parsed: Some(params.edge_type.clone()),
weight: params.weight,
metadata: metadata_json,
content_digest,
recorded_at,
is_invalidated: false,
invalidated_at: None,
invalidation_reason: None,
})
}
pub(crate) fn list_graph_edges_for_node(
conn: &Connection,
node_id: &str,
) -> Result<Vec<StoredGraphEdge>, MemoryError> {
let mut stmt = conn.prepare(
"SELECT id, source, target, edge_type, weight, metadata, content_digest, recorded_at,
is_invalidated, invalidated_at, invalidation_reason
FROM graph_edges
WHERE (source = ?1 OR target = ?1) AND is_invalidated = 0
ORDER BY recorded_at ASC",
)?;
let rows = stmt
.query_map(params![node_id], |row| {
Ok(StoredGraphEdge {
id: row.get(0)?,
source: row.get(1)?,
target: row.get(2)?,
edge_type: row.get(3)?,
edge_type_parsed: None,
weight: row.get(4)?,
metadata: row.get(5)?,
content_digest: row.get(6)?,
recorded_at: row.get(7)?,
is_invalidated: row.get::<_, i64>(8)? != 0,
invalidated_at: row.get(9)?,
invalidation_reason: row.get(10)?,
})
})?
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
}
pub(crate) fn list_all_graph_edges(conn: &Connection) -> Result<Vec<StoredGraphEdge>, MemoryError> {
let mut stmt = conn.prepare(
"SELECT id, source, target, edge_type, weight, metadata, content_digest, recorded_at,
is_invalidated, invalidated_at, invalidation_reason
FROM graph_edges
WHERE is_invalidated = 0
ORDER BY recorded_at ASC",
)?;
let rows = stmt
.query_map([], |row| {
Ok(StoredGraphEdge {
id: row.get(0)?,
source: row.get(1)?,
target: row.get(2)?,
edge_type: row.get(3)?,
edge_type_parsed: None,
weight: row.get(4)?,
metadata: row.get(5)?,
content_digest: row.get(6)?,
recorded_at: row.get(7)?,
is_invalidated: row.get::<_, i64>(8)? != 0,
invalidated_at: row.get(9)?,
invalidation_reason: row.get(10)?,
})
})?
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
}
pub(crate) fn list_graph_edges_for_neighborhood(
conn: &Connection,
seed_ids: &[String],
max_hops: usize,
max_nodes: usize,
) -> Result<Vec<StoredGraphEdge>, MemoryError> {
if seed_ids.is_empty() || max_hops == 0 {
return Ok(Vec::new());
}
let mut visited: HashSet<String> = seed_ids.iter().cloned().collect();
let mut all_edges: Vec<StoredGraphEdge> = Vec::new();
let mut frontier: Vec<String> = seed_ids.iter().cloned().collect();
for _hop in 0..max_hops {
if frontier.is_empty() || visited.len() >= max_nodes {
break;
}
let mut next_frontier: Vec<String> = Vec::new();
for node_id in &frontier {
let edges = list_graph_edges_for_node(conn, node_id)?;
for edge in edges {
let other = if edge.source == *node_id {
&edge.target
} else {
&edge.source
};
if !visited.contains(other) {
visited.insert(other.clone());
next_frontier.push(other.clone());
}
if !all_edges.iter().any(|e: &StoredGraphEdge| e.id == edge.id) {
all_edges.push(edge);
}
}
}
frontier = next_frontier;
if visited.len() >= max_nodes {
break;
}
}
all_edges.sort_by(|a, b| a.recorded_at.cmp(&b.recorded_at));
Ok(all_edges)
}
pub(crate) fn invalidate_graph_edge(
conn: &Connection,
edge_id: &str,
reason: &str,
) -> Result<(), MemoryError> {
let invalidated_at = Utc::now().format("%Y-%m-%d %H:%M:%S%.6f").to_string();
let count = conn.execute(
"UPDATE graph_edges SET is_invalidated = 1, invalidated_at = ?1, invalidation_reason = ?2
WHERE id = ?3 AND is_invalidated = 0",
params![&invalidated_at, reason, edge_id],
)?;
if count == 0 {
return Err(MemoryError::Other(format!(
"graph edge {} not found or already invalidated",
edge_id
)));
}
Ok(())
}
#[allow(dead_code)] pub(crate) fn stored_edges_for_node(
conn: &Connection,
node_id: &str,
) -> Result<Vec<GraphEdge>, MemoryError> {
let rows = list_graph_edges_for_node(conn, node_id)?;
let mut edges = Vec::new();
for row in rows {
let edge_type: GraphEdgeType = serde_json::from_str(&row.edge_type)
.map_err(|e| MemoryError::Other(format!("failed to deserialize edge_type: {e}")))?;
let metadata: Option<serde_json::Value> = match &row.metadata {
Some(s) => Some(serde_json::from_str(s).map_err(|e| {
MemoryError::Other(format!("failed to deserialize metadata: {e}"))
})?),
None => None,
};
edges.push(GraphEdge {
source: row.source,
target: row.target,
edge_type,
weight: row.weight,
metadata,
});
}
Ok(edges)
}
pub(crate) fn stored_outgoing_edges(
conn: &Connection,
node_id: &str,
) -> Result<Vec<GraphEdge>, MemoryError> {
let mut stmt = conn.prepare(
"SELECT id, source, target, edge_type, weight, metadata, content_digest, recorded_at,
is_invalidated, invalidated_at, invalidation_reason
FROM graph_edges
WHERE source = ?1 AND is_invalidated = 0
ORDER BY recorded_at ASC",
)?;
let rows = stmt
.query_map(params![node_id], |row| {
Ok(StoredGraphEdge {
id: row.get(0)?,
source: row.get(1)?,
target: row.get(2)?,
edge_type: row.get(3)?,
edge_type_parsed: None,
weight: row.get(4)?,
metadata: row.get(5)?,
content_digest: row.get(6)?,
recorded_at: row.get(7)?,
is_invalidated: row.get::<_, i64>(8)? != 0,
invalidated_at: row.get(9)?,
invalidation_reason: row.get(10)?,
})
})?
.collect::<Result<Vec<_>, _>>()?;
rows_to_graph_edges(rows)
}
pub(crate) fn stored_incoming_edges(
conn: &Connection,
node_id: &str,
) -> Result<Vec<GraphEdge>, MemoryError> {
let mut stmt = conn.prepare(
"SELECT id, source, target, edge_type, weight, metadata, content_digest, recorded_at,
is_invalidated, invalidated_at, invalidation_reason
FROM graph_edges
WHERE target = ?1 AND is_invalidated = 0
ORDER BY recorded_at ASC",
)?;
let rows = stmt
.query_map(params![node_id], |row| {
Ok(StoredGraphEdge {
id: row.get(0)?,
source: row.get(1)?,
target: row.get(2)?,
edge_type: row.get(3)?,
edge_type_parsed: None,
weight: row.get(4)?,
metadata: row.get(5)?,
content_digest: row.get(6)?,
recorded_at: row.get(7)?,
is_invalidated: row.get::<_, i64>(8)? != 0,
invalidated_at: row.get(9)?,
invalidation_reason: row.get(10)?,
})
})?
.collect::<Result<Vec<_>, _>>()?;
rows_to_graph_edges(rows)
}
fn rows_to_graph_edges(rows: Vec<StoredGraphEdge>) -> Result<Vec<GraphEdge>, MemoryError> {
let mut edges = Vec::new();
for row in rows {
let edge_type: GraphEdgeType = serde_json::from_str(&row.edge_type)
.map_err(|e| MemoryError::Other(format!("failed to deserialize edge_type: {e}")))?;
let metadata: Option<serde_json::Value> = match &row.metadata {
Some(s) => Some(serde_json::from_str(s).map_err(|e| {
MemoryError::Other(format!("failed to deserialize metadata: {e}"))
})?),
None => None,
};
edges.push(GraphEdge {
source: row.source,
target: row.target,
edge_type,
weight: row.weight,
metadata,
});
}
Ok(edges)
}
pub(crate) fn count_graph_edges(conn: &Connection) -> Result<usize, MemoryError> {
let count: i64 = conn
.query_row(
"SELECT COUNT(*) FROM graph_edges WHERE is_invalidated = 0",
[],
|row| row.get(0),
)
.map_err(|e| MemoryError::Database(e))?;
Ok(count as usize)
}
fn compute_edge_digest(
source: &str,
target: &str,
edge_type_json: &str,
weight: f64,
metadata_json: &Option<String>,
) -> String {
let mut builder = stack_ids::DigestBuilder::new();
builder.update(source.as_bytes());
builder.update(target.as_bytes());
builder.update(edge_type_json.as_bytes());
builder.update(&weight.to_le_bytes());
if let Some(meta) = metadata_json {
builder.update(meta.as_bytes());
}
builder.finalize().0
}