Skip to content

Commit b3b1130

Browse files
ymc9ymc9
authored
fix(zod): "strip" mode causes create payload fields to be accidentally dropped (#1747)
1 parent 1197c70 commit b3b1130

File tree

7 files changed

+455
-32
lines changed

7 files changed

+455
-32
lines changed

‎packages/runtime/package.json

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,10 @@
7676
"./models": {
7777
"types": "./models.d.ts"
7878
},
79+
"./zod-utils": {
80+
"types": "./zod-utils.d.ts",
81+
"default": "./zod-utils.js"
82+
},
7983
"./package.json": {
8084
"default": "./package.json"
8185
}

‎packages/runtime/src/zod-utils.ts

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/* eslint-disable @typescript-eslint/no-explicit-any */
2+
import { z as Z } from 'zod';
3+
4+
/**
5+
* A smarter version of `z.union` that decide which candidate to use based on how few unrecognized keys it has.
6+
*
7+
* The helper is used to deal with ambiguity in union generated for Prisma inputs when the zod schemas are configured
8+
* to run in "strip" object parsing mode. Since "strip" automatically drops unrecognized keys, it may result in
9+
* accidentally matching a less-ideal schema candidate.
10+
*
11+
* The helper uses a custom schema to find the candidate that results in the fewest unrecognized keys when parsing the data.
12+
*/
13+
export function smartUnion(z: typeof Z, candidates: Z.ZodSchema[]) {
14+
// strip `z.lazy`
15+
const processedCandidates = candidates.map((candidate) => unwrapLazy(z, candidate));
16+
17+
if (processedCandidates.some((c) => !(c instanceof z.ZodObject || c instanceof z.ZodArray))) {
18+
// fall back to plain union if not all candidates are objects or arrays
19+
return z.union(candidates as any);
20+
}
21+
22+
let resultData: any;
23+
24+
return z
25+
.custom((data) => {
26+
if (Array.isArray(data)) {
27+
const { data: result, success } = smartArrayUnion(
28+
z,
29+
processedCandidates.filter((c) => c instanceof z.ZodArray),
30+
data
31+
);
32+
if (success) {
33+
resultData = result;
34+
}
35+
return success;
36+
} else {
37+
const { data: result, success } = smartObjectUnion(
38+
z,
39+
processedCandidates.filter((c) => c instanceof z.ZodObject),
40+
data
41+
);
42+
if (success) {
43+
resultData = result;
44+
}
45+
return success;
46+
}
47+
})
48+
.transform(() => {
49+
// return the parsed data
50+
return resultData;
51+
});
52+
}
53+
54+
function smartArrayUnion(z: typeof Z, candidates: Array<Z.ZodArray<Z.ZodObject<Z.ZodRawShape>>>, data: any) {
55+
if (candidates.length === 0) {
56+
return { data: undefined, success: false };
57+
}
58+
59+
if (!Array.isArray(data)) {
60+
return { data: undefined, success: false };
61+
}
62+
63+
if (data.length === 0) {
64+
return { data, success: true };
65+
}
66+
67+
// use the first element to identify the candidate schema to use
68+
const item = data[0];
69+
const itemSchema = identifyCandidate(
70+
z,
71+
candidates.map((candidate) => candidate.element),
72+
item
73+
);
74+
75+
// find the matching schema and re-parse the data
76+
const schema = candidates.find((candidate) => candidate.element === itemSchema);
77+
return schema!.safeParse(data);
78+
}
79+
80+
function smartObjectUnion(z: typeof Z, candidates: Z.ZodObject<Z.ZodRawShape>[], data: any) {
81+
if (candidates.length === 0) {
82+
return { data: undefined, success: false };
83+
}
84+
const schema = identifyCandidate(z, candidates, data);
85+
return schema.safeParse(data);
86+
}
87+
88+
function identifyCandidate(
89+
z: typeof Z,
90+
candidates: Array<Z.ZodObject<Z.ZodRawShape> | Z.ZodLazy<Z.ZodObject<Z.ZodRawShape>>>,
91+
data: any
92+
) {
93+
const strictResults = candidates.map((candidate) => {
94+
// make sure to strip `z.lazy` before parsing
95+
const unwrapped = unwrapLazy(z, candidate);
96+
return {
97+
schema: candidate,
98+
// force object schema to run in strict mode to capture unrecognized keys
99+
result: unwrapped.strict().safeParse(data),
100+
};
101+
});
102+
103+
// find the schema with the fewest unrecognized keys
104+
const { schema } = strictResults.sort((a, b) => {
105+
const aCount = countUnrecognizedKeys(a.result.error?.issues ?? []);
106+
const bCount = countUnrecognizedKeys(b.result.error?.issues ?? []);
107+
return aCount - bCount;
108+
})[0];
109+
return schema;
110+
}
111+
112+
function countUnrecognizedKeys(issues: Z.ZodIssue[]) {
113+
return issues
114+
.filter((issue) => issue.code === 'unrecognized_keys')
115+
.map((issue) => issue.keys.length)
116+
.reduce((a, b) => a + b, 0);
117+
}
118+
119+
function unwrapLazy<T extends Z.ZodSchema>(z: typeof Z, schema: T | Z.ZodLazy<T>): T {
120+
return schema instanceof z.ZodLazy ? schema.schema : schema;
121+
}

‎packages/runtime/tests/zod/smart-union.test.ts

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
import { z } from 'zod';
2+
import { smartUnion } from '../../src/zod-utils';
3+
4+
describe('Zod smart union', () => {
5+
it('should work with scalar union', () => {
6+
const schema = smartUnion(z, [z.string(), z.number()]);
7+
expect(schema.safeParse('test')).toMatchObject({ success: true, data: 'test' });
8+
expect(schema.safeParse(1)).toMatchObject({ success: true, data: 1 });
9+
expect(schema.safeParse(true)).toMatchObject({ success: false });
10+
});
11+
12+
it('should work with non-ambiguous object union', () => {
13+
const schema = smartUnion(z, [z.object({ a: z.string() }), z.object({ b: z.number() }).strict()]);
14+
expect(schema.safeParse({ a: 'test' })).toMatchObject({ success: true, data: { a: 'test' } });
15+
expect(schema.safeParse({ b: 1 })).toMatchObject({ success: true, data: { b: 1 } });
16+
expect(schema.safeParse({ a: 'test', b: 1 })).toMatchObject({ success: true });
17+
expect(schema.safeParse({ b: 1, c: 'test' })).toMatchObject({ success: false });
18+
expect(schema.safeParse({ c: 'test' })).toMatchObject({ success: false });
19+
});
20+
21+
it('should work with ambiguous object union', () => {
22+
const schema = smartUnion(z, [
23+
z.object({ a: z.string(), b: z.number() }),
24+
z.object({ a: z.string(), c: z.boolean() }),
25+
]);
26+
expect(schema.safeParse({ a: 'test', b: 1 })).toMatchObject({ success: true, data: { a: 'test', b: 1 } });
27+
expect(schema.safeParse({ a: 'test', c: true })).toMatchObject({ success: true, data: { a: 'test', c: true } });
28+
expect(schema.safeParse({ a: 'test', b: 1, z: 'z' })).toMatchObject({
29+
success: true,
30+
data: { a: 'test', b: 1 },
31+
});
32+
expect(schema.safeParse({ a: 'test', c: true, z: 'z' })).toMatchObject({
33+
success: true,
34+
data: { a: 'test', c: true },
35+
});
36+
expect(schema.safeParse({ c: 'test' })).toMatchObject({ success: false });
37+
});
38+
39+
it('should work with non-ambiguous array union', () => {
40+
const schema = smartUnion(z, [
41+
z.object({ a: z.string() }).array(),
42+
z.object({ b: z.number() }).strict().array(),
43+
]);
44+
45+
expect(schema.safeParse([{ a: 'test' }])).toMatchObject({ success: true, data: [{ a: 'test' }] });
46+
expect(schema.safeParse([{ a: 'test' }, { a: 'test1' }])).toMatchObject({
47+
success: true,
48+
data: [{ a: 'test' }, { a: 'test1' }],
49+
});
50+
51+
expect(schema.safeParse([{ b: 1 }])).toMatchObject({ success: true, data: [{ b: 1 }] });
52+
expect(schema.safeParse([{ a: 'test', b: 1 }])).toMatchObject({ success: true });
53+
expect(schema.safeParse([{ b: 1, c: 'test' }])).toMatchObject({ success: false });
54+
expect(schema.safeParse([{ c: 'test' }])).toMatchObject({ success: false });
55+
56+
// all items must match the same candidate
57+
expect(schema.safeParse([{ a: 'test' }, { b: 1 }])).toMatchObject({ success: false });
58+
});
59+
60+
it('should work with ambiguous array union', () => {
61+
const schema = smartUnion(z, [
62+
z.object({ a: z.string(), b: z.number() }).array(),
63+
z.object({ a: z.string(), c: z.boolean() }).array(),
64+
]);
65+
66+
expect(schema.safeParse([{ a: 'test', b: 1 }])).toMatchObject({ success: true, data: [{ a: 'test', b: 1 }] });
67+
expect(schema.safeParse([{ a: 'test', c: true }])).toMatchObject({
68+
success: true,
69+
data: [{ a: 'test', c: true }],
70+
});
71+
expect(schema.safeParse([{ a: 'test', b: 1, z: 'z' }])).toMatchObject({
72+
success: true,
73+
data: [{ a: 'test', b: 1 }],
74+
});
75+
expect(schema.safeParse([{ a: 'test', c: true, z: 'z' }])).toMatchObject({
76+
success: true,
77+
data: [{ a: 'test', c: true }],
78+
});
79+
expect(schema.safeParse([{ c: 'test' }])).toMatchObject({ success: false });
80+
81+
// all items must match the same candidate
82+
expect(schema.safeParse([{ a: 'test' }, { c: true }])).toMatchObject({ success: false });
83+
});
84+
85+
it('should work with lazy schemas', () => {
86+
const schema = smartUnion(z, [
87+
z.lazy(() => z.object({ a: z.string(), b: z.number() })),
88+
z.lazy(() => z.object({ a: z.string(), c: z.boolean() })),
89+
]);
90+
expect(schema.safeParse({ a: 'test', b: 1 })).toMatchObject({ success: true, data: { a: 'test', b: 1 } });
91+
expect(schema.safeParse({ a: 'test', c: true })).toMatchObject({ success: true, data: { a: 'test', c: true } });
92+
expect(schema.safeParse({ a: 'test', b: 1, z: 'z' })).toMatchObject({
93+
success: true,
94+
data: { a: 'test', b: 1 },
95+
});
96+
});
97+
98+
it('should work with mixed object and array unions', () => {
99+
const schema = smartUnion(z, [
100+
z.object({ a: z.string() }).strict(),
101+
z.object({ b: z.number() }).strict().array(),
102+
]);
103+
104+
expect(schema.safeParse({ a: 'test' })).toMatchObject({ success: true, data: { a: 'test' } });
105+
expect(schema.safeParse([{ b: 1 }])).toMatchObject({ success: true, data: [{ b: 1 }] });
106+
expect(schema.safeParse({ a: 'test', b: 1 })).toMatchObject({ success: false });
107+
expect(schema.safeParse([{ a: 'test' }])).toMatchObject({ success: false });
108+
});
109+
});

‎packages/schema/src/plugins/zod/generator.ts

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,13 @@ import { upperCaseFirst } from 'upper-case-first';
2323
import { name } from '.';
2424
import { getDefaultOutputFolder } from '../plugin-utils';
2525
import Transformer from './transformer';
26+
import { ObjectMode } from './types';
2627
import { makeFieldSchema, makeValidationRefinements } from './utils/schema-gen';
2728

2829
export class ZodSchemaGenerator {
2930
private readonly sourceFiles: SourceFile[] = [];
3031
private readonly globalOptions: PluginGlobalOptions;
32+
private readonly mode: ObjectMode;
3133

3234
constructor(
3335
private readonly model: Model,
@@ -39,6 +41,19 @@ export class ZodSchemaGenerator {
3941
throw new Error('Global options are required');
4042
}
4143
this.globalOptions = globalOptions;
44+
45+
// options validation
46+
if (
47+
this.options.mode &&
48+
(typeof this.options.mode !== 'string' || !['strip', 'strict', 'passthrough'].includes(this.options.mode))
49+
) {
50+
throw new PluginError(
51+
name,
52+
`Invalid mode option: "${this.options.mode}". Must be one of 'strip', 'strict', or 'passthrough'.`
53+
);
54+
}
55+
56+
this.mode = (this.options.mode ?? 'strict') as ObjectMode;
4257
}
4358

4459
async generate() {
@@ -55,17 +70,6 @@ export class ZodSchemaGenerator {
5570
ensureEmptyDir(output);
5671
Transformer.setOutputPath(output);
5772

58-
// options validation
59-
if (
60-
this.options.mode &&
61-
(typeof this.options.mode !== 'string' || !['strip', 'strict', 'passthrough'].includes(this.options.mode))
62-
) {
63-
throw new PluginError(
64-
name,
65-
`Invalid mode option: "${this.options.mode}". Must be one of 'strip', 'strict', or 'passthrough'.`
66-
);
67-
}
68-
6973
// calculate the models to be excluded
7074
const excludeModels = this.getExcludedModels();
7175

@@ -120,6 +124,7 @@ export class ZodSchemaGenerator {
120124
project: this.project,
121125
inputObjectTypes,
122126
zmodel: this.model,
127+
mode: this.mode,
123128
});
124129
await transformer.generateInputSchemas(this.options, this.model);
125130
this.sourceFiles.push(...transformer.sourceFiles);
@@ -215,6 +220,7 @@ export class ZodSchemaGenerator {
215220
project: this.project,
216221
inputObjectTypes: [],
217222
zmodel: this.model,
223+
mode: this.mode,
218224
});
219225
await transformer.generateEnumSchemas();
220226
this.sourceFiles.push(...transformer.sourceFiles);
@@ -243,6 +249,7 @@ export class ZodSchemaGenerator {
243249
project: this.project,
244250
inputObjectTypes,
245251
zmodel: this.model,
252+
mode: this.mode,
246253
});
247254
const moduleName = transformer.generateObjectSchema(generateUnchecked, this.options);
248255
moduleNames.push(moduleName);

0 commit comments

Comments
 (0)