feat: cli plugin support (#181)

* feat: cli plugin support

* update

* update

* more fixes

Author: ymc9

.vscode/launch.json

@@ -10,7 +10,8 @@
             "request": "launch",
             "skipFiles": ["<node_internals>/**"],
             "type": "node",
-            "args": ["generate", "--schema", "${workspaceFolder}/samples/blog/zenstack/schema.zmodel"]
+            "args": ["generate"],
+            "cwd": "${workspaceFolder}/samples/blog/zenstack"
         },
         {
             "name": "Debug with TSX",

package.json

@@ -21,7 +21,7 @@
     "license": "MIT",
     "devDependencies": {
         "@eslint/js": "^9.29.0",
-        "@types/node": "^20.17.24",
+        "@types/node": "catalog:",
         "eslint": "~9.29.0",
         "glob": "^11.0.2",
         "prettier": "^3.5.3",

packages/cli/src/actions/generate.ts

@@ -1,16 +1,18 @@
 import { invariant } from '@zenstackhq/common-helpers';
-import { isPlugin, LiteralExpr, type Model } from '@zenstackhq/language/ast';
-import { PrismaSchemaGenerator, TsSchemaGenerator, type CliGenerator } from '@zenstackhq/sdk';
+import { isPlugin, LiteralExpr, Plugin, type Model } from '@zenstackhq/language/ast';
+import { getLiteral, getLiteralArray } from '@zenstackhq/language/utils';
+import { type CliPlugin } from '@zenstackhq/sdk';
 import colors from 'colors';
-import fs from 'node:fs';
 import path from 'node:path';
+import ora from 'ora';
+import { CliError } from '../cli-error';
+import * as corePlugins from '../plugins';
 import { getPkgJsonConfig, getSchemaFile, loadSchemaDocument } from './action-utils';
 
 type Options = {
     schema?: string;
     output?: string;
     silent?: boolean;
-    savePrismaSchema?: string | boolean;
 };
 
 /**
@@ -24,25 +26,10 @@ export async function run(options: Options) {
     const model = await loadSchemaDocument(schemaFile);
     const outputPath = getOutputPath(options, schemaFile);
 
-    // generate TS schema
-    const tsSchemaFile = path.join(outputPath, 'schema.ts');
-    await new TsSchemaGenerator().generate(schemaFile, [], outputPath);
-
-    await runPlugins(model, outputPath, tsSchemaFile);
-
-    // generate Prisma schema
-    if (options.savePrismaSchema) {
-        const prismaSchema = await new PrismaSchemaGenerator(model).generate();
-        let prismaSchemaFile = path.join(outputPath, 'schema.prisma');
-        if (typeof options.savePrismaSchema === 'string') {
-            prismaSchemaFile = path.resolve(outputPath, options.savePrismaSchema);
-            fs.mkdirSync(path.dirname(prismaSchemaFile), { recursive: true });
-        }
-        fs.writeFileSync(prismaSchemaFile, prismaSchema);
-    }
+    await runPlugins(schemaFile, model, outputPath);
 
     if (!options.silent) {
-        console.log(colors.green(`Generation completed successfully in ${Date.now() - start}ms.`));
+        console.log(colors.green(`Generation completed successfully in ${Date.now() - start}ms.\n`));
         console.log(`You can now create a ZenStack client with it.
 
 \`\`\`ts
@@ -68,18 +55,84 @@ function getOutputPath(options: Options, schemaFile: string) {
     }
 }
 
-async function runPlugins(model: Model, outputPath: string, tsSchemaFile: string) {
+async function runPlugins(schemaFile: string, model: Model, outputPath: string) {
     const plugins = model.declarations.filter(isPlugin);
+    const processedPlugins: { cliPlugin: CliPlugin; pluginOptions: Record<string, unknown> }[] = [];
+
     for (const plugin of plugins) {
-        const providerField = plugin.fields.find((f) => f.name === 'provider');
-        invariant(providerField, `Plugin ${plugin.name} does not have a provider field`);
-        const provider = (providerField.value as LiteralExpr).value as string;
-        let useProvider = provider;
-        if (useProvider.startsWith('@core/')) {
-            useProvider = `@zenstackhq/runtime/plugins/${useProvider.slice(6)}`;
+        const provider = getPluginProvider(plugin);
+
+        let cliPlugin: CliPlugin;
+        if (provider.startsWith('@core/')) {
+            cliPlugin = (corePlugins as any)[provider.slice('@core/'.length)];
+            if (!cliPlugin) {
+                throw new CliError(`Unknown core plugin: ${provider}`);
+            }
+        } else {
+            let moduleSpec = provider;
+            if (moduleSpec.startsWith('.')) {
+                // relative to schema's path
+                moduleSpec = path.resolve(path.dirname(schemaFile), moduleSpec);
+            }
+            try {
+                cliPlugin = (await import(moduleSpec)).default as CliPlugin;
+            } catch (error) {
+                throw new CliError(`Failed to load plugin ${provider}: ${error}`);
+            }
+        }
+
+        processedPlugins.push({ cliPlugin, pluginOptions: getPluginOptions(plugin) });
+    }
+
+    const defaultPlugins = [corePlugins['typescript']].reverse();
+    defaultPlugins.forEach((d) => {
+        if (!processedPlugins.some((p) => p.cliPlugin === d)) {
+            processedPlugins.push({ cliPlugin: d, pluginOptions: {} });
+        }
+    });
+
+    for (const { cliPlugin, pluginOptions } of processedPlugins) {
+        invariant(
+            typeof cliPlugin.generate === 'function',
+            `Plugin ${cliPlugin.name} does not have a generate function`,
+        );
+
+        // run plugin generator
+        const spinner = ora(cliPlugin.statusText ?? `Running plugin ${cliPlugin.name}`).start();
+        try {
+            await cliPlugin.generate({
+                schemaFile,
+                model,
+                defaultOutputPath: outputPath,
+                pluginOptions,
+            });
+            spinner.succeed();
+        } catch (err) {
+            spinner.fail();
+            console.error(err);
+        }
+    }
+}
+
+function getPluginProvider(plugin: Plugin) {
+    const providerField = plugin.fields.find((f) => f.name === 'provider');
+    invariant(providerField, `Plugin ${plugin.name} does not have a provider field`);
+    const provider = (providerField.value as LiteralExpr).value as string;
+    return provider;
+}
+
+function getPluginOptions(plugin: Plugin): Record<string, unknown> {
+    const result: Record<string, unknown> = {};
+    for (const field of plugin.fields) {
+        if (field.name === 'provider') {
+            continue; // skip provider
+        }
+        const value = getLiteral(field.value) ?? getLiteralArray(field.value);
+        if (value === undefined) {
+            console.warn(`Plugin "${plugin.name}" option "${field.name}" has unsupported value, skipping`);
+            continue;
         }
-        const generator = (await import(useProvider)).default as CliGenerator;
-        console.log('Running generator:', provider);
-        await generator({ model, outputPath, tsSchemaFile });
+        result[field.name] = value;
     }
+    return result;
 }

packages/cli/src/index.ts

@@ -55,12 +55,6 @@ export function createProgram() {
         .description('Run code generation.')
         .addOption(schemaOption)
         .addOption(new Option('--silent', 'do not print any output'))
-        .addOption(
-            new Option(
-                '--save-prisma-schema [path]',
-                'save a Prisma schema file, by default into the output directory',
-            ),
-        )
         .addOption(new Option('-o, --output <path>', 'default output directory for core plugins'))
         .action(generateAction);
 

packages/cli/src/plugins/index.ts

@@ -0,0 +1,2 @@
+export { default as prisma } from './prisma';
+export { default as typescript } from './typescript';

packages/cli/src/plugins/prisma.ts

@@ -0,0 +1,21 @@
+import { PrismaSchemaGenerator, type CliPlugin } from '@zenstackhq/sdk';
+import fs from 'node:fs';
+import path from 'node:path';
+
+const plugin: CliPlugin = {
+    name: 'Prisma Schema Generator',
+    statusText: 'Generating Prisma schema',
+    async generate({ model, defaultOutputPath, pluginOptions }) {
+        let outDir = defaultOutputPath;
+        if (typeof pluginOptions['output'] === 'string') {
+            outDir = path.resolve(defaultOutputPath, pluginOptions['output']);
+            if (!fs.existsSync(outDir)) {
+                fs.mkdirSync(outDir, { recursive: true });
+            }
+        }
+        const prismaSchema = await new PrismaSchemaGenerator(model).generate();
+        fs.writeFileSync(path.join(outDir, 'schema.prisma'), prismaSchema);
+    },
+};
+
+export default plugin;

packages/cli/src/plugins/typescript.ts

@@ -0,0 +1,21 @@
+import type { CliPlugin } from '@zenstackhq/sdk';
+import { TsSchemaGenerator } from '@zenstackhq/sdk';
+import fs from 'node:fs';
+import path from 'node:path';
+
+const plugin: CliPlugin = {
+    name: 'TypeScript Schema Generator',
+    statusText: 'Generating TypeScript schema',
+    async generate({ model, defaultOutputPath, pluginOptions }) {
+        let outDir = defaultOutputPath;
+        if (typeof pluginOptions['output'] === 'string') {
+            outDir = path.resolve(defaultOutputPath, pluginOptions['output']);
+            if (!fs.existsSync(outDir)) {
+                fs.mkdirSync(outDir, { recursive: true });
+            }
+        }
+        await new TsSchemaGenerator().generate(model, outDir);
+    },
+};
+
+export default plugin;

packages/cli/test/generate.test.ts

@@ -30,18 +30,6 @@ describe('CLI generate command test', () => {
         expect(fs.existsSync(path.join(workDir, 'zenstack/schema.ts'))).toBe(true);
     });
 
-    it('should respect save prisma schema option', () => {
-        const workDir = createProject(model);
-        runCli('generate --save-prisma-schema', workDir);
-        expect(fs.existsSync(path.join(workDir, 'zenstack/schema.prisma'))).toBe(true);
-    });
-
-    it('should respect save prisma schema custom path option', () => {
-        const workDir = createProject(model);
-        runCli('generate --save-prisma-schema "../prisma/schema.prisma"', workDir);
-        expect(fs.existsSync(path.join(workDir, 'prisma/schema.prisma'))).toBe(true);
-    });
-
     it('should respect package.json config', () => {
         const workDir = createProject(model);
         fs.mkdirSync(path.join(workDir, 'foo'));

packages/cli/test/plugins/custom-plugin.test.ts

@@ -0,0 +1,50 @@
+import fs from 'node:fs';
+import path from 'node:path';
+import { describe, expect, it } from 'vitest';
+import { createProject, runCli } from '../utils';
+import { execSync } from 'node:child_process';
+
+describe('Custom plugins tests', () => {
+    it('runs custom plugin generator', () => {
+        const workDir = createProject(`
+plugin custom {
+    provider = '../my-plugin.js'
+    output = '../custom-output'
+}
+
+model User {
+    id String @id @default(cuid())
+}
+`);
+
+        fs.writeFileSync(
+            path.join(workDir, 'my-plugin.ts'),
+            `
+import type { CliPlugin } from '@zenstackhq/sdk';
+import fs from 'node:fs';
+import path from 'node:path';
+
+const plugin: CliPlugin = {
+    name: 'Custom Generator',
+    statusText: 'Generating foo.txt',
+    async generate({ model, defaultOutputPath, pluginOptions }) {
+        let outDir = defaultOutputPath;
+        if (typeof pluginOptions['output'] === 'string') {
+            outDir = path.resolve(defaultOutputPath, pluginOptions['output']);
+            if (!fs.existsSync(outDir)) {
+                fs.mkdirSync(outDir, { recursive: true });
+            }
+        }
+        fs.writeFileSync(path.join(outDir, 'foo.txt'), 'from my plugin');
+    },
+};
+
+export default plugin;
+`,
+        );
+
+        execSync('npx tsc', { cwd: workDir });
+        runCli('generate', workDir);
+        expect(fs.existsSync(path.join(workDir, 'custom-output/foo.txt'))).toBe(true);
+    });
+});

packages/cli/test/plugins/prisma-plugin.test.ts

@@ -0,0 +1,60 @@
+import fs from 'node:fs';
+import path from 'node:path';
+import { describe, expect, it } from 'vitest';
+import { createProject, runCli } from '../utils';
+
+describe('Core plugins tests', () => {
+    it('can automatically generate a TypeScript schema with default output', () => {
+        const workDir = createProject(`
+model User {
+    id String @id @default(cuid())
+}
+`);
+        runCli('generate', workDir);
+        expect(fs.existsSync(path.join(workDir, 'zenstack/schema.ts'))).toBe(true);
+    });
+
+    it('can automatically generate a TypeScript schema with custom output', () => {
+        const workDir = createProject(`
+plugin typescript {
+    provider = '@core/typescript'
+    output = '../generated-schema'
+}
+
+model User {
+    id String @id @default(cuid())
+}
+`);
+        runCli('generate', workDir);
+        expect(fs.existsSync(path.join(workDir, 'generated-schema/schema.ts'))).toBe(true);
+    });
+
+    it('can generate a Prisma schema with default output', () => {
+        const workDir = createProject(`
+plugin prisma {
+    provider = '@core/prisma'
+}
+
+model User {
+    id String @id @default(cuid())
+}
+`);
+        runCli('generate', workDir);
+        expect(fs.existsSync(path.join(workDir, 'zenstack/schema.prisma'))).toBe(true);
+    });
+
+    it('can generate a Prisma schema with custom output', () => {
+        const workDir = createProject(`
+plugin prisma {
+    provider = '@core/prisma'
+    output = './prisma'
+}
+
+model User {
+    id String @id @default(cuid())
+}
+`);
+        runCli('generate', workDir);
+        expect(fs.existsSync(path.join(workDir, 'zenstack/prisma/schema.prisma'))).toBe(true);
+    });
+});