Render enums using the datamodel renderer

.github/workflows/formatting.yml

@@ -28,7 +28,7 @@ jobs:
       - uses: actions-rs/clippy-check@v1
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
-          args: --all-features
+          args: --all-features --all-targets
   format:
     runs-on: ubuntu-latest
     steps:

introspection-engine/connectors/sql-introspection-connector/src/introspection.rs

@@ -185,17 +185,62 @@ pub(crate) fn introspect(ctx: &Context, warnings: &mut Vec<Warning>) -> Result<(
     // if based on a previous Prisma version add id default opinionations
     add_prisma_1_id_defaults(&version, &mut datamodel, schema, warnings, ctx);
 
-    let data_model = if ctx.render_config {
-        format!(
-            "{}\n{}",
-            render_configuration(ctx.config, schema),
-            psl::render_datamodel_to_string(&datamodel, Some(ctx.config))
-        )
+    let config = if ctx.render_config {
+        render_configuration(ctx.config, schema).to_string()
     } else {
-        psl::render_datamodel_to_string(&datamodel, Some(ctx.config))
+        String::new()
     };
 
-    Ok((version, psl::reformat(&data_model, 2).unwrap(), datamodel.is_empty()))
+    let rendered = format!(
+        "{}\n{}\n{}",
+        config,
+        psl::render_datamodel_to_string(&datamodel, Some(ctx.config)),
+        render_datamodel(&datamodel),
+    );
+
+    Ok((version, psl::reformat(&rendered, 2).unwrap(), datamodel.is_empty()))
+}
+
+fn render_datamodel(dml: &Datamodel) -> render::Datamodel<'_> {
+    let mut data_model = render::Datamodel::new();
+
+    for dml_enum in dml.enums() {
+        let mut r#enum = render::Enum::new(&dml_enum.name);
+
+        if let Some(ref docs) = dml_enum.documentation {
+            r#enum.documentation(docs);
+        }
+
+        if let Some(ref schema) = dml_enum.schema {
+            r#enum.schema(schema);
+        }
+
+        if let Some(ref map) = dml_enum.database_name {
+            r#enum.map(map);
+        }
+
+        for dml_variant in dml_enum.values.iter() {
+            let mut variant = render::EnumVariant::new(&dml_variant.name);
+
+            if dml_variant.commented_out {
+                variant = variant.into_commented_out();
+            }
+
+            if let Some(ref map) = dml_variant.database_name {
+                variant.map(map);
+            }
+
+            if let Some(ref docs) = dml_variant.documentation {
+                variant.documentation(docs);
+            }
+
+            r#enum.push_variant(variant);
+        }
+
+        data_model.push_enum(r#enum);
+    }
+
+    data_model
 }
 
 fn render_configuration<'a>(config: &'a Configuration, schema: &'a SqlSchema) -> render::Configuration<'a> {

introspection-engine/connectors/sql-introspection-connector/src/introspection/postgres.rs

@@ -2,8 +2,6 @@ use datamodel_renderer as render;
 use psl::{builtin_connectors::PostgresDatasourceProperties, Configuration, PreviewFeature};
 use sql_schema_describer::{postgres::PostgresSchemaExt, SqlSchema};
 
-use crate::sanitize_datamodel_names::sanitize_string;
-
 const EXTENSION_ALLOW_LIST: &[&str] = &["citext", "postgis", "pg_crypto", "uuid-ossp"];
 
 pub(super) fn add_extensions<'a>(
@@ -26,14 +24,7 @@ pub(super) fn add_extensions<'a>(
     let mut next_extensions = render::Array::new();
 
     for ext in pg_schema_ext.extension_walkers() {
-        let sanitized_name = sanitize_string(ext.name());
-        let is_sanitized = ext.name() != sanitized_name;
-
-        let mut next_extension = render::Function::new(sanitized_name);
-
-        if is_sanitized {
-            next_extension.push_param(("map", render::Text(ext.name())));
-        }
+        let mut next_extension = render::Function::new(ext.name());
 
         match previous_extensions.and_then(|e| e.find_by_name(ext.name())) {
             Some(prev) => {

introspection-engine/connectors/sql-introspection-connector/src/sanitize_datamodel_names.rs

@@ -9,16 +9,15 @@ use psl::{
 };
 use quaint::prelude::SqlFamily;
 use regex::Regex;
-use std::collections::HashMap;
 
 static EMPTY_ENUM_PLACEHOLDER: &str = "EMPTY_ENUM_VALUE";
 
 static RE_START: Lazy<Regex> = Lazy::new(|| Regex::new("^[^a-zA-Z]+").unwrap());
 static RE: Lazy<Regex> = Lazy::new(|| Regex::new("[^_a-zA-Z0-9]").unwrap());
 
 pub(crate) fn sanitize_datamodel_names(ctx: &Context, datamodel: &mut Datamodel) {
-    let enum_renames = sanitize_models(ctx, datamodel);
-    sanitize_enums(&enum_renames, datamodel);
+    sanitize_models(ctx, datamodel);
+    sanitize_enums(datamodel);
 }
 
 // if after opionated renames we have duplicated names, e.g. a database with
@@ -42,8 +41,7 @@ pub fn sanitization_leads_to_duplicate_names(datamodel: &Datamodel) -> bool {
 }
 
 // Todo: Sanitizing might need to be adjusted to also change the fields in the RelationInfo
-fn sanitize_models(ctx: &Context, datamodel: &mut Datamodel) -> HashMap<String, (String, Option<String>)> {
-    let mut enum_renames = HashMap::new();
+fn sanitize_models(ctx: &Context, datamodel: &mut Datamodel) {
     let sql_family = ctx.sql_family();
 
     for model in datamodel.models_mut() {
@@ -76,26 +74,16 @@ fn sanitize_models(ctx: &Context, datamodel: &mut Datamodel) -> HashMap<String,
                         // Enums in MySQL are defined on the column and do not have a separate name.
                         // Introspection generates an enum name for MySQL as `<model_name>_<field_type>`.
                         // If the sanitization changes the enum name, we need to make sure it's changed everywhere.
-                        let (sanitized_enum_name, db_name) = if let SqlFamily::Mysql = sql_family {
+                        let sanitized_enum_name = if let SqlFamily::Mysql = sql_family {
                             if model_db_name.is_none() && sf.database_name.is_none() {
-                                (enum_name.to_owned(), None)
+                                enum_name.to_owned()
                             } else {
-                                (format!("{}_{}", model_name, sf.name()), Some(enum_name.to_owned()))
+                                format!("{}_{}", model_name, sf.name())
                             }
                         } else {
-                            let sanitized = sanitize_string(enum_name);
-
-                            if &sanitized != enum_name {
-                                (sanitized, Some(enum_name.to_owned()))
-                            } else {
-                                (sanitized, None)
-                            }
+                            sanitize_string(enum_name)
                         };
 
-                        if db_name.is_some() {
-                            enum_renames.insert(enum_name.to_owned(), (sanitized_enum_name.clone(), db_name));
-                        }
-
                         sf.field_type = FieldType::Enum(sanitized_enum_name);
 
                         // If the field also has an associated default enum value, we need to sanitize that enum value.
@@ -128,29 +116,14 @@ fn sanitize_models(ctx: &Context, datamodel: &mut Datamodel) -> HashMap<String,
             sanitize_index_field_names(&mut index.fields);
         }
     }
-
-    enum_renames
 }
 
-fn sanitize_enums(enum_renames: &HashMap<String, (String, Option<String>)>, datamodel: &mut Datamodel) {
+fn sanitize_enums(datamodel: &mut Datamodel) {
     for enm in datamodel.enums_mut() {
-        if let Some((sanitized_name, db_name)) = enum_renames.get(&enm.name) {
-            if enm.database_name().is_none() {
-                enm.set_database_name(db_name.clone());
-            }
-
-            enm.set_name(sanitized_name);
-        } else {
-            sanitize_name(enm);
-        }
+        sanitize_name(enm);
 
         for enum_value in enm.values_mut() {
-            if enum_value.name.is_empty() {
-                enum_value.name = EMPTY_ENUM_PLACEHOLDER.to_string();
-                enum_value.database_name = Some("".to_string());
-            } else {
-                sanitize_name(enum_value);
-            }
+            sanitize_name(enum_value);
         }
     }
 }

introspection-engine/datamodel-renderer/Cargo.toml

@@ -6,8 +6,10 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+once_cell = "1.15.0"
 psl.workspace = true
+regex = "1.6.0"
 
 [dev-dependencies]
 expect-test = "1.4.0"
-indoc = "1.0.7"
+indoc = "1.0.7"

introspection-engine/datamodel-renderer/src/datamodel.rs

@@ -0,0 +1,76 @@
+use std::fmt;
+
+use crate::Enum;
+
+/// The PSL data model declaration.
+#[derive(Default, Debug)]
+pub struct Datamodel<'a> {
+    enums: Vec<Enum<'a>>,
+}
+
+impl<'a> Datamodel<'a> {
+    /// Create a new empty data model.
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Add an enum block to the data model.
+    ///
+    /// ```ignore
+    /// enum Foo { // <
+    ///   Bar      // < this
+    /// }          // <
+    /// ```
+    pub fn push_enum(&mut self, r#enum: Enum<'a>) {
+        self.enums.push(r#enum);
+    }
+}
+
+impl<'a> fmt::Display for Datamodel<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for r#enum in self.enums.iter() {
+            writeln!(f, "{enum}")?;
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::*;
+    use expect_test::expect;
+
+    #[test]
+    fn simple_data_model() {
+        let mut traffic_light = Enum::new("TrafficLight");
+
+        traffic_light.push_variant("Red");
+        traffic_light.push_variant("Yellow");
+        traffic_light.push_variant("Green");
+
+        let mut cat = Enum::new("Cat");
+        cat.push_variant("Asleep");
+        cat.push_variant("Hungry");
+
+        let mut data_model = Datamodel::new();
+        data_model.push_enum(traffic_light);
+        data_model.push_enum(cat);
+
+        let expected = expect![[r#"
+            enum TrafficLight {
+              Red
+              Yellow
+              Green
+            }
+
+            enum Cat {
+              Asleep
+              Hungry
+            }
+        "#]];
+
+        let rendered = psl::reformat(&format!("{data_model}"), 2).unwrap();
+        expected.assert_eq(&rendered);
+    }
+}

introspection-engine/datamodel-renderer/src/datasource.rs

@@ -3,7 +3,7 @@ use std::default::Default;
 
 use psl::datamodel_connector::RelationMode;
 
-use crate::{Array, Commented, Env, Text, Value};
+use crate::{Array, Documentation, Env, Text, Value};
 
 /// The datasource block in a PSL file.
 #[derive(Debug)]
@@ -14,7 +14,7 @@ pub struct Datasource<'a> {
     shadow_database_url: Option<Env<'a>>,
     relation_mode: Option<RelationMode>,
     custom_properties: Vec<(&'a str, Value<'a>)>,
-    documentation: Option<Commented<'a>>,
+    documentation: Option<Documentation<'a>>,
     namespaces: Array<Text<&'a str>>,
 }
 
@@ -102,22 +102,23 @@ impl<'a> Datasource<'a> {
     /// }
     /// ```
     pub fn documentation(&mut self, documentation: &'a str) {
-        self.documentation = Some(Commented::Documentation(documentation));
+        self.documentation = Some(Documentation(documentation));
     }
 
     /// Create a rendering from a PSL datasource.
     pub fn from_psl(psl_ds: &'a psl::Datasource) -> Self {
         let shadow_database_url = psl_ds.shadow_database_url.as_ref().map(|(url, _)| Env::from(url));
+        let namespaces: Vec<Text<_>> = psl_ds.namespaces.iter().map(|(ns, _)| Text(ns.as_str())).collect();
 
         Self {
             name: &psl_ds.name,
             provider: Text(&psl_ds.provider),
             url: Env::from(&psl_ds.url),
             shadow_database_url,
             relation_mode: psl_ds.relation_mode,
-            documentation: psl_ds.documentation.as_deref().map(Commented::Documentation),
+            documentation: psl_ds.documentation.as_deref().map(Documentation),
             custom_properties: Default::default(),
-            namespaces: Array(psl_ds.namespaces.iter().map(|(ns, _)| Text(ns.as_str())).collect()),
+            namespaces: Array::from(namespaces),
         }
     }
 }
@@ -168,12 +169,9 @@ mod tests {
         datasource.shadow_database_url(Env::variable("SHADOW_DATABASE_URL"));
         datasource.relation_mode(RelationMode::ForeignKeys);
 
-        let mut w = Function::new("uuid_ossp");
-        w.push_param(("map", Text("uuid-ossp")));
-
         let mut extensions = Array::new();
         extensions.push(Function::new("postgis"));
-        extensions.push(w);
+        extensions.push(Function::new("uuid-ossp"));
 
         datasource.push_custom_property("extensions", extensions);
 
@@ -192,6 +190,6 @@ mod tests {
         "#]];
 
         let rendered = psl::reformat(&format!("{datasource}"), 2).unwrap();
-        expected.assert_eq(&rendered)
+        expected.assert_eq(&rendered);
     }
 }