use std::collections::HashMap;
use sqlx::AnyPool;
use crate::db::dialect;
use crate::db::DbBackend;
use crate::error::McpSqlError;
pub async fn generate_mermaid_er(
pool: &AnyPool,
backend: DbBackend,
) -> Result<String, McpSqlError> {
let table_rows = dialect::list_tables(pool, backend).await?;
let table_names: Vec<String> = table_rows
.iter()
.filter_map(|r| r.get("table_name").and_then(|v| v.as_str()).map(String::from))
.collect();
if table_names.is_empty() {
return Ok("erDiagram\n %% No tables found".to_string());
}
let mut diagram = String::from("erDiagram\n");
let mut relationships: Vec<(String, String, String, String)> = Vec::new();
for table in &table_names {
let columns = dialect::describe_table(pool, backend, table).await?;
diagram.push_str(&format!(" {} {{\n", table));
for col in &columns {
let name = col
.get("name")
.and_then(|v| v.as_str())
.or_else(|| col.get("column_name").and_then(|v| v.as_str()))
.unwrap_or("?");
let dtype = col
.get("type")
.and_then(|v| v.as_str())
.or_else(|| col.get("data_type").and_then(|v| v.as_str()))
.unwrap_or("?");
let is_pk = col
.get("primary_key")
.and_then(|v| v.as_str())
.map(|s| s == "YES")
.or_else(|| col.get("is_primary_key").and_then(|v| v.as_bool()))
.unwrap_or(false);
let fk = col.get("foreign_key").and_then(|v| v.as_str());
let mut suffix = String::new();
if is_pk {
suffix.push_str(" PK");
}
if fk.is_some() {
suffix.push_str(" FK");
}
diagram.push_str(&format!(
" {} {}{}\n",
dtype.to_uppercase().replace(' ', "_"),
name,
suffix
));
if let Some(fk_ref) = fk {
if let Some((ref_table, ref_col)) = fk_ref.split_once('.') {
relationships.push((
table.clone(),
ref_table.to_string(),
name.to_string(),
ref_col.to_string(),
));
}
}
}
diagram.push_str(" }\n");
}
let mut seen: HashMap<(String, String), String> = HashMap::new();
for (from, to, from_col, _to_col) in &relationships {
seen.entry((from.clone(), to.clone()))
.or_insert_with(|| from_col.clone());
}
for (from, to) in seen.keys() {
diagram.push_str(&format!(" {} ||--o{{ {} : \"\"\n", to, from));
}
Ok(diagram)
}