fix: $containerIndex is not updated after merging imported declarations

packages/plugins/trpc/tests/projects/t3-trpc-v10/package.json

@@ -13,7 +13,7 @@
     "start": "next start"
   },
   "dependencies": {
-    "@prisma/client": "^5.6.0",
+    "@prisma/client": "5.12.0",
     "@t3-oss/env-nextjs": "^0.7.1",
     "@tanstack/react-query": "^4.36.1",
     "@trpc/client": "^10.43.6",
@@ -35,7 +35,7 @@
     "@typescript-eslint/parser": "^6.11.0",
     "eslint": "^8.54.0",
     "eslint-config-next": "^14.0.4",
-    "prisma": "^5.6.0",
+    "prisma": "5.12.0",
     "typescript": "^5.1.6"
   },
   "ct3aMetadata": {

packages/plugins/trpc/tests/projects/t3-trpc-v10/src/server/api/routers/generated/routers/Post.router.ts

@@ -12,6 +12,8 @@ export default function createRouter<Config extends BaseConfig>(router: RouterFa
 
         aggregate: procedure.input($Schema.PostInputSchema.aggregate).query(({ ctx, input }) => checkRead(db(ctx).post.aggregate(input as any))),
 
+        createMany: procedure.input($Schema.PostInputSchema.createMany).mutation(async ({ ctx, input }) => checkMutate(db(ctx).post.createMany(input as any))),
+
         create: procedure.input($Schema.PostInputSchema.create).mutation(async ({ ctx, input }) => checkMutate(db(ctx).post.create(input as any))),
 
         deleteMany: procedure.input($Schema.PostInputSchema.deleteMany).mutation(async ({ ctx, input }) => checkMutate(db(ctx).post.deleteMany(input as any))),
@@ -60,6 +62,20 @@ export interface ClientType<AppRouter extends AnyRouter, Context = AppRouter['_d
             TRPCClientErrorLike<AppRouter>
         >;
 
+    };
+    createMany: {
+
+        useMutation: <T extends Prisma.PostCreateManyArgs>(opts?: UseTRPCMutationOptions<
+            Prisma.PostCreateManyArgs,
+            TRPCClientErrorLike<AppRouter>,
+            Prisma.BatchPayload,
+            Context
+        >,) =>
+            Omit<UseTRPCMutationResult<Prisma.BatchPayload, TRPCClientErrorLike<AppRouter>, Prisma.SelectSubset<T, Prisma.PostCreateManyArgs>, Context>, 'mutateAsync'> & {
+                mutateAsync:
+                <T extends Prisma.PostCreateManyArgs>(variables: T, opts?: UseTRPCMutationOptions<T, TRPCClientErrorLike<AppRouter>, Prisma.BatchPayload, Context>) => Promise<Prisma.BatchPayload>
+            };
+
     };
     create: {
 

packages/plugins/trpc/tests/projects/t3-trpc-v10/src/server/api/routers/generated/routers/User.router.ts

@@ -12,6 +12,8 @@ export default function createRouter<Config extends BaseConfig>(router: RouterFa
 
         aggregate: procedure.input($Schema.UserInputSchema.aggregate).query(({ ctx, input }) => checkRead(db(ctx).user.aggregate(input as any))),
 
+        createMany: procedure.input($Schema.UserInputSchema.createMany).mutation(async ({ ctx, input }) => checkMutate(db(ctx).user.createMany(input as any))),
+
         create: procedure.input($Schema.UserInputSchema.create).mutation(async ({ ctx, input }) => checkMutate(db(ctx).user.create(input as any))),
 
         deleteMany: procedure.input($Schema.UserInputSchema.deleteMany).mutation(async ({ ctx, input }) => checkMutate(db(ctx).user.deleteMany(input as any))),
@@ -60,6 +62,20 @@ export interface ClientType<AppRouter extends AnyRouter, Context = AppRouter['_d
             TRPCClientErrorLike<AppRouter>
         >;
 
+    };
+    createMany: {
+
+        useMutation: <T extends Prisma.UserCreateManyArgs>(opts?: UseTRPCMutationOptions<
+            Prisma.UserCreateManyArgs,
+            TRPCClientErrorLike<AppRouter>,
+            Prisma.BatchPayload,
+            Context
+        >,) =>
+            Omit<UseTRPCMutationResult<Prisma.BatchPayload, TRPCClientErrorLike<AppRouter>, Prisma.SelectSubset<T, Prisma.UserCreateManyArgs>, Context>, 'mutateAsync'> & {
+                mutateAsync:
+                <T extends Prisma.UserCreateManyArgs>(variables: T, opts?: UseTRPCMutationOptions<T, TRPCClientErrorLike<AppRouter>, Prisma.BatchPayload, Context>) => Promise<Prisma.BatchPayload>
+            };
+
     };
     create: {
 

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 }
+        );
+    });
+});