fix(delegate): support _count select of base fields

packages/runtime/src/enhancements/node/delegate.ts

@@ -180,47 +180,102 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
             return;
         }
 
-        for (const kind of ['select', 'include'] as const) {
-            if (args[kind] && typeof args[kind] === 'object') {
-                for (const [field, value] of Object.entries<any>(args[kind])) {
-                    const fieldInfo = resolveField(this.options.modelMeta, model, field);
-                    if (!fieldInfo) {
-                        continue;
-                    }
+        // there're two cases where we need to inject polymorphic base hierarchy for fields
+        // defined in base models
+        // 1. base fields mentioned in select/include clause
+        //     { select: { fieldFromBase: true } } => { select: { delegate_aux_[Base]: { fieldFromBase: true } } }
+        // 2. base fields mentioned in _count select/include clause
+        //     { select: { _count: { select: { fieldFromBase: true } } } } => { select: { delegate_aux_[Base]: { select: { _count: { select: { fieldFromBase: true } } } } } }
+        //
+        // Note that although structurally similar, we need to correctly deal with different injection location of the "delegate_aux" hierarchy
+
+        // selectors for the above two cases
+        const selectors = [
+            // regular select: { select: { field: true } }
+            (payload: any) => ({ data: payload.select, kind: 'select' as const, isCount: false }),
+            // regular include: { include: { field: true } }
+            (payload: any) => ({ data: payload.include, kind: 'include' as const, isCount: false }),
+            // select _count: { select: { _count: { select: { field: true } } } }
+            (payload: any) => ({
+                data: payload.select?._count?.select,
+                kind: 'select' as const,
+                isCount: true,
+            }),
+            // include _count: { include: { _count: { select: { field: true } } } }
+            (payload: any) => ({
+                data: payload.include?._count?.select,
+                kind: 'include' as const,
+                isCount: true,
+            }),
+        ];
+
+        for (const selector of selectors) {
+            const { data, kind, isCount } = selector(args);
+            if (!data || typeof data !== 'object') {
+                continue;
+            }
 
-                    if (this.isDelegateOrDescendantOfDelegate(fieldInfo?.type) && value) {
-                        // delegate model, recursively inject hierarchy
-                        if (args[kind][field]) {
-                            if (args[kind][field] === true) {
-                                // make sure the payload is an object
-                                args[kind][field] = {};
-                            }
-                            await this.injectSelectIncludeHierarchy(fieldInfo.type, args[kind][field]);
+            for (const [field, value] of Object.entries<any>(data)) {
+                const fieldInfo = resolveField(this.options.modelMeta, model, field);
+                if (!fieldInfo) {
+                    continue;
+                }
+
+                if (this.isDelegateOrDescendantOfDelegate(fieldInfo?.type) && value) {
+                    // delegate model, recursively inject hierarchy
+                    if (data[field]) {
+                        if (data[field] === true) {
+                            // make sure the payload is an object
+                            data[field] = {};
                         }
+                        await this.injectSelectIncludeHierarchy(fieldInfo.type, data[field]);
                     }
+                }
 
-                    // refetch the field select/include value because it may have been
-                    // updated during injection
-                    const fieldValue = args[kind][field];
+                // refetch the field select/include value because it may have been
+                // updated during injection
+                const fieldValue = data[field];
 
-                    if (fieldValue !== undefined) {
-                        if (fieldValue.orderBy) {
-                            // `orderBy` may contain fields from base types
-                            enumerate(fieldValue.orderBy).forEach((item) =>
-                                this.injectWhereHierarchy(fieldInfo.type, item)
-                            );
-                        }
+                if (fieldValue !== undefined) {
+                    if (fieldValue.orderBy) {
+                        // `orderBy` may contain fields from base types
+                        enumerate(fieldValue.orderBy).forEach((item) =>
+                            this.injectWhereHierarchy(fieldInfo.type, item)
+                        );
+                    }
 
-                        if (this.injectBaseFieldSelect(model, field, fieldValue, args, kind)) {
-                            delete args[kind][field];
-                        } else if (fieldInfo.isDataModel) {
-                            let nextValue = fieldValue;
-                            if (nextValue === true) {
-                                // make sure the payload is an object
-                                args[kind][field] = nextValue = {};
+                    let injected = false;
+                    if (!isCount) {
+                        // regular select/include injection
+                        injected = await this.injectBaseFieldSelect(model, field, fieldValue, args, kind);
+                        if (injected) {
+                            // if injected, remove the field from the original payload
+                            delete data[field];
+                        }
+                    } else {
+                        // _count select/include injection, inject into an empty payload and then merge to the proper location
+                        const injectTarget = { [kind]: {} };
+                        injected = await this.injectBaseFieldSelect(model, field, fieldValue, injectTarget, kind, true);
+                        if (injected) {
+                            // if injected, remove the field from the original payload
+                            delete data[field];
+                            if (Object.keys(data).length === 0) {
+                                // if the original "_count" payload becomes empty, remove it
+                                delete args[kind]['_count'];
                             }
-                            await this.injectSelectIncludeHierarchy(fieldInfo.type, nextValue);
+                            // finally merge the injection into the original payload
+                            const merged = deepmerge(args[kind], injectTarget[kind]);
+                            args[kind] = merged;
+                        }
+                    }
+
+                    if (!injected && fieldInfo.isDataModel) {
+                        let nextValue = fieldValue;
+                        if (nextValue === true) {
+                            // make sure the payload is an object
+                            data[field] = nextValue = {};
                         }
+                        await this.injectSelectIncludeHierarchy(fieldInfo.type, nextValue);
                     }
                 }
             }
@@ -272,7 +327,8 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
         field: string,
         value: any,
         selectInclude: any,
-        context: 'select' | 'include'
+        context: 'select' | 'include',
+        forCount = false // if the injection is for a "{ _count: { select: { field: true } } }" payload
     ) {
         const fieldInfo = resolveField(this.options.modelMeta, model, field);
         if (!fieldInfo?.inheritedFrom) {
@@ -286,24 +342,35 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
             const baseRelationName = this.makeAuxRelationName(base);
 
             // prepare base layer select/include
-            // let selectOrInclude = 'select';
             let thisLayer: any;
             if (target.include) {
-                // selectOrInclude = 'include';
                 thisLayer = target.include;
             } else if (target.select) {
-                // selectOrInclude = 'select';
                 thisLayer = target.select;
             } else {
-                // selectInclude = 'include';
                 thisLayer = target.select = {};
             }
 
             if (base.name === fieldInfo.inheritedFrom) {
                 if (!thisLayer[baseRelationName]) {
                     thisLayer[baseRelationName] = { [context]: {} };
                 }
-                thisLayer[baseRelationName][context][field] = value;
+                if (forCount) {
+                    // { _count: { select: { field: true } } } => { delegate_aux_[Base]: { select: { _count: { select: { field: true } } } } }
+                    if (
+                        !thisLayer[baseRelationName][context]['_count'] ||
+                        typeof thisLayer[baseRelationName][context] !== 'object'
+                    ) {
+                        thisLayer[baseRelationName][context]['_count'] = {};
+                    }
+                    thisLayer[baseRelationName][context]['_count'] = deepmerge(
+                        thisLayer[baseRelationName][context]['_count'],
+                        { select: { [field]: value } }
+                    );
+                } else {
+                    // { select: { field: true } } => { delegate_aux_[Base]: { select: { field: true } } }
+                    thisLayer[baseRelationName][context][field] = value;
+                }
                 break;
             } else {
                 if (!thisLayer[baseRelationName]) {

tests/integration/tests/enhancements/with-delegate/enhanced-client.test.ts

@@ -378,6 +378,66 @@ describe('Polymorphism Test', () => {
         ).resolves.toHaveLength(1);
     });
 
+    it('read with counting relation defined in base', async () => {
+        const { enhance } = await loadSchema(
+            `
+
+        model A {
+            id Int @id @default(autoincrement())
+            type String
+            bs B[]
+            cs C[]
+            @@delegate(type)
+        }
+
+        model A1 extends A {
+            a1 Int
+            type1 String
+            @@delegate(type1)
+        }
+
+        model A2 extends A1 {
+            a2 Int
+        }
+
+        model B {
+            id Int @id @default(autoincrement())
+            a A @relation(fields: [aId], references: [id])
+            aId Int
+            b Int
+        }
+
+        model C {
+            id Int @id @default(autoincrement())
+            a A @relation(fields: [aId], references: [id])
+            aId Int
+            c Int
+        }
+            `,
+            { enhancements: ['delegate'] }
+        );
+        const db = enhance();
+
+        const a2 = await db.a2.create({
+            data: { a1: 1, a2: 2, bs: { create: [{ b: 1 }, { b: 2 }] }, cs: { create: [{ c: 1 }] } },
+            include: { _count: { select: { bs: true } } },
+        });
+        expect(a2).toMatchObject({ a1: 1, a2: 2, _count: { bs: 2 } });
+
+        await expect(
+            db.a2.findFirst({ select: { a1: true, _count: { select: { bs: true } } } })
+        ).resolves.toStrictEqual({
+            a1: 1,
+            _count: { bs: 2 },
+        });
+
+        await expect(db.a.findFirst({ select: { _count: { select: { bs: true, cs: true } } } })).resolves.toMatchObject(
+            {
+                _count: { bs: 2, cs: 1 },
+            }
+        );
+    });
+
     it('order by base fields', async () => {
         const { db, user } = await setup();
 

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

@@ -0,0 +1,51 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 1467', () => {
+    it('regression', async () => {
+        const { enhance } = await loadSchema(
+            `
+    model User {
+        id   Int    @id @default(autoincrement())
+        type String
+        @@allow('all', true)
+    }
+
+    model Container {
+        id    Int    @id @default(autoincrement())
+        drink Drink @relation(fields: [drinkId], references: [id])
+        drinkId Int
+        @@allow('all', true)
+    }
+
+    model Drink {
+        id                Int   @id @default(autoincrement())
+        name              String @unique
+        containers        Container[]
+        type              String
+
+        @@delegate(type)
+        @@allow('all', true)
+    }
+
+    model Beer extends Drink {
+        @@allow('all', true)
+    }
+            `
+        );
+
+        const db = enhance();
+
+        await db.beer.create({
+            data: { id: 1, name: 'Beer1' },
+        });
+
+        await db.container.create({ data: { drink: { connect: { id: 1 } } } });
+        await db.container.create({ data: { drink: { connect: { id: 1 } } } });
+
+        const beers = await db.beer.findFirst({
+            select: { id: true, name: true, _count: { select: { containers: true } } },
+            orderBy: { name: 'asc' },
+        });
+        expect(beers).toMatchObject({ _count: { containers: 2 } });
+    });
+});