fix: properly set fields with default auth value to optional in creat…

…e input
zenstackhq · Apr 13, 2024 · f53d629 · f53d629
1 parent 70903d3
commit f53d629
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 35 deletions.
diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts
@@ -6,7 +6,9 @@ import {
     getAuthModel,
     getDMMF,
     getDataModels,
+    getForeignKeyFields,
     getPrismaClientImportSpec,
+    hasAttribute,
     isDelegateModel,
     type DMMF,
     type PluginOptions,
@@ -245,22 +247,17 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
                 ]);
             });
 
+        // transform index.d.ts and save it into a new file (better perf than in-line editing)
+
         const sfNew = project.createSourceFile(path.join(prismaClientDir, 'index-fixed.d.ts'), undefined, {
             overwrite: true,
         });
-
-        if (delegateInfo.length > 0) {
-            // transform types for delegated models
-            this.transformDelegate(sf, sfNew, delegateInfo);
-            sfNew.formatText();
-        } else {
-            // just copy
-            sfNew.replaceWithText(sf.getFullText());
-        }
+        this.transform(sf, sfNew, delegateInfo);
+        sfNew.formatText();
         await sfNew.save();
     }
 
-    private transformDelegate(sf: SourceFile, sfNew: SourceFile, delegateInfo: DelegateInfo) {
+    private transform(sf: SourceFile, sfNew: SourceFile, delegateInfo: DelegateInfo) {
         // copy toplevel imports
         sfNew.addImportDeclarations(sf.getImportDeclarations().map((n) => n.getStructure()));
 
@@ -379,22 +376,74 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         // remove aux fields
         source = this.removeAuxFieldsFromTypeAlias(typeAlias, source);
 
-        // remove discriminator field from concrete input types
-        source = this.removeDiscriminatorFromConcreteInput(typeAlias, delegateInfo, source);
+        if (delegateInfo.length > 0) {
+            // remove discriminator field from concrete input types
+            source = this.removeDiscriminatorFromConcreteInput(typeAlias, delegateInfo, source);
+
+            // remove create/connectOrCreate/upsert fields from delegate's input types
+            source = this.removeCreateFromDelegateInput(typeAlias, delegateInfo, source);
 
-        // remove create/connectOrCreate/upsert fields from delegate's input types
-        source = this.removeCreateFromDelegateInput(typeAlias, delegateInfo, source);
+            // remove delegate fields from nested mutation input types
+            source = this.removeDelegateFieldsFromNestedMutationInput(typeAlias, delegateInfo, source);
 
-        // remove delegate fields from nested mutation input types
-        source = this.removeDelegateFieldsFromNestedMutationInput(typeAlias, delegateInfo, source);
+            // fix delegate payload union type
+            source = this.fixDelegatePayloadType(typeAlias, delegateInfo, source);
+        }
 
-        // fix delegate payload union type
-        source = this.fixDelegatePayloadType(typeAlias, delegateInfo, source);
+        // fix the optionality of input args for fields with `auth()` in `@default`
+        source = this.fixFieldsWithDefaultAuth(typeAlias, source);
 
         structure.type = source;
         return structure;
     }
 
+    private fixFieldsWithDefaultAuth(typeAlias: TypeAliasDeclaration, source: string) {
+        const modelsWithDefaultAuth = this.model.declarations.filter(
+            (d): d is DataModel =>
+                isDataModel(d) && d.fields.some((f) => !f.type.optional && f.attributes.some(isDefaultWithAuth))
+        );
+        if (modelsWithDefaultAuth.length === 0) {
+            return source;
+        }
+
+        // set fields optional for create input types
+        const typeName = typeAlias.getName();
+        const createInputRegex = new RegExp(
+            `(${modelsWithDefaultAuth.map((model) => model.name).join('|')})(Unchecked)?Create.*Input`
+        );
+        const match = typeName.match(createInputRegex);
+        if (!match) {
+            return source;
+        }
+
+        const model = modelsWithDefaultAuth.find((model) => model.name === match[1]);
+        if (!model) {
+            return source;
+        }
+
+        const fieldsWithDefaultAuth = model.fields.filter(
+            (f) => !f.type.optional && f.attributes.some(isDefaultWithAuth)
+        );
+        fieldsWithDefaultAuth.forEach((field) => {
+            // mark the field optional
+            const fieldDef = this.findNamedProperty(typeAlias, field.name);
+            if (fieldDef) {
+                source = source.replace(`${field.name}:`, `${field.name}?:`);
+            }
+        });
+
+        const relationFields = model.fields.filter((f) => !f.type.optional && hasAttribute(f, '@relation'));
+        relationFields.forEach((field) => {
+            // if the relation field's all fk fields are optional or have default value, make the relation optional
+            const fkFields = getForeignKeyFields(field);
+            if (fkFields.every((fk) => fk.type.optional || hasAttribute(fk, '@default'))) {
+                source = source.replace(`${field.name}:`, `${field.name}?:`);
+            }
+        });
+
+        return source;
+    }
+
     private fixDelegatePayloadType(typeAlias: TypeAliasDeclaration, delegateInfo: DelegateInfo, source: string) {
         // change the type of `$<DelegateModel>Payload` type of delegate model to a union of concrete types
         const typeName = typeAlias.getName();

diff --git a/packages/schema/src/plugins/prisma/schema-generator.ts b/packages/schema/src/plugins/prisma/schema-generator.ts
@@ -587,23 +587,6 @@ export class PrismaSchemaGenerator {
 
         const type = new ModelFieldType(fieldType, field.type.array, field.type.optional);
 
-        if (this.mode === 'logical') {
-            if (field.attributes.some((attr) => isDefaultWithAuth(attr))) {
-                // field has `@default` with `auth()`, it should be set optional, and the
-                // default value setting is handled outside Prisma
-                type.optional = true;
-            }
-
-            if (isRelationshipField(field)) {
-                // if foreign key field has `@default` with `auth()`, the relation
-                // field should be set optional
-                const foreignKeyFields = getForeignKeyFields(field);
-                if (foreignKeyFields.some((fkField) => fkField.attributes.some((attr) => isDefaultWithAuth(attr)))) {
-                    type.optional = true;
-                }
-            }
-        }
-
         const attributes = field.attributes
             .filter((attr) => this.isPrismaAttribute(attr))
             // `@default` with `auth()` is handled outside Prisma