fix(zmodel): prefer to use triple-slash comments as ZModel documentation (#1817)

Author: ymc9

packages/schema/src/language-server/zmodel-documentation-provider.ts

@@ -0,0 +1,16 @@
+import { AstNode, JSDocDocumentationProvider } from 'langium';
+
+/**
+ * Documentation provider that first tries to use triple-slash comments and falls back to JSDoc comments.
+ */
+export class ZModelDocumentationProvider extends JSDocDocumentationProvider {
+    getDocumentation(node: AstNode): string | undefined {
+        // prefer to use triple-slash comments
+        if ('comments' in node && Array.isArray(node.comments) && node.comments.length > 0) {
+            return node.comments.map((c: string) => c.replace(/^[/]*\s*/, '')).join('\n');
+        }
+
+        // fall back to JSDoc comments
+        return super.getDocumentation(node);
+    }
+}

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

@@ -27,13 +27,14 @@ import { ZModelValidationRegistry, ZModelValidator } from './validator/zmodel-va
 import { ZModelCodeActionProvider } from './zmodel-code-action';
 import { ZModelCompletionProvider } from './zmodel-completion-provider';
 import { ZModelDefinitionProvider } from './zmodel-definition';
+import { ZModelDocumentationProvider } from './zmodel-documentation-provider';
 import { ZModelFormatter } from './zmodel-formatter';
 import { ZModelHighlightProvider } from './zmodel-highlight';
 import { ZModelHoverProvider } from './zmodel-hover';
 import { ZModelLinker } from './zmodel-linker';
 import { ZModelScopeComputation, ZModelScopeProvider } from './zmodel-scope';
 import { ZModelSemanticTokenProvider } from './zmodel-semantic';
-import ZModelWorkspaceManager from './zmodel-workspace-manager';
+import { ZModelWorkspaceManager } from './zmodel-workspace-manager';
 
 /**
  * Declaration of custom services - add your own service classes here.
@@ -77,6 +78,9 @@ export const ZModelModule: Module<ZModelServices, PartialLangiumServices & ZMode
     parser: {
         GrammarConfig: (services) => createGrammarConfig(services),
     },
+    documentation: {
+        DocumentationProvider: (services) => new ZModelDocumentationProvider(services),
+    },
 };
 
 // this duplicates createDefaultSharedModule except that a custom WorkspaceManager is used

packages/schema/src/language-server/zmodel-workspace-manager.ts

@@ -9,7 +9,7 @@ import { PLUGIN_MODULE_NAME, STD_LIB_MODULE_NAME } from './constants';
 /**
  * Custom Langium WorkspaceManager implementation which automatically loads stdlib.zmodel
  */
-export default class ZModelWorkspaceManager extends DefaultWorkspaceManager {
+export class ZModelWorkspaceManager extends DefaultWorkspaceManager {
     public pluginModels = new Set<string>();
 
     protected async loadAdditionalDocuments(

packages/schema/src/plugins/prisma/schema-generator.ts

@@ -100,6 +100,7 @@ export class PrismaSchemaGenerator {
 `;
 
     private mode: 'logical' | 'physical' = 'physical';
+    private customAttributesAsComments = false;
 
     // a mapping from full names to shortened names
     private shortNameMap = new Map<string, string>();
@@ -117,6 +118,14 @@ export class PrismaSchemaGenerator {
             this.mode = options.mode as 'logical' | 'physical';
         }
 
+        if (
+            options.customAttributesAsComments !== undefined &&
+            typeof options.customAttributesAsComments !== 'boolean'
+        ) {
+            throw new PluginError(name, 'option "customAttributesAsComments" must be a boolean');
+        }
+        this.customAttributesAsComments = options.customAttributesAsComments === true;
+
         const prismaVersion = getPrismaVersion();
         if (prismaVersion && semver.lt(prismaVersion, PRISMA_MINIMUM_VERSION)) {
             warnings.push(
@@ -282,12 +291,9 @@ export class PrismaSchemaGenerator {
             this.generateContainerAttribute(model, attr);
         }
 
-        decl.attributes
-            .filter((attr) => attr.decl.ref && !this.isPrismaAttribute(attr))
-            .forEach((attr) => model.addComment('/// ' + this.zModelGenerator.generate(attr)));
-
         // user defined comments pass-through
         decl.comments.forEach((c) => model.addComment(c));
+        this.getCustomAttributesAsComments(decl).forEach((c) => model.addComment(c));
 
         // generate relation fields on base models linking to concrete models
         this.generateDelegateRelationForBase(model, decl);
@@ -763,11 +769,9 @@ export class PrismaSchemaGenerator {
             )
             .map((attr) => this.makeFieldAttribute(attr));
 
-        const nonPrismaAttributes = field.attributes.filter((attr) => attr.decl.ref && !this.isPrismaAttribute(attr));
-
-        const documentations = nonPrismaAttributes.map((attr) => '/// ' + this.zModelGenerator.generate(attr));
-
-        const result = model.addField(field.name, type, attributes, documentations, addToFront);
+        // user defined comments pass-through
+        const docs = [...field.comments, ...this.getCustomAttributesAsComments(field)];
+        const result = model.addField(field.name, type, attributes, docs, addToFront);
 
         if (this.mode === 'logical') {
             if (field.attributes.some((attr) => isDefaultWithAuth(attr))) {
@@ -777,8 +781,6 @@ export class PrismaSchemaGenerator {
             }
         }
 
-        // user defined comments pass-through
-        field.comments.forEach((c) => result.addComment(c));
         return result;
     }
 
@@ -898,23 +900,28 @@ export class PrismaSchemaGenerator {
             this.generateContainerAttribute(_enum, attr);
         }
 
-        decl.attributes
-            .filter((attr) => attr.decl.ref && !this.isPrismaAttribute(attr))
-            .forEach((attr) => _enum.addComment('/// ' + this.zModelGenerator.generate(attr)));
-
         // user defined comments pass-through
         decl.comments.forEach((c) => _enum.addComment(c));
+        this.getCustomAttributesAsComments(decl).forEach((c) => _enum.addComment(c));
     }
 
     private generateEnumField(_enum: PrismaEnum, field: EnumField) {
         const attributes = field.attributes
             .filter((attr) => this.isPrismaAttribute(attr))
             .map((attr) => this.makeFieldAttribute(attr));
 
-        const nonPrismaAttributes = field.attributes.filter((attr) => attr.decl.ref && !this.isPrismaAttribute(attr));
+        const docs = [...field.comments, ...this.getCustomAttributesAsComments(field)];
+        _enum.addField(field.name, attributes, docs);
+    }
 
-        const documentations = nonPrismaAttributes.map((attr) => '/// ' + this.zModelGenerator.generate(attr));
-        _enum.addField(field.name, attributes, documentations.concat(field.comments));
+    private getCustomAttributesAsComments(decl: DataModel | DataModelField | Enum | EnumField) {
+        if (!this.customAttributesAsComments) {
+            return [];
+        } else {
+            return decl.attributes
+                .filter((attr) => attr.decl.ref && !this.isPrismaAttribute(attr))
+                .map((attr) => `/// ${this.zModelGenerator.generate(attr)}`);
+        }
     }
 }
 

packages/schema/src/res/starter.zmodel

@@ -1,8 +1,6 @@
 // This is a sample model to get you started.
 
-/**
- * A sample data source using local sqlite db.
- */
+/// A sample data source using local sqlite db.
 datasource db {
     provider = 'sqlite'
     url = 'file:./dev.db'
@@ -12,9 +10,7 @@ generator client {
     provider = "prisma-client-js"
 }
 
-/**
- * User model
- */
+/// User model
 model User {
     id       String @id @default(cuid())
     email    String @unique @email @length(6, 32)
@@ -28,9 +24,7 @@ model User {
     @@allow('all', auth() == this)
 }
 
-/**
- * Post model
- */
+/// Post model
 model Post {
     id        String   @id @default(cuid())
     createdAt DateTime @default(now())

packages/schema/tests/generator/prisma-generator.test.ts

@@ -47,10 +47,35 @@ describe('Prisma generator test', () => {
                 provider = '@core/prisma'
             }
 
+            /// User roles
+            enum Role {
+                /// Admin role
+                ADMIN
+                /// Regular role
+                USER
+
+                @@schema("auth")
+            }
+
+            /// My user model
+            /// defined here
             model User {
-                id String @id
+                /// the id field
+                id String @id @allow('read', this == auth())
+                role Role
 
                 @@schema("auth")
+                @@allow('all', true)
+                @@deny('update', this != auth())
+            }
+
+            /**
+             * My post model
+             * defined here
+             */
+            model Post {
+                id String @id
+                @@schema("public")
             }
         `);
 
@@ -60,6 +85,7 @@ describe('Prisma generator test', () => {
             schemaPath: 'schema.zmodel',
             output: 'schema.prisma',
             format: false,
+            customAttributesAsComments: true,
         });
 
         const content = fs.readFileSync('schema.prisma', 'utf-8');
@@ -70,6 +96,14 @@ describe('Prisma generator test', () => {
             'extensions = [pg_trgm, postgis(version: "3.3.2"), uuid_ossp(map: "uuid-ossp", schema: "extensions")]'
         );
         expect(content).toContain('schemas = ["auth", "public"]');
+        expect(content).toContain('/// My user model');
+        expect(content).toContain(`/// @@allow('all', true)`);
+        expect(content).toContain(`/// the id field`);
+        expect(content).toContain(`/// @allow('read', this == auth())`);
+        expect(content).not.toContain('/// My post model');
+        expect(content).toContain('/// User roles');
+        expect(content).toContain('/// Admin role');
+        expect(content).toContain('/// Regular role');
         await getDMMF({ datamodel: content });
     });
 
@@ -172,6 +206,7 @@ describe('Prisma generator test', () => {
             provider: '@core/prisma',
             schemaPath: 'schema.zmodel',
             output: name,
+            customAttributesAsComments: true,
         });
 
         const content = fs.readFileSync(name, 'utf-8');
@@ -204,6 +239,7 @@ describe('Prisma generator test', () => {
             provider: '@core/prisma',
             schemaPath: 'schema.zmodel',
             output: name,
+            customAttributesAsComments: true,
         });
 
         const content = fs.readFileSync(name, 'utf-8');
@@ -397,6 +433,7 @@ describe('Prisma generator test', () => {
             schemaPath: 'schema.zmodel',
             output: name,
             generateClient: false,
+            customAttributesAsComments: true,
         });
 
         const content = fs.readFileSync(name, 'utf-8');
@@ -447,6 +484,7 @@ describe('Prisma generator test', () => {
             schemaPath: 'schema.zmodel',
             output: name,
             format: true,
+            customAttributesAsComments: true,
         });
 
         const content = fs.readFileSync(name, 'utf-8');
@@ -478,6 +516,7 @@ describe('Prisma generator test', () => {
             schemaPath: 'schema.zmodel',
             output: name,
             format: true,
+            customAttributesAsComments: true,
         });
 
         const content = fs.readFileSync(name, 'utf-8');