fix: $containerIndex is not updated after merging imported declarations

Author: ymc9

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

@@ -3,7 +3,7 @@ import { getDataModels, getLiteral, hasAttribute } from '@zenstackhq/sdk';
 import colors from 'colors';
 import fs from 'fs';
 import getLatestVersion from 'get-latest-version';
-import { AstNode, getDocument, LangiumDocument, LangiumDocuments, Mutable } from 'langium';
+import { getDocument, LangiumDocument, LangiumDocuments, linkContentToContainer } from 'langium';
 import { NodeFileSystem } from 'langium/node';
 import path from 'path';
 import semver from 'semver';
@@ -153,19 +153,14 @@ export function mergeImportsDeclarations(documents: LangiumDocuments, model: Mod
     const importedModels = resolveTransitiveImports(documents, model);
 
     const importedDeclarations = importedModels.flatMap((m) => m.declarations);
-
-    importedDeclarations.forEach((d) => {
-        const mutable = d as Mutable<AstNode>;
-        // Plugins 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);
 
     // remove import directives
     model.imports = [];
 
+    // fix $containerIndex
+    linkContentToContainer(model);
+
     return importedModels;
 }
 

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

@@ -24,6 +24,7 @@ import {
     getContainerOfType,
     getDocument,
     LangiumDocuments,
+    linkContentToContainer,
     Linker,
     Mutable,
     Reference,
@@ -72,6 +73,9 @@ export function mergeBaseModel(model: Model, linker: Linker) {
                 .filter((attr) => attr.decl.$refText !== '@@map')
                 .map((attr) => cloneAst(attr, dataModel, buildReference))
                 .concat(dataModel.attributes);
+
+            // fix $containerIndex
+            linkContentToContainer(dataModel);
         }
 
         dataModel.$baseMerged = true;
@@ -89,8 +93,6 @@ function cloneAst<T extends InheritableNode>(
 ): Mutable<T> {
     const clone = copyAstNode(node, buildReference) as Mutable<T>;
     clone.$container = newContainer;
-    clone.$containerProperty = node.$containerProperty;
-    clone.$containerIndex = node.$containerIndex;
     clone.$inheritedFrom = node.$inheritedFrom ?? getContainerOfType(node, isDataModel);
     return clone;
 }

tests/integration/tests/regression/issue-1210.test.ts

@@ -0,0 +1,92 @@
+import { FILE_SPLITTER, loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 1210', () => {
+    it('regression', async () => {
+        await loadSchema(
+            `schema.zmodel
+            import "./user"
+            import "./tokens"
+            
+            generator client {
+              provider        = "prisma-client-js"
+              binaryTargets   = ["native"]
+              previewFeatures = ["postgresqlExtensions"]
+            }
+            
+            datasource db {
+              provider   = "postgresql"
+              extensions = [citext]
+            
+              url       = env("DATABASE_URL")
+            }
+            
+            plugin zod {
+              provider = '@core/zod'
+            }
+            
+            ${FILE_SPLITTER}base.zmodel
+            abstract model Base {
+              id String @id @default(uuid())
+              createdAt DateTime @default(now())
+              updatedAt DateTime @updatedAt
+              deletedAt DateTime? @omit
+            
+              @@deny('read', deletedAt != null)
+              @@deny('delete', true)
+            }
+                    
+            ${FILE_SPLITTER}tokens.zmodel
+            import "base"
+            import "user"
+
+            model Session extends Base {
+              expiresAt DateTime
+              userId String
+              user User @relation(references: [id], fields: [userId], onDelete: Cascade)
+            }
+
+            ${FILE_SPLITTER}user.zmodel
+            import "base"
+            import "tokens"
+            enum UserRole {
+              USER
+              ADMIN
+            }
+            
+            model User extends Base {
+              email String @unique @db.Citext @email @trim @lower
+              role UserRole @default(USER) @deny('read,update', auth().role != ADMIN)
+            
+              sessions Session[]
+              posts Post[]
+            
+              @@allow('read,create', auth() == this)
+              @@allow('all', auth().role == ADMIN)
+            }
+            
+            abstract model UserEntity extends Base {
+              userId String
+              user User @relation(fields: [userId], references: [id])
+            
+              @@allow('create', userId == auth().id)
+              @@allow('update', userId == auth().id && future().userId == auth().id)
+            
+              @@allow('all', auth().role == ADMIN)
+            }
+            
+            abstract model PrivateUserEntity extends UserEntity {
+              @@allow('read', userId == auth().id)
+            }
+            
+            abstract model PublicUserEntity extends UserEntity {
+              @@allow('read', true)
+            }
+            
+            model Post extends PublicUserEntity {
+              title String
+            }            
+            `,
+            { addPrelude: false, pushDb: false }
+        );
+    });
+});