From 99a9835e52f6d7e6133cfdd4e79dbfd882aa6480 Mon Sep 17 00:00:00 2001 From: ymc9 <104139426+ymc9@users.noreply.github.com> Date: Sun, 5 Jan 2025 21:39:50 +0800 Subject: [PATCH 1/3] fix(delegate): support _count select of base fields --- .../runtime/src/enhancements/node/delegate.ts | 123 ++++++++++++------ tests/regression/tests/issue-1467.test.ts | 51 ++++++++ 2 files changed, 135 insertions(+), 39 deletions(-) create mode 100644 tests/regression/tests/issue-1467.test.ts diff --git a/packages/runtime/src/enhancements/node/delegate.ts b/packages/runtime/src/enhancements/node/delegate.ts index 80fd09f17..9b66da31b 100644 --- a/packages/runtime/src/enhancements/node/delegate.ts +++ b/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(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(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,16 +322,12 @@ 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 = {}; } @@ -303,7 +335,20 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler { 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]) { diff --git a/tests/regression/tests/issue-1467.test.ts b/tests/regression/tests/issue-1467.test.ts new file mode 100644 index 000000000..374313e45 --- /dev/null +++ b/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 } }); + }); +}); From 045803722f44be0c4bac01c5c0a3d8af356509f1 Mon Sep 17 00:00:00 2001 From: ymc9 <104139426+ymc9@users.noreply.github.com> Date: Mon, 6 Jan 2025 14:32:23 +0800 Subject: [PATCH 2/3] more tests --- .../with-delegate/enhanced-client.test.ts | 60 +++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/tests/integration/tests/enhancements/with-delegate/enhanced-client.test.ts b/tests/integration/tests/enhancements/with-delegate/enhanced-client.test.ts index 91a385db0..7a555e0cd 100644 --- a/tests/integration/tests/enhancements/with-delegate/enhanced-client.test.ts +++ b/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(); From 98e28cf0bfed8b85791ec58e7d528dd45675ffbe Mon Sep 17 00:00:00 2001 From: ymc9 <104139426+ymc9@users.noreply.github.com> Date: Mon, 6 Jan 2025 15:12:53 +0800 Subject: [PATCH 3/3] more comments --- .../runtime/src/enhancements/node/delegate.ts | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/packages/runtime/src/enhancements/node/delegate.ts b/packages/runtime/src/enhancements/node/delegate.ts index 9b66da31b..06c1526e5 100644 --- a/packages/runtime/src/enhancements/node/delegate.ts +++ b/packages/runtime/src/enhancements/node/delegate.ts @@ -180,14 +180,28 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler { return; } + // 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, @@ -232,18 +246,24 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler { 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']; } + // finally merge the injection into the original payload const merged = deepmerge(args[kind], injectTarget[kind]); args[kind] = merged; } @@ -308,7 +328,7 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler { value: any, selectInclude: any, context: 'select' | 'include', - forCount = false + forCount = false // if the injection is for a "{ _count: { select: { field: true } } }" payload ) { const fieldInfo = resolveField(this.options.modelMeta, model, field); if (!fieldInfo?.inheritedFrom) { @@ -336,6 +356,7 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler { thisLayer[baseRelationName] = { [context]: {} }; } if (forCount) { + // { _count: { select: { field: true } } } => { delegate_aux_[Base]: { select: { _count: { select: { field: true } } } } } if ( !thisLayer[baseRelationName][context]['_count'] || typeof thisLayer[baseRelationName][context] !== 'object' @@ -347,6 +368,7 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler { { select: { [field]: value } } ); } else { + // { select: { field: true } } => { delegate_aux_[Base]: { select: { field: true } } } thisLayer[baseRelationName][context][field] = value; } break;