fix: make sure both fk and relation fields are optional in create input types

ymc9 · ymc9 · commit adcbf4bda2d7 · 2024-11-14T20:06:19.000-08:00
fixes #1843
diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts
@@ -7,6 +7,7 @@ import {
     getDataModelAndTypeDefs,
     getDataModels,
     getLiteral,
+    getRelationField,
     isDelegateModel,
     isDiscriminatorField,
     normalizedRelative,
@@ -55,12 +56,23 @@ type DelegateInfo = [DataModel, DataModel[]][];
 const LOGICAL_CLIENT_GENERATION_PATH = './.logical-prisma-client';
 
 export class EnhancerGenerator {
+    // regex for matching "ModelCreateXXXInput" and "ModelUncheckedCreateXXXInput" type
+    // names for models that use `auth()` in `@default` attribute
+    private readonly modelsWithAuthInDefaultCreateInputPattern: RegExp;
+
     constructor(
         private readonly model: Model,
         private readonly options: PluginOptions,
         private readonly project: Project,
         private readonly outDir: string
-    ) {}
+    ) {
+        const modelsWithAuthInDefault = this.model.declarations.filter(
+            (d): d is DataModel => isDataModel(d) && d.fields.some((f) => f.attributes.some(isDefaultWithAuth))
+        );
+        this.modelsWithAuthInDefaultCreateInputPattern = new RegExp(
+            `^(${modelsWithAuthInDefault.map((m) => m.name).join('|')})(Unchecked)?Create.*?Input$`
+        );
+    }
 
     async generate(): Promise<{ dmmf: DMMF.Document | undefined; newPrismaClientDtsPath: string | undefined }> {
         let dmmf: DMMF.Document | undefined;
@@ -69,7 +81,7 @@ export class EnhancerGenerator {
         let prismaTypesFixed = false;
         let resultPrismaImport = prismaImport;
 
-        if (this.needsLogicalClient || this.needsPrismaClientTypeFixes) {
+        if (this.needsLogicalClient) {
             prismaTypesFixed = true;
             resultPrismaImport = `${LOGICAL_CLIENT_GENERATION_PATH}/index-fixed`;
             const result = await this.generateLogicalPrisma();
@@ -230,11 +242,7 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
     }
 
     private get needsLogicalClient() {
-        return this.hasDelegateModel(this.model) || this.hasAuthInDefault(this.model);
-    }
-
-    private get needsPrismaClientTypeFixes() {
-        return this.hasTypeDef(this.model);
+        return this.hasDelegateModel(this.model) || this.hasAuthInDefault(this.model) || this.hasTypeDef(this.model);
     }
 
     private hasDelegateModel(model: Model) {
@@ -449,11 +457,13 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         const auxFields = this.findAuxDecls(variable);
         if (auxFields.length > 0) {
             structure.declarations.forEach((variable) => {
-                let source = variable.type?.toString();
-                auxFields.forEach((f) => {
-                    source = source?.replace(f.getText(), '');
-                });
-                variable.type = source;
+                if (variable.type) {
+                    let source = variable.type.toString();
+                    auxFields.forEach((f) => {
+                        source = this.removeFromSource(source, f.getText());
+                    });
+                    variable.type = source;
+                }
             });
         }
 
@@ -498,72 +508,16 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         // fix delegate payload union type
         source = this.fixDelegatePayloadType(typeAlias, delegateInfo, source);
 
+        // fix fk and relation fields related to using `auth()` in `@default`
+        source = this.fixDefaultAuthType(typeAlias, source);
+
         // fix json field type
         source = this.fixJsonFieldType(typeAlias, source);
 
         structure.type = source;
         return structure;
     }
 
-    private fixJsonFieldType(typeAlias: TypeAliasDeclaration, source: string) {
-        const modelsWithTypeField = this.model.declarations.filter(
-            (d): d is DataModel => isDataModel(d) && d.fields.some((f) => isTypeDef(f.type.reference?.ref))
-        );
-        const typeName = typeAlias.getName();
-
-        const getTypedJsonFields = (model: DataModel) => {
-            return model.fields.filter((f) => isTypeDef(f.type.reference?.ref));
-        };
-
-        const replacePrismaJson = (source: string, field: DataModelField) => {
-            return source.replace(
-                new RegExp(`(${field.name}\\??\\s*):[^\\n]+`),
-                `$1: ${field.type.reference!.$refText}${field.type.array ? '[]' : ''}${
-                    field.type.optional ? ' | null' : ''
-                }`
-            );
-        };
-
-        // fix "$[Model]Payload" type
-        const payloadModelMatch = modelsWithTypeField.find((m) => `$${m.name}Payload` === typeName);
-        if (payloadModelMatch) {
-            const scalars = typeAlias
-                .getDescendantsOfKind(SyntaxKind.PropertySignature)
-                .find((p) => p.getName() === 'scalars');
-            if (!scalars) {
-                return source;
-            }
-
-            const fieldsToFix = getTypedJsonFields(payloadModelMatch);
-            for (const field of fieldsToFix) {
-                source = replacePrismaJson(source, field);
-            }
-        }
-
-        // fix input/output types, "[Model]CreateInput", etc.
-        const inputOutputModelMatch = modelsWithTypeField.find((m) => typeName.startsWith(m.name));
-        if (inputOutputModelMatch) {
-            const relevantTypePatterns = [
-                'GroupByOutputType',
-                '(Unchecked)?Create(\\S+?)?Input',
-                '(Unchecked)?Update(\\S+?)?Input',
-                'CreateManyInput',
-                '(Unchecked)?UpdateMany(Mutation)?Input',
-            ];
-            const typeRegex = modelsWithTypeField.map(
-                (m) => new RegExp(`^(${m.name})(${relevantTypePatterns.join('|')})$`)
-            );
-            if (typeRegex.some((r) => r.test(typeName))) {
-                const fieldsToFix = getTypedJsonFields(inputOutputModelMatch);
-                for (const field of fieldsToFix) {
-                    source = replacePrismaJson(source, field);
-                }
-            }
-        }
-
-        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();
@@ -595,7 +549,7 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
                 .getDescendantsOfKind(SyntaxKind.PropertySignature)
                 .filter((p) => ['create', 'createMany', 'connectOrCreate', 'upsert'].includes(p.getName()));
             toRemove.forEach((r) => {
-                source = source.replace(r.getText(), '');
+                this.removeFromSource(source, r.getText());
             });
         }
         return source;
@@ -633,7 +587,7 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
                 if (isDiscriminatorField(field)) {
                     const fieldDef = this.findNamedProperty(typeAlias, field.name);
                     if (fieldDef) {
-                        source = source.replace(fieldDef.getText(), '');
+                        source = this.removeFromSource(source, fieldDef.getText());
                     }
                 }
             }
@@ -646,7 +600,7 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         const auxDecls = this.findAuxDecls(typeAlias);
         if (auxDecls.length > 0) {
             auxDecls.forEach((d) => {
-                source = source.replace(d.getText(), '');
+                source = this.removeFromSource(source, d.getText());
             });
         }
         return source;
@@ -677,7 +631,7 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         const fieldDef = this.findNamedProperty(typeAlias, relationFieldName);
         if (fieldDef) {
             // remove relation field of delegate type, e.g., `asset`
-            source = source.replace(fieldDef.getText(), '');
+            source = this.removeFromSource(source, fieldDef.getText());
         }
 
         // remove fk fields related to the delegate type relation, e.g., `assetId`
@@ -709,13 +663,103 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         fkFields.forEach((fkField) => {
             const fieldDef = this.findNamedProperty(typeAlias, fkField);
             if (fieldDef) {
-                source = source.replace(fieldDef.getText(), '');
+                source = this.removeFromSource(source, fieldDef.getText());
             }
         });
 
         return source;
     }
 
+    private fixDefaultAuthType(typeAlias: TypeAliasDeclaration, source: string) {
+        const match = typeAlias.getName().match(this.modelsWithAuthInDefaultCreateInputPattern);
+        if (!match) {
+            return source;
+        }
+
+        const modelName = match[1];
+        const dataModel = this.model.declarations.find((d): d is DataModel => isDataModel(d) && d.name === modelName);
+        if (dataModel) {
+            for (const fkField of dataModel.fields.filter((f) => f.attributes.some(isDefaultWithAuth))) {
+                // change fk field to optional since it has a default
+                source = source.replace(new RegExp(`^(\\s*${fkField.name}\\s*):`, 'm'), `$1?:`);
+
+                const relationField = getRelationField(fkField);
+                if (relationField) {
+                    // change relation field to optional since its fk has a default
+                    source = source.replace(new RegExp(`^(\\s*${relationField.name}\\s*):`, 'm'), `$1?:`);
+                }
+            }
+        }
+        return source;
+    }
+
+    private fixJsonFieldType(typeAlias: TypeAliasDeclaration, source: string) {
+        const modelsWithTypeField = this.model.declarations.filter(
+            (d): d is DataModel => isDataModel(d) && d.fields.some((f) => isTypeDef(f.type.reference?.ref))
+        );
+        const typeName = typeAlias.getName();
+
+        const getTypedJsonFields = (model: DataModel) => {
+            return model.fields.filter((f) => isTypeDef(f.type.reference?.ref));
+        };
+
+        const replacePrismaJson = (source: string, field: DataModelField) => {
+            return source.replace(
+                new RegExp(`(${field.name}\\??\\s*):[^\\n]+`),
+                `$1: ${field.type.reference!.$refText}${field.type.array ? '[]' : ''}${
+                    field.type.optional ? ' | null' : ''
+                }`
+            );
+        };
+
+        // fix "$[Model]Payload" type
+        const payloadModelMatch = modelsWithTypeField.find((m) => `$${m.name}Payload` === typeName);
+        if (payloadModelMatch) {
+            const scalars = typeAlias
+                .getDescendantsOfKind(SyntaxKind.PropertySignature)
+                .find((p) => p.getName() === 'scalars');
+            if (!scalars) {
+                return source;
+            }
+
+            const fieldsToFix = getTypedJsonFields(payloadModelMatch);
+            for (const field of fieldsToFix) {
+                source = replacePrismaJson(source, field);
+            }
+        }
+
+        // fix input/output types, "[Model]CreateInput", etc.
+        const inputOutputModelMatch = modelsWithTypeField.find((m) => typeName.startsWith(m.name));
+        if (inputOutputModelMatch) {
+            const relevantTypePatterns = [
+                'GroupByOutputType',
+                '(Unchecked)?Create(\\S+?)?Input',
+                '(Unchecked)?Update(\\S+?)?Input',
+                'CreateManyInput',
+                '(Unchecked)?UpdateMany(Mutation)?Input',
+            ];
+            const typeRegex = modelsWithTypeField.map(
+                (m) => new RegExp(`^(${m.name})(${relevantTypePatterns.join('|')})$`)
+            );
+            if (typeRegex.some((r) => r.test(typeName))) {
+                const fieldsToFix = getTypedJsonFields(inputOutputModelMatch);
+                for (const field of fieldsToFix) {
+                    source = replacePrismaJson(source, field);
+                }
+            }
+        }
+
+        return source;
+    }
+
+    private async generateExtraTypes(sf: SourceFile) {
+        for (const decl of this.model.declarations) {
+            if (isTypeDef(decl)) {
+                generateTypeDefType(sf, decl);
+            }
+        }
+    }
+
     private findNamedProperty(typeAlias: TypeAliasDeclaration, name: string) {
         return typeAlias.getFirstDescendant((d) => d.isKind(SyntaxKind.PropertySignature) && d.getName() === name);
     }
@@ -745,11 +789,12 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         return this.options.generatePermissionChecker === true;
     }
 
-    private async generateExtraTypes(sf: SourceFile) {
-        for (const decl of this.model.declarations) {
-            if (isTypeDef(decl)) {
-                generateTypeDefType(sf, decl);
-            }
-        }
+    private removeFromSource(source: string, text: string) {
+        source = source.replace(text, '');
+        return this.trimEmptyLines(source);
+    }
+
+    private trimEmptyLines(source: string): string {
+        return source.replace(/^\s*[\r\n]/gm, '');
     }
 }
diff --git a/packages/schema/src/plugins/prisma/schema-generator.ts b/packages/schema/src/plugins/prisma/schema-generator.ts
@@ -31,7 +31,7 @@ import {
     StringLiteral,
 } from '@zenstackhq/language/ast';
 import { getPrismaVersion } from '@zenstackhq/sdk/prisma';
-import { match, P } from 'ts-pattern';
+import { match } from 'ts-pattern';
 import { getIdFields } from '../../utils/ast-utils';
 
 import { DELEGATE_AUX_RELATION_PREFIX, PRISMA_MINIMUM_VERSION } from '@zenstackhq/runtime';
@@ -838,14 +838,6 @@ export class PrismaSchemaGenerator {
         const docs = [...field.comments, ...this.getCustomAttributesAsComments(field)];
         const result = model.addField(field.name, type, attributes, docs, addToFront);
 
-        if (this.mode === 'logical') {
-            if (field.attributes.some((attr) => isDefaultWithAuth(attr))) {
-                // field has `@default` with `auth()`, turn it into a dummy default value, and the
-                // real default value setting is handled outside Prisma
-                this.setDummyDefault(result, field);
-            }
-        }
-
         return result;
     }
 
@@ -856,23 +848,6 @@ export class PrismaSchemaGenerator {
         }
     }
 
-    private setDummyDefault(result: ModelField, field: DataModelField) {
-        const dummyDefaultValue = match(field.type.type)
-            .with('String', () => new AttributeArgValue('String', ''))
-            .with(P.union('Int', 'BigInt', 'Float', 'Decimal'), () => new AttributeArgValue('Number', '0'))
-            .with('Boolean', () => new AttributeArgValue('Boolean', 'false'))
-            .with('DateTime', () => new AttributeArgValue('FunctionCall', new PrismaFunctionCall('now')))
-            .with('Json', () => new AttributeArgValue('String', '{}'))
-            .with('Bytes', () => new AttributeArgValue('String', ''))
-            .otherwise(() => {
-                throw new PluginError(name, `Unsupported field type with default value: ${field.type.type}`);
-            });
-
-        result.attributes.push(
-            new PrismaFieldAttribute('@default', [new PrismaAttributeArg(undefined, dummyDefaultValue)])
-        );
-    }
-
     private isInheritedFromDelegate(field: DataModelField) {
         return field.$inheritedFrom && isDelegateModel(field.$inheritedFrom);
     }
diff --git a/tests/regression/tests/issue-1843.test.ts b/tests/regression/tests/issue-1843.test.ts
