Skip to main content

linkrs_query/executor/expression/evaluation_context/
graph_storage.rs

1use crate::storage::StorageReader;
2use linkrs_core::types::VertexId;
3use linkrs_core::{Edge, Vertex};
4use parking_lot::RwLock;
5use std::sync::Arc;
6
7pub struct GraphStorageRef {
8    pub storage: Arc<RwLock<dyn StorageReader>>,
9    pub space: String,
10}
11
12impl GraphStorageRef {
13    pub fn new(storage: Arc<RwLock<dyn StorageReader>>, space: String) -> Self {
14        Self { storage, space }
15    }
16
17    pub fn get_neighbors(&self, node_id: &VertexId) -> Result<Vec<(VertexId, Edge)>, String> {
18        use linkrs_core::types::EdgeDirection;
19        let reader = self.storage.read();
20        let edges = reader
21            .get_node_edges(&self.space, node_id, EdgeDirection::Both, &[])
22            .map_err(|e| format!("Storage error: {}", e))?;
23        let neighbors: Vec<(VertexId, Edge)> = edges
24            .into_iter()
25            .map(|e| {
26                let neighbor_id = if e.src == *node_id { e.dst } else { e.src };
27                (neighbor_id, e)
28            })
29            .collect();
30        Ok(neighbors)
31    }
32
33    pub fn get_vertex(&self, tag: &str, id: &VertexId) -> Result<Option<Vertex>, String> {
34        let reader = self.storage.read();
35        reader
36            .get_vertex(&self.space, tag, id)
37            .map_err(|e| format!("Storage error: {}", e))
38    }
39}