feat: Add option to hide enums (#201)

* feat: Add option to hide enums * refactor: Change option name to ignoreEnums * docs: Add description of ignoreEnums * refactor: Lint by prettier
keonik · May 3, 2023 · 8b3bf99 · 8b3bf99
1 parent e6a3838
commit 8b3bf99
Show file tree

Hide file tree

Showing 4 changed files with 101 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -141,6 +141,18 @@ generator erd {
 }
 ```
 
+### Ignore enums
+
+If you enable this option, enum entities will be hidden.
+This is useful if you want to reduce the number of entities and focus on the tables and their columns and relationships.
+
+```prisma
+generator erd {
+  provider = "prisma-erd-generator"
+  ignoreEnums = true
+}
+```
+
 ### Include relation from field
 
 By default this module skips relation fields in the result diagram. For example fields `userId` and `productId` will not be generated from this prisma schema.

diff --git a/__tests__/ignoreEnums.test.ts b/__tests__/ignoreEnums.test.ts
@@ -0,0 +1,36 @@
+import * as child_process from 'child_process';
+
+test('ignore-enums.prisma', async () => {
+    const fileName = 'ignoreEnums.svg';
+    const folderName = '__tests__';
+    child_process.execSync(`rm -f ${folderName}/${fileName}`);
+    child_process.execSync(
+        `prisma generate --schema ./prisma/ignore-enums.prisma`
+    );
+    const listFile = child_process.execSync(`ls -la ${folderName}/${fileName}`);
+    // did it generate a file
+    expect(listFile.toString()).toContain(fileName);
+
+    const svgAsString = child_process
+        .execSync(`cat ${folderName}/${fileName}`)
+        .toString();
+
+    // did it generate a file without enum
+    expect(svgAsString).toContain(`<svg`);
+    // include tables
+    expect(svgAsString).toContain(`Booking`);
+    expect(svgAsString).toContain(`Event`);
+    // incluse table columns
+    expect(svgAsString).toContain(`name`);
+    expect(svgAsString).toContain(`startDate`);
+    expect(svgAsString).toContain(`status`);
+    expect(svgAsString).toContain(`inviteeEmail`);
+    expect(svgAsString).toContain(`startDateUTC`);
+    expect(svgAsString).toContain(`cancelCode`);
+    // include enum names
+    expect(svgAsString).toContain(`Status`);
+    // exclude enums
+    expect(svgAsString).not.toContain(`PENDING`);
+    expect(svgAsString).not.toContain(`CONFIRMED`);
+    expect(svgAsString).not.toContain(`CANCELLED`);
+});
diff --git a/prisma/ignore-enums.prisma b/prisma/ignore-enums.prisma
@@ -0,0 +1,36 @@
+// This is your Prisma schema file,
+// learn more about it in the docs: https://pris.ly/d/prisma-schema
+
+datasource db {
+    provider = "postgresql"
+    url      = env("DATABASE_URL")
+}
+
+generator erd {
+    provider  = "node ./dist/index.js"
+    output    = "../__tests__/ignoreEnums.svg"
+    theme     = "forest"
+    ignoreEnums = true
+}
+
+model Booking {
+    id           Int      @id @default(autoincrement())
+    inviteeEmail String
+    startDateUTC DateTime
+    cancelCode   String
+    events       Event[]
+}
+
+model Event {
+    id        Int       @id @default(autoincrement())
+    name      String
+    startDate DateTime
+    bookings  Booking[]
+    status    Status    @default(PENDING)
+}
+
+enum Status {
+    PENDING
+    CANCELLED
+    CONFIRMED
+}
diff --git a/src/generate.ts b/src/generate.ts
@@ -24,6 +24,7 @@ export interface DMLModel {
 
 export interface DMLRendererOptions {
     tableOnly?: boolean;
+    ignoreEnums?: boolean;
     includeRelationFromFields?: boolean;
 }
 
@@ -133,19 +134,23 @@ export async function parseDatamodel(
 }
 
 function renderDml(dml: DML, options?: DMLRendererOptions) {
-    const { tableOnly = false, includeRelationFromFields = false } =
-        options ?? {};
+    const {
+        tableOnly = false,
+        ignoreEnums = false,
+        includeRelationFromFields = false,
+    } = options ?? {};
 
     const diagram = 'erDiagram';
 
     // Combine Models and Types as they are pretty similar
     const modellikes = dml.models.concat(dml.types);
 
-    const enums = tableOnly
-        ? ''
-        : dml.enums
-              .map(
-                  (model: DMLEnum) => `
+    const enums =
+        tableOnly || ignoreEnums
+            ? ''
+            : dml.enums
+                  .map(
+                      (model: DMLEnum) => `
         ${model.dbName || model.name} {
             ${model.values
                 .map(
@@ -157,8 +162,8 @@ function renderDml(dml: DML, options?: DMLRendererOptions) {
                 .join('\n')}
         }
     `
-              )
-              .join('\n\n');
+                  )
+                  .join('\n\n');
 
     const classes = modellikes
         .map(
@@ -192,7 +197,7 @@ ${
     for (const model of modellikes) {
         for (const field of model.fields) {
             const isEnum = field.kind === 'enum';
-            if (tableOnly && isEnum) {
+            if (isEnum && (tableOnly || ignoreEnums)) {
                 continue;
             }
 
@@ -364,6 +369,7 @@ export default async (options: GeneratorOptions) => {
             path.join(config.mmdcPath || 'node_modules/.bin', 'mmdc')
         );
         const tableOnly = config.tableOnly === 'true';
+        const ignoreEnums = config.ignoreEnums === 'true';
         const includeRelationFromFields =
             config.includeRelationFromFields === 'true';
         const disabled = Boolean(process.env.DISABLE_ERD);
@@ -418,6 +424,7 @@ export default async (options: GeneratorOptions) => {
 
         const mermaid = renderDml(dml, {
             tableOnly,
+            ignoreEnums,
             includeRelationFromFields,
         });
         if (debug && mermaid) {