fix(delegate): clean up generated zod schemas for delegate auxiliary fields

packages/schema/src/plugins/zod/generator.ts

@@ -1,3 +1,4 @@
+import { DELEGATE_AUX_RELATION_PREFIX } from '@zenstackhq/runtime';
 import {
     ExpressionContext,
     PluginError,
@@ -88,12 +89,18 @@ export class ZodSchemaGenerator {
             (o) => !excludeModels.find((e) => e === o.model)
         );
 
-        // TODO: better way of filtering than string startsWith?
         const inputObjectTypes = prismaClientDmmf.schema.inputObjectTypes.prisma.filter(
-            (type) => !excludeModels.find((e) => type.name.toLowerCase().startsWith(e.toLocaleLowerCase()))
+            (type) =>
+                !excludeModels.some((e) => type.name.toLowerCase().startsWith(e.toLocaleLowerCase())) &&
+                // exclude delegate aux related types
+                !type.name.toLowerCase().includes(DELEGATE_AUX_RELATION_PREFIX)
         );
+
         const outputObjectTypes = prismaClientDmmf.schema.outputObjectTypes.prisma.filter(
-            (type) => !excludeModels.find((e) => type.name.toLowerCase().startsWith(e.toLowerCase()))
+            (type) =>
+                !excludeModels.some((e) => type.name.toLowerCase().startsWith(e.toLowerCase())) &&
+                // exclude delegate aux related types
+                !type.name.toLowerCase().includes(DELEGATE_AUX_RELATION_PREFIX)
         );
 
         const models: DMMF.Model[] = prismaClientDmmf.datamodel.models.filter(
@@ -236,7 +243,8 @@ export class ZodSchemaGenerator {
 
         const moduleNames: string[] = [];
         for (let i = 0; i < inputObjectTypes.length; i += 1) {
-            const fields = inputObjectTypes[i]?.fields;
+            // exclude delegate aux fields
+            const fields = inputObjectTypes[i]?.fields?.filter((f) => !f.name.startsWith(DELEGATE_AUX_RELATION_PREFIX));
             const name = inputObjectTypes[i]?.name;
 
             if (!generateUnchecked && name.includes('Unchecked')) {

packages/schema/src/plugins/zod/transformer.ts

@@ -1,8 +1,11 @@
 /* eslint-disable @typescript-eslint/ban-ts-comment */
+import { DELEGATE_AUX_RELATION_PREFIX } from '@zenstackhq/runtime';
 import {
     getForeignKeyFields,
+    getRelationBackLink,
     hasAttribute,
     indentString,
+    isDelegateModel,
     isDiscriminatorField,
     type PluginOptions,
 } from '@zenstackhq/sdk';
@@ -67,7 +70,11 @@ export default class Transformer {
             const filePath = path.join(Transformer.outputPath, `enums/${name}.schema.ts`);
             const content = `${this.generateImportZodStatement()}\n${this.generateExportSchemaStatement(
                 `${name}`,
-                `z.enum(${JSON.stringify(enumType.values)})`
+                `z.enum(${JSON.stringify(
+                    enumType.values
+                        // exclude fields generated for delegate models
+                        .filter((v) => !v.startsWith(DELEGATE_AUX_RELATION_PREFIX))
+                )})`
             )}`;
             this.sourceFiles.push(this.project.createSourceFile(filePath, content, { overwrite: true }));
             generated.push(enumType.name);
@@ -243,12 +250,19 @@ export default class Transformer {
                         !isFieldRef &&
                         (inputType.namespace === 'prisma' || isEnum)
                     ) {
-                        if (inputType.type !== this.originalName && typeof inputType.type === 'string') {
-                            this.addSchemaImport(inputType.type);
+                        // reduce concrete input types to their delegate base types
+                        // e.g.: "UserCreateNestedOneWithoutDelegate_aux_PostInput" => "UserCreateWithoutAssetInput"
+                        let mappedInputType = inputType;
+                        if (contextDataModel) {
+                            mappedInputType = this.mapDelegateInputType(inputType, contextDataModel, field.name);
+                        }
+
+                        if (mappedInputType.type !== this.originalName && typeof mappedInputType.type === 'string') {
+                            this.addSchemaImport(mappedInputType.type);
                         }
 
                         const contextField = contextDataModel?.fields.find((f) => f.name === field.name);
-                        result.push(this.generatePrismaStringLine(field, inputType, lines.length, contextField));
+                        result.push(this.generatePrismaStringLine(field, mappedInputType, lines.length, contextField));
                     }
                 }
 
@@ -289,6 +303,46 @@ export default class Transformer {
         return [[`  ${fieldName} ${resString} `, field, true]];
     }
 
+    private mapDelegateInputType(
+        inputType: PrismaDMMF.InputTypeRef,
+        contextDataModel: DataModel,
+        contextFieldName: string
+    ) {
+        // input type mapping is only relevant for relation inherited from delegate models
+        const contextField = contextDataModel.fields.find((f) => f.name === contextFieldName);
+        if (!contextField || !isDataModel(contextField.type.reference?.ref)) {
+            return inputType;
+        }
+
+        if (!contextField.$inheritedFrom || !isDelegateModel(contextField.$inheritedFrom)) {
+            return inputType;
+        }
+
+        let processedInputType = inputType;
+
+        // captures: model name and operation, "Without" part that references a concrete model,
+        // and the "Input" or "NestedInput" suffix
+        const match = inputType.type.match(/^(\S+?)((NestedOne)?WithoutDelegate_aux\S+?)((Nested)?Input)$/);
+        if (match) {
+            let mappedInputTypeName = match[1];
+
+            if (contextDataModel) {
+                // get the opposite side of the relation field, which should be of the proper
+                // delegate base type
+                const oppositeRelationField = getRelationBackLink(contextField);
+                if (oppositeRelationField) {
+                    mappedInputTypeName += `Without${upperCaseFirst(oppositeRelationField.name)}`;
+                }
+            }
+
+            // "Input" or "NestedInput" suffix
+            mappedInputTypeName += match[4];
+
+            processedInputType = { ...inputType, type: mappedInputTypeName };
+        }
+        return processedInputType;
+    }
+
     wrapWithZodValidators(
         mainValidators: string | string[],
         field: PrismaDMMF.SchemaArg,

packages/sdk/src/model-meta-generator.ts

@@ -24,16 +24,15 @@ import {
     ExpressionContext,
     getAttribute,
     getAttributeArg,
-    getAttributeArgLiteral,
     getAttributeArgs,
     getAuthDecl,
     getDataModels,
     getInheritedFromDelegate,
     getLiteral,
+    getRelationBackLink,
     getRelationField,
     hasAttribute,
     isAuthInvocation,
-    isDelegateModel,
     isEnumFieldReference,
     isForeignKeyField,
     isIdField,
@@ -289,7 +288,7 @@ function writeFields(
             if (dmField) {
                 // metadata specific to DataModelField
 
-                const backlink = getBackLink(dmField);
+                const backlink = getRelationBackLink(dmField);
                 const fkMapping = generateForeignKeyMapping(dmField);
 
                 if (backlink) {
@@ -336,51 +335,6 @@ function writeFields(
     writer.write(',');
 }
 
-function getBackLink(field: DataModelField) {
-    if (!field.type.reference?.ref || !isDataModel(field.type.reference?.ref)) {
-        return undefined;
-    }
-
-    const relName = getRelationName(field);
-
-    let sourceModel: DataModel;
-    if (field.$inheritedFrom && isDelegateModel(field.$inheritedFrom)) {
-        // field is inherited from a delegate model, use it as the source
-        sourceModel = field.$inheritedFrom;
-    } else {
-        // otherwise use the field's container model as the source
-        sourceModel = field.$container as DataModel;
-    }
-
-    const targetModel = field.type.reference.ref as DataModel;
-
-    for (const otherField of targetModel.fields) {
-        if (otherField === field) {
-            // backlink field is never self
-            continue;
-        }
-        if (otherField.type.reference?.ref === sourceModel) {
-            if (relName) {
-                const otherRelName = getRelationName(otherField);
-                if (relName === otherRelName) {
-                    return otherField;
-                }
-            } else {
-                return otherField;
-            }
-        }
-    }
-    return undefined;
-}
-
-function getRelationName(field: DataModelField) {
-    const relAttr = getAttribute(field, '@relation');
-    if (!relAttr) {
-        return undefined;
-    }
-    return getAttributeArgLiteral(relAttr, 'name');
-}
-
 function getAttributes(target: DataModelField | DataModel | TypeDefField): RuntimeAttribute[] {
     return target.attributes
         .map((attr) => {

packages/sdk/src/utils.ts

@@ -632,3 +632,54 @@ export function getInheritanceChain(from: DataModel, to: DataModel): DataModel[]
 
     return undefined;
 }
+
+/**
+ * Get the opposite side of a relation field.
+ */
+export function getRelationBackLink(field: DataModelField) {
+    if (!field.type.reference?.ref || !isDataModel(field.type.reference?.ref)) {
+        return undefined;
+    }
+
+    const relName = getRelationName(field);
+
+    let sourceModel: DataModel;
+    if (field.$inheritedFrom && isDelegateModel(field.$inheritedFrom)) {
+        // field is inherited from a delegate model, use it as the source
+        sourceModel = field.$inheritedFrom;
+    } else {
+        // otherwise use the field's container model as the source
+        sourceModel = field.$container as DataModel;
+    }
+
+    const targetModel = field.type.reference.ref as DataModel;
+
+    for (const otherField of targetModel.fields) {
+        if (otherField === field) {
+            // backlink field is never self
+            continue;
+        }
+        if (otherField.type.reference?.ref === sourceModel) {
+            if (relName) {
+                const otherRelName = getRelationName(otherField);
+                if (relName === otherRelName) {
+                    return otherField;
+                }
+            } else {
+                return otherField;
+            }
+        }
+    }
+    return undefined;
+}
+
+/**
+ * Get the relation name of a relation field.
+ */
+export function getRelationName(field: DataModelField) {
+    const relAttr = getAttribute(field, '@relation');
+    if (!relAttr) {
+        return undefined;
+    }
+    return getAttributeArgLiteral(relAttr, 'name');
+}

tests/regression/tests/issue-1993.test.ts

@@ -0,0 +1,63 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 1993', () => {
+    it('regression', async () => {
+        const { zodSchemas } = await loadSchema(
+            `
+enum UserType {
+    UserLocal
+    UserGoogle
+}
+
+model User {
+    id          String       @id @default(cuid())
+    companyId   String?
+    type        UserType
+
+    @@delegate(type)
+
+    userFolders UserFolder[]
+
+    @@allow('all', true)
+}
+
+model UserLocal extends User {
+    email    String
+    password String
+}
+
+model UserGoogle extends User {
+    googleId String
+}
+
+model UserFolder {
+    id     String @id @default(cuid())
+    userId String
+    path   String
+
+    user   User   @relation(fields: [userId], references: [id])
+
+    @@allow('all', true)
+}            `,
+            { pushDb: false, fullZod: true, compile: true, output: 'lib/zenstack' }
+        );
+
+        expect(
+            zodSchemas.input.UserLocalInputSchema.create.safeParse({
+                data: {
+                    email: '[email protected]',
+                    password: 'password',
+                },
+            })
+        ).toMatchObject({ success: true });
+
+        expect(
+            zodSchemas.input.UserFolderInputSchema.create.safeParse({
+                data: {
+                    path: '/',
+                    userId: '1',
+                },
+            })
+        ).toMatchObject({ success: true });
+    });
+});