From 08d787b2f23283af9f2d61cf49cc375256b0eb19 Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Sat, 15 Nov 2025 14:36:44 -0800
Subject: [PATCH] feat(cli): allow importing with file extension in generated
 schema

fixes #378
---
 packages/cli/src/actions/generate.ts    | 14 ++++++++-
 packages/cli/src/plugins/typescript.ts  | 13 +++++++-
 packages/sdk/src/ts-schema-generator.ts | 40 ++++++++++++++++++++++---
 3 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/packages/cli/src/actions/generate.ts b/packages/cli/src/actions/generate.ts
index b8aedcb5..bd11c622 100644
--- a/packages/cli/src/actions/generate.ts
+++ b/packages/cli/src/actions/generate.ts
@@ -86,7 +86,19 @@ async function runPlugins(schemaFile: string, model: Model, outputPath: string,
         }
 
         if (cliPlugin) {
-            processedPlugins.push({ cliPlugin, pluginOptions: getPluginOptions(plugin) });
+            const pluginOptions = getPluginOptions(plugin);
+
+            // merge CLI options
+            if (provider === '@core/typescript') {
+                if (pluginOptions['lite'] === undefined) {
+                    pluginOptions['lite'] = options.lite;
+                }
+                if (pluginOptions['liteOnly'] === undefined) {
+                    pluginOptions['liteOnly'] = options.liteOnly;
+                }
+            }
+
+            processedPlugins.push({ cliPlugin, pluginOptions });
         }
     }
 
diff --git a/packages/cli/src/plugins/typescript.ts b/packages/cli/src/plugins/typescript.ts
index 25f950ab..8b3465e8 100644
--- a/packages/cli/src/plugins/typescript.ts
+++ b/packages/cli/src/plugins/typescript.ts
@@ -22,7 +22,18 @@ const plugin: CliPlugin = {
         // liteOnly mode
         const liteOnly = pluginOptions['liteOnly'] === true;
 
-        await new TsSchemaGenerator().generate(model, { outDir, lite, liteOnly });
+        // add .js extension when importing
+        const importWithFileExtension = pluginOptions['importWithFileExtension'];
+        if (importWithFileExtension && typeof importWithFileExtension !== 'string') {
+            throw new Error('The "importWithFileExtension" option must be a string if specified.');
+        }
+
+        await new TsSchemaGenerator().generate(model, {
+            outDir,
+            lite,
+            liteOnly,
+            importWithFileExtension: importWithFileExtension as string | undefined,
+        });
     },
 };
 
diff --git a/packages/sdk/src/ts-schema-generator.ts b/packages/sdk/src/ts-schema-generator.ts
index d1d19fe1..3a7e33bf 100644
--- a/packages/sdk/src/ts-schema-generator.ts
+++ b/packages/sdk/src/ts-schema-generator.ts
@@ -55,6 +55,7 @@ export type TsSchemaGeneratorOptions = {
     outDir: string;
     lite?: boolean;
     liteOnly?: boolean;
+    importWithFileExtension?: string;
 };
 
 export class TsSchemaGenerator {
@@ -117,6 +118,7 @@ export class TsSchemaGenerator {
         const schemaObject = this.createSchemaObject(model, lite);
 
         // Now generate the import declaration with the correct imports
+        // import { type SchemaDef, type OperandExpression, ExpressionUtils } from '@zenstackhq/orm/schema';
         const runtimeImportDecl = ts.factory.createImportDeclaration(
             undefined,
             ts.factory.createImportClause(
@@ -1290,7 +1292,15 @@ export class TsSchemaGenerator {
         const statements: ts.Statement[] = [];
 
         // generate: import { schema as $schema, type SchemaType as $Schema } from './schema';
-        statements.push(this.generateSchemaImport(model, true, true, !!(options.lite || options.liteOnly)));
+        statements.push(
+            this.generateSchemaImport(
+                model,
+                true,
+                true,
+                !!(options.lite || options.liteOnly),
+                options.importWithFileExtension,
+            ),
+        );
 
         // generate: import type { ModelResult as $ModelResult } from '@zenstackhq/orm';
         statements.push(
@@ -1416,7 +1426,13 @@ export class TsSchemaGenerator {
         fs.writeFileSync(outputFile, result);
     }
 
-    private generateSchemaImport(model: Model, schemaObject: boolean, schemaType: boolean, useLite: boolean) {
+    private generateSchemaImport(
+        model: Model,
+        schemaObject: boolean,
+        schemaType: boolean,
+        useLite: boolean,
+        importWithFileExtension: string | undefined,
+    ) {
         const importSpecifiers = [];
 
         if (schemaObject) {
@@ -1442,10 +1458,18 @@ export class TsSchemaGenerator {
             );
         }
 
+        let importFrom = useLite ? './schema-lite' : './schema';
+        if (importWithFileExtension) {
+            importFrom += importWithFileExtension.startsWith('.')
+                ? importWithFileExtension
+                : `.${importWithFileExtension}`;
+        }
+
+        // import { schema as $schema, type SchemaType as $Schema } from './schema';
         return ts.factory.createImportDeclaration(
             undefined,
             ts.factory.createImportClause(false, undefined, ts.factory.createNamedImports(importSpecifiers)),
-            ts.factory.createStringLiteral(useLite ? './schema-lite' : './schema'),
+            ts.factory.createStringLiteral(importFrom),
         );
     }
 
@@ -1466,7 +1490,15 @@ export class TsSchemaGenerator {
         const statements: ts.Statement[] = [];
 
         // generate: import { SchemaType as $Schema } from './schema';
-        statements.push(this.generateSchemaImport(model, false, true, !!(options.lite || options.liteOnly)));
+        statements.push(
+            this.generateSchemaImport(
+                model,
+                false,
+                true,
+                !!(options.lite || options.liteOnly),
+                options.importWithFileExtension,
+            ),
+        );
 
         // generate: import { CreateArgs as $CreateArgs, ... } from '@zenstackhq/orm';
         const inputTypes = [