fix(cli): enhancer code fails to compile when generated into a custom folder (#1678)

Author: ymc9

packages/schema/src/plugins/enhancer/enhance/index.ts

@@ -38,6 +38,7 @@ import {
 import { upperCaseFirst } from 'upper-case-first';
 import { name } from '..';
 import { execPackage } from '../../../utils/exec-utils';
+import { CorePlugins, getPluginCustomOutputFolder } from '../../plugin-utils';
 import { trackPrismaSchemaError } from '../../prisma';
 import { PrismaSchemaGenerator } from '../../prisma/schema-generator';
 import { isDefaultWithAuth } from '../enhancer-utils';
@@ -100,7 +101,11 @@ import { type EnhancementContext, type EnhancementOptions, type ZodSchemas, type
 import { createEnhancement } from '@zenstackhq/runtime/enhancements';
 import modelMeta from './model-meta';
 import policy from './policy';
-${this.options.withZodSchemas ? "import * as zodSchemas from './zod';" : 'const zodSchemas = undefined;'}
+${
+    this.options.withZodSchemas
+        ? `import * as zodSchemas from '${this.getZodImport()}';`
+        : 'const zodSchemas = undefined;'
+}
 
 ${
     logicalPrismaClientDir
@@ -126,6 +131,27 @@ ${
         return { dmmf };
     }
 
+    private getZodImport() {
+        const zodCustomOutput = getPluginCustomOutputFolder(this.model, CorePlugins.Zod);
+
+        if (!this.options.output && !zodCustomOutput) {
+            // neither zod or me (enhancer) have custom output, use the default
+            return './zod';
+        }
+
+        if (!zodCustomOutput) {
+            // I have a custom output, but zod doesn't, import from runtime
+            return '@zenstackhq/runtime/zod';
+        }
+
+        // both zod and me have custom output, resolve to relative path and import
+        const schemaDir = path.dirname(this.options.schemaPath);
+        const zodAbsPath = path.isAbsolute(zodCustomOutput)
+            ? zodCustomOutput
+            : path.resolve(schemaDir, zodCustomOutput);
+        return path.relative(this.outDir, zodAbsPath);
+    }
+
     private createSimplePrismaImports(prismaImport: string) {
         return `import { Prisma, type PrismaClient } from '${prismaImport}';
 import type * as _P from '${prismaImport}';

packages/schema/src/plugins/plugin-utils.ts

@@ -1,8 +1,9 @@
 import { DEFAULT_RUNTIME_LOAD_PATH, type PolicyOperationKind } from '@zenstackhq/runtime';
-import { PluginGlobalOptions, ensureEmptyDir } from '@zenstackhq/sdk';
+import { PluginGlobalOptions, ensureEmptyDir, getLiteral } from '@zenstackhq/sdk';
 import fs from 'fs';
 import path from 'path';
 import { PluginRunnerOptions } from '../cli/plugin-runner';
+import { isPlugin, Model, Plugin } from '@zenstackhq/sdk/ast';
 
 export const ALL_OPERATION_KINDS: PolicyOperationKind[] = ['create', 'update', 'postUpdate', 'read', 'delete'];
 
@@ -114,3 +115,18 @@ export enum CorePlugins {
     Zod = '@core/zod',
     Enhancer = '@core/enhancer',
 }
+
+/**
+ * Gets the custom output folder for a plugin.
+ */
+export function getPluginCustomOutputFolder(zmodel: Model, provider: string) {
+    const plugin = zmodel.declarations.find(
+        (d): d is Plugin =>
+            isPlugin(d) && d.fields.some((f) => f.name === 'provider' && getLiteral<string>(f.value) === provider)
+    );
+    if (!plugin) {
+        return undefined;
+    }
+    const output = plugin.fields.find((f) => f.name === 'output');
+    return output && getLiteral<string>(output.value);
+}

packages/testtools/src/schema.ts

@@ -135,6 +135,7 @@ export type SchemaLoadOptions = {
     generatePermissionChecker?: boolean;
     previewFeatures?: string[];
     prismaClientOptions?: object;
+    generateNoCompile?: boolean;
 };
 
 const defaultOptions: SchemaLoadOptions = {
@@ -146,6 +147,7 @@ const defaultOptions: SchemaLoadOptions = {
     logPrismaQuery: false,
     provider: 'sqlite',
     preserveTsFiles: false,
+    generateNoCompile: false,
 };
 
 export async function loadSchemaFromFile(schemaFile: string, options?: SchemaLoadOptions) {
@@ -225,13 +227,20 @@ export async function loadSchema(schema: string, options?: SchemaLoadOptions) {
     }
 
     const outputArg = opt.output ? ` --output ${opt.output}` : '';
+    let otherArgs = '';
+    if (opt.generateNoCompile) {
+        otherArgs = ' --no-compile';
+    }
 
     if (opt.customSchemaFilePath) {
-        run(`npx zenstack generate --no-version-check --schema ${zmodelPath} --no-dependency-check${outputArg}`, {
-            NODE_PATH: './node_modules',
-        });
+        run(
+            `npx zenstack generate --no-version-check --schema ${zmodelPath} --no-dependency-check${outputArg}${otherArgs}`,
+            {
+                NODE_PATH: './node_modules',
+            }
+        );
     } else {
-        run(`npx zenstack generate --no-version-check --no-dependency-check${outputArg}`, {
+        run(`npx zenstack generate --no-version-check --no-dependency-check${outputArg}${otherArgs}`, {
             NODE_PATH: './node_modules',
         });
     }

tests/integration/tests/plugins/zod.test.ts

@@ -497,7 +497,7 @@ describe('Zod plugin tests', () => {
         expect(schema.safeParse({ arr: [4] }).error.issues[1].path).toEqual(['arr', 'every']);
         expect(schema.safeParse({ arr: [4] }).error.issues[2].path).toEqual(['arr', 'some']);
         expect(schema.safeParse({ arr: [1, 2, 3] }).success).toBeTruthy();
-    })
+    });
 
     it('full-text search', async () => {
         const model = `
@@ -788,6 +788,11 @@ describe('Zod plugin tests', () => {
             provider = 'prisma-client-js'
         }
     
+        plugin enhancer {
+            provider = '@core/enhancer'
+            output = "$projectRoot/enhance"
+        }
+    
         plugin zod {
             provider = "@core/zod"
             output = "$projectRoot/zod"
@@ -801,7 +806,7 @@ describe('Zod plugin tests', () => {
             password String @omit
         }
         `,
-            { addPrelude: false, pushDb: false, projectDir }
+            { addPrelude: false, pushDb: false, projectDir, getPrismaOnly: true, generateNoCompile: true }
         );
 
         expect(fs.existsSync(path.join(projectDir, 'zod', 'test.txt'))).toBeFalsy();
@@ -822,6 +827,11 @@ describe('Zod plugin tests', () => {
         generator js {
             provider = 'prisma-client-js'
         }
+
+        plugin enhancer {
+            provider = '@core/enhancer'
+            output = "$projectRoot/enhance"
+        }
     
         plugin zod {
             provider = "@core/zod"
@@ -836,7 +846,7 @@ describe('Zod plugin tests', () => {
             password String @omit
         }
         `,
-                { addPrelude: false, pushDb: false, projectDir }
+                { addPrelude: false, pushDb: false, projectDir, generateNoCompile: true }
             )
         ).rejects.toThrow('already exists and is not a directory');
     });

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

@@ -0,0 +1,60 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 1667', () => {
+    it('custom enhance standard zod output', async () => {
+        await loadSchema(
+            `
+            generator client {
+                provider = "prisma-client-js"
+            }
+
+            datasource db {
+                provider = "sqlite"
+                url      = "file:./dev.db"
+            }            
+
+            plugin enhancer {
+                provider = '@core/enhancer'
+                output = './zen'
+            }
+
+            model User {
+                id Int @id
+                email String @unique @email
+            }
+            `,
+            { addPrelude: false, getPrismaOnly: true, preserveTsFiles: true }
+        );
+    });
+
+    it('custom enhance custom zod output', async () => {
+        await loadSchema(
+            `
+            generator client {
+                provider = "prisma-client-js"
+            }
+
+            datasource db {
+                provider = "sqlite"
+                url      = "file:./dev.db"
+            }            
+
+            plugin enhancer {
+                provider = '@core/enhancer'
+                output = './zen'
+            }
+
+            plugin zod {
+                provider = '@core/zod'
+                output = './myzod'
+            }
+
+            model User {
+                id Int @id
+                email String @unique @email
+            }
+            `,
+            { addPrelude: false, getPrismaOnly: true, generateNoCompile: true, compile: true }
+        );
+    });
+});