feat: implement zmodel import (#126)

* feat: implement zmodel import

* addressing PR comments

* update

Author: ymc9

packages/cli/package.json

@@ -44,12 +44,12 @@
     },
     "devDependencies": {
         "@types/better-sqlite3": "^7.6.13",
-        "@types/tmp": "^0.2.6",
+        "@types/tmp": "catalog:",
         "@zenstackhq/eslint-config": "workspace:*",
         "@zenstackhq/runtime": "workspace:*",
         "@zenstackhq/testtools": "workspace:*",
         "@zenstackhq/typescript-config": "workspace:*",
         "better-sqlite3": "^11.8.1",
-        "tmp": "^0.2.3"
+        "tmp": "catalog:"
     }
 }

packages/cli/src/actions/action-utils.ts

@@ -1,4 +1,5 @@
 import { loadDocument } from '@zenstackhq/language';
+import { isDataSource } from '@zenstackhq/language/ast';
 import { PrismaSchemaGenerator } from '@zenstackhq/sdk';
 import colors from 'colors';
 import fs from 'node:fs';
@@ -41,6 +42,9 @@ export async function loadSchemaDocument(schemaFile: string) {
         });
         throw new CliError('Failed to load schema');
     }
+    loadResult.warnings.forEach((warn) => {
+        console.warn(colors.yellow(warn));
+    });
     return loadResult.model;
 }
 
@@ -54,6 +58,9 @@ export function handleSubProcessError(err: unknown) {
 
 export async function generateTempPrismaSchema(zmodelPath: string, folder?: string) {
     const model = await loadSchemaDocument(zmodelPath);
+    if (!model.declarations.some(isDataSource)) {
+        throw new CliError('Schema must define a datasource');
+    }
     const prismaSchema = await new PrismaSchemaGenerator(model).generate();
     if (!folder) {
         folder = path.dirname(zmodelPath);

packages/cli/test/ts-schema-gen.test.ts

@@ -1,5 +1,7 @@
 import { ExpressionUtils } from '@zenstackhq/runtime/schema';
-import { generateTsSchema } from '@zenstackhq/testtools';
+import { createTestProject, generateTsSchema, generateTsSchemaInPlace } from '@zenstackhq/testtools';
+import fs from 'node:fs';
+import path from 'node:path';
 import { describe, expect, it } from 'vitest';
 
 describe('TypeScript schema generation tests', () => {
@@ -325,4 +327,37 @@ model User extends Base {
             },
         });
     });
+
+    it('merges all declarations from imported modules', async () => {
+        const workDir = createTestProject();
+        fs.writeFileSync(
+            path.join(workDir, 'a.zmodel'),
+            `
+        enum Role {
+          Admin
+          User
+        }
+          `,
+        );
+        fs.writeFileSync(
+            path.join(workDir, 'b.zmodel'),
+            `
+        import './a'
+
+        datasource db {
+            provider = 'sqlite'
+            url = 'file:./test.db'
+        }
+
+        model User {
+          id Int @id
+          role Role
+        }
+        `,
+        );
+
+        const { schema } = await generateTsSchemaInPlace(path.join(workDir, 'b.zmodel'));
+        expect(schema.enums).toMatchObject({ Role: expect.any(Object) });
+        expect(schema.models).toMatchObject({ User: expect.any(Object) });
+    });
 });

packages/language/package.json

@@ -62,7 +62,9 @@
         "@zenstackhq/eslint-config": "workspace:*",
         "@zenstackhq/typescript-config": "workspace:*",
         "@zenstackhq/common-helpers": "workspace:*",
-        "langium-cli": "catalog:"
+        "langium-cli": "catalog:",
+        "tmp": "catalog:",
+        "@types/tmp": "catalog:"
     },
     "volta": {
         "node": "18.19.1",

packages/language/src/index.ts

@@ -1,11 +1,12 @@
-import { URI } from 'langium';
+import { isAstNode, URI, type LangiumDocument, type LangiumDocuments, type Mutable } from 'langium';
 import { NodeFileSystem } from 'langium/node';
 import fs from 'node:fs';
 import path from 'node:path';
 import { fileURLToPath } from 'node:url';
-import type { Model } from './ast';
+import { isDataSource, type AstNode, type Model } from './ast';
 import { STD_LIB_MODULE_NAME } from './constants';
 import { createZModelLanguageServices } from './module';
+import { getDataModelAndTypeDefs, getDocument, hasAttribute, resolveImport, resolveTransitiveImports } from './utils';
 
 export function createZModelServices() {
     return createZModelLanguageServices(NodeFileSystem);
@@ -60,8 +61,15 @@ export async function loadDocument(
     const langiumDocuments = services.shared.workspace.LangiumDocuments;
     const document = await langiumDocuments.getOrCreateDocument(URI.file(path.resolve(fileName)));
 
+    // load imports
+    const importedURIs = await loadImports(document, langiumDocuments);
+    const importedDocuments: LangiumDocument[] = [];
+    for (const uri of importedURIs) {
+        importedDocuments.push(await langiumDocuments.getOrCreateDocument(uri));
+    }
+
     // build the document together with standard library, plugin modules, and imported documents
-    await services.shared.workspace.DocumentBuilder.build([stdLib, ...pluginDocs, document], {
+    await services.shared.workspace.DocumentBuilder.build([stdLib, ...pluginDocs, document, ...importedDocuments], {
         validation: true,
     });
 
@@ -95,11 +103,100 @@ export async function loadDocument(
         };
     }
 
+    const model = document.parseResult.value as Model;
+
+    // merge all declarations into the main document
+    const imported = mergeImportsDeclarations(langiumDocuments, model);
+
+    // remove imported documents
+    imported.forEach((model) => {
+        langiumDocuments.deleteDocument(model.$document!.uri);
+        services.shared.workspace.IndexManager.remove(model.$document!.uri);
+    });
+
+    // extra validation after merging imported declarations
+    const additionalErrors = validationAfterImportMerge(model);
+    if (additionalErrors.length > 0) {
+        return {
+            success: false,
+            errors: additionalErrors,
+            warnings,
+        };
+    }
+
     return {
         success: true,
         model: document.parseResult.value as Model,
         warnings,
     };
 }
 
+async function loadImports(document: LangiumDocument, documents: LangiumDocuments, uris: Set<string> = new Set()) {
+    const uriString = document.uri.toString();
+    if (!uris.has(uriString)) {
+        uris.add(uriString);
+        const model = document.parseResult.value as Model;
+        for (const imp of model.imports) {
+            const importedModel = resolveImport(documents, imp);
+            if (importedModel) {
+                const importedDoc = getDocument(importedModel);
+                await loadImports(importedDoc, documents, uris);
+            }
+        }
+    }
+    return Array.from(uris)
+        .filter((x) => uriString != x)
+        .map((e) => URI.parse(e));
+}
+
+function mergeImportsDeclarations(documents: LangiumDocuments, model: Model) {
+    const importedModels = resolveTransitiveImports(documents, model);
+
+    const importedDeclarations = importedModels.flatMap((m) => m.declarations);
+    model.declarations.push(...importedDeclarations);
+
+    // remove import directives
+    model.imports = [];
+
+    // fix $container, $containerIndex, and $containerProperty
+    linkContentToContainer(model);
+
+    return importedModels;
+}
+
+function linkContentToContainer(node: AstNode): void {
+    for (const [name, value] of Object.entries(node)) {
+        if (!name.startsWith('$')) {
+            if (Array.isArray(value)) {
+                value.forEach((item, index) => {
+                    if (isAstNode(item)) {
+                        (item as Mutable<AstNode>).$container = node;
+                        (item as Mutable<AstNode>).$containerProperty = name;
+                        (item as Mutable<AstNode>).$containerIndex = index;
+                    }
+                });
+            } else if (isAstNode(value)) {
+                (value as Mutable<AstNode>).$container = node;
+                (value as Mutable<AstNode>).$containerProperty = name;
+            }
+        }
+    }
+}
+
+function validationAfterImportMerge(model: Model) {
+    const errors: string[] = [];
+    const dataSources = model.declarations.filter((d) => isDataSource(d));
+    if (dataSources.length > 1) {
+        errors.push('Validation error: Multiple datasource declarations are not allowed');
+    }
+
+    // at most one `@@auth` model
+    const decls = getDataModelAndTypeDefs(model, true);
+    const authDecls = decls.filter((d) => hasAttribute(d, '@@auth'));
+    if (authDecls.length > 1) {
+        errors.push('Validation error: Multiple `@@auth` declarations are not allowed');
+    }
+    return errors;
+}
+
 export * from './module';

packages/language/src/utils.ts

@@ -1,5 +1,5 @@
 import { invariant } from '@zenstackhq/common-helpers';
-import { AstUtils, URI, type AstNode, type LangiumDocuments, type Reference } from 'langium';
+import { AstUtils, URI, type AstNode, type LangiumDocument, type LangiumDocuments, type Reference } from 'langium';
 import fs from 'node:fs';
 import path from 'path';
 import { STD_LIB_MODULE_NAME, type ExpressionContext } from './constants';
@@ -413,38 +413,13 @@ export function resolveImport(documents: LangiumDocuments, imp: ModelImport) {
 }
 
 export function resolveImportUri(imp: ModelImport) {
-    if (!imp.path) return undefined; // This will return true if imp.path is undefined, null, or an empty string ("").
-
-    if (!imp.path.endsWith('.zmodel')) {
-        imp.path += '.zmodel';
-    }
-
-    if (
-        !imp.path.startsWith('.') && // Respect relative paths
-        !path.isAbsolute(imp.path) // Respect Absolute paths
-    ) {
-        // use the current model's path as the search context
-        const contextPath = imp.$container.$document
-            ? path.dirname(imp.$container.$document.uri.fsPath)
-            : process.cwd();
-        imp.path = findNodeModulesFile(imp.path, contextPath) ?? imp.path;
+    if (!imp.path) {
+        return undefined;
     }
-
     const doc = AstUtils.getDocument(imp);
     const dir = path.dirname(doc.uri.fsPath);
-    return URI.file(path.resolve(dir, imp.path));
-}
-
-export function findNodeModulesFile(name: string, cwd: string = process.cwd()) {
-    if (!name) return undefined;
-    try {
-        // Use require.resolve to find the module/file. The paths option allows specifying the directory to start from.
-        const resolvedPath = require.resolve(name, { paths: [cwd] });
-        return resolvedPath;
-    } catch {
-        // If require.resolve fails to find the module/file, it will throw an error.
-        return undefined;
-    }
+    const importPath = imp.path.endsWith('.zmodel') ? imp.path : `${imp.path}.zmodel`;
+    return URI.file(path.resolve(dir, importPath));
 }
 
 /**
@@ -577,3 +552,28 @@ export function getAllAttributes(
     attributes.push(...decl.attributes);
     return attributes;
 }
+
+/**
+ * Retrieve the document in which the given AST node is contained. A reference to the document is
+ * usually held by the root node of the AST.
+ *
+ * @throws an error if the node is not contained in a document.
+ */
+export function getDocument<T extends AstNode = AstNode>(node: AstNode): LangiumDocument<T> {
+    const rootNode = findRootNode(node);
+    const result = rootNode.$document;
+    if (!result) {
+        throw new Error('AST node has no document.');
+    }
+    return result as LangiumDocument<T>;
+}
+
+/**
+ * Returns the root node of the given AST node by following the `$container` references.
+ */
+export function findRootNode(node: AstNode): AstNode {
+    while (node.$container) {
+        node = node.$container;
+    }
+    return node;
+}

packages/language/src/validators/schema-validator.ts

@@ -47,8 +47,8 @@ export default class SchemaValidator implements AstValidator<Model> {
     private validateImports(model: Model, accept: ValidationAcceptor) {
         model.imports.forEach((imp) => {
             const importedModel = resolveImport(this.documents, imp);
-            const importPath = imp.path.endsWith('.zmodel') ? imp.path : `${imp.path}.zmodel`;
             if (!importedModel) {
+                const importPath = imp.path.endsWith('.zmodel') ? imp.path : `${imp.path}.zmodel`;
                 accept('error', `Cannot find model file ${importPath}`, {
                     node: imp,
                 });