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        let src = super::dot_escape(src_id);
79        let tgt = super::dot_escape(tgt_id);
80        writeln!(w, r#"  "{src}" -> "{tgt}" [label="{label}"{style}];"#,)?;
81    }
82
83    writeln!(w, "}}")?;
84    Ok(())
85}
86
87/// Write nodes without grouping (flat list, sorted by unique_id)
88fn write_nodes_flat<W: Write>(w: &mut W, graph: &LineageGraph) -> io::Result<()> {
89    let mut nodes: Vec<_> = graph.node_indices().map(|idx| &graph[idx]).collect();
90    nodes.sort_by_key(|n| &n.unique_id);
91    for node in &nodes {
92        write_node(w, node, "  ")?;
93    }
94    Ok(())
95}
96
97/// Write nodes grouped by node type using DOT subgraph clusters
98fn write_nodes_grouped_by_type<W: Write>(w: &mut W, graph: &LineageGraph) -> io::Result<()> {
99    // Group nodes by NodeType directly (exhaustive match ensures compile error on new variants)
100    let mut groups: BTreeMap<NodeType, Vec<&NodeData>> = BTreeMap::new();
101    for idx in graph.node_indices() {
102        let node = &graph[idx];
103        groups.entry(node.node_type).or_default().push(node);
104    }
105
106    for (node_type, mut group_nodes) in groups {
107        group_nodes.sort_by_key(|n| &n.unique_id);
108        let type_label = node_type.label();
109        let (bg_color, _) = node_colors(node_type);
110        let title = super::capitalize(type_label);
111        writeln!(w, r#"  subgraph cluster_{type_label} {{"#)?;
112        writeln!(w, r#"    label="{title}";"#)?;
113        writeln!(w, "    style=rounded;")?;
114        writeln!(w, r#"    color="{bg_color}";"#)?;
115        writeln!(w)?;
116        for node in &group_nodes {
117            write_node(w, node, "    ")?;
118        }
119        writeln!(w, "  }}")?;
120    }
121    Ok(())
122}
123
124/// Write nodes grouped by file directory using DOT subgraph clusters
125fn write_nodes_grouped_by_directory<W: Write>(w: &mut W, graph: &LineageGraph) -> io::Result<()> {
126    let mut groups: BTreeMap<String, Vec<&NodeData>> = BTreeMap::new();
127    for idx in graph.node_indices() {
128        let node = &graph[idx];
129        let dir = super::directory_label(node);
130        groups.entry(dir).or_default().push(node);
131    }
132
133    for (dir, mut group_nodes) in groups {
134        group_nodes.sort_by_key(|n| &n.unique_id);
135        let cluster_id = super::sanitize_id(&dir);
136        let dir_label = super::dot_escape(&dir);
137        writeln!(w, r#"  subgraph cluster_{cluster_id} {{"#)?;
138        writeln!(w, r#"    label="{dir_label}";"#)?;
139        writeln!(w, "    style=rounded;")?;
140        writeln!(w)?;
141        for node in &group_nodes {
142            write_node(w, node, "    ")?;
143        }
144        writeln!(w, "  }}")?;
145    }
146    Ok(())
147}
148
149/// Write a single node definition
150fn write_node<W: Write>(w: &mut W, node: &NodeData, indent: &str) -> io::Result<()> {
151    let (color, fontcolor) = node_colors(node.node_type);
152    let id = super::dot_escape(&node.unique_id);
153    let label = super::dot_escape(&node.display_name());
154    writeln!(
155        w,
156        r#"{indent}"{id}" [label="{label}", fillcolor="{color}", fontcolor="{fontcolor}"];"#,
157    )
158}
159
160fn node_colors(node_type: NodeType) -> (&'static str, &'static str) {
161    match node_type {
162        NodeType::Model => ("#4A90D9", "white"),
163        NodeType::Source => ("#27AE60", "white"),
164        NodeType::Seed => ("#F39C12", "white"),
165        NodeType::Snapshot => ("#8E44AD", "white"),
166        NodeType::Test => ("#1ABC9C", "white"),
167        NodeType::Exposure => ("#E74C3C", "white"),
168        NodeType::SemanticModel => ("#16A085", "white"),
169        NodeType::Metric => ("#D35400", "white"),
170        NodeType::SavedQuery => ("#2980B9", "white"),
171        NodeType::Phantom => ("#BDC3C7", "black"),
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use crate::render::test_helpers::{make_node, make_node_with_path};
179
180    fn render_to_string(graph: &LineageGraph) -> String {
181        let mut buf = Vec::new();
182        render_dot_to_writer(graph, &mut buf, None, Direction::LR).unwrap();
183        String::from_utf8(buf).unwrap()
184    }
185
186    fn render_to_string_grouped(graph: &LineageGraph) -> String {
187        let mut buf = Vec::new();
188        render_dot_to_writer(graph, &mut buf, Some(GroupBy::NodeType), Direction::LR).unwrap();
189        String::from_utf8(buf).unwrap()
190    }
191
192    #[test]
193    fn test_empty_graph() {
194        let graph = LineageGraph::new();
195        let output = render_to_string(&graph);
196        assert!(output.contains("digraph dbt_lineage {"));
197        assert!(output.contains("}"));
198    }
199
200    #[test]
201    fn test_single_node() {
202        let mut graph = LineageGraph::new();
203        graph.add_node(make_node("model.orders", "orders", NodeType::Model));
204        let output = render_to_string(&graph);
205        assert!(output.contains("\"model.orders\""));
206        assert!(output.contains("label=\"orders\""));
207        assert!(output.contains("fillcolor=\"#4A90D9\""));
208    }
209
210    #[test]
211    fn test_edge_styles() {
212        let mut graph = LineageGraph::new();
213        let a = graph.add_node(make_node(
214            "source.raw.orders",
215            "raw.orders",
216            NodeType::Source,
217        ));
218        let b = graph.add_node(make_node("model.stg_orders", "stg_orders", NodeType::Model));
219        graph.add_edge(a, b, EdgeData::direct(EdgeType::Source));
220
221        let output = render_to_string(&graph);
222        assert!(output.contains("style=dashed"));
223        assert!(output.contains("label=\"source\""));
224    }
225
226    #[test]
227    fn test_all_edge_type_labels() {
228        let types = [
229            (EdgeType::Ref, "ref"),
230            (EdgeType::Source, "source"),
231            (EdgeType::Test, "test"),
232            (EdgeType::Exposure, "exposure"),
233        ];
234        for (et, expected) in types {
235            let ed = EdgeData::direct(et);
236            assert_eq!(ed.edge_type.label(), expected);
237        }
238    }
239
240    #[test]
241    fn test_node_colors_all_types() {
242        let types = [
243            NodeType::Model,
244            NodeType::Source,
245            NodeType::Seed,
246            NodeType::Snapshot,
247            NodeType::Test,
248            NodeType::Exposure,
249            NodeType::Phantom,
250        ];
251        for nt in types {
252            let (color, fontcolor) = node_colors(nt);
253            assert!(
254                color.starts_with('#'),
255                "Color for {:?} should start with #",
256                nt
257            );
258            assert!(!fontcolor.is_empty());
259        }
260    }
261
262    #[test]
263    fn test_multiple_edges_different_styles() {
264        let mut graph = LineageGraph::new();
265        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
266        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
267        let c = graph.add_node(make_node("test.t", "t", NodeType::Test));
268        let d = graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
269
270        graph.add_edge(a, b, EdgeData::direct(EdgeType::Ref));
271        graph.add_edge(b, c, EdgeData::direct(EdgeType::Test));
272        graph.add_edge(b, d, EdgeData::direct(EdgeType::Exposure));
273
274        let output = render_to_string(&graph);
275        // Ref edges have no extra style
276        assert!(output.contains("label=\"ref\""));
277        assert!(output.contains("style=dotted"));
278        assert!(output.contains("style=bold"));
279    }
280
281    #[test]
282    fn test_all_node_types_render() {
283        let mut graph = LineageGraph::new();
284        graph.add_node(make_node("model.m", "m", NodeType::Model));
285        graph.add_node(make_node("source.s", "s", NodeType::Source));
286        graph.add_node(make_node("seed.sd", "sd", NodeType::Seed));
287        graph.add_node(make_node("snapshot.sn", "sn", NodeType::Snapshot));
288        graph.add_node(make_node("test.t", "t", NodeType::Test));
289        graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
290        graph.add_node(make_node("phantom.p", "p", NodeType::Phantom));
291
292        let output = render_to_string(&graph);
293        // Verify all node colors appear
294        assert!(output.contains("#4A90D9")); // Model
295        assert!(output.contains("#27AE60")); // Source
296        assert!(output.contains("#F39C12")); // Seed
297        assert!(output.contains("#8E44AD")); // Snapshot
298        assert!(output.contains("#1ABC9C")); // Test
299        assert!(output.contains("#E74C3C")); // Exposure
300        assert!(output.contains("#BDC3C7")); // Phantom
301        assert!(output.contains("fontcolor=\"black\"")); // Phantom font
302    }
303
304    #[test]
305    fn test_all_four_edge_styles_in_render() {
306        let mut graph = LineageGraph::new();
307        let s = graph.add_node(make_node("source.raw.o", "raw.o", NodeType::Source));
308        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
309        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
310        let t = graph.add_node(make_node("test.t", "t", NodeType::Test));
311        let e = graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
312
313        graph.add_edge(s, a, EdgeData::direct(EdgeType::Source));
314        graph.add_edge(a, b, EdgeData::direct(EdgeType::Ref));
315        graph.add_edge(b, t, EdgeData::direct(EdgeType::Test));
316        graph.add_edge(b, e, EdgeData::direct(EdgeType::Exposure));
317
318        let output = render_to_string(&graph);
319        assert!(output.contains("label=\"source\""));
320        assert!(output.contains("label=\"ref\""));
321        assert!(output.contains("label=\"test\""));
322        assert!(output.contains("label=\"exposure\""));
323        assert!(output.contains("style=dashed"));
324        assert!(output.contains("style=dotted"));
325        assert!(output.contains("style=bold"));
326    }
327
328    #[test]
329    fn test_transitive_ref_edge_style() {
330        let mut graph = LineageGraph::new();
331        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
332        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
333        graph.add_edge(a, b, EdgeData::transitive(EdgeType::Ref, 2));
334
335        let output = render_to_string(&graph);
336        assert!(output.contains(r#"label="ref (via 2)""#));
337        assert!(output.contains("style=dashed"));
338    }
339
340    #[test]
341    fn test_transitive_source_edge_preserves_dashed() {
342        let mut graph = LineageGraph::new();
343        let a = graph.add_node(make_node("source.raw.a", "a", NodeType::Source));
344        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
345        graph.add_edge(a, b, EdgeData::transitive(EdgeType::Source, 3));
346
347        let output = render_to_string(&graph);
348        assert!(output.contains(r#"label="source (via 3)""#));
349        assert!(output.contains(r#"style="dashed,bold""#));
350    }
351
352    #[test]
353    fn test_transitive_exposure_edge_preserves_bold() {
354        let mut graph = LineageGraph::new();
355        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
356        let b = graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
357        graph.add_edge(a, b, EdgeData::transitive(EdgeType::Exposure, 1));
358
359        let output = render_to_string(&graph);
360        assert!(output.contains(r#"label="exposure (via 1)""#));
361        assert!(output.contains(r#"style="bold,dashed""#));
362    }
363
364    #[test]
365    fn test_snapshot_lineage() {
366        let graph = crate::render::test_helpers::make_sample_lineage_graph();
367        let output = render_to_string(&graph);
368        insta::assert_snapshot!(output);
369    }
370
371    #[test]
372    fn test_group_by_node_type() {
373        let graph = crate::render::test_helpers::make_sample_lineage_graph();
374        let output = render_to_string_grouped(&graph);
375        insta::assert_snapshot!(output);
376    }
377
378    #[test]
379    fn test_snapshot_all_edge_types() {
380        let mut graph = LineageGraph::new();
381        let s = graph.add_node(make_node("source.raw.o", "raw.o", NodeType::Source));
382        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
383        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
384        let t = graph.add_node(make_node("test.t", "t", NodeType::Test));
385        let e = graph.add_node(make_node("exposure.e", "e", NodeType::Exposure));
386
387        graph.add_edge(s, a, EdgeData::direct(EdgeType::Source));
388        graph.add_edge(a, b, EdgeData::direct(EdgeType::Ref));
389        graph.add_edge(b, t, EdgeData::direct(EdgeType::Test));
390        graph.add_edge(b, e, EdgeData::direct(EdgeType::Exposure));
391
392        let output = render_to_string(&graph);
393        insta::assert_snapshot!(output);
394    }
395
396    #[test]
397    fn test_snapshot_transitive_edges() {
398        let mut graph = LineageGraph::new();
399        let a = graph.add_node(make_node("source.raw.a", "a", NodeType::Source));
400        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
401        let c = graph.add_node(make_node("model.c", "c", NodeType::Model));
402        graph.add_edge(a, b, EdgeData::direct(EdgeType::Source));
403        graph.add_edge(a, c, EdgeData::transitive(EdgeType::Source, 3));
404        graph.add_edge(b, c, EdgeData::direct(EdgeType::Ref));
405
406        let output = render_to_string(&graph);
407        insta::assert_snapshot!(output);
408    }
409
410    #[test]
411    fn test_group_by_node_type_cluster_structure() {
412        let mut graph = LineageGraph::new();
413        graph.add_node(make_node("model.a", "a", NodeType::Model));
414        graph.add_node(make_node("source.raw.b", "raw.b", NodeType::Source));
415
416        let output = render_to_string_grouped(&graph);
417        assert!(output.contains("subgraph cluster_model {"));
418        assert!(output.contains("subgraph cluster_source {"));
419        assert!(output.contains("label=\"Model\""));
420        assert!(output.contains("label=\"Source\""));
421        assert!(output.contains("style=rounded;"));
422    }
423
424    fn render_to_string_directory(graph: &LineageGraph) -> String {
425        let mut buf = Vec::new();
426        render_dot_to_writer(graph, &mut buf, Some(GroupBy::Directory), Direction::LR).unwrap();
427        String::from_utf8(buf).unwrap()
428    }
429
430    #[test]
431    fn test_group_by_directory_cluster_structure() {
432        let mut graph = LineageGraph::new();
433        graph.add_node(make_node_with_path(
434            "model.stg_orders",
435            "stg_orders",
436            NodeType::Model,
437            "models/staging/stg_orders.sql",
438        ));
439        graph.add_node(make_node_with_path(
440            "model.orders",
441            "orders",
442            NodeType::Model,
443            "models/marts/orders.sql",
444        ));
445        graph.add_node(make_node(
446            "exposure.dashboard",
447            "dashboard",
448            NodeType::Exposure,
449        ));
450
451        let output = render_to_string_directory(&graph);
452        assert!(output.contains("subgraph cluster_models_staging {"));
453        assert!(output.contains(r#"label="models/staging";"#));
454        assert!(output.contains("subgraph cluster_models_marts {"));
455        assert!(output.contains(r#"label="models/marts";"#));
456        assert!(output.contains("subgraph cluster__other_ {"));
457        assert!(output.contains(r#"label="(other)";"#));
458    }
459
460    #[test]
461    fn test_snapshot_group_by_directory() {
462        let mut graph = LineageGraph::new();
463        let src = graph.add_node(make_node_with_path(
464            "source.raw.orders",
465            "raw.orders",
466            NodeType::Source,
467            "models/staging/schema.yml",
468        ));
469        let stg = graph.add_node(make_node_with_path(
470            "model.stg_orders",
471            "stg_orders",
472            NodeType::Model,
473            "models/staging/stg_orders.sql",
474        ));
475        let mart = graph.add_node(make_node_with_path(
476            "model.orders",
477            "orders",
478            NodeType::Model,
479            "models/marts/orders.sql",
480        ));
481        let exp = graph.add_node(make_node(
482            "exposure.dashboard",
483            "dashboard",
484            NodeType::Exposure,
485        ));
486
487        graph.add_edge(src, stg, EdgeData::direct(EdgeType::Source));
488        graph.add_edge(stg, mart, EdgeData::direct(EdgeType::Ref));
489        graph.add_edge(mart, exp, EdgeData::direct(EdgeType::Exposure));
490
491        let output = render_to_string_directory(&graph);
492        insta::assert_snapshot!(output);
493    }
494
495    #[test]
496    fn test_directory_cluster_label_with_special_chars_is_escaped() {
497        // A file_path whose directory contains `"` must be escaped in the
498        // cluster label to produce syntactically valid DOT output.
499        let mut graph = LineageGraph::new();
500        graph.add_node(make_node_with_path(
501            "model.m",
502            "m",
503            NodeType::Model,
504            r#"models/my"dir/m.sql"#,
505        ));
506        let output = render_to_string_directory(&graph);
507        assert!(
508            output.contains(r#"label="models/my\"dir";"#),
509            "directory cluster label not escaped:\n{output}"
510        );
511    }
512
513    #[test]
514    fn test_snapshot_direction_tb() {
515        let graph = crate::render::test_helpers::make_sample_lineage_graph();
516        let mut buf = Vec::new();
517        render_dot_to_writer(&graph, &mut buf, None, Direction::TB).unwrap();
518        let output = String::from_utf8(buf).unwrap();
519        insta::assert_snapshot!(output);
520    }
521
522    #[test]
523    fn test_label_with_quotes_and_backslash_is_escaped() {
524        let mut graph = LineageGraph::new();
525        let mut node = make_node(
526            r#"metric.revenue_"net""#,
527            r#"Revenue "Net" (100%\off)"#,
528            NodeType::Metric,
529        );
530        // unique_id also contains a special char to exercise edge escaping
531        node.unique_id = r#"metric.revenue_"net""#.into();
532        graph.add_node(node);
533        let output = render_to_string(&graph);
534        // display_name() prepends the node-type prefix, e.g. "metric:Revenue ..."
535        // Both " and \ must be backslash-escaped inside the DOT attribute value.
536        assert!(
537            output.contains(r#"label="metric:Revenue \"Net\" (100%\\off)""#),
538            "label not escaped:\n{output}"
539        );
540        // The node identifier in DOT must also be escaped
541        assert!(
542            output.contains(r#""metric.revenue_\"net\"""#),
543            "unique_id not escaped:\n{output}"
544        );
545    }
546
547    #[test]
548    fn test_snapshot_direction_tb_grouped() {
549        let graph = crate::render::test_helpers::make_sample_lineage_graph();
550        let mut buf = Vec::new();
551        render_dot_to_writer(&graph, &mut buf, Some(GroupBy::NodeType), Direction::TB).unwrap();
552        let output = String::from_utf8(buf).unwrap();
553        insta::assert_snapshot!(output);
554    }
555}
dlin_core/render/dot.rs

dlin_core/render/
dot.rs
