402 lines
14 KiB
Rust
402 lines
14 KiB
Rust
//! 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_fill: String,
|
|
pub status_stroke: 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_fill.is_empty() {
|
|
lines.push(format!(" style {} fill:{},stroke:{}", node.id, node.status_fill, node.status_stroke));
|
|
}
|
|
}
|
|
|
|
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_fill(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
|
|
}
|
|
}
|
|
|
|
pub fn status_to_stroke(status_name: &str) -> &'static str {
|
|
match status_name.to_lowercase().as_str() {
|
|
"wantsuccessful" | "partitionlive" | "jobrunsucceeded" => "#22c55e", // green-500
|
|
"wantbuilding" | "partitionbuilding" | "jobrunrunning" => "#8b5cf6", // violet-500
|
|
"wantidle" | "jobrunqueued" => "#6b7280", // gray-500
|
|
"wantfailed" | "partitionfailed" | "jobrunfailed" => "#ef4444", // red-500
|
|
"wantdepmiss" | "jobrundepmiss" => "#f59e0b", // amber-500
|
|
"wantupstreamfailed" | "partitionupstreamfailed" => "#ef4444", // red-500
|
|
"wantcanceled" | "jobruncanceled" => "#4b5563", // gray-600
|
|
"wantupstreambuilding" | "partitionupstreambuilding" => "#7c3aed", // violet-600
|
|
_ => "#9ca3af", // gray-400 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, 50),
|
|
node_type: LineageNodeType::Want,
|
|
status_fill: status_to_fill(status_name).to_string(),
|
|
status_stroke: status_to_stroke(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, 50),
|
|
node_type: LineageNodeType::Partition,
|
|
status_fill: status_to_fill(&partition_status).to_string(),
|
|
status_stroke: status_to_stroke(&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, 50),
|
|
node_type: LineageNodeType::JobRun,
|
|
status_fill: status_to_fill(job_status_name).to_string(),
|
|
status_stroke: status_to_stroke(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_fill: "#22c55e".to_string(),
|
|
status_stroke: "#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_fill: "#22c55e".to_string(),
|
|
status_stroke: "#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_fill: "#f97316".to_string(),
|
|
status_stroke: "#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"));
|
|
}
|
|
}
|