Skip to content

Commit bcbfb9a

Browse files
ymc9ymc9
authored
fix(delegate): support _count select of base fields (#1937)
1 parent 1956bdb commit bcbfb9a

File tree

3 files changed

+217
-39
lines changed

3 files changed

+217
-39
lines changed

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

Lines changed: 106 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -180,47 +180,102 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
180180
return;
181181
}
182182

183-
for (const kind of ['select', 'include'] as const) {
184-
if (args[kind] && typeof args[kind] === 'object') {
185-
for (const [field, value] of Object.entries<any>(args[kind])) {
186-
const fieldInfo = resolveField(this.options.modelMeta, model, field);
187-
if (!fieldInfo) {
188-
continue;
189-
}
183+
// there're two cases where we need to inject polymorphic base hierarchy for fields
184+
// defined in base models
185+
// 1. base fields mentioned in select/include clause
186+
// { select: { fieldFromBase: true } } => { select: { delegate_aux_[Base]: { fieldFromBase: true } } }
187+
// 2. base fields mentioned in _count select/include clause
188+
// { select: { _count: { select: { fieldFromBase: true } } } } => { select: { delegate_aux_[Base]: { select: { _count: { select: { fieldFromBase: true } } } } } }
189+
//
190+
// Note that although structurally similar, we need to correctly deal with different injection location of the "delegate_aux" hierarchy
191+
192+
// selectors for the above two cases
193+
const selectors = [
194+
// regular select: { select: { field: true } }
195+
(payload: any) => ({ data: payload.select, kind: 'select' as const, isCount: false }),
196+
// regular include: { include: { field: true } }
197+
(payload: any) => ({ data: payload.include, kind: 'include' as const, isCount: false }),
198+
// select _count: { select: { _count: { select: { field: true } } } }
199+
(payload: any) => ({
200+
data: payload.select?._count?.select,
201+
kind: 'select' as const,
202+
isCount: true,
203+
}),
204+
// include _count: { include: { _count: { select: { field: true } } } }
205+
(payload: any) => ({
206+
data: payload.include?._count?.select,
207+
kind: 'include' as const,
208+
isCount: true,
209+
}),
210+
];
211+
212+
for (const selector of selectors) {
213+
const { data, kind, isCount } = selector(args);
214+
if (!data || typeof data !== 'object') {
215+
continue;
216+
}
190217

191-
if (this.isDelegateOrDescendantOfDelegate(fieldInfo?.type) && value) {
192-
// delegate model, recursively inject hierarchy
193-
if (args[kind][field]) {
194-
if (args[kind][field] === true) {
195-
// make sure the payload is an object
196-
args[kind][field] = {};
197-
}
198-
await this.injectSelectIncludeHierarchy(fieldInfo.type, args[kind][field]);
218+
for (const [field, value] of Object.entries<any>(data)) {
219+
const fieldInfo = resolveField(this.options.modelMeta, model, field);
220+
if (!fieldInfo) {
221+
continue;
222+
}
223+
224+
if (this.isDelegateOrDescendantOfDelegate(fieldInfo?.type) && value) {
225+
// delegate model, recursively inject hierarchy
226+
if (data[field]) {
227+
if (data[field] === true) {
228+
// make sure the payload is an object
229+
data[field] = {};
199230
}
231+
await this.injectSelectIncludeHierarchy(fieldInfo.type, data[field]);
200232
}
233+
}
201234

202-
// refetch the field select/include value because it may have been
203-
// updated during injection
204-
const fieldValue = args[kind][field];
235+
// refetch the field select/include value because it may have been
236+
// updated during injection
237+
const fieldValue = data[field];
205238

206-
if (fieldValue !== undefined) {
207-
if (fieldValue.orderBy) {
208-
// `orderBy` may contain fields from base types
209-
enumerate(fieldValue.orderBy).forEach((item) =>
210-
this.injectWhereHierarchy(fieldInfo.type, item)
211-
);
212-
}
239+
if (fieldValue !== undefined) {
240+
if (fieldValue.orderBy) {
241+
// `orderBy` may contain fields from base types
242+
enumerate(fieldValue.orderBy).forEach((item) =>
243+
this.injectWhereHierarchy(fieldInfo.type, item)
244+
);
245+
}
213246

214-
if (this.injectBaseFieldSelect(model, field, fieldValue, args, kind)) {
215-
delete args[kind][field];
216-
} else if (fieldInfo.isDataModel) {
217-
let nextValue = fieldValue;
218-
if (nextValue === true) {
219-
// make sure the payload is an object
220-
args[kind][field] = nextValue = {};
247+
let injected = false;
248+
if (!isCount) {
249+
// regular select/include injection
250+
injected = await this.injectBaseFieldSelect(model, field, fieldValue, args, kind);
251+
if (injected) {
252+
// if injected, remove the field from the original payload
253+
delete data[field];
254+
}
255+
} else {
256+
// _count select/include injection, inject into an empty payload and then merge to the proper location
257+
const injectTarget = { [kind]: {} };
258+
injected = await this.injectBaseFieldSelect(model, field, fieldValue, injectTarget, kind, true);
259+
if (injected) {
260+
// if injected, remove the field from the original payload
261+
delete data[field];
262+
if (Object.keys(data).length === 0) {
263+
// if the original "_count" payload becomes empty, remove it
264+
delete args[kind]['_count'];
221265
}
222-
await this.injectSelectIncludeHierarchy(fieldInfo.type, nextValue);
266+
// finally merge the injection into the original payload
267+
const merged = deepmerge(args[kind], injectTarget[kind]);
268+
args[kind] = merged;
269+
}
270+
}
271+
272+
if (!injected && fieldInfo.isDataModel) {
273+
let nextValue = fieldValue;
274+
if (nextValue === true) {
275+
// make sure the payload is an object
276+
data[field] = nextValue = {};
223277
}
278+
await this.injectSelectIncludeHierarchy(fieldInfo.type, nextValue);
224279
}
225280
}
226281
}
@@ -272,7 +327,8 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
272327
field: string,
273328
value: any,
274329
selectInclude: any,
275-
context: 'select' | 'include'
330+
context: 'select' | 'include',
331+
forCount = false // if the injection is for a "{ _count: { select: { field: true } } }" payload
276332
) {
277333
const fieldInfo = resolveField(this.options.modelMeta, model, field);
278334
if (!fieldInfo?.inheritedFrom) {
@@ -286,24 +342,35 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
286342
const baseRelationName = this.makeAuxRelationName(base);
287343

288344
// prepare base layer select/include
289-
// let selectOrInclude = 'select';
290345
let thisLayer: any;
291346
if (target.include) {
292-
// selectOrInclude = 'include';
293347
thisLayer = target.include;
294348
} else if (target.select) {
295-
// selectOrInclude = 'select';
296349
thisLayer = target.select;
297350
} else {
298-
// selectInclude = 'include';
299351
thisLayer = target.select = {};
300352
}
301353

302354
if (base.name === fieldInfo.inheritedFrom) {
303355
if (!thisLayer[baseRelationName]) {
304356
thisLayer[baseRelationName] = { [context]: {} };
305357
}
306-
thisLayer[baseRelationName][context][field] = value;
358+
if (forCount) {
359+
// { _count: { select: { field: true } } } => { delegate_aux_[Base]: { select: { _count: { select: { field: true } } } } }
360+
if (
361+
!thisLayer[baseRelationName][context]['_count'] ||
362+
typeof thisLayer[baseRelationName][context] !== 'object'
363+
) {
364+
thisLayer[baseRelationName][context]['_count'] = {};
365+
}
366+
thisLayer[baseRelationName][context]['_count'] = deepmerge(
367+
thisLayer[baseRelationName][context]['_count'],
368+
{ select: { [field]: value } }
369+
);
370+
} else {
371+
// { select: { field: true } } => { delegate_aux_[Base]: { select: { field: true } } }
372+
thisLayer[baseRelationName][context][field] = value;
373+
}
307374
break;
308375
} else {
309376
if (!thisLayer[baseRelationName]) {

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

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,6 +378,66 @@ describe('Polymorphism Test', () => {
378378
).resolves.toHaveLength(1);
379379
});
380380

381+
it('read with counting relation defined in base', async () => {
382+
const { enhance } = await loadSchema(
383+
`
384+
385+
model A {
386+
id Int @id @default(autoincrement())
387+
type String
388+
bs B[]
389+
cs C[]
390+
@@delegate(type)
391+
}
392+
393+
model A1 extends A {
394+
a1 Int
395+
type1 String
396+
@@delegate(type1)
397+
}
398+
399+
model A2 extends A1 {
400+
a2 Int
401+
}
402+
403+
model B {
404+
id Int @id @default(autoincrement())
405+
a A @relation(fields: [aId], references: [id])
406+
aId Int
407+
b Int
408+
}
409+
410+
model C {
411+
id Int @id @default(autoincrement())
412+
a A @relation(fields: [aId], references: [id])
413+
aId Int
414+
c Int
415+
}
416+
`,
417+
{ enhancements: ['delegate'] }
418+
);
419+
const db = enhance();
420+
421+
const a2 = await db.a2.create({
422+
data: { a1: 1, a2: 2, bs: { create: [{ b: 1 }, { b: 2 }] }, cs: { create: [{ c: 1 }] } },
423+
include: { _count: { select: { bs: true } } },
424+
});
425+
expect(a2).toMatchObject({ a1: 1, a2: 2, _count: { bs: 2 } });
426+
427+
await expect(
428+
db.a2.findFirst({ select: { a1: true, _count: { select: { bs: true } } } })
429+
).resolves.toStrictEqual({
430+
a1: 1,
431+
_count: { bs: 2 },
432+
});
433+
434+
await expect(db.a.findFirst({ select: { _count: { select: { bs: true, cs: true } } } })).resolves.toMatchObject(
435+
{
436+
_count: { bs: 2, cs: 1 },
437+
}
438+
);
439+
});
440+
381441
it('order by base fields', async () => {
382442
const { db, user } = await setup();
383443

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import { loadSchema } from '@zenstackhq/testtools';
2+
3+
describe('issue 1467', () => {
4+
it('regression', async () => {
5+
const { enhance } = await loadSchema(
6+
`
7+
model User {
8+
id Int @id @default(autoincrement())
9+
type String
10+
@@allow('all', true)
11+
}
12+
13+
model Container {
14+
id Int @id @default(autoincrement())
15+
drink Drink @relation(fields: [drinkId], references: [id])
16+
drinkId Int
17+
@@allow('all', true)
18+
}
19+
20+
model Drink {
21+
id Int @id @default(autoincrement())
22+
name String @unique
23+
containers Container[]
24+
type String
25+
26+
@@delegate(type)
27+
@@allow('all', true)
28+
}
29+
30+
model Beer extends Drink {
31+
@@allow('all', true)
32+
}
33+
`
34+
);
35+
36+
const db = enhance();
37+
38+
await db.beer.create({
39+
data: { id: 1, name: 'Beer1' },
40+
});
41+
42+
await db.container.create({ data: { drink: { connect: { id: 1 } } } });
43+
await db.container.create({ data: { drink: { connect: { id: 1 } } } });
44+
45+
const beers = await db.beer.findFirst({
46+
select: { id: true, name: true, _count: { select: { containers: true } } },
47+
orderBy: { name: 'asc' },
48+
});
49+
expect(beers).toMatchObject({ _count: { containers: 2 } });
50+
});
51+
});

0 commit comments

Comments
 (0)