dlin_core/render/
dot.rs

1use std::collections::BTreeMap;
2use std::io::{self, Write};
3
4use petgraph::visit::{EdgeRef, IntoEdgeReferences};
5
6use crate::graph::types::*;
7use crate::{Direction, GroupBy};
8
9/// Render the lineage graph as Graphviz DOT format to stdout
10pub fn render_dot(graph: &LineageGraph, group_by: Option<GroupBy>, direction: Direction) {
11    super::handle_stdout_result(render_dot_to_writer(
12        graph,
13        &mut std::io::stdout().lock(),
14        group_by,
15        direction,
16    ));
17}
18
19fn render_dot_to_writer<W: Write>(
20    graph: &LineageGraph,
21    w: &mut W,
22    group_by: Option<GroupBy>,
23    direction: Direction,
24) -> io::Result<()> {
25    writeln!(w, "digraph dbt_lineage {{")?;
26    writeln!(w, "  rankdir={direction};")?;
27    writeln!(
28        w,
29        r#"  node [shape=box, style=filled, fontname="Helvetica"];"#
30    )?;
31    writeln!(w)?;
32
33    match group_by {
34        Some(GroupBy::NodeType) => write_nodes_grouped_by_type(w, graph)?,
35        Some(GroupBy::Directory) => write_nodes_grouped_by_directory(w, graph)?,
36        None => write_nodes_flat(w, graph)?,
37    }
38
39    writeln!(w)?;
40
41    // Collect and sort edges for deterministic output
42    let mut edges: Vec<_> = graph
43        .edge_references()
44        .map(|edge| {
45            let source = &graph[edge.source()];
46            let target = &graph[edge.target()];
47            let ed = edge.weight();
48            (
49                &source.unique_id,
50                &target.unique_id,
51                ed.edge_type,
52                ed.collapsed_through,
53            )
54        })
55        .collect();
56    edges.sort_by(|a, b| {
57        a.0.cmp(b.0)
58            .then(a.1.cmp(b.1))
59            .then(a.2.cmp(&b.2))
60            .then(a.3.cmp(&b.3))
61    });
62
63    for (src_id, tgt_id, edge_type, collapsed) in &edges {
64        let style = match (edge_type, collapsed.is_some()) {
65            (EdgeType::Ref, false) => "",
66            (EdgeType::Ref, true) => ", style=dashed",
67            (EdgeType::Source, false) => ", style=dashed",
68            (EdgeType::Source, true) => r#", style="dashed,bold""#,
69            (EdgeType::Test, false) => ", style=dotted",
70            (EdgeType::Test, true) => r#", style="dotted,dashed""#,
71            (EdgeType::Exposure, false) => ", style=bold",
72            (EdgeType::Exposure, true) => r#", style="bold,dashed""#,
73        };
74        let label = match collapsed {
75            Some(n) => format!("{} (via {})", edge_type.label(), n),
76            None => edge_type.label().to_string(),
77        };
78        writeln!(w, r#"  "{src_id}" -> "{tgt_id}" [label="{label}"{style}];"#,)?;
79    }
80
81    writeln!(w, "}}")?;
82    Ok(())
83}
84
85/// Write nodes without grouping (flat list, sorted by unique_id)
86fn write_nodes_flat<W: Write>(w: &mut W, graph: &LineageGraph) -> io::Result<()> {
87    let mut nodes: Vec<_> = graph.node_indices().map(|idx| &graph[idx]).collect();
88    nodes.sort_by_key(|n| &n.unique_id);
89    for node in &nodes {
90        write_node(w, node, "  ")?;
91    }
92    Ok(())
93}
94
95/// Write nodes grouped by node type using DOT subgraph clusters
96fn write_nodes_grouped_by_type<W: Write>(w: &mut W, graph: &LineageGraph) -> io::Result<()> {
97    // Group nodes by NodeType directly (exhaustive match ensures compile error on new variants)
98    let mut groups: BTreeMap<NodeType, Vec<&NodeData>> = BTreeMap::new();
99    for idx in graph.node_indices() {
100        let node = &graph[idx];
101        groups.entry(node.node_type).or_default().push(node);
102    }
103
104    for (node_type, mut group_nodes) in groups {
105        group_nodes.sort_by_key(|n| &n.unique_id);
106        let type_label = node_type.label();
107        let (bg_color, _) = node_colors(node_type);
108        let title = super::capitalize(type_label);
109        writeln!(w, r#"  subgraph cluster_{type_label} {{"#)?;
110        writeln!(w, r#"    label="{title}";"#)?;
111        writeln!(w, "    style=rounded;")?;
112        writeln!(w, r#"    color="{bg_color}";"#)?;
113        writeln!(w)?;
114        for node in &group_nodes {
115            write_node(w, node, "    ")?;
116        }
117        writeln!(w, "  }}")?;
118    }
119    Ok(())
120}
121
122/// Write nodes grouped by file directory using DOT subgraph clusters
123fn write_nodes_grouped_by_directory<W: Write>(w: &mut W, graph: &LineageGraph) -> io::Result<()> {
124    let mut groups: BTreeMap<String, Vec<&NodeData>> = BTreeMap::new();
125    for idx in graph.node_indices() {
126        let node = &graph[idx];
127        let dir = super::directory_label(node);
128        groups.entry(dir).or_default().push(node);
129    }
130
131    for (dir, mut group_nodes) in groups {
132        group_nodes.sort_by_key(|n| &n.unique_id);
133        let cluster_id = super::sanitize_id(&dir);
134        writeln!(w, r#"  subgraph cluster_{cluster_id} {{"#)?;
135        writeln!(w, r#"    label="{dir}";"#)?;
136        writeln!(w, "    style=rounded;")?;
137        writeln!(w)?;
138        for node in &group_nodes {
139            write_node(w, node, "    ")?;
140        }
141        writeln!(w, "  }}")?;
142    }
143    Ok(())
144}
145
146/// Write a single node definition
147fn write_node<W: Write>(w: &mut W, node: &NodeData, indent: &str) -> io::Result<()> {
148    let (color, fontcolor) = node_colors(node.node_type);
149    let label = node.display_name();
150    writeln!(
151        w,
152        r#"{indent}"{}" [label="{label}", fillcolor="{color}", fontcolor="{fontcolor}"];"#,
153        node.unique_id,
154    )
155}
156
157fn node_colors(node_type: NodeType) -> (&'static str, &'static str) {
158    match node_type {
159        NodeType::Model => ("#4A90D9", "white"),
160        NodeType::Source => ("#27AE60", "white"),
161        NodeType::Seed => ("#F39C12", "white"),
162        NodeType::Snapshot => ("#8E44AD", "white"),
163        NodeType::Test => ("#1ABC9C", "white"),
164        NodeType::Exposure => ("#E74C3C", "white"),
165        NodeType::Phantom => ("#BDC3C7", "black"),
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172    use crate::render::test_helpers::{make_node, make_node_with_path};
173
174    fn render_to_string(graph: &LineageGraph) -> String {
175        let mut buf = Vec::new();
176        render_dot_to_writer(graph, &mut buf, None, Direction::LR).unwrap();
177        String::from_utf8(buf).unwrap()
178    }
179
180    fn render_to_string_grouped(graph: &LineageGraph) -> String {
181        let mut buf = Vec::new();
182        render_dot_to_writer(graph, &mut buf, Some(GroupBy::NodeType), Direction::LR).unwrap();
183        String::from_utf8(buf).unwrap()
184    }
185
186    #[test]
187    fn test_empty_graph() {
188        let graph = LineageGraph::new();
189        let output = render_to_string(&graph);
190        assert!(output.contains("digraph dbt_lineage {"));
191        assert!(output.contains("}"));
192    }
193
194    #[test]
195    fn test_single_node() {
196        let mut graph = LineageGraph::new();
197        graph.add_node(make_node("model.orders", "orders", NodeType::Model));
198        let output = render_to_string(&graph);
199        assert!(output.contains("\"model.orders\""));
200        assert!(output.contains("label=\"orders\""));
201        assert!(output.contains("fillcolor=\"#4A90D9\""));
202    }
203
204    #[test]
205    fn test_edge_styles() {
206        let mut graph = LineageGraph::new();
207        let a = graph.add_node(make_node(
208            "source.raw.orders",
209            "raw.orders",
210            NodeType::Source,
211        ));
212        let b = graph.add_node(make_node("model.stg_orders", "stg_orders", NodeType::Model));
213        graph.add_edge(a, b, EdgeData::direct(EdgeType::Source));
214
215        let output = render_to_string(&graph);
216        assert!(output.contains("style=dashed"));
217        assert!(output.contains("label=\"source\""));
218    }
219
220    #[test]
221    fn test_all_edge_type_labels() {
222        let types = [
223            (EdgeType::Ref, "ref"),
224            (EdgeType::Source, "source"),
225            (EdgeType::Test, "test"),
226            (EdgeType::Exposure, "exposure"),
227        ];
228        for (et, expected) in types {
229            let ed = EdgeData::direct(et);
230            assert_eq!(ed.edge_type.label(), expected);
231        }
232    }
233
234    #[test]
235    fn test_node_colors_all_types() {
236        let types = [
237            NodeType::Model,
238            NodeType::Source,
239            NodeType::Seed,
240            NodeType::Snapshot,
241            NodeType::Test,
242            NodeType::Exposure,
243            NodeType::Phantom,
244        ];
245        for nt in types {
246            let (color, fontcolor) = node_colors(nt);
247            assert!(
248                color.starts_with('#'),
249                "Color for {:?} should start with #",
250                nt
251            );
252            assert!(!fontcolor.is_empty());
253        }
254    }
255
256    #[test]
257    fn test_multiple_edges_different_styles() {
258        let mut graph = LineageGraph::new();
259        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
260        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
261        let c = graph.add_node(make_node("test.t", "t", NodeType::Test));
262        let d = graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
263
264        graph.add_edge(a, b, EdgeData::direct(EdgeType::Ref));
265        graph.add_edge(b, c, EdgeData::direct(EdgeType::Test));
266        graph.add_edge(b, d, EdgeData::direct(EdgeType::Exposure));
267
268        let output = render_to_string(&graph);
269        // Ref edges have no extra style
270        assert!(output.contains("label=\"ref\""));
271        assert!(output.contains("style=dotted"));
272        assert!(output.contains("style=bold"));
273    }
274
275    #[test]
276    fn test_all_node_types_render() {
277        let mut graph = LineageGraph::new();
278        graph.add_node(make_node("model.m", "m", NodeType::Model));
279        graph.add_node(make_node("source.s", "s", NodeType::Source));
280        graph.add_node(make_node("seed.sd", "sd", NodeType::Seed));
281        graph.add_node(make_node("snapshot.sn", "sn", NodeType::Snapshot));
282        graph.add_node(make_node("test.t", "t", NodeType::Test));
283        graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
284        graph.add_node(make_node("phantom.p", "p", NodeType::Phantom));
285
286        let output = render_to_string(&graph);
287        // Verify all node colors appear
288        assert!(output.contains("#4A90D9")); // Model
289        assert!(output.contains("#27AE60")); // Source
290        assert!(output.contains("#F39C12")); // Seed
291        assert!(output.contains("#8E44AD")); // Snapshot
292        assert!(output.contains("#1ABC9C")); // Test
293        assert!(output.contains("#E74C3C")); // Exposure
294        assert!(output.contains("#BDC3C7")); // Phantom
295        assert!(output.contains("fontcolor=\"black\"")); // Phantom font
296    }
297
298    #[test]
299    fn test_all_four_edge_styles_in_render() {
300        let mut graph = LineageGraph::new();
301        let s = graph.add_node(make_node("source.raw.o", "raw.o", NodeType::Source));
302        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
303        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
304        let t = graph.add_node(make_node("test.t", "t", NodeType::Test));
305        let e = graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
306
307        graph.add_edge(s, a, EdgeData::direct(EdgeType::Source));
308        graph.add_edge(a, b, EdgeData::direct(EdgeType::Ref));
309        graph.add_edge(b, t, EdgeData::direct(EdgeType::Test));
310        graph.add_edge(b, e, EdgeData::direct(EdgeType::Exposure));
311
312        let output = render_to_string(&graph);
313        assert!(output.contains("label=\"source\""));
314        assert!(output.contains("label=\"ref\""));
315        assert!(output.contains("label=\"test\""));
316        assert!(output.contains("label=\"exposure\""));
317        assert!(output.contains("style=dashed"));
318        assert!(output.contains("style=dotted"));
319        assert!(output.contains("style=bold"));
320    }
321
322    #[test]
323    fn test_transitive_ref_edge_style() {
324        let mut graph = LineageGraph::new();
325        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
326        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
327        graph.add_edge(a, b, EdgeData::transitive(EdgeType::Ref, 2));
328
329        let output = render_to_string(&graph);
330        assert!(output.contains(r#"label="ref (via 2)""#));
331        assert!(output.contains("style=dashed"));
332    }
333
334    #[test]
335    fn test_transitive_source_edge_preserves_dashed() {
336        let mut graph = LineageGraph::new();
337        let a = graph.add_node(make_node("source.raw.a", "a", NodeType::Source));
338        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
339        graph.add_edge(a, b, EdgeData::transitive(EdgeType::Source, 3));
340
341        let output = render_to_string(&graph);
342        assert!(output.contains(r#"label="source (via 3)""#));
343        assert!(output.contains(r#"style="dashed,bold""#));
344    }
345
346    #[test]
347    fn test_transitive_exposure_edge_preserves_bold() {
348        let mut graph = LineageGraph::new();
349        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
350        let b = graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
351        graph.add_edge(a, b, EdgeData::transitive(EdgeType::Exposure, 1));
352
353        let output = render_to_string(&graph);
354        assert!(output.contains(r#"label="exposure (via 1)""#));
355        assert!(output.contains(r#"style="bold,dashed""#));
356    }
357
358    #[test]
359    fn test_snapshot_lineage() {
360        let graph = crate::render::test_helpers::make_sample_lineage_graph();
361        let output = render_to_string(&graph);
362        insta::assert_snapshot!(output);
363    }
364
365    #[test]
366    fn test_group_by_node_type() {
367        let graph = crate::render::test_helpers::make_sample_lineage_graph();
368        let output = render_to_string_grouped(&graph);
369        insta::assert_snapshot!(output);
370    }
371
372    #[test]
373    fn test_snapshot_all_edge_types() {
374        let mut graph = LineageGraph::new();
375        let s = graph.add_node(make_node("source.raw.o", "raw.o", NodeType::Source));
376        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
377        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
378        let t = graph.add_node(make_node("test.t", "t", NodeType::Test));
379        let e = graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
380
381        graph.add_edge(s, a, EdgeData::direct(EdgeType::Source));
382        graph.add_edge(a, b, EdgeData::direct(EdgeType::Ref));
383        graph.add_edge(b, t, EdgeData::direct(EdgeType::Test));
384        graph.add_edge(b, e, EdgeData::direct(EdgeType::Exposure));
385
386        let output = render_to_string(&graph);
387        insta::assert_snapshot!(output);
388    }
389
390    #[test]
391    fn test_snapshot_transitive_edges() {
392        let mut graph = LineageGraph::new();
393        let a = graph.add_node(make_node("source.raw.a", "a", NodeType::Source));
394        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
395        let c = graph.add_node(make_node("model.c", "c", NodeType::Model));
396        graph.add_edge(a, b, EdgeData::direct(EdgeType::Source));
397        graph.add_edge(a, c, EdgeData::transitive(EdgeType::Source, 3));
398        graph.add_edge(b, c, EdgeData::direct(EdgeType::Ref));
399
400        let output = render_to_string(&graph);
401        insta::assert_snapshot!(output);
402    }
403
404    #[test]
405    fn test_group_by_node_type_cluster_structure() {
406        let mut graph = LineageGraph::new();
407        graph.add_node(make_node("model.a", "a", NodeType::Model));
408        graph.add_node(make_node("source.raw.b", "raw.b", NodeType::Source));
409
410        let output = render_to_string_grouped(&graph);
411        assert!(output.contains("subgraph cluster_model {"));
412        assert!(output.contains("subgraph cluster_source {"));
413        assert!(output.contains("label=\"Model\""));
414        assert!(output.contains("label=\"Source\""));
415        assert!(output.contains("style=rounded;"));
416    }
417
418    fn render_to_string_directory(graph: &LineageGraph) -> String {
419        let mut buf = Vec::new();
420        render_dot_to_writer(graph, &mut buf, Some(GroupBy::Directory), Direction::LR).unwrap();
421        String::from_utf8(buf).unwrap()
422    }
423
424    #[test]
425    fn test_group_by_directory_cluster_structure() {
426        let mut graph = LineageGraph::new();
427        graph.add_node(make_node_with_path(
428            "model.stg_orders",
429            "stg_orders",
430            NodeType::Model,
431            "models/staging/stg_orders.sql",
432        ));
433        graph.add_node(make_node_with_path(
434            "model.orders",
435            "orders",
436            NodeType::Model,
437            "models/marts/orders.sql",
438        ));
439        graph.add_node(make_node(
440            "exposure.dashboard",
441            "dashboard",
442            NodeType::Exposure,
443        ));
444
445        let output = render_to_string_directory(&graph);
446        assert!(output.contains("subgraph cluster_models_staging {"));
447        assert!(output.contains(r#"label="models/staging";"#));
448        assert!(output.contains("subgraph cluster_models_marts {"));
449        assert!(output.contains(r#"label="models/marts";"#));
450        assert!(output.contains("subgraph cluster__other_ {"));
451        assert!(output.contains(r#"label="(other)";"#));
452    }
453
454    #[test]
455    fn test_snapshot_group_by_directory() {
456        let mut graph = LineageGraph::new();
457        let src = graph.add_node(make_node_with_path(
458            "source.raw.orders",
459            "raw.orders",
460            NodeType::Source,
461            "models/staging/schema.yml",
462        ));
463        let stg = graph.add_node(make_node_with_path(
464            "model.stg_orders",
465            "stg_orders",
466            NodeType::Model,
467            "models/staging/stg_orders.sql",
468        ));
469        let mart = graph.add_node(make_node_with_path(
470            "model.orders",
471            "orders",
472            NodeType::Model,
473            "models/marts/orders.sql",
474        ));
475        let exp = graph.add_node(make_node(
476            "exposure.dashboard",
477            "dashboard",
478            NodeType::Exposure,
479        ));
480
481        graph.add_edge(src, stg, EdgeData::direct(EdgeType::Source));
482        graph.add_edge(stg, mart, EdgeData::direct(EdgeType::Ref));
483        graph.add_edge(mart, exp, EdgeData::direct(EdgeType::Exposure));
484
485        let output = render_to_string_directory(&graph);
486        insta::assert_snapshot!(output);
487    }
488
489    #[test]
490    fn test_snapshot_direction_tb() {
491        let graph = crate::render::test_helpers::make_sample_lineage_graph();
492        let mut buf = Vec::new();
493        render_dot_to_writer(&graph, &mut buf, None, Direction::TB).unwrap();
494        let output = String::from_utf8(buf).unwrap();
495        insta::assert_snapshot!(output);
496    }
497
498    #[test]
499    fn test_snapshot_direction_tb_grouped() {
500        let graph = crate::render::test_helpers::make_sample_lineage_graph();
501        let mut buf = Vec::new();
502        render_dot_to_writer(&graph, &mut buf, Some(GroupBy::NodeType), Direction::TB).unwrap();
503        let output = String::from_utf8(buf).unwrap();
504        insta::assert_snapshot!(output);
505    }
506}
dlin_core/render/dot.rs

dlin_core/render/
dot.rs
