use std::collections::{HashSet, VecDeque};
use crate::error::Result;
use crate::graph::{Edge, GraphStore, Node, NodeId};
#[derive(Debug, Clone, Default)]
pub struct NodeQuery {
pub label: Option<String>,
pub property_eq: Vec<(String, serde_json::Value)>,
}
impl NodeQuery {
pub fn label(label: impl Into<String>) -> Self {
Self {
label: Some(label.into()),
..Self::default()
}
}
#[must_use]
pub fn property_eq(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
self.property_eq.push((key.into(), value));
self
}
pub async fn run(&self, store: &dyn GraphStore) -> Result<Vec<Node>> {
let label = self.label.as_deref().unwrap_or("");
let candidates: Vec<Node> = if label.is_empty() {
return Ok(Vec::new()); } else {
store.nodes_by_label(label).await?
};
Ok(candidates
.into_iter()
.filter(|n| {
self.property_eq
.iter()
.all(|(k, v)| n.properties.get(k) == Some(v))
})
.collect())
}
}
#[derive(Debug, Clone, Default)]
pub struct EdgeQuery {
pub label: Option<String>,
pub direction: EdgeDirection,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum EdgeDirection {
#[default]
Out,
In,
Either,
}
impl EdgeQuery {
pub fn outbound(label: impl Into<String>) -> Self {
Self {
label: Some(label.into()),
direction: EdgeDirection::Out,
}
}
#[must_use]
pub fn direction(mut self, d: EdgeDirection) -> Self {
self.direction = d;
self
}
pub async fn run(&self, store: &dyn GraphStore, node: &NodeId) -> Result<Vec<Edge>> {
let label = self.label.as_deref();
match self.direction {
EdgeDirection::Out => store.edges_from(node, label).await,
EdgeDirection::In => store.edges_to(node, label).await,
EdgeDirection::Either => {
let mut out = store.edges_from(node, label).await?;
out.extend(store.edges_to(node, label).await?);
Ok(out)
}
}
}
}
pub async fn traverse(
store: &dyn GraphStore,
start: &NodeId,
edge_label: Option<&str>,
max_depth: usize,
) -> Result<Vec<Node>> {
let mut seen: HashSet<NodeId> = HashSet::new();
let mut queue: VecDeque<(NodeId, usize)> = VecDeque::new();
queue.push_back((start.clone(), 0));
seen.insert(start.clone());
let mut out = Vec::new();
while let Some((current, depth)) = queue.pop_front() {
if let Some(node) = store.get_node(¤t).await? {
out.push(node);
}
if depth >= max_depth {
continue;
}
let edges = store.edges_from(¤t, edge_label).await?;
for e in edges {
if seen.insert(e.to.clone()) {
queue.push_back((e.to.clone(), depth + 1));
}
}
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::graph::{Edge, InMemoryGraph, Node};
use serde_json::json;
async fn fixture() -> InMemoryGraph {
let g = InMemoryGraph::new();
g.upsert_node(
Node::new("pet:1", "pet")
.with_property("name", json!("Rex"))
.with_property("status", json!("available")),
)
.await
.unwrap();
g.upsert_node(
Node::new("pet:2", "pet")
.with_property("name", json!("Buddy"))
.with_property("status", json!("sold")),
)
.await
.unwrap();
g.upsert_node(Node::new("owner:1", "owner")).await.unwrap();
g.upsert_node(Node::new("owner:2", "owner")).await.unwrap();
g.add_edge(Edge::new("owner:1", "pet:1", "owns"))
.await
.unwrap();
g.add_edge(Edge::new("owner:2", "pet:2", "owns"))
.await
.unwrap();
g.add_edge(Edge::new("owner:1", "owner:2", "knows"))
.await
.unwrap();
g
}
#[tokio::test]
async fn node_query_filters_by_property() {
let g = fixture().await;
let q = NodeQuery::label("pet").property_eq("status", json!("available"));
let results = q.run(&g).await.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].id, "pet:1");
}
#[tokio::test]
async fn edge_query_directions() {
let g = fixture().await;
let owns_out = EdgeQuery::outbound("owns")
.run(&g, &"owner:1".into())
.await
.unwrap();
assert_eq!(owns_out.len(), 1);
let either = EdgeQuery {
label: None,
direction: EdgeDirection::Either,
}
.run(&g, &"owner:1".into())
.await
.unwrap();
assert_eq!(either.len(), 2);
}
#[tokio::test]
async fn traverse_walks_depth() {
let g = fixture().await;
let nodes = traverse(&g, &"owner:1".into(), None, 2).await.unwrap();
let ids: Vec<_> = nodes.iter().map(|n| n.id.as_str()).collect();
assert!(ids.contains(&"owner:1"));
assert!(ids.contains(&"pet:1"));
assert!(ids.contains(&"owner:2"));
assert!(ids.contains(&"pet:2"));
}
}