From b8059ee10f3f693dd72cb43179a769f262ef8056 Mon Sep 17 00:00:00 2001
From: Dylan Lundy <4567380+diesal11@users.noreply.github.com>
Date: Thu, 6 Mar 2025 17:31:38 +1030
Subject: [PATCH] fix(zod): delegate aux fields appearing in WhereUniqueInput
 schemas

---
 .../src/plugins/enhancer/enhance/index.ts     |   6 +-
 packages/schema/src/plugins/zod/generator.ts  |   2 +-
 tests/regression/tests/issue-2028.test.ts     | 116 ++++++++++++++++++
 3 files changed, 120 insertions(+), 4 deletions(-)
 create mode 100644 tests/regression/tests/issue-2028.test.ts

diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts
index 1cd49c0a4..c0331e61e 100644
--- a/packages/schema/src/plugins/enhancer/enhance/index.ts
+++ b/packages/schema/src/plugins/enhancer/enhance/index.ts
@@ -565,10 +565,10 @@ export type Enhanced<Client> =
         const structure = iface.getStructure();
 
         // filter out aux fields
-        structure.properties = structure.properties?.filter((p) => !p.name.startsWith(DELEGATE_AUX_RELATION_PREFIX));
+        structure.properties = structure.properties?.filter((p) => !p.name.includes(DELEGATE_AUX_RELATION_PREFIX));
 
         // filter out aux methods
-        structure.methods = structure.methods?.filter((m) => !m.name.startsWith(DELEGATE_AUX_RELATION_PREFIX));
+        structure.methods = structure.methods?.filter((m) => !m.name.includes(DELEGATE_AUX_RELATION_PREFIX));
 
         if (delegateInfo.some(([delegate]) => `${delegate.name}Delegate` === iface.getName())) {
             // delegate models cannot be created directly, remove create/createMany/upsert
@@ -850,7 +850,7 @@ export type Enhanced<Client> =
     private findAuxDecls(node: Node) {
         return node
             .getDescendantsOfKind(SyntaxKind.PropertySignature)
-            .filter((n) => n.getName().startsWith(DELEGATE_AUX_RELATION_PREFIX));
+            .filter((n) => n.getName().includes(DELEGATE_AUX_RELATION_PREFIX));
     }
 
     private saveSourceFile(sf: SourceFile) {
diff --git a/packages/schema/src/plugins/zod/generator.ts b/packages/schema/src/plugins/zod/generator.ts
index 341b8cae5..9a5cd53b4 100644
--- a/packages/schema/src/plugins/zod/generator.ts
+++ b/packages/schema/src/plugins/zod/generator.ts
@@ -244,7 +244,7 @@ export class ZodSchemaGenerator {
         const moduleNames: string[] = [];
         for (let i = 0; i < inputObjectTypes.length; i += 1) {
             // exclude delegate aux fields
-            const fields = inputObjectTypes[i]?.fields?.filter((f) => !f.name.startsWith(DELEGATE_AUX_RELATION_PREFIX));
+            const fields = inputObjectTypes[i]?.fields?.filter((f) => !f.name.includes(DELEGATE_AUX_RELATION_PREFIX));
             const name = inputObjectTypes[i]?.name;
 
             if (!generateUnchecked && name.includes('Unchecked')) {
diff --git a/tests/regression/tests/issue-2028.test.ts b/tests/regression/tests/issue-2028.test.ts
new file mode 100644
index 000000000..a98f0af7d
--- /dev/null
+++ b/tests/regression/tests/issue-2028.test.ts
@@ -0,0 +1,116 @@
+import { createPostgresDb, loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 2028', () => {
+    it('regression', async () => {
+        const dbUrl = await createPostgresDb('issue-2028');
+        const { enhance, zodSchemas } = await loadSchema(
+            `
+enum FooType {
+    Bar
+    Baz
+}
+
+model User {
+    id          String       @id @default(cuid())
+    userFolders UserFolder[]
+    @@allow('all', true)
+}
+
+model Foo {
+    id          String       @id @default(cuid())
+    type        FooType
+
+    userFolders UserFolder[]
+
+    @@delegate(type)
+    @@allow('all', true)
+}
+
+model Bar extends Foo {
+    name String
+}
+
+model Baz extends Foo {
+    age Int
+}
+
+model UserFolder {
+    id     String @id @default(cuid())
+    userId String
+    fooId  String
+
+    user   User   @relation(fields: [userId], references: [id])
+    foo    Foo    @relation(fields: [fooId], references: [id])
+
+    @@unique([userId, fooId])
+    @@allow('all', true)
+}
+            `,
+            {
+                fullZod: true,
+                provider: 'postgresql',
+                dbUrl,
+            }
+        );
+        // Ensure Zod Schemas don't include the delegate fields
+        expect(
+            zodSchemas.objects.UserFolderWhereUniqueInputObjectSchema.safeParse({
+                userId_delegate_aux_UserFolder_fooId_Bar: {
+                    userId: '1',
+                    fooId: '2',
+                },
+            }).success
+        ).toBeFalsy();
+
+        expect(
+            zodSchemas.objects.UserFolderWhereUniqueInputObjectSchema.safeParse({
+                userId_delegate_aux_UserFolder_fooId_Baz: {
+                    userId: '1',
+                    fooId: '2',
+                },
+            }).success
+        ).toBeFalsy();
+
+        // Ensure we can query by the CompoundUniqueInput
+        const db = enhance();
+        const user = await db.user.create({ data: {} });
+        const bar = await db.bar.create({ data: { name: 'bar' } });
+        const baz = await db.baz.create({ data: { age: 1 } });
+
+        const userFolderA = await db.userFolder.create({
+            data: {
+                userId: user.id,
+                fooId: bar.id,
+            },
+        });
+
+        const userFolderB = await db.userFolder.create({
+            data: {
+                userId: user.id,
+                fooId: baz.id,
+            },
+        });
+
+        await expect(
+            db.userFolder.findUnique({
+                where: {
+                    userId_fooId: {
+                        userId: user.id,
+                        fooId: bar.id,
+                    },
+                },
+            })
+        ).resolves.toMatchObject(userFolderA);
+
+        await expect(
+            db.userFolder.findUnique({
+                where: {
+                    userId_fooId: {
+                        userId: user.id,
+                        fooId: baz.id,
+                    },
+                },
+            })
+        ).resolves.toMatchObject(userFolderB);
+    });
+});