feat: count and aggregate for delegate models

packages/runtime/src/client/crud/dialects/base.ts

@@ -818,20 +818,21 @@ export abstract class BaseCrudDialect<Schema extends SchemaDef> {
         return result;
     }
 
-    buildSelectField(query: SelectQueryBuilder<any, any, any>, model: string, modelAlias: string, field: string) {
+    buildSelectField(
+        query: SelectQueryBuilder<any, any, any>,
+        model: string,
+        modelAlias: string,
+        field: string,
+    ): SelectQueryBuilder<any, any, any> {
         const fieldDef = requireField(this.schema, model, field);
-
         if (fieldDef.computed) {
             // TODO: computed field from delegate base?
             return query.select((eb) => buildFieldRef(this.schema, model, field, this.options, eb).as(field));
         } else if (!fieldDef.originModel) {
             // regular field
             return query.select(sql.ref(`${modelAlias}.${field}`).as(field));
         } else {
-            // field from delegate base, build a join
-            let result = query;
-            result = this.buildSelectField(result, fieldDef.originModel, fieldDef.originModel, field);
-            return result;
+            return this.buildSelectField(query, fieldDef.originModel, fieldDef.originModel, field);
         }
     }
 

packages/runtime/src/client/crud/operations/aggregate.ts

@@ -1,3 +1,4 @@
+import type { ExpressionBuilder } from 'kysely';
 import { sql } from 'kysely';
 import { match } from 'ts-pattern';
 import type { SchemaDef } from '../../../schema';
@@ -15,12 +16,33 @@ export class AggregateOperationHandler<Schema extends SchemaDef> extends BaseOpe
         let query = this.kysely.selectFrom((eb) => {
             // nested query for filtering and pagination
 
-            // where
-            let subQuery = eb
-                .selectFrom(this.model)
-                .selectAll(this.model as any) // TODO: check typing
+            // table and where
+            let subQuery = this.dialect
+                .buildSelectModel(eb as ExpressionBuilder<any, any>, this.model)
                 .where((eb1) => this.dialect.buildFilter(eb1, this.model, this.model, parsedArgs?.where));
 
+            // select fields: collect fields from aggregation body
+            const selectedFields: string[] = [];
+            for (const [key, value] of Object.entries(parsedArgs)) {
+                if (key.startsWith('_') && value && typeof value === 'object') {
+                    // select fields
+                    Object.entries(value)
+                        .filter(([field]) => field !== '_all')
+                        .filter(([, val]) => val === true)
+                        .forEach(([field]) => {
+                            if (!selectedFields.includes(field)) selectedFields.push(field);
+                        });
+                }
+            }
+            if (selectedFields.length > 0) {
+                for (const field of selectedFields) {
+                    subQuery = this.dialect.buildSelectField(subQuery, this.model, this.model, field);
+                }
+            } else {
+                // if no field is explicitly selected, just do a `select 1` so `_count` works
+                subQuery = subQuery.select(() => eb.lit(1).as('_all'));
+            }
+
             // skip & take
             const skip = parsedArgs?.skip;
             let take = parsedArgs?.take;

packages/runtime/src/client/crud/operations/count.ts

@@ -1,3 +1,4 @@
+import type { ExpressionBuilder } from 'kysely';
 import { sql } from 'kysely';
 import type { SchemaDef } from '../../../schema';
 import { BaseOperationHandler } from './base';
@@ -9,15 +10,29 @@ export class CountOperationHandler<Schema extends SchemaDef> extends BaseOperati
 
         // parse args
         const parsedArgs = this.inputValidator.validateCountArgs(this.model, normalizedArgs);
+        const subQueryName = '$sub';
 
         let query = this.kysely.selectFrom((eb) => {
             // nested query for filtering and pagination
-            let subQuery = eb
-                .selectFrom(this.model)
-                .selectAll()
+
+            let subQuery = this.dialect
+                .buildSelectModel(eb as ExpressionBuilder<any, any>, this.model)
                 .where((eb1) => this.dialect.buildFilter(eb1, this.model, this.model, parsedArgs?.where));
+
+            if (parsedArgs?.select && typeof parsedArgs.select === 'object') {
+                // select fields
+                for (const [key, value] of Object.entries(parsedArgs.select)) {
+                    if (key !== '_all' && value === true) {
+                        subQuery = this.dialect.buildSelectField(subQuery, this.model, this.model, key);
+                    }
+                }
+            } else {
+                // no field selection, just build a `select 1`
+                subQuery = subQuery.select(() => eb.lit(1).as('_all'));
+            }
+
             subQuery = this.dialect.buildSkipTake(subQuery, parsedArgs?.skip, parsedArgs?.take);
-            return subQuery.as('$sub');
+            return subQuery.as(subQueryName);
         });
 
         if (parsedArgs?.select && typeof parsedArgs.select === 'object') {
@@ -26,7 +41,7 @@ export class CountOperationHandler<Schema extends SchemaDef> extends BaseOperati
                 Object.keys(parsedArgs.select!).map((key) =>
                     key === '_all'
                         ? eb.cast(eb.fn.countAll(), 'integer').as('_all')
-                        : eb.cast(eb.fn.count(sql.ref(`$sub.${key}`)), 'integer').as(key),
+                        : eb.cast(eb.fn.count(sql.ref(`${subQueryName}.${key}`)), 'integer').as(key),
                 ),
             );
 

packages/runtime/src/client/executor/name-mapper.ts

@@ -263,7 +263,7 @@ export class QueryNameMapper extends OperationNodeTransformer {
         model = model ?? this.currentModel;
         const modelDef = requireModel(this.schema, model!);
         const scalarFields = Object.entries(modelDef.fields)
-            .filter(([, fieldDef]) => !fieldDef.relation && !fieldDef.computed)
+            .filter(([, fieldDef]) => !fieldDef.relation && !fieldDef.computed && !fieldDef.originModel)
             .map(([fieldName]) => fieldName);
         return scalarFields;
     }

packages/runtime/test/client-api/delegate.test.ts

@@ -1070,5 +1070,146 @@ model Gallery {
                 await expect(client.asset.findMany()).toResolveWithLength(1);
             });
         });
+
+        describe('Delegate aggregation tests', () => {
+            beforeEach(async () => {
+                const u = await client.user.create({
+                    data: {
+                        id: 1,
+                        email: 'u1@example.com',
+                    },
+                });
+                await client.ratedVideo.create({
+                    data: {
+                        id: 1,
+                        viewCount: 0,
+                        duration: 100,
+                        url: 'v1',
+                        rating: 5,
+                        owner: { connect: { id: u.id } },
+                        user: { connect: { id: u.id } },
+                        comments: { create: [{ content: 'c1' }, { content: 'c2' }] },
+                    },
+                });
+                await client.ratedVideo.create({
+                    data: {
+                        id: 2,
+                        viewCount: 2,
+                        duration: 200,
+                        url: 'v2',
+                        rating: 3,
+                    },
+                });
+            });
+
+            it('works with count', async () => {
+                await expect(
+                    client.ratedVideo.count({
+                        where: { rating: 5 },
+                    }),
+                ).resolves.toEqual(1);
+                await expect(
+                    client.ratedVideo.count({
+                        where: { duration: 100 },
+                    }),
+                ).resolves.toEqual(1);
+                await expect(
+                    client.ratedVideo.count({
+                        where: { viewCount: 2 },
+                    }),
+                ).resolves.toEqual(1);
+
+                await expect(
+                    client.video.count({
+                        where: { duration: 100 },
+                    }),
+                ).resolves.toEqual(1);
+                await expect(
+                    client.asset.count({
+                        where: { viewCount: { gt: 0 } },
+                    }),
+                ).resolves.toEqual(1);
+
+                // field selection
+                await expect(
+                    client.ratedVideo.count({
+                        select: { _all: true, viewCount: true, url: true, rating: true },
+                    }),
+                ).resolves.toMatchObject({
+                    _all: 2,
+                    viewCount: 2,
+                    url: 2,
+                    rating: 2,
+                });
+                await expect(
+                    client.video.count({
+                        select: { _all: true, viewCount: true, url: true },
+                    }),
+                ).resolves.toMatchObject({
+                    _all: 2,
+                    viewCount: 2,
+                    url: 2,
+                });
+                await expect(
+                    client.asset.count({
+                        select: { _all: true, viewCount: true },
+                    }),
+                ).resolves.toMatchObject({
+                    _all: 2,
+                    viewCount: 2,
+                });
+            });
+
+            it('works with aggregate', async () => {
+                await expect(
+                    client.ratedVideo.aggregate({
+                        where: { viewCount: { gte: 0 }, duration: { gt: 0 }, rating: { gt: 0 } },
+                        _avg: { viewCount: true, duration: true, rating: true },
+                        _count: true,
+                    }),
+                ).resolves.toMatchObject({
+                    _avg: {
+                        viewCount: 1,
+                        duration: 150,
+                        rating: 4,
+                    },
+                    _count: 2,
+                });
+                await expect(
+                    client.video.aggregate({
+                        where: { viewCount: { gte: 0 }, duration: { gt: 0 } },
+                        _avg: { viewCount: true, duration: true },
+                        _count: true,
+                    }),
+                ).resolves.toMatchObject({
+                    _avg: {
+                        viewCount: 1,
+                        duration: 150,
+                    },
+                    _count: 2,
+                });
+                await expect(
+                    client.asset.aggregate({
+                        where: { viewCount: { gte: 0 } },
+                        _avg: { viewCount: true },
+                        _count: true,
+                    }),
+                ).resolves.toMatchObject({
+                    _avg: {
+                        viewCount: 1,
+                    },
+                    _count: 2,
+                });
+
+                // just count
+                await expect(
+                    client.ratedVideo.aggregate({
+                        _count: true,
+                    }),
+                ).resolves.toMatchObject({
+                    _count: 2,
+                });
+            });
+        });
     },
 );