mcp_sql/
schema.rs

1use std::collections::HashMap;
2
3use sqlx::AnyPool;
4
5use crate::db::dialect;
6use crate::db::DbBackend;
7use crate::error::McpSqlError;
8
9/// Generate a Mermaid ER diagram for all tables in a database.
10pub async fn generate_mermaid_er(
11    pool: &AnyPool,
12    backend: DbBackend,
13) -> Result<String, McpSqlError> {
14    // Get all tables
15    let table_rows = dialect::list_tables(pool, backend).await?;
16    let table_names: Vec<String> = table_rows
17        .iter()
18        .filter_map(|r| r.get("table_name").and_then(|v| v.as_str()).map(String::from))
19        .collect();
20
21    if table_names.is_empty() {
22        return Ok("erDiagram\n    %% No tables found".to_string());
23    }
24
25    let mut diagram = String::from("erDiagram\n");
26
27    // FK relationships: (from_table, to_table, from_col, to_col)
28    let mut relationships: Vec<(String, String, String, String)> = Vec::new();
29
30    // Describe each table
31    for table in &table_names {
32        let columns = dialect::describe_table(pool, backend, table).await?;
33        diagram.push_str(&format!("    {} {{\n", table));
34        for col in &columns {
35            let name = col
36                .get("name")
37                .and_then(|v| v.as_str())
38                .or_else(|| col.get("column_name").and_then(|v| v.as_str()))
39                .unwrap_or("?");
40            let dtype = col
41                .get("type")
42                .and_then(|v| v.as_str())
43                .or_else(|| col.get("data_type").and_then(|v| v.as_str()))
44                .unwrap_or("?");
45            let is_pk = col
46                .get("primary_key")
47                .and_then(|v| v.as_str())
48                .map(|s| s == "YES")
49                .or_else(|| col.get("is_primary_key").and_then(|v| v.as_bool()))
50                .unwrap_or(false);
51            let fk = col.get("foreign_key").and_then(|v| v.as_str());
52
53            let mut suffix = String::new();
54            if is_pk {
55                suffix.push_str(" PK");
56            }
57            if fk.is_some() {
58                suffix.push_str(" FK");
59            }
60            // Mermaid ER format: TYPE name CONSTRAINT
61            // Type names cannot contain spaces, so replace spaces with underscores.
62            diagram.push_str(&format!(
63                "        {} {}{}\n",
64                dtype.to_uppercase().replace(' ', "_"),
65                name,
66                suffix
67            ));
68
69            // Track FK relationships
70            if let Some(fk_ref) = fk {
71                if let Some((ref_table, ref_col)) = fk_ref.split_once('.') {
72                    relationships.push((
73                        table.clone(),
74                        ref_table.to_string(),
75                        name.to_string(),
76                        ref_col.to_string(),
77                    ));
78                }
79            }
80        }
81        diagram.push_str("    }\n");
82    }
83
84    // Add relationships
85    // Deduplicate by (from_table, to_table) pair
86    let mut seen: HashMap<(String, String), String> = HashMap::new();
87    for (from, to, from_col, _to_col) in &relationships {
88        seen.entry((from.clone(), to.clone()))
89            .or_insert_with(|| from_col.clone());
90    }
91    for (from, to) in seen.keys() {
92        diagram.push_str(&format!("    {} ||--o{{ {} : \"\"\n", to, from));
93    }
94
95    Ok(diagram)
96}
mcp_sql/schema.rs

mcp_sql/
schema.rs
