fix(zod): "strip" mode causes create payload fields to be accidentally dropped

packages/runtime/package.json

@@ -76,6 +76,10 @@
         "./models": {
             "types": "./models.d.ts"
         },
+        "./zod-utils": {
+            "types": "./zod-utils.d.ts",
+            "default": "./zod-utils.js"
+        },
         "./package.json": {
             "default": "./package.json"
         }

packages/runtime/src/zod-utils.ts

@@ -0,0 +1,121 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { z as Z } from 'zod';
+
+/**
+ * A smarter version of `z.union` that decide which candidate to use based on how few unrecognized keys it has.
+ *
+ * The helper is used to deal with ambiguity in union generated for Prisma inputs when the zod schemas are configured
+ * to run in "strip" object parsing mode. Since "strip" automatically drops unrecognized keys, it may result in
+ * accidentally matching a less-ideal schema candidate.
+ *
+ * The helper uses a custom schema to find the candidate that results in the fewest unrecognized keys when parsing the data.
+ */
+export function smartUnion(z: typeof Z, candidates: Z.ZodSchema[]) {
+    // strip `z.lazy`
+    const processedCandidates = candidates.map((candidate) => unwrapLazy(z, candidate));
+
+    if (processedCandidates.some((c) => !(c instanceof z.ZodObject || c instanceof z.ZodArray))) {
+        // fall back to plain union if not all candidates are objects or arrays
+        return z.union(candidates as any);
+    }
+
+    let resultData: any;
+
+    return z
+        .custom((data) => {
+            if (Array.isArray(data)) {
+                const { data: result, success } = smartArrayUnion(
+                    z,
+                    processedCandidates.filter((c) => c instanceof z.ZodArray),
+                    data
+                );
+                if (success) {
+                    resultData = result;
+                }
+                return success;
+            } else {
+                const { data: result, success } = smartObjectUnion(
+                    z,
+                    processedCandidates.filter((c) => c instanceof z.ZodObject),
+                    data
+                );
+                if (success) {
+                    resultData = result;
+                }
+                return success;
+            }
+        })
+        .transform(() => {
+            // return the parsed data
+            return resultData;
+        });
+}
+
+function smartArrayUnion(z: typeof Z, candidates: Array<Z.ZodArray<Z.ZodObject<Z.ZodRawShape>>>, data: any) {
+    if (candidates.length === 0) {
+        return { data: undefined, success: false };
+    }
+
+    if (!Array.isArray(data)) {
+        return { data: undefined, success: false };
+    }
+
+    if (data.length === 0) {
+        return { data, success: true };
+    }
+
+    // use the first element to identify the candidate schema to use
+    const item = data[0];
+    const itemSchema = identifyCandidate(
+        z,
+        candidates.map((candidate) => candidate.element),
+        item
+    );
+
+    // find the matching schema and re-parse the data
+    const schema = candidates.find((candidate) => candidate.element === itemSchema);
+    return schema!.safeParse(data);
+}
+
+function smartObjectUnion(z: typeof Z, candidates: Z.ZodObject<Z.ZodRawShape>[], data: any) {
+    if (candidates.length === 0) {
+        return { data: undefined, success: false };
+    }
+    const schema = identifyCandidate(z, candidates, data);
+    return schema.safeParse(data);
+}
+
+function identifyCandidate(
+    z: typeof Z,
+    candidates: Array<Z.ZodObject<Z.ZodRawShape> | Z.ZodLazy<Z.ZodObject<Z.ZodRawShape>>>,
+    data: any
+) {
+    const strictResults = candidates.map((candidate) => {
+        // make sure to strip `z.lazy` before parsing
+        const unwrapped = unwrapLazy(z, candidate);
+        return {
+            schema: candidate,
+            // force object schema to run in strict mode to capture unrecognized keys
+            result: unwrapped.strict().safeParse(data),
+        };
+    });
+
+    // find the schema with the fewest unrecognized keys
+    const { schema } = strictResults.sort((a, b) => {
+        const aCount = countUnrecognizedKeys(a.result.error?.issues ?? []);
+        const bCount = countUnrecognizedKeys(b.result.error?.issues ?? []);
+        return aCount - bCount;
+    })[0];
+    return schema;
+}
+
+function countUnrecognizedKeys(issues: Z.ZodIssue[]) {
+    return issues
+        .filter((issue) => issue.code === 'unrecognized_keys')
+        .map((issue) => issue.keys.length)
+        .reduce((a, b) => a + b, 0);
+}
+
+function unwrapLazy<T extends Z.ZodSchema>(z: typeof Z, schema: T | Z.ZodLazy<T>): T {
+    return schema instanceof z.ZodLazy ? schema.schema : schema;
+}

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

@@ -0,0 +1,109 @@
+import { z } from 'zod';
+import { smartUnion } from '../../src/zod-utils';
+
+describe('Zod smart union', () => {
+    it('should work with scalar union', () => {
+        const schema = smartUnion(z, [z.string(), z.number()]);
+        expect(schema.safeParse('test')).toMatchObject({ success: true, data: 'test' });
+        expect(schema.safeParse(1)).toMatchObject({ success: true, data: 1 });
+        expect(schema.safeParse(true)).toMatchObject({ success: false });
+    });
+
+    it('should work with non-ambiguous object union', () => {
+        const schema = smartUnion(z, [z.object({ a: z.string() }), z.object({ b: z.number() }).strict()]);
+        expect(schema.safeParse({ a: 'test' })).toMatchObject({ success: true, data: { a: 'test' } });
+        expect(schema.safeParse({ b: 1 })).toMatchObject({ success: true, data: { b: 1 } });
+        expect(schema.safeParse({ a: 'test', b: 1 })).toMatchObject({ success: true });
+        expect(schema.safeParse({ b: 1, c: 'test' })).toMatchObject({ success: false });
+        expect(schema.safeParse({ c: 'test' })).toMatchObject({ success: false });
+    });
+
+    it('should work with ambiguous object union', () => {
+        const schema = smartUnion(z, [
+            z.object({ a: z.string(), b: z.number() }),
+            z.object({ a: z.string(), c: z.boolean() }),
+        ]);
+        expect(schema.safeParse({ a: 'test', b: 1 })).toMatchObject({ success: true, data: { a: 'test', b: 1 } });
+        expect(schema.safeParse({ a: 'test', c: true })).toMatchObject({ success: true, data: { a: 'test', c: true } });
+        expect(schema.safeParse({ a: 'test', b: 1, z: 'z' })).toMatchObject({
+            success: true,
+            data: { a: 'test', b: 1 },
+        });
+        expect(schema.safeParse({ a: 'test', c: true, z: 'z' })).toMatchObject({
+            success: true,
+            data: { a: 'test', c: true },
+        });
+        expect(schema.safeParse({ c: 'test' })).toMatchObject({ success: false });
+    });
+
+    it('should work with non-ambiguous array union', () => {
+        const schema = smartUnion(z, [
+            z.object({ a: z.string() }).array(),
+            z.object({ b: z.number() }).strict().array(),
+        ]);
+
+        expect(schema.safeParse([{ a: 'test' }])).toMatchObject({ success: true, data: [{ a: 'test' }] });
+        expect(schema.safeParse([{ a: 'test' }, { a: 'test1' }])).toMatchObject({
+            success: true,
+            data: [{ a: 'test' }, { a: 'test1' }],
+        });
+
+        expect(schema.safeParse([{ b: 1 }])).toMatchObject({ success: true, data: [{ b: 1 }] });
+        expect(schema.safeParse([{ a: 'test', b: 1 }])).toMatchObject({ success: true });
+        expect(schema.safeParse([{ b: 1, c: 'test' }])).toMatchObject({ success: false });
+        expect(schema.safeParse([{ c: 'test' }])).toMatchObject({ success: false });
+
+        // all items must match the same candidate
+        expect(schema.safeParse([{ a: 'test' }, { b: 1 }])).toMatchObject({ success: false });
+    });
+
+    it('should work with ambiguous array union', () => {
+        const schema = smartUnion(z, [
+            z.object({ a: z.string(), b: z.number() }).array(),
+            z.object({ a: z.string(), c: z.boolean() }).array(),
+        ]);
+
+        expect(schema.safeParse([{ a: 'test', b: 1 }])).toMatchObject({ success: true, data: [{ a: 'test', b: 1 }] });
+        expect(schema.safeParse([{ a: 'test', c: true }])).toMatchObject({
+            success: true,
+            data: [{ a: 'test', c: true }],
+        });
+        expect(schema.safeParse([{ a: 'test', b: 1, z: 'z' }])).toMatchObject({
+            success: true,
+            data: [{ a: 'test', b: 1 }],
+        });
+        expect(schema.safeParse([{ a: 'test', c: true, z: 'z' }])).toMatchObject({
+            success: true,
+            data: [{ a: 'test', c: true }],
+        });
+        expect(schema.safeParse([{ c: 'test' }])).toMatchObject({ success: false });
+
+        // all items must match the same candidate
+        expect(schema.safeParse([{ a: 'test' }, { c: true }])).toMatchObject({ success: false });
+    });
+
+    it('should work with lazy schemas', () => {
+        const schema = smartUnion(z, [
+            z.lazy(() => z.object({ a: z.string(), b: z.number() })),
+            z.lazy(() => z.object({ a: z.string(), c: z.boolean() })),
+        ]);
+        expect(schema.safeParse({ a: 'test', b: 1 })).toMatchObject({ success: true, data: { a: 'test', b: 1 } });
+        expect(schema.safeParse({ a: 'test', c: true })).toMatchObject({ success: true, data: { a: 'test', c: true } });
+        expect(schema.safeParse({ a: 'test', b: 1, z: 'z' })).toMatchObject({
+            success: true,
+            data: { a: 'test', b: 1 },
+        });
+    });
+
+    it('should work with mixed object and array unions', () => {
+        const schema = smartUnion(z, [
+            z.object({ a: z.string() }).strict(),
+            z.object({ b: z.number() }).strict().array(),
+        ]);
+
+        expect(schema.safeParse({ a: 'test' })).toMatchObject({ success: true, data: { a: 'test' } });
+        expect(schema.safeParse([{ b: 1 }])).toMatchObject({ success: true, data: [{ b: 1 }] });
+        expect(schema.safeParse({ a: 'test', b: 1 })).toMatchObject({ success: false });
+        expect(schema.safeParse([{ a: 'test' }])).toMatchObject({ success: false });
+    });
+});

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

@@ -23,11 +23,13 @@ import { upperCaseFirst } from 'upper-case-first';
 import { name } from '.';
 import { getDefaultOutputFolder } from '../plugin-utils';
 import Transformer from './transformer';
+import { ObjectMode } from './types';
 import { makeFieldSchema, makeValidationRefinements } from './utils/schema-gen';
 
 export class ZodSchemaGenerator {
     private readonly sourceFiles: SourceFile[] = [];
     private readonly globalOptions: PluginGlobalOptions;
+    private readonly mode: ObjectMode;
 
     constructor(
         private readonly model: Model,
@@ -39,6 +41,19 @@ export class ZodSchemaGenerator {
             throw new Error('Global options are required');
         }
         this.globalOptions = globalOptions;
+
+        // options validation
+        if (
+            this.options.mode &&
+            (typeof this.options.mode !== 'string' || !['strip', 'strict', 'passthrough'].includes(this.options.mode))
+        ) {
+            throw new PluginError(
+                name,
+                `Invalid mode option: "${this.options.mode}". Must be one of 'strip', 'strict', or 'passthrough'.`
+            );
+        }
+
+        this.mode = (this.options.mode ?? 'strict') as ObjectMode;
     }
 
     async generate() {
@@ -55,17 +70,6 @@ export class ZodSchemaGenerator {
         ensureEmptyDir(output);
         Transformer.setOutputPath(output);
 
-        // options validation
-        if (
-            this.options.mode &&
-            (typeof this.options.mode !== 'string' || !['strip', 'strict', 'passthrough'].includes(this.options.mode))
-        ) {
-            throw new PluginError(
-                name,
-                `Invalid mode option: "${this.options.mode}". Must be one of 'strip', 'strict', or 'passthrough'.`
-            );
-        }
-
         // calculate the models to be excluded
         const excludeModels = this.getExcludedModels();
 
@@ -120,6 +124,7 @@ export class ZodSchemaGenerator {
                 project: this.project,
                 inputObjectTypes,
                 zmodel: this.model,
+                mode: this.mode,
             });
             await transformer.generateInputSchemas(this.options, this.model);
             this.sourceFiles.push(...transformer.sourceFiles);
@@ -215,6 +220,7 @@ export class ZodSchemaGenerator {
             project: this.project,
             inputObjectTypes: [],
             zmodel: this.model,
+            mode: this.mode,
         });
         await transformer.generateEnumSchemas();
         this.sourceFiles.push(...transformer.sourceFiles);
@@ -243,6 +249,7 @@ export class ZodSchemaGenerator {
                 project: this.project,
                 inputObjectTypes,
                 zmodel: this.model,
+                mode: this.mode,
             });
             const moduleName = transformer.generateObjectSchema(generateUnchecked, this.options);
             moduleNames.push(moduleName);