add mermaid lineage chart for wants

This commit is contained in:
Stuart Axelbrooke 2025-12-01 05:11:46 +08:00
parent 8176a8261e
commit 21633c69c3
5 changed files with 474 additions and 7 deletions

View file

@ -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(),

View file

@ -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;

381
databuild/lineage.rs Normal file
View file

@ -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<LineageNode>,
pub edges: Vec<LineageEdge>,
}
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<String> = HashSet::new();
fn add_want_to_graph(
build_state: &BuildState,
graph: &mut LineageGraph,
visited_wants: &mut HashSet<String>,
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"));
}
}

View file

@ -126,12 +126,18 @@ pub struct WantDetailView {
pub derivative_want_ids: Vec<String>,
pub job_runs: Vec<JobRunDetailView>,
pub derivative_wants: Vec<DerivativeWantView>,
/// 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<DerivativeWantView>) -> 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<DerivativeWantView>,
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<WantDetail> 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"
);
}
}
}

View file

@ -54,6 +54,17 @@
</ul>
</div>
{% if !want.lineage_mermaid.is_empty() %}
<div class="detail-section">
<h2>Lineage Graph</h2>
<div class="mermaid-container" style="background: var(--color-bg-secondary); border-radius: 8px; padding: 16px; overflow-x: auto; display: flex; justify-content: center;">
<pre class="mermaid" style="margin: 0;">
{{ want.lineage_mermaid }}
</pre>
</div>
</div>
{% endif %}
{% if !want.upstreams.is_empty() %}
<div class="detail-section">
<h2>Upstream Dependencies ({{ want.upstreams.len() }})</h2>
@ -117,4 +128,20 @@
</div>
{% endif %}
{% if !want.lineage_mermaid.is_empty() %}
<script type="module">
import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@11/dist/mermaid.esm.min.mjs';
mermaid.initialize({
startOnLoad: true,
theme: 'neutral',
flowchart: {
useMaxWidth: true,
htmlLabels: true,
curve: 'basis'
},
securityLevel: 'loose'
});
</script>
{% endif %}
{% call base::footer() %}