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

ymc9 · ymc9 · commit c68df7030423 · 2024-09-26T23:16:13.000-07:00
Fixes #1746
diff --git a/packages/runtime/package.json b/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"
         }
diff --git a/packages/runtime/src/zod-utils.ts b/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;
+}
diff --git a/packages/schema/src/plugins/zod/generator.ts b/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);
diff --git a/packages/schema/src/plugins/zod/transformer.ts b/packages/schema/src/plugins/zod/transformer.ts
@@ -7,7 +7,7 @@ import path from 'path';
 import type { Project, SourceFile } from 'ts-morph';
 import { upperCaseFirst } from 'upper-case-first';
 import { computePrismaClientImport } from './generator';
-import { AggregateOperationSupport, TransformerParams } from './types';
+import { AggregateOperationSupport, ObjectMode, TransformerParams } from './types';
 
 export default class Transformer {
     name: string;
@@ -28,6 +28,7 @@ export default class Transformer {
     private inputObjectTypes: PrismaDMMF.InputType[];
     public sourceFiles: SourceFile[] = [];
     private zmodel: Model;
+    private mode: ObjectMode;
 
     constructor(params: TransformerParams) {
         this.originalName = params.name ?? '';
@@ -40,6 +41,7 @@ export default class Transformer {
         this.project = params.project;
         this.inputObjectTypes = params.inputObjectTypes;
         this.zmodel = params.zmodel;
+        this.mode = params.mode;
     }
 
     static setOutputPath(outPath: string) {
@@ -73,7 +75,12 @@ export default class Transformer {
     }
 
     generateImportZodStatement() {
-        return "import { z } from 'zod';\n";
+        let r = "import { z } from 'zod';\n";
+        if (this.mode === 'strip') {
+            // import the additional `smartUnion` helper
+            r += `import { smartUnion } from '@zenstackhq/runtime/zod-utils';\n`;
+        }
+        return r;
     }
 
     generateExportSchemaStatement(name: string, schema: string) {
@@ -210,8 +217,19 @@ export default class Transformer {
 
         const opt = !field.isRequired ? '.optional()' : '';
 
-        let resString =
-            alternatives.length === 1 ? alternatives.join(',\r\n') : `z.union([${alternatives.join(',\r\n')}])${opt}`;
+        let resString: string;
+
+        if (alternatives.length === 1) {
+            resString = alternatives.join(',\r\n');
+        } else {
+            if (alternatives.some((alt) => alt.includes('Unchecked'))) {
+                // if the union is for combining checked and unchecked input types, use `smartUnion`
+                // to parse with the best candidate at runtime
+                resString = this.wrapWithSmartUnion(...alternatives) + `${opt}`;
+            } else {
+                resString = `z.union([${alternatives.join(',\r\n')}])${opt}`;
+            }
+        }
 
         if (field.isNullable) {
             resString += '.nullable()';
@@ -391,17 +409,6 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
         return `${modelName}InputSchema.${queryName}`;
     }
 
-    wrapWithZodUnion(zodStringFields: string[]) {
-        let wrapped = '';
-
-        wrapped += 'z.union([';
-        wrapped += '\n';
-        wrapped += '  ' + zodStringFields.join(',');
-        wrapped += '\n';
-        wrapped += '])';
-        return wrapped;
-    }
-
     wrapWithZodObject(zodStringFields: string | string[], mode = 'strict') {
         let wrapped = '';
 
@@ -425,6 +432,14 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
         return wrapped;
     }
 
+    wrapWithSmartUnion(...schemas: string[]) {
+        if (this.mode === 'strip') {
+            return `smartUnion(z, [${schemas.join(', ')}])`;
+        } else {
+            return `z.union([${schemas.join(', ')}])`;
+        }
+    }
+
     async generateInputSchemas(options: PluginOptions, zmodel: Model) {
         const globalExports: string[] = [];
 
@@ -464,7 +479,7 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
                 this.resolveSelectIncludeImportAndZodSchemaLine(model);
 
             let imports = [
-                `import { z } from 'zod'`,
+                this.generateImportZodStatement(),
                 this.generateImportPrismaStatement(options),
                 selectImport,
                 includeImport,
@@ -523,7 +538,10 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
                     );
                 }
                 const dataSchema = generateUnchecked
-                    ? `z.union([${modelName}CreateInputObjectSchema, ${modelName}UncheckedCreateInputObjectSchema])`
+                    ? this.wrapWithSmartUnion(
+                          `${modelName}CreateInputObjectSchema`,
+                          `${modelName}UncheckedCreateInputObjectSchema`
+                      )
                     : `${modelName}CreateInputObjectSchema`;
                 const fields = `${selectZodSchemaLineLazy} ${includeZodSchemaLineLazy} data: ${dataSchema}`;
                 codeBody += `create: ${this.wrapWithZodObject(fields, mode)},`;
@@ -568,7 +586,10 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
                     );
                 }
                 const dataSchema = generateUnchecked
-                    ? `z.union([${modelName}UpdateInputObjectSchema, ${modelName}UncheckedUpdateInputObjectSchema])`
+                    ? this.wrapWithSmartUnion(
+                          `${modelName}UpdateInputObjectSchema`,
+                          `${modelName}UncheckedUpdateInputObjectSchema`
+                      )
                     : `${modelName}UpdateInputObjectSchema`;
                 const fields = `${selectZodSchemaLineLazy} ${includeZodSchemaLineLazy} data: ${dataSchema}, where: ${modelName}WhereUniqueInputObjectSchema`;
                 codeBody += `update: ${this.wrapWithZodObject(fields, mode)},`;
@@ -586,7 +607,10 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
                     );
                 }
                 const dataSchema = generateUnchecked
-                    ? `z.union([${modelName}UpdateManyMutationInputObjectSchema, ${modelName}UncheckedUpdateManyInputObjectSchema])`
+                    ? this.wrapWithSmartUnion(
+                          `${modelName}UpdateManyMutationInputObjectSchema`,
+                          `${modelName}UncheckedUpdateManyInputObjectSchema`
+                      )
                     : `${modelName}UpdateManyMutationInputObjectSchema`;
                 const fields = `data: ${dataSchema}, where: ${modelName}WhereInputObjectSchema.optional()`;
                 codeBody += `updateMany: ${this.wrapWithZodObject(fields, mode)},`;
@@ -606,10 +630,16 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
                     );
                 }
                 const createSchema = generateUnchecked
-                    ? `z.union([${modelName}CreateInputObjectSchema, ${modelName}UncheckedCreateInputObjectSchema])`
+                    ? this.wrapWithSmartUnion(
+                          `${modelName}CreateInputObjectSchema`,
+                          `${modelName}UncheckedCreateInputObjectSchema`
+                      )
                     : `${modelName}CreateInputObjectSchema`;
                 const updateSchema = generateUnchecked
-                    ? `z.union([${modelName}UpdateInputObjectSchema, ${modelName}UncheckedUpdateInputObjectSchema])`
+                    ? this.wrapWithSmartUnion(
+                          `${modelName}UpdateInputObjectSchema`,
+                          `${modelName}UncheckedUpdateInputObjectSchema`
+                      )
                     : `${modelName}UpdateInputObjectSchema`;
                 const fields = `${selectZodSchemaLineLazy} ${includeZodSchemaLineLazy} where: ${modelName}WhereUniqueInputObjectSchema, create: ${createSchema}, update: ${updateSchema}`;
                 codeBody += `upsert: ${this.wrapWithZodObject(fields, mode)},`;
diff --git a/packages/schema/src/plugins/zod/types.ts b/packages/schema/src/plugins/zod/types.ts
@@ -14,6 +14,7 @@ export type TransformerParams = {
     project: Project;
     inputObjectTypes: PrismaDMMF.InputType[];
     zmodel: Model;
+    mode: ObjectMode;
 };
 
 export type AggregateOperationSupport = {
@@ -25,3 +26,5 @@ export type AggregateOperationSupport = {
         avg?: boolean;
     };
 };
+
+export type ObjectMode = 'strict' | 'strip' | 'passthrough';
diff --git a/tests/regression/tests/issue-1746.test.ts b/tests/regression/tests/issue-1746.test.ts
