From 21633c69c3c56a4603cb3bed17f498a8bf19ee37 Mon Sep 17 00:00:00 2001 From: Stuart Axelbrooke Date: Mon, 1 Dec 2025 05:11:46 +0800 Subject: [PATCH] add mermaid lineage chart for wants --- databuild/http_server.rs | 6 +- databuild/lib.rs | 1 + databuild/lineage.rs | 381 ++++++++++++++++++++++ databuild/web/templates.rs | 66 +++- databuild/web/templates/wants/detail.html | 27 ++ 5 files changed, 474 insertions(+), 7 deletions(-) create mode 100644 databuild/lineage.rs diff --git a/databuild/http_server.rs b/databuild/http_server.rs index d920edd..72ba7dd 100644 --- a/databuild/http_server.rs +++ b/databuild/http_server.rs @@ -1,6 +1,7 @@ use crate::build_event_log::BELStorage; use crate::build_state::BuildState; use crate::commands::Command; +use crate::lineage::build_lineage_graph; use crate::web::templates::{ BaseContext, DerivativeWantView, HomePage, JobRunDetailPage, JobRunDetailView, JobRunsListPage, PartitionDetailPage, PartitionDetailView, PartitionsListPage, WantCreatePage, WantDetailPage, @@ -268,9 +269,12 @@ async fn want_detail_page( .map(|w| DerivativeWantView::from(&w)) .collect(); + // Build lineage graph (up to 3 generations) + let lineage_mermaid = build_lineage_graph(&build_state, &want.want_id, 3).to_mermaid(); + let template = WantDetailPage { base: BaseContext::default(), - want: WantDetailView::new(&want, derivative_wants), + want: WantDetailView::new(&want, derivative_wants, lineage_mermaid), }; match template.render() { Ok(html) => Html(html).into_response(), diff --git a/databuild/lib.rs b/databuild/lib.rs index dc76cab..896fe53 100644 --- a/databuild/lib.rs +++ b/databuild/lib.rs @@ -9,6 +9,7 @@ pub mod http_server; mod job; mod job_run; mod job_run_state; +pub mod lineage; mod mock_job_run; pub mod orchestrator; mod partition_state; diff --git a/databuild/lineage.rs b/databuild/lineage.rs new file mode 100644 index 0000000..63bd9c8 --- /dev/null +++ b/databuild/lineage.rs @@ -0,0 +1,381 @@ +//! Lineage graph generation for visualizing want → partition → job run dependencies. +//! +//! This module provides functionality for building Mermaid flowcharts that show +//! the dependency relationships between wants, partitions, and job runs. + +use crate::build_state::BuildState; +use std::collections::HashSet; + +// ============================================================================= +// Lineage Graph Data Structures +// ============================================================================= + +/// Node types in the lineage graph +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LineageNodeType { + Want, + Partition, + JobRun, +} + +/// A node in the lineage graph +#[derive(Debug, Clone)] +pub struct LineageNode { + pub id: String, + pub label: String, + pub node_type: LineageNodeType, + pub status_color: String, + pub url: String, +} + +/// An edge in the lineage graph +#[derive(Debug, Clone)] +pub struct LineageEdge { + pub from: String, + pub to: String, + pub dashed: bool, +} + +/// A directed graph representing want/partition/job run lineage +#[derive(Debug, Clone, Default)] +pub struct LineageGraph { + pub nodes: Vec, + pub edges: Vec, +} + +impl LineageGraph { + pub fn new() -> Self { + Self::default() + } + + pub fn add_node(&mut self, node: LineageNode) { + // Only add if not already present + if !self.nodes.iter().any(|n| n.id == node.id) { + self.nodes.push(node); + } + } + + pub fn add_edge(&mut self, from: String, to: String, dashed: bool) { + let edge = LineageEdge { from, to, dashed }; + // Only add if not already present + if !self + .edges + .iter() + .any(|e| e.from == edge.from && e.to == edge.to) + { + self.edges.push(edge); + } + } + + /// Generate Mermaid flowchart syntax + pub fn to_mermaid(&self) -> String { + let mut lines = vec!["flowchart TD".to_string()]; + + // Add nodes with labels and shapes + for node in &self.nodes { + let shape = match node.node_type { + LineageNodeType::Want => format!("{}[\"🎯 {}\"]", node.id, node.label), + LineageNodeType::Partition => format!("{}[/\"📦 {}\"/]", node.id, node.label), + LineageNodeType::JobRun => format!("{}([\"⚙️ {}\"])", node.id, node.label), + }; + lines.push(format!(" {}", shape)); + } + + lines.push(String::new()); + + // Add edges + for edge in &self.edges { + let arrow = if edge.dashed { "-.->" } else { "-->" }; + lines.push(format!(" {} {} {}", edge.from, arrow, edge.to)); + } + + lines.push(String::new()); + + // Add styles for status colors + for node in &self.nodes { + if !node.status_color.is_empty() { + lines.push(format!(" style {} fill:{}", node.id, node.status_color)); + } + } + + lines.push(String::new()); + + // Add click handlers for navigation + for node in &self.nodes { + if !node.url.is_empty() { + lines.push(format!(" click {} \"{}\"", node.id, node.url)); + } + } + + lines.join("\n") + } +} + +// ============================================================================= +// Status Color Mapping +// ============================================================================= + +/// Map status names to colors for the lineage graph. +/// Status names are matched case-insensitively against proto-generated enum names. +pub fn status_to_color(status_name: &str) -> &'static str { + match status_name.to_lowercase().as_str() { + "wantsuccessful" | "partitionlive" | "jobrunsucceeded" => "#dcfce7", // green + "wantbuilding" | "partitionbuilding" | "jobrunrunning" => "#ede9fe", // blue + "wantidle" | "jobrunqueued" => "#9ca3af", // gray + "wantfailed" | "partitionfailed" | "jobrunfailed" => "#fee2e2;", // red + "wantdepmiss" | "jobrundepmiss" => "#fef3e2", // orange + "wantupstreamfailed" | "partitionupstreamfailed" => "#fee2e2", // red + "wantupstreambuilding" | "partitionupstreambuilding" => "#a855f7", // purple + "wantcanceled" | "jobruncanceled" => "#6b7280", // dark gray + _ => "#e5e7eb", // light gray default + } +} + +// ============================================================================= +// Graph Builder +// ============================================================================= + +/// Build a lineage graph starting from a want, traversing up to `max_depth` generations. +/// The graph shows: Want -> Partitions -> JobRuns -> DerivativeWants -> ... +pub fn build_lineage_graph( + build_state: &BuildState, + want_id: &str, + max_depth: usize, +) -> LineageGraph { + let mut graph = LineageGraph::new(); + let mut visited_wants: HashSet = HashSet::new(); + + fn add_want_to_graph( + build_state: &BuildState, + graph: &mut LineageGraph, + visited_wants: &mut HashSet, + want_id: &str, + depth: usize, + max_depth: usize, + ) { + if depth > max_depth || visited_wants.contains(want_id) { + return; + } + visited_wants.insert(want_id.to_string()); + + let Some(want) = build_state.get_want(want_id) else { + return; + }; + + // Create node ID (sanitize for mermaid) + let want_node_id = format!("W{}", sanitize_node_id(want_id)); + let status_name = want.status.as_ref().map(|s| s.name.as_str()).unwrap_or(""); + + // Add want node + graph.add_node(LineageNode { + id: want_node_id.clone(), + label: truncate_label(want_id, 30), + node_type: LineageNodeType::Want, + status_color: status_to_color(status_name).to_string(), + url: format!("/wants/{}", want_id), + }); + + // Add partition nodes and edges from want to partitions + for partition in &want.partitions { + let partition_ref = &partition.r#ref; + let partition_node_id = format!("P{}", sanitize_node_id(partition_ref)); + + // Get partition status if available + let partition_status = build_state + .get_partition(partition_ref) + .and_then(|p| p.status) + .map(|s| s.name) + .unwrap_or_default(); + + graph.add_node(LineageNode { + id: partition_node_id.clone(), + label: truncate_label(partition_ref, 25), + node_type: LineageNodeType::Partition, + status_color: status_to_color(&partition_status).to_string(), + url: format!("/partitions/{}", urlencoding::encode(partition_ref)), + }); + + // Want -> Partition edge + graph.add_edge(want_node_id.clone(), partition_node_id.clone(), false); + } + + // Add job run nodes and edges from partitions to job runs + for job_run in &want.job_runs { + let job_run_node_id = format!("JR{}", sanitize_node_id(&job_run.id)); + let job_status_name = job_run + .status + .as_ref() + .map(|s| s.name.as_str()) + .unwrap_or(""); + + // Use job_label if available, otherwise job_run_id + let job_label = if job_run.job_label.is_empty() { + &job_run.id + } else { + &job_run.job_label + }; + + graph.add_node(LineageNode { + id: job_run_node_id.clone(), + label: truncate_label(job_label, 25), + node_type: LineageNodeType::JobRun, + status_color: status_to_color(job_status_name).to_string(), + url: format!("/job_runs/{}", job_run.id), + }); + + // Connect partitions being built to this job run (dashed = building relationship) + for partition in &job_run.building_partitions { + let partition_node_id = format!("P{}", sanitize_node_id(&partition.r#ref)); + graph.add_edge(partition_node_id, job_run_node_id.clone(), true); + } + + // Recurse into derivative wants + for derivative_want_id in &job_run.derivative_want_ids { + add_want_to_graph( + build_state, + graph, + visited_wants, + derivative_want_id, + depth + 1, + max_depth, + ); + + // Add edge from job run to derivative want + let derivative_want_node_id = format!("W{}", sanitize_node_id(derivative_want_id)); + graph.add_edge(job_run_node_id.clone(), derivative_want_node_id, false); + } + } + } + + add_want_to_graph( + build_state, + &mut graph, + &mut visited_wants, + want_id, + 0, + max_depth, + ); + + graph +} + +/// Sanitize a string to be a valid mermaid node ID (alphanumeric + underscore only) +pub fn sanitize_node_id(s: &str) -> String { + s.chars() + .map(|c| if c.is_alphanumeric() { c } else { '_' }) + .collect() +} + +/// Truncate a label to fit in the graph, adding ellipsis if needed +pub fn truncate_label(s: &str, max_len: usize) -> String { + if s.len() <= max_len { + s.to_string() + } else { + format!("{}...", &s[..max_len - 3]) + } +} + +// ============================================================================= +// Tests +// ============================================================================= + +#[cfg(test)] +mod tests { + use super::*; + use crate::build_state::BuildState; + use crate::util::test_scenarios::multihop_scenario; + + /// Helper to replay events into a fresh BuildState + fn build_state_from_events(events: &[crate::data_build_event::Event]) -> BuildState { + let mut state = BuildState::default(); + for event in events { + state.handle_event(event); + } + state + } + + #[test] + fn test_lineage_graph_mermaid_generation() { + let mut graph = LineageGraph::new(); + + graph.add_node(LineageNode { + id: "W_beta_want".to_string(), + label: "beta-want".to_string(), + node_type: LineageNodeType::Want, + status_color: "#22c55e".to_string(), + url: "/wants/beta-want".to_string(), + }); + + graph.add_node(LineageNode { + id: "P_data_beta".to_string(), + label: "data/beta".to_string(), + node_type: LineageNodeType::Partition, + status_color: "#22c55e".to_string(), + url: "/partitions/data%2Fbeta".to_string(), + }); + + graph.add_node(LineageNode { + id: "JR_beta_job".to_string(), + label: "//job_beta".to_string(), + node_type: LineageNodeType::JobRun, + status_color: "#f97316".to_string(), + url: "/job_runs/beta-job".to_string(), + }); + + graph.add_edge("W_beta_want".to_string(), "P_data_beta".to_string(), false); + graph.add_edge("P_data_beta".to_string(), "JR_beta_job".to_string(), true); + + let mermaid = graph.to_mermaid(); + + assert!( + mermaid.contains("flowchart TD"), + "Should have flowchart header" + ); + assert!(mermaid.contains("W_beta_want"), "Should have want node"); + assert!( + mermaid.contains("P_data_beta"), + "Should have partition node" + ); + assert!(mermaid.contains("JR_beta_job"), "Should have job run node"); + assert!(mermaid.contains("-->"), "Should have solid edge"); + assert!(mermaid.contains("-.->"), "Should have dashed edge"); + assert!( + mermaid.contains("style W_beta_want fill:#22c55e"), + "Should have status color" + ); + assert!( + mermaid.contains("click W_beta_want"), + "Should have click handler" + ); + } + + #[test] + fn test_sanitize_node_id() { + assert_eq!(sanitize_node_id("data/beta"), "data_beta"); + assert_eq!(sanitize_node_id("job-run-123"), "job_run_123"); + assert_eq!(sanitize_node_id("simple"), "simple"); + } + + #[test] + fn test_truncate_label() { + assert_eq!(truncate_label("short", 10), "short"); + assert_eq!(truncate_label("this is very long", 10), "this is..."); + } + + #[test] + fn test_build_lineage_graph_from_multihop_scenario() { + let (events, ids) = multihop_scenario(); + let state = build_state_from_events(&events); + + let graph = build_lineage_graph(&state, &ids.beta_want_id, 3); + + // Should have nodes for beta want, alpha want, partitions, and job runs + assert!(!graph.nodes.is_empty(), "Graph should have nodes"); + assert!(!graph.edges.is_empty(), "Graph should have edges"); + + // Check that the mermaid output is valid + let mermaid = graph.to_mermaid(); + assert!(mermaid.contains("flowchart TD")); + assert!(mermaid.contains("beta_want") || mermaid.contains("beta")); + } +} diff --git a/databuild/web/templates.rs b/databuild/web/templates.rs index f3b33e8..7c38cca 100644 --- a/databuild/web/templates.rs +++ b/databuild/web/templates.rs @@ -126,12 +126,18 @@ pub struct WantDetailView { pub derivative_want_ids: Vec, pub job_runs: Vec, pub derivative_wants: Vec, + /// Mermaid flowchart showing want -> partition -> job run -> derivative want lineage + pub lineage_mermaid: String, } impl WantDetailView { - /// Create a WantDetailView with derivative wants populated. - /// Use this for the detail page where derivative wants need to be shown. - pub fn new(w: &WantDetail, derivative_wants: Vec) -> Self { + /// Create a WantDetailView with derivative wants and lineage graph populated. + /// Use this for the detail page where full lineage information is shown. + pub fn new( + w: &WantDetail, + derivative_wants: Vec, + lineage_mermaid: String, + ) -> Self { use crate::want_detail::Lifetime; let (lifetime, data_timestamp, ttl_seconds, sla_seconds, is_ephemeral, source_job_run_id) = @@ -179,14 +185,15 @@ impl WantDetailView { derivative_want_ids: w.derivative_want_ids.clone(), job_runs: w.job_runs.iter().map(JobRunDetailView::from).collect(), derivative_wants, + lineage_mermaid, } } } -/// For list pages where derivative wants aren't needed +/// For list pages where derivative wants and lineage aren't needed impl From for WantDetailView { fn from(w: WantDetail) -> Self { - Self::new(&w, vec![]) + Self::new(&w, vec![], String::new()) } } @@ -590,7 +597,7 @@ mod tests { let template = WantDetailPage { base: BaseContext::default(), - want: WantDetailView::new(&want_detail, derivative_wants), + want: WantDetailView::new(&want_detail, derivative_wants, String::new()), }; let html = template.render().expect("template should render"); @@ -601,5 +608,52 @@ mod tests { html ); } + + /// Given: A want detail view with lineage mermaid content + /// When: We render the template + /// Then: The HTML should contain the mermaid diagram + #[test] + fn test_want_detail_page_shows_lineage_mermaid() { + let (events, ids) = multihop_scenario(); + let state = build_state_from_events(&events); + + let want_detail = state + .get_want(&ids.beta_want_id) + .expect("beta want should exist"); + + let derivative_wants: Vec<_> = want_detail + .derivative_want_ids + .iter() + .filter_map(|id| state.get_want(id)) + .map(|w| DerivativeWantView::from(&w)) + .collect(); + + // Create a simple lineage mermaid string + let lineage_mermaid = "flowchart TD\n W[Want]\n P[Partition]".to_string(); + + let template = WantDetailPage { + base: BaseContext::default(), + want: WantDetailView::new(&want_detail, derivative_wants, lineage_mermaid), + }; + let html = template.render().expect("template should render"); + + // Verify lineage graph section is present + assert!( + html.contains("Lineage Graph"), + "Should have Lineage Graph section header" + ); + assert!( + html.contains("class=\"mermaid\""), + "Should have mermaid pre element" + ); + assert!( + html.contains("flowchart TD"), + "Should contain the mermaid content" + ); + assert!( + html.contains("mermaid.esm.min.mjs"), + "Should include mermaid script" + ); + } } } diff --git a/databuild/web/templates/wants/detail.html b/databuild/web/templates/wants/detail.html index 1fa655a..89b0087 100644 --- a/databuild/web/templates/wants/detail.html +++ b/databuild/web/templates/wants/detail.html @@ -54,6 +54,17 @@ +{% if !want.lineage_mermaid.is_empty() %} +
+

Lineage Graph

+
+
+{{ want.lineage_mermaid }}
+        
+
+
+{% endif %} + {% if !want.upstreams.is_empty() %}

Upstream Dependencies ({{ want.upstreams.len() }})

@@ -117,4 +128,20 @@
{% endif %} +{% if !want.lineage_mermaid.is_empty() %} + +{% endif %} + {% call base::footer() %}