fix: duplicated zod schema imported when there're multiple fields with an enum type

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

@@ -220,10 +220,14 @@ async function generateModelSchema(model: DataModel, project: Project, output: s
         }
 
         // import enum schemas
+        const importedEnumSchemas = new Set<string>();
         for (const field of fields) {
             if (field.type.reference?.ref && isEnum(field.type.reference?.ref)) {
                 const name = upperCaseFirst(field.type.reference?.ref.name);
-                writer.writeLine(`import { ${name}Schema } from '../enums/${name}.schema';`);
+                if (!importedEnumSchemas.has(name)) {
+                    writer.writeLine(`import { ${name}Schema } from '../enums/${name}.schema';`);
+                    importedEnumSchemas.add(name);
+                }
             }
         }
 

packages/testtools/package.json

@@ -23,11 +23,13 @@
     "@zenstackhq/runtime": "workspace:*",
     "@zenstackhq/sdk": "workspace:*",
     "json5": "^2.2.3",
+    "pg": "^8.11.1",
     "tmp": "^0.2.1",
     "zenstack": "workspace:*"
   },
   "devDependencies": {
     "@types/node": "^18.0.0",
+    "@types/pg": "^8.10.2",
     "@types/tmp": "^0.2.3",
     "copyfiles": "^2.4.1",
     "rimraf": "^3.0.2",

packages/testtools/src/db.ts

@@ -0,0 +1,16 @@
+import { Pool } from 'pg';
+
+const USERNAME = 'postgres';
+const PASSWORD = 'abc123';
+
+export async function createPostgresDb(db: string) {
+    const pool = new Pool({ user: USERNAME, password: PASSWORD });
+    await pool.query(`DROP DATABASE IF EXISTS "${db}";`);
+    await pool.query(`CREATE DATABASE "${db}";`);
+    return `postgresql://${USERNAME}:${PASSWORD}@localhost:5432/${db}`;
+}
+
+export async function dropPostgresDb(db: string) {
+    const pool = new Pool({ user: USERNAME, password: PASSWORD });
+    await pool.query(`DROP DATABASE IF EXISTS "${db}";`);
+}

packages/testtools/src/index.ts

@@ -1 +1,2 @@
 export * from './schema';
+export * from './db';

packages/testtools/src/schema.ts

@@ -70,10 +70,11 @@ export function getWorkspaceNpmCacheFolder(start: string) {
     return root ? path.join(root, '.npmcache') : './.npmcache';
 }
 
-const MODEL_PRELUDE = `
+function makePrelude(options: SchemaLoadOptions) {
+    return `
 datasource db {
-    provider = 'sqlite'
-    url = 'file:./test.db'
+    provider = '${options.provider}'
+    url = '${options.dbUrl}'
 }
 
 generator js {
@@ -83,26 +84,10 @@ generator js {
 
 plugin zod {
     provider = '@core/zod'
-    modelOnly = true
+    modelOnly = ${!options.fullZod}
 }
 `;
-
-const MODEL_PRELUDE_FULL_ZOD = `
-datasource db {
-    provider = 'sqlite'
-    url = 'file:./test.db'
-}
-
-generator js {
-    provider = 'prisma-client-js'
-    previewFeatures = ['clientExtensions']
-}
-
-plugin zod {
-    provider = '@core/zod'
-    modelOnly = false
 }
-`;
 
 export type SchemaLoadOptions = {
     addPrelude?: boolean;
@@ -112,6 +97,8 @@ export type SchemaLoadOptions = {
     compile?: boolean;
     customSchemaFilePath?: string;
     logPrismaQuery?: boolean;
+    provider?: 'sqlite' | 'postgresql';
+    dbUrl?: string;
 };
 
 const defaultOptions: SchemaLoadOptions = {
@@ -121,6 +108,8 @@ const defaultOptions: SchemaLoadOptions = {
     extraDependencies: [],
     compile: false,
     logPrismaQuery: false,
+    provider: 'sqlite',
+    dbUrl: 'file:./test.db',
 };
 
 export async function loadSchemaFromFile(schemaFile: string, options?: SchemaLoadOptions) {
@@ -164,7 +153,7 @@ export async function loadSchema(schema: string, options?: SchemaLoadOptions) {
                 zmodelPath = path.join(projectRoot, fileName);
                 if (opt.addPrelude) {
                     // plugin need to be added after import statement
-                    fileContent = `${fileContent}\n${opt.fullZod ? MODEL_PRELUDE_FULL_ZOD : MODEL_PRELUDE}`;
+                    fileContent = `${fileContent}\n${makePrelude(opt)}`;
                 }
             }
 
@@ -174,7 +163,7 @@ export async function loadSchema(schema: string, options?: SchemaLoadOptions) {
         });
     } else {
         schema = schema.replaceAll('$projectRoot', projectRoot);
-        const content = opt.addPrelude ? `${opt.fullZod ? MODEL_PRELUDE_FULL_ZOD : MODEL_PRELUDE}\n${schema}` : schema;
+        const content = opt.addPrelude ? `${makePrelude(opt)}\n${schema}` : schema;
         if (opt.customSchemaFilePath) {
             zmodelPath = path.join(projectRoot, opt.customSchemaFilePath);
             fs.mkdirSync(path.dirname(zmodelPath), { recursive: true });

tests/integration/tests/enhancements/with-policy/field-comparison.test.ts

@@ -1,44 +1,33 @@
-import { loadSchema } from '@zenstackhq/testtools';
+import { loadSchema, createPostgresDb, dropPostgresDb } from '@zenstackhq/testtools';
 import path from 'path';
-import { Pool } from 'pg';
 
 const DB_NAME = 'field-comparison';
 
 describe('WithPolicy: field comparison tests', () => {
     let origDir: string;
+    let dbUrl: string;
     let prisma: any;
 
-    const pool = new Pool({ user: 'postgres', password: 'abc123' });
-
     beforeAll(async () => {
         origDir = path.resolve('.');
     });
 
     beforeEach(async () => {
-        await pool.query(`DROP DATABASE IF EXISTS "${DB_NAME}";`);
-        await pool.query(`CREATE DATABASE "${DB_NAME}";`);
+        dbUrl = await createPostgresDb(DB_NAME);
     });
 
     afterEach(async () => {
-        process.chdir(origDir);
         if (prisma) {
             await prisma.$disconnect();
+            prisma = undefined;
         }
-        await pool.query(`DROP DATABASE IF EXISTS "${DB_NAME}";`);
+        await dropPostgresDb(DB_NAME);
+        process.chdir(origDir);
     });
 
     it('field comparison success with input check', async () => {
         const r = await loadSchema(
             `
-        datasource db {
-            provider = 'postgresql'
-            url = 'postgres://postgres:abc123@localhost:5432/${DB_NAME}'
-        }
-
-        generator js {
-            provider = 'prisma-client-js'
-        }
-
         model Model {
             id String @id @default(uuid())
             x Int
@@ -48,7 +37,7 @@ describe('WithPolicy: field comparison tests', () => {
             @@allow('read', true)
         }
         `,
-            { addPrelude: false }
+            { provider: 'postgresql', dbUrl }
         );
 
         prisma = r.prisma;
@@ -60,15 +49,6 @@ describe('WithPolicy: field comparison tests', () => {
     it('field comparison success with policy check', async () => {
         const r = await loadSchema(
             `
-        datasource db {
-            provider = 'postgresql'
-            url = 'postgres://postgres:abc123@localhost:5432/${DB_NAME}'
-        }
-
-        generator js {
-            provider = 'prisma-client-js'
-        }
-
         model Model {
             id String @id @default(uuid())
             x Int @default(0)
@@ -78,7 +58,7 @@ describe('WithPolicy: field comparison tests', () => {
             @@allow('read', true)
         }
         `,
-            { addPrelude: false }
+            { provider: 'postgresql', dbUrl }
         );
 
         prisma = r.prisma;
@@ -90,15 +70,6 @@ describe('WithPolicy: field comparison tests', () => {
     it('field in operator success with input check', async () => {
         const r = await loadSchema(
             `
-        datasource db {
-            provider = 'postgresql'
-            url = 'postgres://postgres:abc123@localhost:5432/${DB_NAME}'
-        }
-
-        generator js {
-            provider = 'prisma-client-js'
-        }
-
         model Model {
             id String @id @default(uuid())
             x String
@@ -108,7 +79,7 @@ describe('WithPolicy: field comparison tests', () => {
             @@allow('read', x in y)
         }
         `,
-            { addPrelude: false }
+            { provider: 'postgresql', dbUrl }
         );
 
         prisma = r.prisma;
@@ -120,15 +91,6 @@ describe('WithPolicy: field comparison tests', () => {
     it('field in operator success with policy check', async () => {
         const r = await loadSchema(
             `
-        datasource db {
-            provider = 'postgresql'
-            url = 'postgres://postgres:abc123@localhost:5432/${DB_NAME}'
-        }
-
-        generator js {
-            provider = 'prisma-client-js'
-        }
-
         model Model {
             id String @id @default(uuid())
             x String @default('x')
@@ -138,7 +100,7 @@ describe('WithPolicy: field comparison tests', () => {
             @@allow('read', x in y)
         }
         `,
-            { addPrelude: false }
+            { provider: 'postgresql', dbUrl }
         );
 
         prisma = r.prisma;