From f5007f063b638a41d52c65ac2e5397d9a85ff23b Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Sun, 11 Feb 2024 23:04:09 +0800
Subject: [PATCH] fix: @omit doesn't remove fields inside to-many relation

---
 packages/runtime/src/enhancements/omit.ts     | 15 ++++--
 .../enhancements/with-omit/with-omit.test.ts  | 48 +++++++++++++++++++
 .../tests/regression/issue-992.test.ts        | 45 +++++++++++++++++
 3 files changed, 104 insertions(+), 4 deletions(-)
 create mode 100644 tests/integration/tests/regression/issue-992.test.ts

diff --git a/packages/runtime/src/enhancements/omit.ts b/packages/runtime/src/enhancements/omit.ts
index 8b2937845..2df81b40a 100644
--- a/packages/runtime/src/enhancements/omit.ts
+++ b/packages/runtime/src/enhancements/omit.ts
@@ -54,11 +54,18 @@ class OmitHandler extends DefaultPrismaProxyHandler {
                 continue;
             }
 
-            if (fieldInfo.attributes?.find((attr) => attr.name === '@omit')) {
+            const shouldOmit = fieldInfo.attributes?.find((attr) => attr.name === '@omit');
+            if (shouldOmit) {
                 delete entityData[field];
-            } else if (fieldInfo.isDataModel) {
-                // recurse
-                await this.doPostProcess(entityData[field], fieldInfo.type);
+            }
+
+            if (fieldInfo.isDataModel) {
+                const items =
+                    fieldInfo.isArray && Array.isArray(entityData[field]) ? entityData[field] : [entityData[field]];
+                for (const item of items) {
+                    // recurse
+                    await this.doPostProcess(item, fieldInfo.type);
+                }
             }
         }
     }
diff --git a/tests/integration/tests/enhancements/with-omit/with-omit.test.ts b/tests/integration/tests/enhancements/with-omit/with-omit.test.ts
index 61d44b440..e2db0a3f8 100644
--- a/tests/integration/tests/enhancements/with-omit/with-omit.test.ts
+++ b/tests/integration/tests/enhancements/with-omit/with-omit.test.ts
@@ -105,4 +105,52 @@ describe('Omit test', () => {
         expect(r1.password).toBeUndefined();
         expect(r1.profile.image).toBeUndefined();
     });
+
+    it('to-many', async () => {
+        const { withOmit } = await loadSchema(
+            `
+            model User {
+                id String @id @default(cuid())
+                posts Post[]
+            
+                @@allow('all', true)
+            }
+            
+            model Post {
+                id String @id @default(cuid())
+                user User @relation(fields: [userId], references: [id])
+                userId String
+                images Image[]
+            
+                @@allow('all', true)
+            }
+
+            model Image {
+                id String @id @default(cuid())
+                post Post @relation(fields: [postId], references: [id])
+                postId String
+                url String @omit
+
+                @@allow('all', true)
+            }
+            `
+        );
+
+        const db = withOmit();
+        const r = await db.user.create({
+            include: { posts: { include: { images: true } } },
+            data: {
+                posts: {
+                    create: [
+                        { images: { create: { url: 'img1' } } },
+                        { images: { create: [{ url: 'img2' }, { url: 'img3' }] } },
+                    ],
+                },
+            },
+        });
+
+        expect(r.posts[0].images[0].url).toBeUndefined();
+        expect(r.posts[1].images[0].url).toBeUndefined();
+        expect(r.posts[1].images[1].url).toBeUndefined();
+    });
 });
diff --git a/tests/integration/tests/regression/issue-992.test.ts b/tests/integration/tests/regression/issue-992.test.ts
new file mode 100644
index 000000000..40a1aac47
--- /dev/null
+++ b/tests/integration/tests/regression/issue-992.test.ts
@@ -0,0 +1,45 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('Regression: issue 992', () => {
+    it('regression', async () => {
+        const { enhance, prisma } = await loadSchema(
+            `
+            model Product {
+                id String @id @default(cuid())
+                category Category @relation(fields: [categoryId], references: [id])
+                categoryId String
+              
+                deleted Int @default(0) @omit
+                @@deny('read', deleted != 0)
+                @@allow('all', true)
+            }
+              
+            model Category {
+                id String @id @default(cuid())
+                products Product[]
+                @@allow('all', true)
+            }
+            `
+        );
+
+        await prisma.category.create({
+            data: {
+                products: {
+                    create: [
+                        {
+                            deleted: 0,
+                        },
+                        {
+                            deleted: 0,
+                        },
+                    ],
+                },
+            },
+        });
+
+        const db = enhance();
+        const category = await db.category.findFirst({ include: { products: true } });
+        expect(category.products[0].deleted).toBeUndefined();
+        expect(category.products[1].deleted).toBeUndefined();
+    });
+});