rich_data/
er.rs

1//! ER diagrams from the schema model (#247), behind the `er` feature.
2//!
3//! [`model`] turns a [`Schema`] (from SQL DDL through
4//! [`rich_ext::schema::sql`], from JSON Schema, Arrow or code) into
5//! `rs-rich-diagram`'s [`ErModel`]: a box per table with its columns and
6//! keys, and an edge per foreign key. [`from_sql`] does both steps for DDL.
7//! The two crates do not depend on each other, so the bridge lives here.
8//!
9//! - A schema with tables gives an entity per table; one without gives a
10//!   single entity, named after the schema (`schema` when it has no name).
11//! - A column is a primary key when the table's key names it, a foreign key
12//!   when a key starting from it does, unique when it is marked so, and
13//!   nullable unless it is required.
14//! - A foreign key is drawn from the referring table to the one it names,
15//!   `N:1`, or `1:1` when its columns are the referring table's whole primary
16//!   key or a unique column. A table named `users` matches `public.users`
17//!   when that is the only table so named; otherwise the diagram notes that
18//!   the table is missing.
19//!
20//! ```
21//! use rich::Console;
22//! use rich_data::er;
23//!
24//! let ddl = "CREATE TABLE users (id INT PRIMARY KEY, email TEXT UNIQUE);
25//!            CREATE TABLE posts (id INT PRIMARY KEY,
26//!                                author INT NOT NULL REFERENCES users (id));";
27//! let (diagram, notes) = er::from_sql(ddl).unwrap();
28//! assert!(notes.is_empty());
29//! let model = diagram.model();
30//! assert_eq!(model.entities.len(), 2);
31//! assert_eq!(model.relationships[0].from, "posts");
32//! assert_eq!(model.relationships[0].to, "users");
33//!
34//! let console = Console::builder().width(80).color_system(None).build();
35//! let out = console.render_to_string(&diagram);
36//! assert!(out.contains("users") && out.contains("author"));
37//! ```
38
39use std::collections::{HashMap, HashSet};
40
41use rich_diagram::er::{Cardinality, Column, Entity, Relationship};
42use rich_diagram::{ErDiagram, ErModel};
43use rich_ext::schema::sql::{self, Note, SqlError};
44use rich_ext::schema::{ForeignKey, Schema};
45
46/// The ER model of `schema`. See the [module docs](self).
47pub fn model(schema: &Schema) -> ErModel {
48    let tables: Vec<(String, &Schema)> = if schema.tables().is_empty() {
49        vec![(schema.name().unwrap_or("schema").to_string(), schema)]
50    } else {
51        schema
52            .tables()
53            .iter()
54            .enumerate()
55            .map(|(index, table)| {
56                let name = table
57                    .name()
58                    .map_or_else(|| format!("table {}", index + 1), str::to_string);
59                (name, table)
60            })
61            .collect()
62    };
63    let names = Names::new(tables.iter().map(|(name, _)| name.as_str()));
64    let mut er = ErModel::new();
65    for (name, table) in &tables {
66        let keys = table.foreign_keys();
67        // Looked up by name, so wide tables with many keys stay linear.
68        let marks = Marks::new(table, &keys);
69        let entity = Entity::new(name.as_str()).columns(table.fields().iter().map(|field| {
70            let mut column = Column::new(field.name()).data_type(field.type_label());
71            if marks.primary.contains(field.name()) {
72                column = column.primary_key();
73            }
74            if marks.foreign.contains(field.name()) {
75                column = column.foreign_key();
76            }
77            if field.is_unique() {
78                column = column.unique();
79            }
80            if !field.is_required() && !field.is_primary_key() {
81                column = column.nullable();
82            }
83            column
84        }));
85        er = er.entity(entity);
86        for key in &keys {
87            er = er.relationship(relationship(name, &marks, key, &names));
88        }
89    }
90    er
91}
92
93/// Read SQL DDL and draw it: the diagram, and the notes on what the reader
94/// skipped. See [`rich_ext::schema::sql::parse`] for what it reads.
95pub fn from_sql(ddl: &str) -> Result<(ErDiagram, Vec<Note>), SqlError> {
96    let parsed = sql::parse(ddl)?;
97    Ok((ErDiagram::new(model(&parsed.schema)), parsed.notes))
98}
99
100/// A table's primary key, foreign-key and unique columns, by name.
101struct Marks<'a> {
102    /// The primary key's columns, in order (a name may repeat).
103    key: Vec<&'a str>,
104    primary: HashSet<&'a str>,
105    foreign: HashSet<&'a str>,
106    unique: HashSet<&'a str>,
107}
108
109impl<'a> Marks<'a> {
110    fn new(table: &'a Schema, keys: &'a [ForeignKey]) -> Self {
111        let key = table.primary_key();
112        Marks {
113            primary: key.iter().copied().collect(),
114            key,
115            foreign: keys
116                .iter()
117                .flat_map(|k| k.columns.iter().map(String::as_str))
118                .collect(),
119            unique: table
120                .fields()
121                .iter()
122                .filter(|f| f.is_unique())
123                .map(|f| f.name())
124                .collect(),
125        }
126    }
127}
128
129/// The edge for one foreign key of table `from`.
130fn relationship(from: &str, marks: &Marks, key: &ForeignKey, names: &Names) -> Relationship {
131    let to = names.resolve(&key.table);
132    let one = (!marks.key.is_empty()
133        && marks.key.len() == key.columns.len()
134        && key
135            .columns
136            .iter()
137            .all(|c| marks.primary.contains(c.as_str())))
138        || (key.columns.len() == 1 && marks.unique.contains(key.columns[0].as_str()));
139    let mut edge = Relationship::new(from, to).cardinality(if one {
140        Cardinality::OneToOne
141    } else {
142        Cardinality::ManyToOne
143    });
144    if !key.columns.is_empty() {
145        edge = edge.columns(key.columns.clone(), key.references.clone());
146    }
147    edge
148}
149
150/// The tables' names, to resolve the table a key names without a search.
151struct Names<'a> {
152    all: HashSet<&'a str>,
153    /// By last dotted part: the one table, or `None` when several share it.
154    last: HashMap<&'a str, Option<&'a str>>,
155}
156
157/// `users` of `public.users`.
158fn last_part(name: &str) -> &str {
159    name.rsplit('.').next().unwrap_or(name)
160}
161
162impl<'a> Names<'a> {
163    fn new(names: impl IntoIterator<Item = &'a str>) -> Self {
164        let mut all = HashSet::new();
165        let mut last: HashMap<&str, Option<&str>> = HashMap::new();
166        for name in names {
167            all.insert(name);
168            last.entry(last_part(name))
169                .and_modify(|only| *only = None)
170                .or_insert(Some(name));
171        }
172        Names { all, last }
173    }
174
175    /// The table a key names: as written when a table has that name, else
176    /// the one table whose last dotted part matches it, else as written (and
177    /// the diagram notes that it is missing).
178    fn resolve(&self, name: &str) -> String {
179        if self.all.contains(name) {
180            return name.to_string();
181        }
182        match self.last.get(last_part(name)) {
183            Some(Some(only)) => (*only).to_string(),
184            _ => name.to_string(),
185        }
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn keys_and_cardinality_come_from_the_ddl() {
195        let (diagram, notes) = from_sql(
196            "CREATE TABLE public.users (id INT PRIMARY KEY, email TEXT UNIQUE NOT NULL, bio TEXT);
197             CREATE TABLE profiles (user_id INT PRIMARY KEY REFERENCES users (id));
198             CREATE TABLE posts (id INT, author INT, PRIMARY KEY (id),
199                                 FOREIGN KEY (author) REFERENCES public.users (id));
200             CREATE INDEX posts_author ON posts (author);",
201        )
202        .unwrap();
203        assert_eq!(notes.len(), 1, "{notes:?}");
204        let model = diagram.model();
205        let users = &model.entities[0];
206        assert_eq!(users.name, "public.users");
207        assert!(users.columns[0].primary_key && !users.columns[0].nullable);
208        assert!(users.columns[1].unique && !users.columns[1].nullable);
209        assert!(users.columns[2].nullable);
210        let edges: Vec<_> = model
211            .relationships
212            .iter()
213            .map(|r| (r.from.as_str(), r.to.as_str(), r.cardinality))
214            .collect();
215        // `users` resolves to the one table so named, `public.users`.
216        assert_eq!(
217            edges,
218            [
219                ("profiles", "public.users", Some(Cardinality::OneToOne)),
220                ("posts", "public.users", Some(Cardinality::ManyToOne)),
221            ]
222        );
223        assert!(model.entities[2].columns[1].foreign_key);
224    }
225
226    #[test]
227    fn a_schema_without_tables_is_one_entity() {
228        let schema = Schema::new(vec![rich_ext::schema::Field::new(
229            "id",
230            rich_ext::schema::DataType::Integer,
231        )]);
232        let er = model(&schema);
233        assert_eq!(er.entities.len(), 1);
234        assert_eq!(er.entities[0].name, "schema");
235        assert!(er.relationships.is_empty());
236    }
237
238    /// Each foreign key's table was looked up among all the tables' names,
239    /// split anew each time: 50,000 keys to a missing table took minutes.
240    #[test]
241    fn many_keys_resolve_in_linear_time() {
242        let mut ddl = String::new();
243        for i in 0..50_000 {
244            ddl.push_str(&format!("CREATE TABLE s.t{i} (a INT REFERENCES ghost);\n"));
245        }
246        ddl.push_str("CREATE TABLE u (b INT REFERENCES t7);");
247        let schema = sql::parse(&ddl).unwrap().schema;
248        let started = std::time::Instant::now();
249        let er = model(&schema);
250        assert!(started.elapsed() < std::time::Duration::from_secs(20));
251        assert_eq!(er.relationships.len(), 50_001);
252        assert_eq!(er.relationships[0].to, "ghost");
253        assert_eq!(er.relationships[50_000].to, "s.t7");
254    }
255
256    /// Every column searched every foreign key, and every key the table's
257    /// fields: tables of 2,000 columns and 2,000 keys took 16 s (debug).
258    #[test]
259    fn wide_tables_with_many_keys_are_linear() {
260        let columns: Vec<String> = (0..2_000).map(|i| format!("c{i} INT UNIQUE")).collect();
261        let keys: Vec<String> = (0..2_000)
262            .map(|i| format!("FOREIGN KEY (c{i}) REFERENCES t{}", i % 7))
263            .collect();
264        let table = format!("({}, {})", columns.join(", "), keys.join(", "));
265        let mut ddl = String::new();
266        for t in 0..100 {
267            ddl.push_str(&format!("CREATE TABLE t{t} {table};\n"));
268        }
269        let schema = sql::parse(&ddl).unwrap().schema;
270        let started = std::time::Instant::now();
271        let er = model(&schema);
272        assert!(started.elapsed() < std::time::Duration::from_secs(5));
273        assert_eq!(er.relationships.len(), 200_000);
274        assert_eq!(er.relationships[0].cardinality, Some(Cardinality::OneToOne));
275        let column = &er.entities[99].columns[1_999];
276        assert!(column.foreign_key && column.unique && column.nullable);
277    }
278
279    #[test]
280    fn an_ambiguous_or_missing_table_is_kept_as_written() {
281        let resolve = |name: &str, names: &[&str]| Names::new(names.iter().copied()).resolve(name);
282        assert_eq!(resolve("users", &["a.users", "b.users"]), "users");
283        assert_eq!(resolve("ghost", &["users"]), "ghost");
284        assert_eq!(resolve("users", &["users", "a.users"]), "users");
285        assert_eq!(resolve("x.users", &["a.users"]), "a.users");
286    }
287}
rich_data/er.rs

rich_data/
er.rs
