Render enums using the datamodel renderer

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(dbg!(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/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);
     }
 }