Support enums, smallintegers and add around mysql column names

jhoobergs · Jul 9, 2024 · 044dc52 · 044dc52
1 parent 813575e
commit 044dc52
Show file tree

Hide file tree

Showing 7 changed files with 225 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/erd/Cargo.toml b/erd/Cargo.toml
@@ -13,3 +13,4 @@ pest = "2.1"
 pest_derive = "2.1"
 serde = { version = "1.0", features = ["derive"] }
 clap = { version = "4.5.8", features = ["derive"] }
+lazy_static = "1.5.0"
diff --git a/erd/src/ast.rs b/erd/src/ast.rs
@@ -172,6 +172,7 @@ pub enum Expr {
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum DataType {
     Integer,
+    SmallInteger,
     AutoIncrement,
     Float,
     Boolean,
@@ -187,6 +188,7 @@ pub enum DataType {
     /// M is the precision: total number of digits (. and - not counted)
     /// D is the scale: the number of digits after the decimal point
     Decimal(usize, usize),
+    Enum(Vec<String>),
 }
 
 impl std::convert::From<String> for DataType {
@@ -202,12 +204,17 @@ impl std::convert::From<String> for DataType {
             let d = part2[..(part2.len() - 1)].trim().parse().unwrap();
 
             Self::Decimal(m, d)
+        } else if s.starts_with("enum(") {
+            let part = &s["enum(".len()..(s.len() - 1)];
+            let elements: Vec<String> = part.split(",").map(|x| x.trim().to_string()).collect();
+            Self::Enum(elements)
         } else {
             match &s[..] {
                 "uuid" => Self::Uuid,
                 "text" => Self::Text,
                 "blob" => Self::Blob,
                 "integer" => Self::Integer,
+                "smallinteger" => Self::SmallInteger,
                 "autoincrement" => Self::AutoIncrement,
                 "float" => Self::Float,
                 "boolean" => Self::Boolean,
@@ -224,6 +231,7 @@ impl DataType {
     pub fn foreign_key_type(&self) -> DataType {
         match self {
             Self::Integer => Self::Integer,
+            Self::SmallInteger => Self::SmallInteger,
             Self::AutoIncrement => Self::Integer,
             Self::Float => Self::Float,
             Self::Boolean => Self::Boolean,
@@ -236,6 +244,7 @@ impl DataType {
             Self::Text => Self::Text,
             Self::Uuid => Self::Uuid,
             Self::Decimal(m, d) => Self::Decimal(*m, *d),
+            Self::Enum(m) => Self::Enum(m.to_owned()),
         }
     }
 }
diff --git a/erd/src/erd.pest b/erd/src/erd.pest
@@ -4,7 +4,26 @@ COMMENT = _{ "//" ~ (!"\n" ~ ANY)* }
 
 ident = @{ ASCII_ALPHA ~ (ASCII_ALPHANUMERIC | "_")* }
 
-datatype = { "integer" | "autoincrement" | "float" | "boolean" | (!"datetime" ~ "date") | "time" | "datetime" | "blob" | "text" | "uuid" | "varchar(" ~ (!"0" ~ ASCII_DIGIT) ~ ASCII_DIGIT* ~ ")" | "varbinary(" ~ (!"0" ~ ASCII_DIGIT) ~ ASCII_DIGIT* ~ ")" | "decimal(" ~ (!"0" ~ ASCII_DIGIT) ~ ASCII_DIGIT* ~ "," ~ (!"0" ~ ASCII_DIGIT) ~ ASCII_DIGIT* ~ ")" }
+datatype = { 
+    "integer" |
+    "smallinteger" | 
+    "autoincrement" | 
+    "float" | 
+    "boolean" | 
+    (!"datetime" ~ "date") | 
+    "time" | 
+    "datetime" | 
+    "blob" | 
+    "text" | 
+    "uuid" | 
+    "varchar(" ~ (!"0" ~ ASCII_DIGIT) ~ ASCII_DIGIT* ~ ")" | 
+    "varbinary(" ~ (!"0" ~ ASCII_DIGIT) ~ ASCII_DIGIT* ~ ")" | 
+    "decimal(" ~ (!"0" ~ ASCII_DIGIT) ~ ASCII_DIGIT* ~ "," ~ (!"0" ~ ASCII_DIGIT) ~ ASCII_DIGIT* ~ ")" | 
+    "enum(" ~ enum_item ~ ("," ~ enum_item)*  ~ ")"
+}
+not_comma_or_space = { !("," | " " | ")") ~ ANY }
+enum_item = { not_comma_or_space+ }
+
 attribute_prefix = { "attribute" | "id" }
 attribute = { attribute_prefix ~ ident ~ ("type" ~ datatype)? } 
 entity = { "entity" ~ ident ~ (!"\n\n" ~ "\n" ~ attribute)* }

diff --git a/erd/src/physical.rs b/erd/src/physical.rs
@@ -168,10 +168,15 @@ pub struct Table {
 
 impl Table {
     fn write_sql_create(&self, s: &mut String, sql: SQL) {
+        let mut additional_definitions = Vec::new();
+
         write!(s, "CREATE TABLE {} (\n", self.name);
         for col in self.columns.iter() {
             col.write_sql_create_lines(s, sql);
             write!(s, "\n");
+            if let Some(x) = sql.to_additional_definitions(&col.datatype) {
+                additional_definitions.push(x);
+            }
         }
         write!(
             s,
@@ -183,6 +188,12 @@ impl Table {
                 .join(","),
         );
         write!(s, ");");
+        if !additional_definitions.is_empty() {
+            for def in additional_definitions {
+                writeln!(s, "");
+                write!(s, "{}", def);
+            }
+        }
     }
 }