use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::sync::RwLock;
use serde::{Serialize, Deserialize};
use crate::error::Result;
use crate::graph::Graph;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Default)]
pub struct SchemaInfo {
pub node_labels: Vec<String>,
pub edge_types: Vec<String>,
pub node_properties: HashMap<String, HashSet<String>>,
pub edge_properties: HashMap<String, HashSet<String>>,
pub node_counts: HashMap<String, usize>,
pub edge_counts: HashMap<String, usize>,
pub total_nodes: usize,
pub total_edges: usize,
}
impl SchemaInfo {
pub fn new() -> Self {
Self::default()
}
pub fn add_node_label(&mut self, label: String) {
if !self.node_labels.contains(&label) {
self.node_labels.push(label.clone());
self.node_properties.insert(label.clone(), HashSet::new());
self.node_counts.insert(label, 0);
}
}
pub fn add_node_property(&mut self, label: &str, property: String) {
self.node_properties
.entry(label.to_string())
.or_default()
.insert(property);
}
pub fn add_edge_type(&mut self, edge_type: String) {
if !self.edge_types.contains(&edge_type) {
self.edge_types.push(edge_type.clone());
self.edge_properties.insert(edge_type.clone(), HashSet::new());
self.edge_counts.insert(edge_type, 0);
}
}
pub fn add_edge_property(&mut self, edge_type: &str, property: String) {
self.edge_properties
.entry(edge_type.to_string())
.or_default()
.insert(property);
}
pub fn increment_node_count(&mut self, label: &str) {
*self.node_counts.entry(label.to_string()).or_insert(0) += 1;
self.total_nodes += 1;
}
pub fn increment_edge_count(&mut self, edge_type: &str) {
*self.edge_counts.entry(edge_type.to_string()).or_insert(0) += 1;
self.total_edges += 1;
}
}
pub struct SchemaManager {
info: Arc<RwLock<SchemaInfo>>,
dirty: Arc<AtomicBool>,
}
impl SchemaManager {
pub fn new() -> Self {
Self {
info: Arc::new(RwLock::new(SchemaInfo::new())),
dirty: Arc::new(AtomicBool::new(true)),
}
}
pub fn is_dirty(&self) -> bool {
self.dirty.load(Ordering::SeqCst)
}
pub fn mark_dirty(&self) {
self.dirty.store(true, Ordering::SeqCst);
}
pub async fn rebuild(&self, graph: &Graph) -> Result<()> {
log::info!("Rebuilding schema from graph...");
let mut info = SchemaInfo::new();
let nodes = graph.get_all_nodes().await?;
log::debug!("Scanning {} nodes for schema", nodes.len());
for node in nodes {
let label = node.label.clone();
info.add_node_label(label.clone());
info.increment_node_count(&label);
for key in node.properties.keys() {
info.add_node_property(&label, key.clone());
}
}
let edges = graph.get_all_edges().await?;
log::debug!("Scanning {} edges for schema", edges.len());
for edge in edges {
let edge_type = edge.edge_type.clone();
info.add_edge_type(edge_type.clone());
info.increment_edge_count(&edge_type);
for key in edge.properties.keys() {
info.add_edge_property(&edge_type, key.clone());
}
}
*self.info.write().await = info;
self.dirty.store(false, Ordering::SeqCst);
log::info!("Schema rebuilt successfully");
Ok(())
}
pub async fn get_info(&self, graph: &Graph) -> Result<SchemaInfo> {
if self.is_dirty() {
self.rebuild(graph).await?;
}
Ok(self.info.read().await.clone())
}
pub async fn get_info_cached(&self) -> SchemaInfo {
self.info.read().await.clone()
}
}
impl Default for SchemaManager {
fn default() -> Self {
Self::new()
}
}