feat: Added possibility to generate diagram with relation fields and …

…Fixed generation of ERD with composite keys.(Closes #135) (#136) * Fixed composite primary keys * Added possibility to show relation fields in the result diagram
keonik · Sep 10, 2022 · f36325b · f36325b
1 parent 7b1e27c
commit f36325b
Show file tree

Hide file tree

Showing 7 changed files with 212 additions and 49 deletions.
diff --git a/.gitignore b/.gitignore
@@ -5,4 +5,6 @@ dist
 prisma.mmd
 prisma/debug
 coverage
-!__test__/*.ts
+!__test__/*.ts
+__tests__/*.svg
+__tests__/*.png
diff --git a/README.md b/README.md
@@ -1,7 +1,9 @@
 # Prisma Entity Relationship Diagram Generator
 
 <!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->
+
 [![All Contributors](https://img.shields.io/badge/all_contributors-10-orange.svg?style=flat-square)](#contributors-)
+
 <!-- ALL-CONTRIBUTORS-BADGE:END -->
 
 Prisma generator to create an ER Diagram every time you generate your prisma client.
@@ -119,6 +121,44 @@ generator erd {
 }
 ```
 
+### 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.
+
+```prisma
+model User {
+  id            String         @id
+  email         String
+  favoriteProducts  FavoriteProducts[]
+}
+
+
+model Product {
+  id              String        @id
+  title           String
+  inFavorites  FavoriteProducts[]
+}
+
+model FavoriteProducts {
+  userId      String
+  user        User    @relation(fields: [userId], references: [id])
+  productId   String
+  product     Product @relation(fields: [productId], references: [id])
+
+  @@id([userId, productId])
+}
+
+```
+
+It can be useful to show them when working with RDBMS. To show them use `includeRelationFromFields = true`
+
+```prisma
+generator erd {
+  provider = "prisma-erd-generator"
+  includeRelationFromFields = true
+}
+```
+
 ## Contributors ✨
 
 Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):

diff --git a/__tests__/compositePk.test.ts b/__tests__/compositePk.test.ts
@@ -0,0 +1,24 @@
+import * as child_process from 'child_process';
+
+test('composite-pk.prisma', async () => {
+    const fileName = 'CompositePk.svg';
+    const folderName = '__tests__';
+    child_process.execSync(`rm -f ${folderName}/${fileName}`);
+    child_process.execSync(
+        `npx prisma generate --schema ./prisma/composite-pk.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();
+
+    const pks = svgAsString.match(/PK/g);
+
+    // did it generate a file with the correct content
+    expect(svgAsString).toContain(`<svg`);
+    expect(svgAsString).toContain(`Booking`);
+    expect(pks).toHaveLength(2);
+});
diff --git a/__tests__/includeRelationFromFields.test.ts b/__tests__/includeRelationFromFields.test.ts
@@ -0,0 +1,25 @@
+import * as child_process from 'child_process';
+
+test('include-relation-from-fields.prisma', async () => {
+    const fileName = 'includeRelationFromFields.svg';
+    const folderName = '__tests__';
+    child_process.execSync(`rm -f ${folderName}/${fileName}`);
+    child_process.execSync(
+        `npx prisma generate --schema ./prisma/include-relation-from-fields.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 with the correct content
+    expect(svgAsString).toContain(`<svg`);
+    expect(svgAsString).toContain(`User`);
+    expect(svgAsString).toContain(`Product`);
+    expect(svgAsString).toContain(`FavoriteProducts`);
+    expect(svgAsString).toContain(`userId`);
+    expect(svgAsString).toContain(`productId`);
+});
diff --git a/prisma/composite-pk.prisma b/prisma/composite-pk.prisma
@@ -0,0 +1,23 @@
+// 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__/CompositePk.svg"
+    theme    = "forest"
+}
+
+model Booking {
+    id1          Int
+    id2          Int
+    inviteeEmail String
+    startDateUTC DateTime
+    cancelCode   String
+
+    @@id([id1, id2])
+}
diff --git a/prisma/include-relation-from-fields.prisma b/prisma/include-relation-from-fields.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__/includeRelationFromFields.svg"
+    theme    = "forest"
+    includeRelationFromFields = true
+}
+
+model User {
+  id            String         @id
+  email         String
+  favoriteProducts  FavoriteProducts[]
+}
+
+
+model Product {
+  id              String        @id
+  title           String
+  inFavorites  FavoriteProducts[]
+}
+
+model FavoriteProducts {
+  userId      String
+  user        User    @relation(fields: [userId], references: [id])
+  productId   String
+  product     Product @relation(fields: [productId], references: [id])
+
+  @@id([userId, productId])
+}
diff --git a/src/generate.ts b/src/generate.ts
@@ -11,31 +11,20 @@ export interface DMLModel {
     name: string;
     isEmbedded: boolean;
     dbName: string | null;
-    fields: {
-        name: string;
-        hasDefaultValue: boolean;
-        isGenerated: boolean;
-        isId: boolean;
-        isList: boolean;
-        isReadOnly: boolean;
-        isRequired: boolean;
-        isUnique: boolean;
-        isUpdatedAt: boolean;
-        kind: 'scalar' | 'object' | 'enum';
-        type: string;
-        relationFromFields?: any[];
-        relationName?: string;
-        relationOnDelete?: string;
-        relationToFields?: any[];
-    }[];
+    fields: DMLField[];
     idFields: any[];
     uniqueFields: any[];
     uniqueIndexes: any[];
     isGenerated: boolean;
+    primaryKey: {
+        name: string | null;
+        fields: string[];
+    } | null;
 }
 
 export interface DMLRendererOptions {
     tableOnly?: boolean;
+    includeRelationFromFields?: boolean;
 }
 
 // Copy paste of the DMLModel
@@ -44,27 +33,33 @@ export interface DMLType {
     name: string;
     isEmbedded: boolean;
     dbName: string | null;
-    fields: {
-        name: string;
-        hasDefaultValue: boolean;
-        isGenerated: boolean;
-        isId: boolean;
-        isList: boolean;
-        isReadOnly: boolean;
-        isRequired: boolean;
-        isUnique: boolean;
-        isUpdatedAt: boolean;
-        kind: 'scalar' | 'object' | 'enum';
-        type: string;
-        relationFromFields?: any[];
-        relationName?: string;
-        relationOnDelete?: string;
-        relationToFields?: any[];
-    }[];
+    fields: DMLField[];
     idFields: any[];
     uniqueFields: any[];
     uniqueIndexes: any[];
     isGenerated: boolean;
+    primaryKey: {
+        name: string | null;
+        fields: string[];
+    } | null;
+}
+
+export interface DMLField {
+    name: string;
+    hasDefaultValue: boolean;
+    isGenerated: boolean;
+    isId: boolean;
+    isList: boolean;
+    isReadOnly: boolean;
+    isRequired: boolean;
+    isUnique: boolean;
+    isUpdatedAt: boolean;
+    kind: 'scalar' | 'object' | 'enum';
+    type: string;
+    relationFromFields?: any[];
+    relationName?: string;
+    relationOnDelete?: string;
+    relationToFields?: any[];
 }
 
 export interface DMLEnum {
@@ -138,7 +133,8 @@ export async function parseDatamodel(
 }
 
 function renderDml(dml: DML, options?: DMLRendererOptions) {
-    const { tableOnly = false } = options ?? {};
+    const { tableOnly = false, includeRelationFromFields = false } =
+        options ?? {};
 
     const diagram = 'erDiagram';
 
@@ -172,23 +168,18 @@ ${
     tableOnly
         ? ''
         : model.fields
-              .filter(
-                  (field) =>
-                      field.kind !== 'object' &&
-                      !model.fields.find(
-                          ({ relationFromFields }) =>
-                              relationFromFields &&
-                              relationFromFields.includes(field.name)
-                      )
-              )
+              .filter(isFieldShownInSchema(model, includeRelationFromFields))
               // the replace is a hack to make MongoDB style ID columns like _id valid for Mermaid
               .map((field) => {
                   return `    ${field.type.trimStart()} ${field.name.replace(
                       /^_/,
                       'z_'
-                  )} ${field.isId ? 'PK' : ''} ${
-                      field.isRequired ? '' : '"nullable"'
-                  }`;
+                  )} ${
+                      field.isId ||
+                      model.primaryKey?.fields?.includes(field.name)
+                          ? 'PK'
+                          : ''
+                  } ${field.isRequired ? '' : '"nullable"'}`;
               })
               .join('\n')
 }
@@ -289,6 +280,23 @@ ${
     return diagram + '\n' + enums + '\n' + classes + '\n' + relationships;
 }
 
+const isFieldShownInSchema =
+    (model: DMLModel, includeRelationFromFields: boolean) =>
+    (field: DMLField) => {
+        if (includeRelationFromFields) {
+            return field.kind !== 'object';
+        }
+
+        return (
+            field.kind !== 'object' &&
+            !model.fields.find(
+                ({ relationFromFields }) =>
+                    relationFromFields &&
+                    relationFromFields.includes(field.name)
+            )
+        );
+    };
+
 export const mapPrismaToDb = (dmlModels: DMLModel[], dataModel: string) => {
     const splitDataModel = dataModel
         ?.split('\n')
@@ -336,6 +344,8 @@ export default async (options: GeneratorOptions) => {
         const config = options.generator.config;
         const theme = config.theme || 'forest';
         const tableOnly = config.tableOnly === 'true';
+        const includeRelationFromFields =
+            config.includeRelationFromFields === 'true';
         const disabled = Boolean(process.env.DISABLE_ERD);
         const debug =
             config.erdDebug === 'true' || Boolean(process.env.ERD_DEBUG);
@@ -386,7 +396,10 @@ export default async (options: GeneratorOptions) => {
             console.log(`applied @map to fields written to ${mapAppliedFile}`);
         }
 
-        const mermaid = renderDml(dml, { tableOnly });
+        const mermaid = renderDml(dml, {
+            tableOnly,
+            includeRelationFromFields,
+        });
         if (debug && mermaid) {
             const mermaidFile = path.resolve('prisma/debug/3-mermaid.mmd');
             fs.writeFileSync(mermaidFile, mermaid);