fix(zmodel): resolve auth() from all loaded and reachable documents (#1428)

Author: ymc9

packages/schema/src/language-server/utils.ts

@@ -1,4 +1,10 @@
-import { DataModel, DataModelField, isArrayExpr, isReferenceExpr, ReferenceExpr } from '@zenstackhq/language/ast';
+import {
+    isArrayExpr,
+    isReferenceExpr,
+    type DataModel,
+    type DataModelField,
+    type ReferenceExpr,
+} from '@zenstackhq/language/ast';
 import { resolved } from '@zenstackhq/sdk';
 
 /**

packages/schema/src/language-server/zmodel-linker.ts

@@ -35,13 +35,7 @@ import {
     isReferenceExpr,
     isStringLiteral,
 } from '@zenstackhq/language/ast';
-import {
-    getAuthModel,
-    getContainingModel,
-    getModelFieldsWithBases,
-    isAuthInvocation,
-    isFutureExpr,
-} from '@zenstackhq/sdk';
+import { getAuthModel, getModelFieldsWithBases, isAuthInvocation, isFutureExpr } from '@zenstackhq/sdk';
 import {
     AstNode,
     AstNodeDescription,
@@ -52,13 +46,14 @@ import {
     LangiumServices,
     LinkingError,
     Reference,
+    getContainerOfType,
     interruptAndCheck,
     isReference,
     streamContents,
 } from 'langium';
 import { match } from 'ts-pattern';
 import { CancellationToken } from 'vscode-jsonrpc';
-import { getAllDataModelsIncludingImports, getContainingDataModel } from '../utils/ast-utils';
+import { getAllLoadedAndReachableDataModels, getContainingDataModel } from '../utils/ast-utils';
 import { mapBuiltinTypeToExpressionType } from './validator/utils';
 
 interface DefaultReference extends Reference {
@@ -283,15 +278,17 @@ export class ZModelLinker extends DefaultLinker {
             // eslint-disable-next-line @typescript-eslint/ban-types
             const funcDecl = node.function.ref as FunctionDecl;
             if (isAuthInvocation(node)) {
-                // auth() function is resolved to User model in the current document
-                const model = getContainingModel(node);
-
-                if (model) {
-                    const allDataModels = getAllDataModelsIncludingImports(this.langiumDocuments(), model);
-                    const authModel = getAuthModel(allDataModels);
-                    if (authModel) {
-                        node.$resolvedType = { decl: authModel, nullable: true };
-                    }
+                // auth() function is resolved against all loaded and reachable documents
+
+                // get all data models from loaded and reachable documents
+                const allDataModels = getAllLoadedAndReachableDataModels(
+                    this.langiumDocuments(),
+                    getContainerOfType(node, isDataModel)
+                );
+
+                const authModel = getAuthModel(allDataModels);
+                if (authModel) {
+                    node.$resolvedType = { decl: authModel, nullable: true };
                 }
             } else if (isFutureExpr(node)) {
                 // future() function is resolved to current model

packages/schema/src/language-server/zmodel-scope.ts

@@ -32,7 +32,7 @@ import {
 import { match } from 'ts-pattern';
 import { CancellationToken } from 'vscode-jsonrpc';
 import {
-    getAllDataModelsIncludingImports,
+    getAllLoadedAndReachableDataModels,
     isCollectionPredicate,
     isFutureInvocation,
     resolveImportUri,
@@ -219,18 +219,18 @@ export class ZModelScopeProvider extends DefaultScopeProvider {
     }
 
     private createScopeForAuthModel(node: AstNode, globalScope: Scope) {
-        const model = getContainerOfType(node, isModel);
-        if (model) {
-            const allDataModels = getAllDataModelsIncludingImports(
-                this.services.shared.workspace.LangiumDocuments,
-                model
-            );
-            const authModel = getAuthModel(allDataModels);
-            if (authModel) {
-                return this.createScopeForModel(authModel, globalScope);
-            }
+        // get all data models from loaded and reachable documents
+        const allDataModels = getAllLoadedAndReachableDataModels(
+            this.services.shared.workspace.LangiumDocuments,
+            getContainerOfType(node, isDataModel)
+        );
+
+        const authModel = getAuthModel(allDataModels);
+        if (authModel) {
+            return this.createScopeForModel(authModel, globalScope);
+        } else {
+            return EMPTY_SCOPE;
         }
-        return EMPTY_SCOPE;
     }
 }
 

packages/schema/src/utils/ast-utils.ts

@@ -280,3 +280,36 @@ export function findUpAst(node: AstNode, predicate: (node: AstNode) => boolean):
     }
     return undefined;
 }
+
+/**
+ * Gets all data models from all loaded documents
+ */
+export function getAllLoadedDataModels(langiumDocuments: LangiumDocuments) {
+    return langiumDocuments.all
+        .map((doc) => doc.parseResult.value as Model)
+        .flatMap((model) => model.declarations.filter(isDataModel))
+        .toArray();
+}
+
+/**
+ * Gets all data models from loaded and reachable documents
+ */
+export function getAllLoadedAndReachableDataModels(langiumDocuments: LangiumDocuments, fromModel?: DataModel) {
+    // get all data models from loaded documents
+    const allDataModels = getAllLoadedDataModels(langiumDocuments);
+
+    if (fromModel) {
+        // merge data models transitively reached from the current model
+        const model = getContainerOfType(fromModel, isModel);
+        if (model) {
+            const transitiveDataModels = getAllDataModelsIncludingImports(langiumDocuments, model);
+            transitiveDataModels.forEach((dm) => {
+                if (!allDataModels.includes(dm)) {
+                    allDataModels.push(dm);
+                }
+            });
+        }
+    }
+
+    return allDataModels;
+}

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

@@ -0,0 +1,26 @@
+import { FILE_SPLITTER, loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 1388', () => {
+    it('regression', async () => {
+        await loadSchema(
+            `schema.zmodel
+            import './auth'
+            import './post'
+
+            ${FILE_SPLITTER}auth.zmodel
+            model User {
+              id   String @id @default(cuid())
+              role String
+            }
+            
+            ${FILE_SPLITTER}post.zmodel
+            model Post {
+              id        String  @id @default(nanoid(6))
+              title String
+              @@deny('all', auth() == null)
+              @@allow('all', auth().id == 'user1')
+            }
+            `
+        );
+    });
+});