Fix generators not support string ids (#2267)

Co-authored-by: Brandon Bayer <[email protected]> (patch)
blitz-js · Apr 28, 2021 · 4bb6ded · 4bb6ded
1 parent 1d8f69d
commit 4bb6ded
Show file tree

Hide file tree

Showing 15 changed files with 356 additions and 238 deletions.
diff --git a/packages/generator/package.json b/packages/generator/package.json
@@ -34,6 +34,7 @@
     "@babel/core": "7.12.10",
     "@babel/plugin-transform-typescript": "7.12.1",
     "@blitzjs/display": "0.35.0-canary.2",
+    "@mrleebo/prisma-ast": "^0.2.3",
     "@types/jscodeshift": "0.7.2",
     "chalk": "^4.1.0",
     "cross-spawn": "7.0.3",

diff --git a/packages/generator/src/generators/model-generator.ts b/packages/generator/src/generators/model-generator.ts
@@ -1,11 +1,11 @@
 import {log} from "@blitzjs/display"
+import * as ast from "@mrleebo/prisma-ast"
 import {spawn} from "cross-spawn"
 import which from "npm-which"
 import path from "path"
 import {Generator, GeneratorOptions, SourceRootType} from "../generator"
 import {Field} from "../prisma/field"
 import {Model} from "../prisma/model"
-import {matchBetween} from "../utils/match-between"
 
 export interface ModelGeneratorOptions extends GeneratorOptions {
   modelName: string
@@ -36,55 +36,60 @@ export class ModelGenerator extends Generator<ModelGeneratorOptions> {
 
   // eslint-disable-next-line require-await
   async write() {
+    const schemaPath = path.resolve("db/schema.prisma")
+    if (!this.fs.exists(schemaPath)) {
+      throw new Error("Prisma schema file was not found")
+    }
+
+    let schema: ast.Schema
     try {
-      if (!this.fs.exists(path.resolve("db/schema.prisma"))) {
-        throw new Error("Prisma schema file was not found")
-      }
-      let updatedOrCreated = "created"
+      schema = ast.getSchema(this.fs.read(schemaPath))
+    } catch (err) {
+      log.error("Failed to parse db/schema.prisma file")
+      throw err
+    }
+    const {modelName, extraArgs, dryRun} = this.options
+    let updatedOrCreated = "created"
 
-      const extraArgs =
-        this.options.extraArgs.length === 1 && this.options.extraArgs[0].includes(" ")
-          ? this.options.extraArgs[0].split(" ")
-          : this.options.extraArgs
-      const modelDefinition = new Model(
-        this.options.modelName,
-        extraArgs.flatMap((def) => Field.parse(def)),
-      )
-      if (!this.options.dryRun) {
-        // wrap in newlines to put a space below the previously generated model and
-        // to preserve the EOF newline
-        const schema = this.fs.read(path.resolve("db/schema.prisma"))
+    let fields = (extraArgs.length === 1 && extraArgs[0].includes(" ")
+      ? extraArgs[0].split(" ")
+      : extraArgs
+    ).flatMap((input) => Field.parse(input, schema))
 
-        if (schema.indexOf(`model ${modelDefinition.name}`) === -1) {
-          //model does not exist in schema - just add it
-          this.fs.append(path.resolve("db/schema.prisma"), `\n${modelDefinition.toString()}\n`)
-        } else {
-          const model = matchBetween(schema, `model ${modelDefinition.name}`, "}")
-          if (model) {
-            //filter out all fields that are already defined
-            modelDefinition.fields = modelDefinition.fields.filter((field) => {
-              return model.indexOf(field.name) === -1
-            })
+    const modelDefinition = new Model(modelName, fields)
 
-            //add new fields to the selected model
-            const newModel = model.replace("}", `${modelDefinition.getNewFields()}}`)
-
-            //replace all content with the newly added fields for the already existing model
-            this.fs.write(path.resolve("db/schema.prisma"), schema.replace(model, newModel))
-            updatedOrCreated = "updated"
-          }
+    let model: ast.Model | undefined
+    if (!dryRun) {
+      model = schema.list.find(function (component): component is ast.Model {
+        return component.type === "model" && component.name === modelDefinition.name
+      })
+      try {
+        if (model) {
+          for (const field of fields) field.appendTo(model)
+          this.fs.write(schemaPath, ast.printSchema(schema))
+          updatedOrCreated = "updated"
+        } else {
+          model = modelDefinition.appendTo(schema)
+          this.fs.write(schemaPath, ast.printSchema(schema))
         }
+      } catch (err) {
+        console.error(`Failed to apply changes to model '${modelDefinition.name}'`)
+        throw err
       }
+    }
+
+    if (model) {
       log.newline()
       log.success(
-        `Model for '${this.options.modelName}'${
-          this.options.dryRun ? "" : ` ${updatedOrCreated} in schema.prisma`
+        `Model '${modelDefinition.name}'${
+          dryRun ? "" : ` ${updatedOrCreated} in schema.prisma`
         }:\n`,
       )
-      modelDefinition.toString().split("\n").map(log.progress)
+      ast
+        .printSchema({type: "schema", list: [model]})
+        .split("\n")
+        .map(log.progress)
       log.newline()
-    } catch (error) {
-      throw error
     }
   }
 

diff --git a/packages/generator/src/prisma/field.ts b/packages/generator/src/prisma/field.ts
@@ -1,4 +1,5 @@
 import {log} from "@blitzjs/display"
+import * as ast from "@mrleebo/prisma-ast"
 import {capitalize, singlePascal, uncapitalize} from "../utils/inflector"
 
 export enum FieldType {
@@ -53,10 +54,11 @@ export class Field {
   relationToFields?: string[]
 
   // 'name:type?[]:attribute' => Field
-  static parse(input: string): Field[] {
+  static parse(input: string, schema?: ast.Schema): Field[] {
     const [_fieldName, _fieldType = "String", attribute] = input.split(":")
     let fieldName = uncapitalize(_fieldName)
     let fieldType = capitalize(_fieldType)
+    let isId = fieldName === "id"
     let isRequired = true
     let isList = false
     let isUpdatedAt = false
@@ -91,7 +93,21 @@ export class Field {
           const idFieldName = `${fieldName}Id`
           relationFromFields = [idFieldName]
           relationToFields = ["id"]
-          maybeIdField = new Field(idFieldName, {type: FieldType.Int, isRequired})
+
+          const relationModel = schema?.list.find(function (component): component is ast.Model {
+            return component.type === "model" && component.name === fieldType
+          })
+          const relationField =
+            relationModel &&
+            relationModel.properties.find(function (prop): prop is ast.Field {
+              return prop.type === "field" && prop.name === "id"
+            })
+
+          maybeIdField = new Field(idFieldName, {
+            // find the matching field based on the relation and, if found, match its field type
+            type: relationField ? (relationField.fieldType as FieldType) : FieldType.Int,
+            isRequired,
+          })
           isList = false
           break
       }
@@ -109,14 +125,14 @@ export class Field {
       if (defaultValueTest.test(attribute)) {
         const [, _defaultValue] = attribute.match(defaultValueTest)!
         defaultValue = builtInGenerators.includes(_defaultValue)
-          ? {name: _defaultValue}
+          ? {type: "function", name: _defaultValue, params: []}
           : _defaultValue
       }
     }
     try {
       const parseResult = new Field(fieldName, {
         default: defaultValue,
-        isId: false,
+        isId,
         isList,
         isRequired,
         isUnique,
@@ -168,72 +184,88 @@ export class Field {
     }
   }
 
-  private getDefault() {
-    if (this.default === undefined) return ""
-    let defaultValue: string
-    if (typeof this.default === "object") {
-      // { name: 'fnname' } is based off of the Prisma model definition
-      defaultValue = `${this.default.name}()`
-    } else {
-      defaultValue = String(this.default)
-    }
-    return `@default(${defaultValue})`
-  }
+  appendTo(model: ast.Model) {
+    if (model.properties.some((prop) => prop.type === "field" && prop.name === this.name)) return
 
-  private getId() {
-    return this.isId ? "@id" : ""
-  }
+    const attributes = [
+      this.getId(),
+      this.getIsUnique(),
+      this.getDefault(),
+      this.getIsUpdatedAt(),
+      this.getRelation(),
+    ].filter(Boolean) as ast.Attribute[]
 
-  private getIsUnique() {
-    return this.isUnique ? "@unique" : ""
+    model.properties.push({
+      type: "field",
+      name: this.name,
+      fieldType: this.type,
+      optional: !this.isRequired,
+      array: this.isList,
+      attributes,
+    })
   }
 
-  private getIsUpdatedAt() {
-    return this.isUpdatedAt ? "@updatedAt" : ""
+  private getDefault(): ast.Attribute | undefined {
+    if (this.default == null) return
+
+    return {
+      type: "attribute",
+      kind: "field",
+      name: "default",
+      args: [
+        {
+          type: "attributeArgument",
+          value: typeof this.default === "object" ? `${this.default.name}()` : String(this.default),
+        },
+      ],
+    }
   }
 
-  private getRelation() {
-    if (this.relationFromFields === undefined || this.relationToFields === undefined) return ""
-    const separator =
-      this.relationToFields &&
-      this.relationToFields.length > 0 &&
-      this.relationFromFields &&
-      this.relationFromFields.length
-        ? ", "
-        : ""
-    const fromFields =
-      this.relationFromFields && this.relationFromFields.length > 0
-        ? `fields: [${this.relationFromFields.toString()}]`
-        : ""
-
-    const toFields =
-      this.relationToFields && this.relationToFields.length > 0
-        ? `references: [${this.relationToFields.toString()}]`
-        : ""
-    return `@relation(${fromFields}${separator}${toFields})`
+  private getId(): ast.Attribute | undefined {
+    if (!this.isId) return
+
+    return {
+      type: "attribute",
+      kind: "field",
+      name: "id",
+    }
   }
 
-  private getTypeModifiers() {
-    return `${this.isRequired ? "" : "?"}${this.isList ? "[]" : ""}`
+  private getIsUnique(): ast.Attribute | undefined {
+    if (!this.isUnique) return
+    return {type: "attribute", kind: "field", name: "unique"}
   }
 
-  private getAttributes() {
-    const possibleAttributes = [
-      this.getDefault(),
-      this.getId(),
-      this.getIsUnique(),
-      this.getIsUpdatedAt(),
-      this.getRelation(),
-    ]
-    // filter out any attributes that return ''
-    const attrs = possibleAttributes.filter((attr) => attr)
-    if (attrs.length > 0) {
-      return `  ${attrs.join(" ")}`
-    }
-    return ""
+  private getIsUpdatedAt(): ast.Attribute | undefined {
+    if (!this.isUpdatedAt) return
+    return {type: "attribute", kind: "field", name: "updatedAt"}
   }
 
-  toString() {
-    return `${this.name}  ${this.type}${this.getTypeModifiers()}${this.getAttributes()}`
+  private getRelation(): ast.Attribute | undefined {
+    if (this.relationFromFields == null || this.relationToFields == null) return
+
+    return {
+      type: "attribute",
+      kind: "field",
+      name: "relation",
+      args: [
+        {
+          type: "attributeArgument",
+          value: {
+            type: "keyValue",
+            key: "fields",
+            value: {type: "array", args: this.relationFromFields},
+          },
+        },
+        {
+          type: "attributeArgument",
+          value: {
+            type: "keyValue",
+            key: "references",
+            value: {type: "array", args: this.relationToFields},
+          },
+        },
+      ],
+    }
   }
 }