sql_splitter/graph/format/
dot.rs

1//! Graphviz DOT format output for ERD diagrams.
2
3use crate::graph::format::Layout;
4use crate::graph::view::GraphView;
5
6/// Generate DOT format output with ERD-style tables showing all columns
7pub fn to_dot(view: &GraphView, layout: Layout) -> String {
8    let mut output = String::new();
9
10    // Header
11    output.push_str("digraph ERD {\n");
12    output.push_str("  graph [pad=\"0.5\", nodesep=\"1\", ranksep=\"1.5\"];\n");
13
14    // Layout direction
15    let rankdir = match layout {
16        Layout::LR => "LR",
17        Layout::TB => "TB",
18    };
19    output.push_str(&format!("  rankdir={};\n", rankdir));
20
21    // Node styling for ERD tables
22    output.push_str("  node [shape=none, margin=0];\n");
23    output.push_str("  edge [arrowhead=crow, arrowtail=none, dir=both];\n\n");
24
25    // Generate each table as an HTML-like label
26    for table in view.sorted_tables() {
27        let label = generate_table_label(table);
28        output.push_str(&format!(
29            "  {} [label=<{}>];\n",
30            escape_dot_id(&table.name),
31            label
32        ));
33    }
34
35    if !view.edges.is_empty() {
36        output.push('\n');
37    }
38
39    // Generate edges (FK relationships)
40    for edge in &view.edges {
41        let label = format!("{}→{}", edge.from_column, edge.to_column);
42        output.push_str(&format!(
43            "  {}:{} -> {}:{} [label=\"{}\"];\n",
44            escape_dot_id(&edge.from_table),
45            escape_dot_id(&edge.from_column),
46            escape_dot_id(&edge.to_table),
47            escape_dot_id(&edge.to_column),
48            label
49        ));
50    }
51
52    output.push_str("}\n");
53    output
54}
55
56/// Generate HTML-like table label for DOT
57fn generate_table_label(table: &crate::graph::view::TableInfo) -> String {
58    let mut html = String::new();
59
60    // Table structure with styling
61    html.push_str("<TABLE BORDER=\"0\" CELLBORDER=\"1\" CELLSPACING=\"0\" CELLPADDING=\"4\">");
62
63    // Table header
64    html.push_str(&format!(
65        "<TR><TD BGCOLOR=\"#4a5568\" COLSPAN=\"3\"><FONT COLOR=\"white\"><B>{}</B></FONT></TD></TR>",
66        escape_html(&table.name)
67    ));
68
69    // Column headers
70    html.push_str("<TR>");
71    html.push_str("<TD BGCOLOR=\"#e2e8f0\"><FONT POINT-SIZE=\"10\"><B>Column</B></FONT></TD>");
72    html.push_str("<TD BGCOLOR=\"#e2e8f0\"><FONT POINT-SIZE=\"10\"><B>Type</B></FONT></TD>");
73    html.push_str("<TD BGCOLOR=\"#e2e8f0\"><FONT POINT-SIZE=\"10\"><B>Key</B></FONT></TD>");
74    html.push_str("</TR>");
75
76    // Columns
77    for col in &table.columns {
78        let key_marker = if col.is_primary_key {
79            "🔑 PK"
80        } else if col.is_foreign_key {
81            "🔗 FK"
82        } else {
83            ""
84        };
85
86        let null_marker = if col.is_nullable && !col.is_primary_key {
87            " <FONT COLOR=\"#888888\">NULL</FONT>"
88        } else {
89            ""
90        };
91
92        html.push_str("<TR>");
93        html.push_str(&format!(
94            "<TD ALIGN=\"LEFT\" PORT=\"{}\">{}{}</TD>",
95            escape_html(&col.name),
96            escape_html(&col.name),
97            null_marker
98        ));
99        html.push_str(&format!(
100            "<TD ALIGN=\"LEFT\"><FONT COLOR=\"#666666\">{}</FONT></TD>",
101            escape_html(&col.col_type)
102        ));
103        html.push_str(&format!("<TD ALIGN=\"CENTER\">{}</TD>", key_marker));
104        html.push_str("</TR>");
105    }
106
107    html.push_str("</TABLE>");
108    html
109}
110
111/// Escape a string for use in DOT HTML labels
112fn escape_html(s: &str) -> String {
113    s.replace('&', "&amp;")
114        .replace('<', "&lt;")
115        .replace('>', "&gt;")
116        .replace('"', "&quot;")
117}
118
119/// Escape a string for use as a DOT node ID
120fn escape_dot_id(s: &str) -> String {
121    if s.chars().all(|c| c.is_alphanumeric() || c == '_') && !s.is_empty() {
122        s.to_string()
123    } else {
124        format!("\"{}\"", s.replace('\\', "\\\\").replace('"', "\\\""))
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131    use crate::graph::view::{Cardinality, ColumnInfo, EdgeInfo, TableInfo};
132    use ahash::AHashMap;
133
134    fn create_test_view() -> GraphView {
135        let mut tables = AHashMap::new();
136
137        tables.insert(
138            "users".to_string(),
139            TableInfo {
140                name: "users".to_string(),
141                columns: vec![
142                    ColumnInfo {
143                        name: "id".to_string(),
144                        col_type: "INT".to_string(),
145                        is_primary_key: true,
146                        is_foreign_key: false,
147                        is_nullable: false,
148                        references_table: None,
149                        references_column: None,
150                    },
151                    ColumnInfo {
152                        name: "email".to_string(),
153                        col_type: "VARCHAR(255)".to_string(),
154                        is_primary_key: false,
155                        is_foreign_key: false,
156                        is_nullable: false,
157                        references_table: None,
158                        references_column: None,
159                    },
160                ],
161            },
162        );
163
164        tables.insert(
165            "orders".to_string(),
166            TableInfo {
167                name: "orders".to_string(),
168                columns: vec![
169                    ColumnInfo {
170                        name: "id".to_string(),
171                        col_type: "INT".to_string(),
172                        is_primary_key: true,
173                        is_foreign_key: false,
174                        is_nullable: false,
175                        references_table: None,
176                        references_column: None,
177                    },
178                    ColumnInfo {
179                        name: "user_id".to_string(),
180                        col_type: "INT".to_string(),
181                        is_primary_key: false,
182                        is_foreign_key: true,
183                        is_nullable: false,
184                        references_table: Some("users".to_string()),
185                        references_column: Some("id".to_string()),
186                    },
187                ],
188            },
189        );
190
191        let edges = vec![EdgeInfo {
192            from_table: "orders".to_string(),
193            from_column: "user_id".to_string(),
194            to_table: "users".to_string(),
195            to_column: "id".to_string(),
196            cardinality: Cardinality::ManyToOne,
197        }];
198
199        GraphView { tables, edges }
200    }
201
202    #[test]
203    fn test_dot_contains_table_structure() {
204        let view = create_test_view();
205        let output = to_dot(&view, Layout::LR);
206
207        assert!(output.contains("digraph ERD"));
208        assert!(output.contains("rankdir=LR"));
209        assert!(output.contains("<B>users</B>"));
210        assert!(output.contains("<B>orders</B>"));
211        assert!(output.contains("🔑 PK"));
212        assert!(output.contains("🔗 FK"));
213    }
214
215    #[test]
216    fn test_dot_contains_columns() {
217        let view = create_test_view();
218        let output = to_dot(&view, Layout::LR);
219
220        assert!(output.contains("email"));
221        assert!(output.contains("VARCHAR(255)"));
222        assert!(output.contains("user_id"));
223    }
224
225    #[test]
226    fn test_dot_contains_edges() {
227        let view = create_test_view();
228        let output = to_dot(&view, Layout::LR);
229
230        assert!(output.contains("orders:user_id -> users:id"));
231    }
232}