From ec6494ee590aba5c22d85477c18e10c080f998e9 Mon Sep 17 00:00:00 2001 From: Stuart Axelbrooke Date: Fri, 18 Jul 2025 12:50:15 -0700 Subject: [PATCH] Add edge status colors --- databuild/mermaid_utils.rs | 224 +++++++++++++++++++++++++++++++++++-- 1 file changed, 217 insertions(+), 7 deletions(-) diff --git a/databuild/mermaid_utils.rs b/databuild/mermaid_utils.rs index cd92fa6..80794b5 100644 --- a/databuild/mermaid_utils.rs +++ b/databuild/mermaid_utils.rs @@ -84,6 +84,20 @@ pub fn extract_status_map(events: &[BuildEvent]) -> (HashMap (job_statuses, partition_statuses) } +/// Convert NodeStatus to EdgeStatus for edge coloring +fn map_node_status_to_edge_status(node_status: &NodeStatus) -> EdgeStatus { + match node_status { + NodeStatus::Failed => EdgeStatus::Failed, + NodeStatus::Running => EdgeStatus::Running, + NodeStatus::Completed => EdgeStatus::Completed, + NodeStatus::Available => EdgeStatus::Available, + NodeStatus::Pending => EdgeStatus::Pending, + NodeStatus::Cancelled => EdgeStatus::Failed, // Treat cancelled as failed + NodeStatus::Skipped => EdgeStatus::Pending, // Treat skipped as pending + NodeStatus::Delegated => EdgeStatus::Available, // Treat delegated as available + } +} + /// Encodes ID for safe usage in mermaid graph fn encode_id(id: &str) -> String { id.replace("/", "_").replace("=", "_").replace(":", "_") @@ -197,6 +211,16 @@ enum EdgeType { Dotted, // Weak dependency } +/// Status of an edge for coloring purposes +#[derive(Debug, Clone, PartialEq)] +enum EdgeStatus { + Failed, // Red - critical path issues + Running, // Yellow - actively processing + Completed, // Green - successfully processed + Available, // Light green - data ready + Pending, // Gray - waiting/not started +} + /// Represents an edge between two nodes #[derive(Debug, Clone, PartialEq, Eq, Hash)] struct MermaidEdge { @@ -380,6 +404,16 @@ impl MermaidStyleSheet { } result } + + fn get_edge_color(&self, status: &EdgeStatus) -> &'static str { + match status { + EdgeStatus::Failed => "#ff4444", // Red + EdgeStatus::Running => "#ffaa00", // Orange + EdgeStatus::Completed => "#44aa44", // Green + EdgeStatus::Available => "#88cc88", // Light green + EdgeStatus::Pending => "#888888", // Gray + } + } } /// Builder for constructing Mermaid diagrams @@ -388,6 +422,7 @@ struct MermaidDiagramBuilder { partition_nodes: HashMap, edges: EdgeCollection, output_refs: HashSet, + edge_count: usize, } impl MermaidDiagramBuilder { @@ -397,6 +432,7 @@ impl MermaidDiagramBuilder { partition_nodes: HashMap::new(), edges: EdgeCollection::new(), output_refs: HashSet::new(), + edge_count: 0, } } @@ -423,7 +459,34 @@ impl MermaidDiagramBuilder { self.edges.add(MermaidEdge::new(from_id, to_id, edge_type)); } - fn build(self, statuses: &(HashMap, HashMap), stylesheet: MermaidStyleSheet) -> String { + fn add_edge_with_status(&mut self, from_id: String, to_id: String, edge_type: EdgeType, + edge_status: EdgeStatus, result: &mut String, stylesheet: &MermaidStyleSheet) { + // Create the edge + let edge = MermaidEdge::new(from_id, to_id, edge_type); + + // Check if this edge already exists (for deduplication) + if self.edges.edges.contains(&edge) { + return; // Skip duplicate edge + } + + // Render the edge + result.push_str(&edge.render()); + + // Add edge to collection for deduplication tracking + self.edges.add(edge); + + // Immediately render the linkStyle if status is not pending + if edge_status != EdgeStatus::Pending { + let color = stylesheet.get_edge_color(&edge_status); + result.push_str(&format!(" linkStyle {} stroke:{},stroke-width:2px\n", + self.edge_count, color)); + } + + self.edge_count += 1; + } + + fn build_with_edges(self, statuses: &(HashMap, HashMap), + stylesheet: MermaidStyleSheet, edges_content: String) -> String { let (job_statuses, partition_statuses) = statuses; let mut result = String::from("flowchart TD\n"); @@ -438,8 +501,8 @@ impl MermaidDiagramBuilder { result.push_str(&node.render(status)); } - // Render all edges - result.push_str(&self.edges.render_all()); + // Add the edges content (which includes linkStyle statements) + result.push_str(&edges_content); // Apply styles result.push_str(&stylesheet.render()); @@ -459,11 +522,16 @@ pub fn generate_mermaid_with_status( events: &[BuildEvent], ) -> String { let statuses = extract_status_map(events); + let (job_statuses, partition_statuses) = &statuses; let mut builder = MermaidDiagramBuilder::new(); + let stylesheet = MermaidStyleSheet::default(); // Set output refs for highlighting builder.set_output_refs(&graph.outputs); + // String to accumulate edges with their styles + let mut edges_content = String::new(); + // Process all task nodes for task in &graph.nodes { if let Some(job_node) = MermaidJobNode::from(task) { @@ -480,21 +548,35 @@ pub fn generate_mermaid_with_status( } else { EdgeType::Dotted }; - builder.add_edge(ref_id, job_id.clone(), edge_type); + + // Get partition status for edge coloring + let partition_status = partition_statuses.get(&partition_ref.str) + .unwrap_or(&NodeStatus::Pending); + let edge_status = map_node_status_to_edge_status(partition_status); + + builder.add_edge_with_status(ref_id, job_id.clone(), edge_type, + edge_status, &mut edges_content, &stylesheet); } } // Process outputs for output in &config.outputs { let ref_id = builder.add_partition_node(&output.str); - builder.add_edge(job_id.clone(), ref_id, EdgeType::Solid); + + // Get job status for edge coloring + let job_status = job_statuses.get(&job_id) + .unwrap_or(&NodeStatus::Pending); + let edge_status = map_node_status_to_edge_status(job_status); + + builder.add_edge_with_status(job_id.clone(), ref_id, EdgeType::Solid, + edge_status, &mut edges_content, &stylesheet); } } } } - // Build the diagram with default styling - builder.build(&statuses, MermaidStyleSheet::default()) + // Build the diagram with edges content + builder.build_with_edges(&statuses, stylesheet, edges_content) } #[cfg(test)] @@ -697,4 +779,132 @@ mod tests { assert_eq!(job_statuses.get("same_job___output2"), Some(&NodeStatus::Completed)); assert_eq!(job_statuses.get("same_job"), None, "Should not find job by label alone"); } + + #[test] + fn test_edge_coloring_with_status() { + // Create a simple graph with status + let mut task1 = Task::default(); + task1.job = Some(JobLabel { label: "job1".to_string() }); + task1.config = Some(JobConfig { + inputs: vec![{ + let mut input = DataDep::default(); + input.partition_ref = Some(PartitionRef { str: "input/data".to_string() }); + input.dep_type = 1; // Solid dependency + input + }], + outputs: vec![ + PartitionRef { str: "intermediate/data".to_string() }, + ], + args: vec![], + env: HashMap::new(), + }); + + let mut graph = JobGraph::default(); + graph.nodes = vec![task1]; + graph.outputs = vec![ + PartitionRef { str: "intermediate/data".to_string() }, + ]; + + // Create events to set status + let mut partition_event = BuildEvent::default(); + partition_event.event_type = Some(crate::build_event::EventType::PartitionEvent({ + let mut pe = PartitionEvent::default(); + pe.partition_ref = Some(PartitionRef { str: "input/data".to_string() }); + pe.status = 4; // PARTITION_AVAILABLE + pe + })); + + let mut job_event = BuildEvent::default(); + job_event.event_type = Some(crate::build_event::EventType::JobEvent({ + let mut je = JobEvent::default(); + je.job_label = Some(JobLabel { label: "job1".to_string() }); + je.target_partitions = vec![PartitionRef { str: "intermediate/data".to_string() }]; + je.status = 2; // JOB_RUNNING + je + })); + + let events = vec![partition_event, job_event]; + let mermaid = generate_mermaid_with_status(&graph, &events); + + // Check that linkStyle statements are present + assert!(mermaid.contains("linkStyle"), "Should contain linkStyle statements"); + assert!(mermaid.contains("#88cc88"), "Should contain available edge color (light green)"); + assert!(mermaid.contains("#ffaa00"), "Should contain running edge color (orange)"); + + // Check basic structure is still intact + assert!(mermaid.contains("flowchart TD")); + assert!(mermaid.contains("job1___intermediate_data")); + assert!(mermaid.contains("ref_input_data")); + assert!(mermaid.contains("ref_intermediate_data")); + } + + #[test] + fn test_edge_status_mapping() { + assert_eq!(map_node_status_to_edge_status(&NodeStatus::Failed), EdgeStatus::Failed); + assert_eq!(map_node_status_to_edge_status(&NodeStatus::Running), EdgeStatus::Running); + assert_eq!(map_node_status_to_edge_status(&NodeStatus::Completed), EdgeStatus::Completed); + assert_eq!(map_node_status_to_edge_status(&NodeStatus::Available), EdgeStatus::Available); + assert_eq!(map_node_status_to_edge_status(&NodeStatus::Pending), EdgeStatus::Pending); + assert_eq!(map_node_status_to_edge_status(&NodeStatus::Cancelled), EdgeStatus::Failed); + assert_eq!(map_node_status_to_edge_status(&NodeStatus::Skipped), EdgeStatus::Pending); + assert_eq!(map_node_status_to_edge_status(&NodeStatus::Delegated), EdgeStatus::Available); + } + + #[test] + fn test_edge_deduplication() { + // Create a graph that could potentially have duplicate edges + let mut task1 = Task::default(); + task1.job = Some(JobLabel { label: "job1".to_string() }); + task1.config = Some(JobConfig { + inputs: vec![{ + let mut input = DataDep::default(); + input.partition_ref = Some(PartitionRef { str: "shared_input".to_string() }); + input.dep_type = 1; + input + }], + outputs: vec![ + PartitionRef { str: "output1".to_string() }, + ], + args: vec![], + env: HashMap::new(), + }); + + let mut task2 = Task::default(); + task2.job = Some(JobLabel { label: "job2".to_string() }); + task2.config = Some(JobConfig { + inputs: vec![{ + let mut input = DataDep::default(); + input.partition_ref = Some(PartitionRef { str: "shared_input".to_string() }); + input.dep_type = 1; + input + }], + outputs: vec![ + PartitionRef { str: "output2".to_string() }, + ], + args: vec![], + env: HashMap::new(), + }); + + let mut graph = JobGraph::default(); + graph.nodes = vec![task1, task2]; + graph.outputs = vec![ + PartitionRef { str: "output1".to_string() }, + PartitionRef { str: "output2".to_string() }, + ]; + + let mermaid = generate_mermaid_diagram(&graph); + + // Count how many times the shared edge appears + let shared_edge_count = mermaid.matches("ref_shared_input --> job").count(); + + // Should only appear once per job (2 total), not duplicated + assert_eq!(shared_edge_count, 2, "Should have exactly 2 edges from shared_input (one to each job)"); + + // Verify no duplicate edges in the output + let lines: Vec<&str> = mermaid.lines().collect(); + let edge_lines: Vec<&str> = lines.iter().filter(|line| line.contains("-->") || line.contains("-.->")).cloned().collect(); + let unique_edges: std::collections::HashSet<&str> = edge_lines.iter().cloned().collect(); + + assert_eq!(edge_lines.len(), unique_edges.len(), "Should have no duplicate edges in output"); + } } \ No newline at end of file