From 8be5b35d5692cfae8be6a68d94471b8030814207 Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Mon, 11 Mar 2024 15:13:46 -0700
Subject: [PATCH] fix(polymorphism): wrong query result when selecting "_count"

Fixes #1123
---
 packages/runtime/src/enhancements/delegate.ts | 56 ++++++++++++-------
 .../with-delegate/issue-1123.test.ts          | 47 ++++++++++++++++
 2 files changed, 83 insertions(+), 20 deletions(-)
 create mode 100644 tests/integration/tests/enhancements/with-delegate/issue-1123.test.ts

diff --git a/packages/runtime/src/enhancements/delegate.ts b/packages/runtime/src/enhancements/delegate.ts
index dce180bd6..12c249912 100644
--- a/packages/runtime/src/enhancements/delegate.ts
+++ b/packages/runtime/src/enhancements/delegate.ts
@@ -1106,23 +1106,31 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
 
         const modelInfo = getModelInfo(this.options.modelMeta, model, true);
 
-        for (const field of Object.values(modelInfo.fields)) {
+        for (const [key, value] of Object.entries(entity)) {
+            if (key.startsWith(DELEGATE_AUX_RELATION_PREFIX)) {
+                continue;
+            }
+
+            const field = modelInfo.fields[key];
+            if (!field) {
+                // not a field, could be `_count`, `_sum`, etc.
+                result[key] = value;
+                continue;
+            }
+
             if (field.inheritedFrom) {
                 // already merged from base
                 continue;
             }
 
-            if (field.name in entity) {
-                const fieldValue = entity[field.name];
-                if (field.isDataModel) {
-                    if (Array.isArray(fieldValue)) {
-                        result[field.name] = fieldValue.map((item) => this.assembleUp(field.type, item));
-                    } else {
-                        result[field.name] = this.assembleUp(field.type, fieldValue);
-                    }
+            if (field.isDataModel) {
+                if (Array.isArray(value)) {
+                    result[field.name] = value.map((item) => this.assembleUp(field.type, item));
                 } else {
-                    result[field.name] = fieldValue;
+                    result[field.name] = this.assembleUp(field.type, value);
                 }
+            } else {
+                result[field.name] = value;
             }
         }
 
@@ -1147,18 +1155,26 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
             }
         }
 
-        for (const field of Object.values(modelInfo.fields)) {
-            if (field.name in entity) {
-                const fieldValue = entity[field.name];
-                if (field.isDataModel) {
-                    if (Array.isArray(fieldValue)) {
-                        result[field.name] = fieldValue.map((item) => this.assembleDown(field.type, item));
-                    } else {
-                        result[field.name] = this.assembleDown(field.type, fieldValue);
-                    }
+        for (const [key, value] of Object.entries(entity)) {
+            if (key.startsWith(DELEGATE_AUX_RELATION_PREFIX)) {
+                continue;
+            }
+
+            const field = modelInfo.fields[key];
+            if (!field) {
+                // not a field, could be `_count`, `_sum`, etc.
+                result[key] = value;
+                continue;
+            }
+
+            if (field.isDataModel) {
+                if (Array.isArray(value)) {
+                    result[field.name] = value.map((item) => this.assembleDown(field.type, item));
                 } else {
-                    result[field.name] = fieldValue;
+                    result[field.name] = this.assembleDown(field.type, value);
                 }
+            } else {
+                result[field.name] = value;
             }
         }
 
diff --git a/tests/integration/tests/enhancements/with-delegate/issue-1123.test.ts b/tests/integration/tests/enhancements/with-delegate/issue-1123.test.ts
new file mode 100644
index 000000000..02ee7a983
--- /dev/null
+++ b/tests/integration/tests/enhancements/with-delegate/issue-1123.test.ts
@@ -0,0 +1,47 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('Regression for issue 1123', () => {
+    it('regression', async () => {
+        const { enhance } = await loadSchema(
+            `
+            model Content {
+                id String @id @default(cuid())
+                published Boolean @default(false)
+                contentType String
+                likes Like[]
+                @@delegate(contentType)
+                @@allow('all', true)
+            }
+              
+            model Post extends Content {
+              title String
+            }
+            
+            model Image extends Content {
+              url String
+            }
+            
+            model Like {
+              id String @id @default(cuid())
+              content Content @relation(fields: [contentId], references: [id])
+              contentId String
+              @@allow('all', true)
+            }
+            `
+        );
+
+        const db = enhance();
+        await db.post.create({
+            data: {
+                title: 'a post',
+                likes: { create: {} },
+            },
+        });
+
+        await expect(db.content.findFirst({ include: { _count: { select: { likes: true } } } })).resolves.toMatchObject(
+            {
+                _count: { likes: 1 },
+            }
+        );
+    });
+});