Skip to content

Commit ea6d16b

Browse files
ymc9ymc9
committed
feat: support updateManyAndReturn 
1 parent 502dc83 commit ea6d16b

File tree

3 files changed

+237
-8
lines changed

3 files changed

+237
-8
lines changed

‎packages/runtime/src/enhancements/node/policy/handler.ts‎

Lines changed: 73 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -511,7 +511,7 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
511511
});
512512
}
513513

514-
// throw read-back error if any of create result read-back fails
514+
// throw read-back error if any of the create result read-back fails
515515
const error = result.find((r) => !!r.error)?.error;
516516
if (error) {
517517
throw error;
@@ -1268,6 +1268,14 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
12681268
}
12691269

12701270
updateMany(args: any) {
1271+
return this.doUpdateMany(args, 'updateMany');
1272+
}
1273+
1274+
updateManyAndReturn(args: any): Promise<unknown[]> {
1275+
return this.doUpdateMany(args, 'updateManyAndReturn');
1276+
}
1277+
1278+
private doUpdateMany(args: any, action: 'updateMany' | 'updateManyAndReturn'): Promise<any> {
12711279
if (!args) {
12721280
throw prismaClientValidationError(this.prisma, this.prismaModule, 'query argument is required');
12731281
}
@@ -1279,9 +1287,10 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
12791287
);
12801288
}
12811289

1282-
return createDeferredPromise(() => {
1290+
return createDeferredPromise(async () => {
12831291
this.policyUtils.tryReject(this.prisma, this.model, 'update');
12841292

1293+
const origArgs = args;
12851294
args = this.policyUtils.safeClone(args);
12861295
this.policyUtils.injectAuthGuardAsWhere(this.prisma, args, this.model, 'update');
12871296

@@ -1302,13 +1311,37 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
13021311
if (this.shouldLogQuery) {
13031312
this.logger.info(`[policy] \`updateMany\` ${this.model}: ${formatObject(args)}`);
13041313
}
1305-
return this.modelClient.updateMany(args);
1314+
if (action === 'updateMany') {
1315+
return this.modelClient.updateMany(args);
1316+
} else {
1317+
// make sure only id fields are returned so we can directly use the result
1318+
// for read-back check
1319+
const updatedArg = {
1320+
...args,
1321+
select: this.policyUtils.makeIdSelection(this.model),
1322+
include: undefined,
1323+
};
1324+
const updated = await this.modelClient.updateManyAndReturn(updatedArg);
1325+
// process read-back
1326+
const result = await Promise.all(
1327+
updated.map((item) =>
1328+
this.policyUtils.readBack(this.prisma, this.model, 'update', origArgs, item)
1329+
)
1330+
);
1331+
// throw read-back error if any of create result read-back fails
1332+
const error = result.find((r) => !!r.error)?.error;
1333+
if (error) {
1334+
throw error;
1335+
} else {
1336+
return result.map((r) => r.result);
1337+
}
1338+
}
13061339
}
13071340

13081341
// collect post-update checks
13091342
const postWriteChecks: PostWriteCheckRecord[] = [];
13101343

1311-
return this.queryUtils.transaction(this.prisma, async (tx) => {
1344+
const result = await this.queryUtils.transaction(this.prisma, async (tx) => {
13121345
// collect pre-update values
13131346
let select = this.policyUtils.makeIdSelection(this.model);
13141347
const preValueSelect = this.policyUtils.getPreValueSelect(this.model);
@@ -1352,13 +1385,45 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
13521385
if (this.shouldLogQuery) {
13531386
this.logger.info(`[policy] \`updateMany\` in tx for ${this.model}: ${formatObject(args)}`);
13541387
}
1355-
const result = await tx[this.model].updateMany(args);
13561388

1357-
// run post-write checks
1358-
await this.runPostWriteChecks(postWriteChecks, tx);
1389+
if (action === 'updateMany') {
1390+
const result = await tx[this.model].updateMany(args);
1391+
// run post-write checks
1392+
await this.runPostWriteChecks(postWriteChecks, tx);
1393+
return result;
1394+
} else {
1395+
// make sure only id fields are returned so we can directly use the result
1396+
// for read-back check
1397+
const updatedArg = {
1398+
...args,
1399+
select: this.policyUtils.makeIdSelection(this.model),
1400+
include: undefined,
1401+
};
1402+
const result = await tx[this.model].updateManyAndReturn(updatedArg);
1403+
// run post-write checks
1404+
await this.runPostWriteChecks(postWriteChecks, tx);
1405+
return result;
1406+
}
1407+
});
13591408

1409+
if (action === 'updateMany') {
1410+
// no further processing needed
13601411
return result;
1361-
});
1412+
} else {
1413+
// process read-back
1414+
const readBackResult = await Promise.all(
1415+
(result as unknown[]).map((item) =>
1416+
this.policyUtils.readBack(this.prisma, this.model, 'update', origArgs, item)
1417+
)
1418+
);
1419+
// throw read-back error if any of the update result read-back fails
1420+
const error = readBackResult.find((r) => !!r.error)?.error;
1421+
if (error) {
1422+
throw error;
1423+
} else {
1424+
return readBackResult.map((r) => r.result);
1425+
}
1426+
}
13621427
});
13631428
}
13641429

‎tests/integration/tests/enhancements/with-policy/update-many-and-return.test.ts‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
import { loadSchema } from '@zenstackhq/testtools';
2+
3+
describe('Test API updateManyAndReturn', () => {
4+
it('model-level policies', async () => {
5+
const { prisma, enhance } = await loadSchema(
6+
`
7+
model User {
8+
id Int @id @default(autoincrement())
9+
posts Post[]
10+
level Int
11+
12+
@@allow('read', level > 0)
13+
}
14+
15+
model Post {
16+
id Int @id @default(autoincrement())
17+
title String
18+
published Boolean @default(false)
19+
userId Int
20+
user User @relation(fields: [userId], references: [id])
21+
22+
@@allow('read', published)
23+
@@allow('update', contains(title, 'hello'))
24+
}
25+
`
26+
);
27+
28+
await prisma.user.createMany({
29+
data: [{ id: 1, level: 1 }],
30+
});
31+
await prisma.user.createMany({
32+
data: [{ id: 2, level: 0 }],
33+
});
34+
35+
await prisma.post.createMany({
36+
data: [
37+
{ id: 1, title: 'hello1', userId: 1, published: true },
38+
{ id: 2, title: 'world1', userId: 1, published: false },
39+
],
40+
});
41+
42+
const db = enhance();
43+
44+
// only post#1 is updated
45+
let r = await db.post.updateManyAndReturn({
46+
data: { title: 'foo' },
47+
});
48+
expect(r).toHaveLength(1);
49+
expect(r[0].id).toBe(1);
50+
51+
// post#2 is excluded from update
52+
await expect(
53+
db.post.updateManyAndReturn({
54+
where: { id: 2 },
55+
data: { title: 'foo' },
56+
})
57+
).resolves.toHaveLength(0);
58+
59+
// reset
60+
await prisma.post.update({ where: { id: 1 }, data: { title: 'hello1' } });
61+
62+
// post#1 is updated
63+
await expect(
64+
db.post.updateManyAndReturn({
65+
where: { id: 1 },
66+
data: { title: 'foo' },
67+
})
68+
).resolves.toHaveLength(1);
69+
70+
// reset
71+
await prisma.post.update({ where: { id: 1 }, data: { title: 'hello1' } });
72+
73+
// read-back check
74+
// post#1 updated but can't be read back
75+
await expect(
76+
db.post.updateManyAndReturn({
77+
data: { published: false },
78+
})
79+
).toBeRejectedByPolicy(['result is not allowed to be read back']);
80+
// but the update should have been applied
81+
await expect(prisma.post.findUnique({ where: { id: 1 } })).resolves.toMatchObject({ published: false });
82+
83+
// reset
84+
await prisma.post.update({ where: { id: 1 }, data: { published: true } });
85+
86+
// return relation
87+
r = await db.post.updateManyAndReturn({
88+
include: { user: true },
89+
data: { title: 'hello2' },
90+
});
91+
expect(r[0]).toMatchObject({ user: { id: 1 } });
92+
93+
// relation filtered
94+
await prisma.post.create({ data: { id: 3, title: 'hello3', userId: 2, published: true } });
95+
await expect(
96+
db.post.updateManyAndReturn({
97+
where: { id: 3 },
98+
include: { user: true },
99+
data: { title: 'hello4' },
100+
})
101+
).toBeRejectedByPolicy(['result is not allowed to be read back']);
102+
// update is applied
103+
await expect(prisma.post.findUnique({ where: { id: 3 } })).resolves.toMatchObject({ title: 'hello4' });
104+
});
105+
106+
it('field-level policies', async () => {
107+
const { prisma, enhance } = await loadSchema(
108+
`
109+
model Post {
110+
id Int @id @default(autoincrement())
111+
title String @allow('read', published)
112+
published Boolean @default(false)
113+
114+
@@allow('all', true)
115+
}
116+
`
117+
);
118+
119+
const db = enhance();
120+
121+
// update should succeed but one result's title field can't be read back
122+
await prisma.post.createMany({
123+
data: [
124+
{ id: 1, title: 'post1', published: true },
125+
{ id: 2, title: 'post2', published: false },
126+
],
127+
});
128+
129+
const r = await db.post.updateManyAndReturn({
130+
data: { title: 'foo' },
131+
});
132+
133+
expect(r.length).toBe(2);
134+
expect(r[0].title).toBeTruthy();
135+
expect(r[1].title).toBeUndefined();
136+
137+
// check posts are updated
138+
await expect(prisma.post.findMany({ where: { title: 'foo' } })).resolves.toHaveLength(2);
139+
});
140+
});

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ describe('issue 1955', () => {
2121
_prisma = prisma;
2222

2323
const db = enhance();
24+
2425
await expect(
2526
db.post.createManyAndReturn({
2627
data: [
@@ -38,6 +39,17 @@ describe('issue 1955', () => {
3839
expect.objectContaining({ name: 'blu' }),
3940
])
4041
);
42+
43+
await expect(
44+
db.post.updateManyAndReturn({
45+
data: { name: 'foo' },
46+
})
47+
).resolves.toEqual(
48+
expect.arrayContaining([
49+
expect.objectContaining({ name: 'foo' }),
50+
expect.objectContaining({ name: 'foo' }),
51+
])
52+
);
4153
} finally {
4254
await _prisma.$disconnect();
4355
await dropPostgresDb('issue-1955-1');
@@ -72,6 +84,7 @@ describe('issue 1955', () => {
7284
_prisma = prisma;
7385

7486
const db = enhance();
87+
7588
await expect(
7689
db.post.createManyAndReturn({
7790
data: [
@@ -89,6 +102,17 @@ describe('issue 1955', () => {
89102
expect.objectContaining({ name: 'blu' }),
90103
])
91104
);
105+
106+
await expect(
107+
db.post.updateManyAndReturn({
108+
data: { name: 'foo' },
109+
})
110+
).resolves.toEqual(
111+
expect.arrayContaining([
112+
expect.objectContaining({ name: 'foo' }),
113+
expect.objectContaining({ name: 'foo' }),
114+
])
115+
);
92116
} finally {
93117
await _prisma.$disconnect();
94118
await dropPostgresDb('issue-1955-2');

0 commit comments

Comments
 (0)