fix(zod): avoid importing Prisma enum, recognize enum fields with default (#2307)

Author: ymc9

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

@@ -340,7 +340,7 @@ export class ZodSchemaGenerator {
         });
         this.sourceFiles.push(sf);
         sf.replaceWithText((writer) => {
-            this.addPreludeAndImports(typeDef, writer, output);
+            this.addPreludeAndImports(typeDef, writer);
 
             writer.write(`const baseSchema = z.object(`);
             writer.inlineBlock(() => {
@@ -383,7 +383,7 @@ export const ${typeDef.name}Schema = ${refineFuncName}(${noRefineSchema});
         return schemaName;
     }
 
-    private addPreludeAndImports(decl: DataModel | TypeDef, writer: CodeBlockWriter, output: string) {
+    private addPreludeAndImports(decl: DataModel | TypeDef, writer: CodeBlockWriter) {
         writer.writeLine(`import { z } from 'zod/${this.zodVersion}';`);
 
         // import user-defined enums from Prisma as they might be referenced in the expressions
@@ -396,10 +396,6 @@ export const ${typeDef.name}Schema = ${refineFuncName}(${noRefineSchema});
                 }
             }
         }
-        if (importEnums.size > 0) {
-            const prismaImport = computePrismaClientImport(path.join(output, 'models'), this.options);
-            writer.writeLine(`import { ${[...importEnums].join(', ')} } from '${prismaImport}';`);
-        }
 
         // import enum schemas
         const importedEnumSchemas = new Set<string>();
@@ -448,7 +444,7 @@ export const ${typeDef.name}Schema = ${refineFuncName}(${noRefineSchema});
             const relations = model.fields.filter((field) => isDataModel(field.type.reference?.ref));
             const fkFields = model.fields.filter((field) => isForeignKeyField(field));
 
-            this.addPreludeAndImports(model, writer, output);
+            this.addPreludeAndImports(model, writer);
 
             // base schema - including all scalar fields, with optionality following the schema
             this.createModelBaseSchema('baseSchema', writer, scalarFields, true);
@@ -730,9 +726,7 @@ export const ${upperCaseFirst(model.name)}UpdateSchema = ${updateSchema};
     /**
     * Schema refinement function for applying \`@@validate\` rules.
     */
-    export function ${refineFuncName}<T>(schema: z.ZodType<T>) { return schema${refinements.join(
-                    '\n'
-                )};
+    export function ${refineFuncName}<T>(schema: z.ZodType<T>) { return schema${refinements.join('\n')};
     }
     `
             );
@@ -766,6 +760,7 @@ export const ${upperCaseFirst(model.name)}UpdateSchema = ${updateSchema};
                     let expr = new TypeScriptExpressionTransformer({
                         context: ExpressionContext.ValidationRule,
                         fieldReferenceContext: 'value',
+                        useLiteralEnum: true,
                     }).transform(valueArg);
 
                     if (isDataModelFieldReference(valueArg)) {

packages/schema/src/plugins/zod/utils/schema-gen.ts

@@ -1,5 +1,5 @@
 import { upperCaseFirst } from '@zenstackhq/runtime/local-helpers';
-import { getLiteral, hasAttribute, isFromStdlib } from '@zenstackhq/sdk';
+import { getLiteral, hasAttribute, isEnumFieldReference, isFromStdlib } from '@zenstackhq/sdk';
 import {
     DataModelField,
     DataModelFieldAttribute,
@@ -246,6 +246,8 @@ export function getFieldSchemaDefault(field: DataModelField | TypeDefField) {
             return arg.value.value;
         } else if (isBooleanLiteral(arg.value)) {
             return arg.value.value;
+        } else if (isEnumFieldReference(arg.value) && arg.value.target.ref) {
+            return JSON.stringify(arg.value.target.ref.name);
         } else if (
             isInvocationExpr(arg.value) &&
             isFromStdlib(arg.value.function.ref!) &&

packages/sdk/src/typescript-expression-transformer.ts

@@ -40,6 +40,7 @@ type Options = {
     futureRefContext?: string;
     context: ExpressionContext;
     operationContext?: 'read' | 'create' | 'update' | 'postUpdate' | 'delete';
+    useLiteralEnum?: boolean;
 };
 
 type Casing = 'original' | 'upper' | 'lower' | 'capitalize' | 'uncapitalize';
@@ -392,7 +393,9 @@ export class TypeScriptExpressionTransformer {
         }
 
         if (isEnumField(expr.target.ref)) {
-            return `${expr.target.ref.$container.name}.${expr.target.ref.name}`;
+            return this.options.useLiteralEnum
+                ? JSON.stringify(expr.target.ref.name)
+                : `${expr.target.ref.$container.name}.${expr.target.ref.name}`;
         } else {
             if (this.options?.isPostGuard) {
                 // if we're processing post-update, any direct field access should be

tests/regression/tests/issue-2291.test.ts

@@ -0,0 +1,23 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('Issue 2291', () => {
+    it('should work', async () => {
+        const { zodSchemas } = await loadSchema(
+            `
+enum SomeEnum {
+    Ex1
+    Ex2
+}
+
+/// Post model
+model Post {
+    id        String   @id @default(cuid())
+    e SomeEnum @default(Ex1)
+}
+`,
+            { fullZod: true }
+        );
+
+        expect(zodSchemas.models.PostSchema.parse({ id: '1' })).toEqual({ id: '1', e: 'Ex1' });
+    });
+});