Skip to content

Commit

Permalink
fix(delegate): support _count select of base fields (#1937)
Browse files Browse the repository at this point in the history
  • Loading branch information
@ymc9
ymc9 authored Jan 6, 2025
1 parent 1956bdb commit bcbfb9a
Show file tree
Hide file tree
Showing 3 changed files with 217 additions and 39 deletions.
145 changes: 106 additions & 39 deletions packages/runtime/src/enhancements/node/delegate.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}
Expand Down Expand Up @@ -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) {
Expand All @@ -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]) {
Expand Down
60 changes: 60 additions & 0 deletions tests/integration/tests/enhancements/with-delegate/enhanced-client.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand Down
51 changes: 51 additions & 0 deletions tests/regression/tests/issue-1467.test.ts
View file
Original file line number Diff line number Diff line change
@@ -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 } });
});
});

0 comments on commit bcbfb9a

Please sign in to comment.