fix: User model not found when using policy in the imported model (#457)

Author: jiashengguo

packages/schema/src/cli/cli-util.ts

@@ -3,7 +3,7 @@ import { getLiteral, PluginError } from '@zenstackhq/sdk';
 import colors from 'colors';
 import fs from 'fs';
 import getLatestVersion from 'get-latest-version';
-import { getDocument, LangiumDocument, LangiumDocuments } from 'langium';
+import { AstNode, getDocument, LangiumDocument, LangiumDocuments, Mutable } from 'langium';
 import { NodeFileSystem } from 'langium/node';
 import ora from 'ora';
 import path from 'path';
@@ -193,7 +193,17 @@ export function eagerLoadAllImports(
 
 export function mergeImportsDeclarations(documents: LangiumDocuments, model: Model) {
     const importedModels = resolveTransitiveImports(documents, model);
-    model.declarations.push(...importedModels.flatMap((m) => m.declarations));
+
+    const importedDeclarations = importedModels.flatMap((m) => m.declarations);
+
+    importedDeclarations.forEach((d) => {
+        const mutable = d as Mutable<AstNode>;
+        // The plugin might use $container to access the model
+        // need to make sure it is always resolved to the main model
+        mutable.$container = model;
+    });
+
+    model.declarations.push(...importedDeclarations);
 }
 
 export async function getPluginDocuments(services: ZModelServices, fileName: string): Promise<LangiumDocument[]> {

packages/testtools/src/schema.ts

@@ -11,6 +11,18 @@ import tmp from 'tmp';
 import { loadDocument } from 'zenstack/cli/cli-util';
 import prismaPlugin from 'zenstack/plugins/prisma';
 
+/** 
+ * Use it to represent multiple files in a single string like this
+   `schema.zmodel
+    import "user"
+    ${FILE_SPLITTER}user.zmodel
+    import "schema"
+    model User {
+    ...
+    }`
+*/
+export const FILE_SPLITTER = '#FILE_SPLITTER#';
+
 export type WeakDbOperations = {
     [key in keyof DbOperations]: (...args: any[]) => Promise<any>;
 };
@@ -112,17 +124,42 @@ export async function loadSchema(
     console.log('Workdir:', projectRoot);
     process.chdir(projectRoot);
 
-    schema = schema.replaceAll('$projectRoot', projectRoot);
-
     let zmodelPath = path.join(projectRoot, 'schema.zmodel');
-    const content = addPrelude ? `${MODEL_PRELUDE}\n${schema}` : schema;
-    if (customSchemaFilePath) {
-        zmodelPath = path.join(projectRoot, customSchemaFilePath);
-        fs.mkdirSync(path.dirname(zmodelPath), { recursive: true });
-        fs.writeFileSync(zmodelPath, content);
+
+    const files = schema.split(FILE_SPLITTER);
+
+    if (files.length > 1) {
+        // multiple files
+        files.forEach((file, index) => {
+            //first line is the file name
+            const firstLine = file.indexOf('\n');
+            const fileName = file.substring(0, firstLine).trim();
+            let fileContent = file.substring(firstLine + 1);
+            if (index === 0) {
+                // The first file is the main schema file
+                zmodelPath = path.join(projectRoot, fileName);
+                if (addPrelude) {
+                    // plugin need to be added after import statement
+                    fileContent = `${fileContent}\n${MODEL_PRELUDE}`;
+                }
+            }
+
+            fileContent = fileContent.replaceAll('$projectRoot', projectRoot);
+            const filePath = path.join(projectRoot, fileName);
+            fs.writeFileSync(filePath, fileContent);
+        });
     } else {
-        fs.writeFileSync('schema.zmodel', content);
+        schema = schema.replaceAll('$projectRoot', projectRoot);
+        const content = addPrelude ? `${MODEL_PRELUDE}\n${schema}` : schema;
+        if (customSchemaFilePath) {
+            zmodelPath = path.join(projectRoot, customSchemaFilePath);
+            fs.mkdirSync(path.dirname(zmodelPath), { recursive: true });
+            fs.writeFileSync(zmodelPath, content);
+        } else {
+            fs.writeFileSync('schema.zmodel', content);
+        }
     }
+
     run('npm install');
 
     if (customSchemaFilePath) {

tests/integration/tests/with-policy/multi-file.test.ts

@@ -0,0 +1,32 @@
+import { loadSchema, FILE_SPLITTER } from '@zenstackhq/testtools';
+
+describe('Multiple file test', () => {
+    it('model loading', async () => {
+        await loadSchema(
+            `schema.zmodel
+            import "post"
+             ${FILE_SPLITTER}post.zmodel
+             import "user"
+             model Post {
+                id Int @id @default(autoincrement())
+                owner User @relation(fields: [ownerId], references: [id])
+                ownerId Int
+
+                // require login
+                @@deny('all', auth() == null)
+
+                // can be read by owner or space members (only if not private) 
+                @@allow('all', owner == auth())
+                }
+            ${FILE_SPLITTER}user.zmodel
+            import "post"
+            model User {
+                id Int @id @default(autoincrement())
+                name String
+                email String @unique
+                posts Post[]
+            }
+            `
+        );
+    });
+});