fix(delegate): support _count select of base fields

Author: ymc9

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

@@ -180,47 +180,82 @@ 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;
-                    }
+        const selectors = [
+            (payload: any) => ({ data: payload.select, kind: 'select' as const, isCount: false }),
+            (payload: any) => ({ data: payload.include, kind: 'include' as const, isCount: false }),
+            (payload: any) => ({
+                data: payload.select?._count?.select,
+                kind: 'select' as const,
+                isCount: 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) {
+                        injected = await this.injectBaseFieldSelect(model, field, fieldValue, args, kind);
+                        if (injected) {
+                            delete data[field];
+                        }
+                    } else {
+                        const injectTarget = { [kind]: {} };
+                        injected = await this.injectBaseFieldSelect(model, field, fieldValue, injectTarget, kind, true);
+                        if (injected) {
+                            delete data[field];
+                            if (Object.keys(data).length === 0) {
+                                delete args[kind]['_count'];
                             }
-                            await this.injectSelectIncludeHierarchy(fieldInfo.type, nextValue);
+                            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 +307,8 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
         field: string,
         value: any,
         selectInclude: any,
-        context: 'select' | 'include'
+        context: 'select' | 'include',
+        forCount = false
     ) {
         const fieldInfo = resolveField(this.options.modelMeta, model, field);
         if (!fieldInfo?.inheritedFrom) {
@@ -286,24 +322,33 @@ 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) {
+                    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 {
+                    thisLayer[baseRelationName][context][field] = value;
+                }
                 break;
             } else {
                 if (!thisLayer[baseRelationName]) {

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 } });
+    });
+});