fix(zmodel): enum is resolved to wrong element after merging base models

Author: ymc9

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

@@ -12,7 +12,7 @@ import { URI } from 'vscode-uri';
 import { PLUGIN_MODULE_NAME, STD_LIB_MODULE_NAME } from '../language-server/constants';
 import { ZModelFormatter } from '../language-server/zmodel-formatter';
 import { createZModelServices, ZModelServices } from '../language-server/zmodel-module';
-import { mergeBaseModel, resolveImport, resolveTransitiveImports } from '../utils/ast-utils';
+import { mergeBaseModels, resolveImport, resolveTransitiveImports } from '../utils/ast-utils';
 import { findUp } from '../utils/pkg-utils';
 import { getVersion } from '../utils/version-utils';
 import { CliError } from './cli-error';
@@ -101,10 +101,10 @@ export async function loadDocument(fileName: string): Promise<Model> {
         imported.map((m) => m.$document!.uri)
     );
 
-    validationAfterMerge(model);
+    validationAfterImportMerge(model);
 
     // merge fields and attributes from base models
-    mergeBaseModel(model, services.references.Linker);
+    mergeBaseModels(model, services.references.Linker);
 
     // finally relink all references
     const relinkedModel = await relinkAll(model, services);
@@ -113,7 +113,7 @@ export async function loadDocument(fileName: string): Promise<Model> {
 }
 
 // check global unique thing after merge imports
-function validationAfterMerge(model: Model) {
+function validationAfterImportMerge(model: Model) {
     const dataSources = model.declarations.filter((d) => isDataSource(d));
     if (dataSources.length == 0) {
         console.error(colors.red('Validation error: Model must define a datasource'));

packages/schema/src/language-server/validator/datamodel-validator.ts

@@ -2,14 +2,19 @@ import {
     ArrayExpr,
     DataModel,
     DataModelField,
-    isDataModel,
-    isStringLiteral,
     ReferenceExpr,
+    isDataModel,
     isEnum,
+    isStringLiteral,
 } from '@zenstackhq/language/ast';
-import { getLiteral, getModelIdFields, getModelUniqueFields, isDelegateModel } from '@zenstackhq/sdk';
-import { AstNode, DiagnosticInfo, getDocument, ValidationAcceptor } from 'langium';
-import { getModelFieldsWithBases } from '../../utils/ast-utils';
+import {
+    getLiteral,
+    getModelFieldsWithBases,
+    getModelIdFields,
+    getModelUniqueFields,
+    isDelegateModel,
+} from '@zenstackhq/sdk';
+import { AstNode, DiagnosticInfo, ValidationAcceptor, getDocument } from 'langium';
 import { IssueCodes, SCALAR_TYPES } from '../constants';
 import { AstValidator } from '../types';
 import { getUniqueFields } from '../utils';

packages/schema/src/language-server/zmodel-code-action.ts

@@ -10,8 +10,8 @@ import {
     getDocument,
 } from 'langium';
 
+import { getModelFieldsWithBases } from '@zenstackhq/sdk';
 import { CodeAction, CodeActionKind, CodeActionParams, Command, Diagnostic } from 'vscode-languageserver';
-import { getModelFieldsWithBases } from '../utils/ast-utils';
 import { IssueCodes } from './constants';
 import { MissingOppositeRelationData } from './validator/datamodel-validator';
 import { ZModelFormatter } from './zmodel-formatter';

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

@@ -17,7 +17,7 @@ import {
     ModelImport,
     ReferenceExpr,
 } from '@zenstackhq/language/ast';
-import { isDelegateModel, isFromStdlib } from '@zenstackhq/sdk';
+import { getModelFieldsWithBases, getRecursiveBases, isDelegateModel, isFromStdlib } from '@zenstackhq/sdk';
 import {
     AstNode,
     copyAstNode,
@@ -47,16 +47,13 @@ type BuildReference = (
     refText: string
 ) => Reference<AstNode>;
 
-export function mergeBaseModel(model: Model, linker: Linker) {
+export function mergeBaseModels(model: Model, linker: Linker) {
     const buildReference = linker.buildReference.bind(linker);
 
-    model.declarations.filter(isDataModel).forEach((decl) => {
-        const dataModel = decl as DataModel;
-
+    model.declarations.filter(isDataModel).forEach((dataModel) => {
         const bases = getRecursiveBases(dataModel).reverse();
         if (bases.length > 0) {
             dataModel.fields = bases
-                // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
                 .flatMap((base) => base.fields)
                 // don't inherit skip-level fields
                 .filter((f) => !f.$inheritedFrom)
@@ -67,16 +64,25 @@ export function mergeBaseModel(model: Model, linker: Linker) {
                 .flatMap((base) => base.attributes.filter((attr) => filterBaseAttribute(base, attr)))
                 .map((attr) => cloneAst(attr, dataModel, buildReference))
                 .concat(dataModel.attributes);
-
-            // fix $containerIndex
-            linkContentToContainer(dataModel);
         }
 
+        // mark base merged
         dataModel.$baseMerged = true;
     });
 
     // remove abstract models
     model.declarations = model.declarations.filter((x) => !(isDataModel(x) && x.isAbstract));
+
+    model.declarations.filter(isDataModel).forEach((dm) => {
+        // remove abstract super types
+        dm.superTypes = dm.superTypes.filter((t) => t.ref && isDelegateModel(t.ref));
+
+        // fix $containerIndex
+        linkContentToContainer(dm);
+    });
+
+    // fix $containerIndex after deleting abstract models
+    linkContentToContainer(model);
 }
 
 function filterBaseAttribute(base: DataModel, attr: DataModelAttribute) {
@@ -244,29 +250,6 @@ export function getContainingDataModel(node: Expression): DataModel | undefined
     return undefined;
 }
 
-export function getModelFieldsWithBases(model: DataModel, includeDelegate = true) {
-    if (model.$baseMerged) {
-        return model.fields;
-    } else {
-        return [...model.fields, ...getRecursiveBases(model, includeDelegate).flatMap((base) => base.fields)];
-    }
-}
-
-export function getRecursiveBases(dataModel: DataModel, includeDelegate = true): DataModel[] {
-    const result: DataModel[] = [];
-    dataModel.superTypes.forEach((superType) => {
-        const baseDecl = superType.ref;
-        if (baseDecl) {
-            if (!includeDelegate && isDelegateModel(baseDecl)) {
-                return;
-            }
-            result.push(baseDecl);
-            result.push(...getRecursiveBases(baseDecl));
-        }
-    });
-    return result;
-}
-
 /**
  * Walk upward from the current AST node to find the first node that satisfies the predicate.
  */

packages/schema/tests/utils.ts

@@ -5,7 +5,7 @@ import * as path from 'path';
 import * as tmp from 'tmp';
 import { URI } from 'vscode-uri';
 import { createZModelServices } from '../src/language-server/zmodel-module';
-import { mergeBaseModel } from '../src/utils/ast-utils';
+import { mergeBaseModels } from '../src/utils/ast-utils';
 
 tmp.setGracefulCleanup();
 
@@ -68,7 +68,7 @@ export async function loadModel(content: string, validate = true, verbose = true
     const model = (await doc.parseResult.value) as Model;
 
     if (mergeBase) {
-        mergeBaseModel(model, ZModel.references.Linker);
+        mergeBaseModels(model, ZModel.references.Linker);
     }
 
     return model;
@@ -87,13 +87,13 @@ export async function loadModelWithError(content: string, verbose = false) {
 }
 
 export async function safelyLoadModel(content: string, validate = true, verbose = false) {
-    const [ result ] = await Promise.allSettled([ loadModel(content, validate, verbose) ]);
+    const [result] = await Promise.allSettled([loadModel(content, validate, verbose)]);
 
-    return result
+    return result;
 }
 
 export const errorLike = (msg: string) => ({
     reason: {
-        message: expect.stringContaining(msg)
+        message: expect.stringContaining(msg),
     },
-})
+});

packages/sdk/src/utils.ts

@@ -504,17 +504,24 @@ export function getDataModelFieldReference(expr: Expression): DataModelField | u
     }
 }
 
-export function getModelFieldsWithBases(model: DataModel) {
-    return [...model.fields, ...getRecursiveBases(model).flatMap((base) => base.fields)];
+export function getModelFieldsWithBases(model: DataModel, includeDelegate = true) {
+    if (model.$baseMerged) {
+        return model.fields;
+    } else {
+        return [...model.fields, ...getRecursiveBases(model, includeDelegate).flatMap((base) => base.fields)];
+    }
 }
 
-export function getRecursiveBases(dataModel: DataModel): DataModel[] {
+export function getRecursiveBases(dataModel: DataModel, includeDelegate = true): DataModel[] {
     const result: DataModel[] = [];
     dataModel.superTypes.forEach((superType) => {
         const baseDecl = superType.ref;
         if (baseDecl) {
+            if (!includeDelegate && isDelegateModel(baseDecl)) {
+                return;
+            }
             result.push(baseDecl);
-            result.push(...getRecursiveBases(baseDecl));
+            result.push(...getRecursiveBases(baseDecl, includeDelegate));
         }
     });
     return result;

packages/testtools/src/model.ts

@@ -5,7 +5,7 @@ import * as path from 'path';
 import * as tmp from 'tmp';
 import { URI } from 'vscode-uri';
 import { createZModelServices } from 'zenstack/language-server/zmodel-module';
-import { mergeBaseModel } from 'zenstack/utils/ast-utils';
+import { mergeBaseModels } from 'zenstack/utils/ast-utils';
 
 tmp.setGracefulCleanup();
 
@@ -53,7 +53,7 @@ export async function loadModel(content: string, validate = true, verbose = true
 
     const model = (await doc.parseResult.value) as Model;
 
-    mergeBaseModel(model, ZModel.references.Linker);
+    mergeBaseModels(model, ZModel.references.Linker);
 
     return model;
 }

tests/integration/tsconfig.json

@@ -8,5 +8,5 @@
         "skipLibCheck": true,
         "experimentalDecorators": true
     },
-    "include": ["**/*.ts", "**/*.d.ts", "../regression/tests/issue-177.test.ts", "../regression/tests/issue-416.test.ts", "../regression/tests/issue-646.test.ts", "../regression/tests/issue-657.test.ts", "../regression/tests/issue-665.test.ts", "../regression/tests/issue-674.test.ts", "../regression/tests/issue-689.test.ts", "../regression/tests/issue-703.test.ts", "../regression/tests/issue-714.test.ts", "../regression/tests/issue-724.test.ts", "../regression/tests/issue-735.test.ts", "../regression/tests/issue-744.test.ts", "../regression/tests/issue-756.test.ts", "../regression/tests/issue-764.test.ts", "../regression/tests/issue-765.test.ts", "../regression/tests/issue-804.test.ts", "../regression/tests/issue-811.test.ts", "../regression/tests/issue-814.test.ts", "../regression/tests/issue-825.test.ts", "../regression/tests/issue-864.test.ts", "../regression/tests/issue-886.test.ts", "../regression/tests/issue-925.test.ts", "../regression/tests/issue-947.test.ts", "../regression/tests/issue-961.test.ts", "../regression/tests/issue-965.test.ts", "../regression/tests/issue-971.test.ts", "../regression/tests/issue-992.test.ts", "../regression/tests/issue-1014.test.ts", "../regression/tests/issue-1078.test.ts", "../regression/tests/issue-1080.test.ts", "../regression/tests/issue-1129.test.ts", "../regression/tests/issue-1167.test.ts", "../regression/tests/issue-1179.test.ts", "../regression/tests/issue-1186.test.ts", "../regression/tests/issue-1210.test.ts", "../regression/tests/issue-1235.test.ts", "../regression/tests/issue-1241.test.ts", "../regression/tests/issue-1257.test.ts", "../regression/tests/issue-1265.test.ts", "../regression/tests/issues.test.ts"]
+    "include": ["tests/**/*.ts"]
 }

tests/integration/tests/enhancements/with-delegate/issue-1058.test.ts renamed to tests/regression/tests/issue-1058.test.ts

@@ -0,0 +1,119 @@
+import { createPostgresDb, dropPostgresDb, loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 1435', () => {
+    it('regression', async () => {
+        let prisma: any;
+        let dbUrl: string;
+
+        try {
+            dbUrl = await createPostgresDb('issue-1435');
+            const r = await loadSchema(
+                `
+            /* Interfaces */
+            abstract model IBase {
+              updatedAt DateTime @updatedAt
+              createdAt DateTime @default(now())
+            }
+            
+            abstract model IAuth extends IBase {
+              user   User   @relation(fields: [userId], references: [id], onDelete: Cascade)
+              userId String @unique
+            
+              @@allow('create', true)
+              @@allow('all', auth() == user)
+            }
+            
+            abstract model IIntegration extends IBase {
+              organization   Organization @relation(fields: [organizationId], references: [id], onDelete: Cascade)
+              organizationId String       @unique
+            
+              @@allow('all', organization.members?[user == auth() && type == OWNER])
+              @@allow('read', organization.members?[user == auth()])
+            }
+            
+            /* Auth Stuff */
+            model User extends IBase {
+              id          String      @id @default(cuid())
+              firstName   String
+              lastName    String
+              google      GoogleAuth?
+              memberships Member[]
+            
+              @@allow('create', true)
+              @@allow('all', auth() == this)
+            }
+            
+            model GoogleAuth extends IAuth {
+              reference    String @id
+              refreshToken String
+            }
+            
+            /* Org Stuff */
+            enum MemberType {
+              OWNER
+              MEMBER
+            }
+            
+            model Organization extends IBase {
+              id      String             @id @default(cuid())
+              name    String
+              members Member[]
+              google  GoogleIntegration?
+            
+              @@allow('create', true)
+              @@allow('all', members?[user == auth() && type == OWNER])
+              @@allow('read', members?[user == auth()])
+            }
+            
+            
+            model Member extends IBase {
+              type           MemberType   @default(MEMBER)
+              organization   Organization @relation(fields: [organizationId], references: [id], onDelete: Cascade)
+              organizationId String
+              user           User         @relation(fields: [userId], references: [id], onDelete: Cascade)
+              userId         String
+            
+              @@id([organizationId, userId])
+              @@allow('all', organization.members?[user == auth() && type == OWNER])
+              @@allow('read', user == auth())
+            }
+            
+            /* Google Stuff */
+            model GoogleIntegration extends IIntegration {
+              reference String @id
+            }
+            `,
+                { provider: 'postgresql', dbUrl, logPrismaQuery: true }
+            );
+
+            prisma = r.prisma;
+            const enhance = r.enhance;
+
+            await prisma.organization.create({
+                data: {
+                    name: 'My Organization',
+                    members: {
+                        create: {
+                            type: 'OWNER',
+                            user: {
+                                create: {
+                                    id: '1',
+                                    firstName: 'John',
+                                    lastName: 'Doe',
+                                },
+                            },
+                        },
+                    },
+                },
+            });
+
+            const db = enhance({ id: '1' });
+            await expect(db.organization.findMany()).resolves.toHaveLength(1);
+        } finally {
+            if (prisma) {
+                await prisma.$disconnect();
+            }
+            await dropPostgresDb('issue-1435');
+        }
+    });
+});