Merge pull request #27 from JackCme/master

Fix/Relation model class and Self relation class missing imports

Author: omar-dulaimi

src/generate-class.ts

@@ -215,15 +215,31 @@ function generateRelationClass(
  }
 
  const relationImports = new Set();
+ let hasSelfRelation = false;
+
  relationFields.forEach((field) => {
- if (field.relationName && model.name !== field.type) {
- relationImports.add(field.type);
+ if (field.relationName) {
+ if (model.name !== field.type) {
+ relationImports.add(field.type);
+ } else {
+ hasSelfRelation = true;
+ }
  }
  });
 
- generateRelationImportsImport(sourceFile, [
- ...relationImports,
- ] as Array<string>);
+ // For self-relations in the Relations class, import the combined model class
+ if (hasSelfRelation) {
+ sourceFile.addImportDeclaration({
+ moduleSpecifier: `./${model.name}.model`,
+ namedImports: [model.name],
+ });
+ }
+
+ if (relationImports.size > 0) {
+ generateRelationImportsImport(sourceFile, [
+ ...relationImports,
+ ] as Array<string>);
+ }
 
  sourceFile.addClass({
  name: `${model.name}Relations`,
@@ -264,6 +280,31 @@ function generateCombinedClass(
  namedImports: [`${model.name}Base`],
  });
 
+ // Add class validator imports for relation fields
+ const validatorImports = [
+ ...new Set(
+ relationFields
+ .map((field) => getDecoratorsImportsByType(field))
+ .flatMap((item) => item),
+ ),
+ ];
+
+ if (validatorImports.length > 0) {
+ generateClassValidatorImport(sourceFile, validatorImports as Array<string>);
+ }
+
+ // Add Swagger imports if enabled
+ if (
+ config.swagger &&
+ relationFields.length > 0 &&
+ shouldImportSwagger(relationFields as PrismaDMMF.Field[])
+ ) {
+ const swaggerImports = getSwaggerImportsByType(
+ relationFields as PrismaDMMF.Field[],
+ );
+ generateSwaggerImport(sourceFile, swaggerImports);
+ }
+
  // Import relation types for the combined class
  const relationImports = new Set();
  relationFields.forEach((field) => {
@@ -293,6 +334,7 @@ function generateCombinedClass(
  hasExclamationToken: field.isRequired,
  hasQuestionToken: !field.isRequired,
  trailingTrivia: '\r\n',
+ decorators: getDecoratorsByFieldType(field, config.swagger),
  };
  },
  ),

tests/schemas/self-relation.prisma

@@ -0,0 +1,27 @@
+generator client {
+ provider = "prisma-client-js"
+}
+
+generator class_validator {
+ provider = "node ./lib/generator.js"
+ output = "../generated/self-relation"
+ swagger = "true"
+ separateRelationFields = "true"
+}
+
+datasource db {
+ provider = "sqlite"
+ url = "file:./test.db"
+}
+
+model User {
+ id Int @id @default(autoincrement())
+ email String @unique
+ name String?
+ createdAt DateTime @default(now())
+
+ // Self-relation: User can have a mentor (one-to-many)
+ mentorId Int?
+ mentor User? @relation("UserMentor", fields: [mentorId], references: [id])
+ mentees User[] @relation("UserMentor")
+}

tests/self-relation.test.ts

@@ -0,0 +1,130 @@
+import { exec } from 'child_process';
+import { promisify } from 'util';
+import { existsSync, readFileSync } from 'fs';
+import { describe, it, expect, beforeAll } from 'vitest';
+import path from 'path';
+
+const execAsync = promisify(exec);
+
+describe('Self-Relation Generation', () => {
+ beforeAll(async () => {
+ // Build the generator first
+ await execAsync('npm run build');
+
+ // Generate models for self-relation schema
+ const schemaPath = path.resolve(__dirname, 'schemas/self-relation.prisma');
+ await execAsync(`npx prisma generate --schema="${schemaPath}"`);
+ }, 60000);
+
+ it('should generate UserBase class without self-relations', () => {
+ const outputPath = path.resolve(__dirname, 'generated/self-relation');
+ const userBasePath = path.join(outputPath, 'models', 'UserBase.model.ts');
+
+ expect(existsSync(userBasePath)).toBe(true);
+ const userBase = readFileSync(userBasePath, 'utf-8');
+
+ // Should have scalar fields
+ expect(userBase).toContain('id!: number');
+ expect(userBase).toContain('email!: string');
+ expect(userBase).toContain('name?: string');
+ expect(userBase).toContain('createdAt!: Date');
+ expect(userBase).toContain('mentorId?: number');
+
+ // Should have decorators
+ expect(userBase).toContain('@IsInt()');
+ expect(userBase).toContain('@IsString()');
+ expect(userBase).toContain('@IsDate()');
+ expect(userBase).toContain('@IsDefined()');
+ expect(userBase).toContain('@IsOptional()');
+
+ // Should NOT have self-relation fields
+ expect(userBase).not.toContain('mentor?');
+ expect(userBase).not.toContain('mentees');
+ expect(userBase).not.toContain('User[]');
+ });
+
+ it('should generate UserRelations class with only self-relation fields', () => {
+ const outputPath = path.resolve(__dirname, 'generated/self-relation');
+ const userRelationsPath = path.join(
+ outputPath,
+ 'models',
+ 'UserRelations.model.ts',
+ );
+
+ expect(existsSync(userRelationsPath)).toBe(true);
+ const userRelations = readFileSync(userRelationsPath, 'utf-8');
+
+ // Should have self-relation fields
+ expect(userRelations).toContain('mentor?: User');
+ expect(userRelations).toContain('mentees!: User[]');
+
+ // Should have class-validator decorators
+ expect(userRelations).toContain('@IsOptional()');
+ expect(userRelations).toContain('@IsDefined()');
+
+ // Should have Swagger decorators
+ expect(userRelations).toContain('@ApiProperty({');
+ expect(userRelations).toContain('type: () => User');
+ expect(userRelations).toContain('required: false');
+ expect(userRelations).toContain('isArray: true');
+
+ // For self-relations in Relations class, should import from the combined model
+ expect(userRelations).toContain('import { User } from "./User.model"');
+
+ // Should NOT have scalar fields
+ expect(userRelations).not.toContain('id!: number');
+ expect(userRelations).not.toContain('email!: string');
+ expect(userRelations).not.toContain('mentorId');
+ });
+
+ it('should generate combined User class with self-relations', () => {
+ const outputPath = path.resolve(__dirname, 'generated/self-relation');
+ const userModelPath = path.join(outputPath, 'models', 'User.model.ts');
+
+ expect(existsSync(userModelPath)).toBe(true);
+ const userModel = readFileSync(userModelPath, 'utf-8');
+
+ // Should import from UserBase
+ expect(userModel).toContain('import { UserBase } from "./UserBase.model"');
+
+ // Should extend UserBase
+ expect(userModel).toContain('export class User extends UserBase');
+
+ // Should have self-relation fields with decorators
+ expect(userModel).toContain('mentor?: User');
+ expect(userModel).toContain('mentees!: User[]');
+
+ // Should have class-validator imports
+ expect(userModel).toContain(
+ 'import { IsOptional, IsDefined } from "class-validator"',
+ );
+
+ // Should have Swagger imports
+ expect(userModel).toContain(
+ 'import { ApiProperty } from "@nestjs/swagger"',
+ );
+
+ // Should have decorators on relation fields
+ expect(userModel).toContain('@IsOptional()');
+ expect(userModel).toContain('@IsDefined()');
+ expect(userModel).toContain('@ApiProperty({');
+
+ // Should NOT import User from itself (no circular import)
+ expect(userModel).not.toContain('import { User } from "./"');
+ expect(userModel).not.toContain('import { User } from "./User.model"');
+ });
+
+ it('should handle self-relations without circular import issues', () => {
+ const outputPath = path.resolve(__dirname, 'generated/self-relation');
+
+ // Check that the index file exports User
+ const indexPath = path.join(outputPath, 'models', 'index.ts');
+ expect(existsSync(indexPath)).toBe(true);
+ const index = readFileSync(indexPath, 'utf-8');
+
+ // When separateRelationFields is enabled, index should export all classes
+ expect(index).toContain('export { User } from "./User.model"');
+ // Note: UserBase and UserRelations might not be in index if only User is exported
+ // This depends on the generator's index generation logic
+ });
+});

tests/swagger-generation.test.ts

@@ -70,3 +70,124 @@ describe('Swagger Generation', () => {
  expect(userModel).toContain('@ApiProperty({');
  });
 });
+
+describe('Combined Class with separateRelationFields', () => {
+ beforeAll(async () => {
+ // Build the generator first
+ await execAsync('npm run build');
+
+ // Generate models for full-features schema (has both swagger and separateRelationFields)
+ const schemaPath = path.resolve(__dirname, 'schemas/full-features.prisma');
+ await execAsync(`npx prisma generate --schema="${schemaPath}"`);
+ }, 60000);
+
+ it('should generate combined class with class-validator and Swagger decorators for relations', () => {
+ const outputPath = path.resolve(__dirname, 'generated/full-features');
+
+ // Check User combined class
+ const userModelPath = path.join(outputPath, 'models', 'User.model.ts');
+ const userModel = readFileSync(userModelPath, 'utf-8');
+
+ // Check that it imports from UserBase
+ expect(userModel).toContain('import { UserBase } from "./UserBase.model"');
+
+ // Check for class-validator imports
+ expect(userModel).toContain('import { IsDefined } from "class-validator"');
+
+ // Check for Swagger imports
+ expect(userModel).toContain(
+ 'import { ApiProperty } from "@nestjs/swagger"',
+ );
+
+ // Check for relation type imports (from index file)
+ expect(userModel).toContain('import { Post } from "./"');
+
+ // Check that class extends UserBase
+ expect(userModel).toContain('export class User extends UserBase');
+
+ // Check for relation field with decorators
+ expect(userModel).toContain('posts!: Post[]');
+ expect(userModel).toContain('@IsDefined()');
+ expect(userModel).toContain(
+ '@ApiProperty({ isArray: true, type: () => Post })',
+ );
+ });
+
+ it('should generate Post combined class with decorators for author relation', () => {
+ const outputPath = path.resolve(__dirname, 'generated/full-features');
+
+ // Check Post combined class
+ const postModelPath = path.join(outputPath, 'models', 'Post.model.ts');
+ const postModel = readFileSync(postModelPath, 'utf-8');
+
+ // Check that it imports from PostBase
+ expect(postModel).toContain('import { PostBase } from "./PostBase.model"');
+
+ // Check for class-validator imports for optional relation
+ expect(postModel).toContain('import { IsOptional } from "class-validator"');
+
+ // Check for Swagger imports
+ expect(postModel).toContain(
+ 'import { ApiProperty } from "@nestjs/swagger"',
+ );
+
+ // Check for relation type imports (from index file)
+ expect(postModel).toContain('import { User } from "./"');
+
+ // Check that class extends PostBase
+ expect(postModel).toContain('export class Post extends PostBase');
+
+ // Check for optional relation field with decorators
+ expect(postModel).toContain('author?: User');
+ expect(postModel).toContain('@IsOptional()');
+ expect(postModel).toContain(
+ '@ApiProperty({ required: false, type: () => User })',
+ );
+ });
+
+ it('should generate base classes without relation fields', () => {
+ const outputPath = path.resolve(__dirname, 'generated/full-features');
+
+ // Check UserBase class
+ const userBasePath = path.join(outputPath, 'models', 'UserBase.model.ts');
+ const userBase = readFileSync(userBasePath, 'utf-8');
+
+ // Should have scalar fields with decorators
+ expect(userBase).toContain('@IsInt()');
+ expect(userBase).toContain('@IsString()');
+ expect(userBase).toContain('@ApiProperty({');
+ expect(userBase).toContain('id!: number');
+ expect(userBase).toContain('email!: string');
+ expect(userBase).toContain('name?: string');
+
+ // Should NOT have relation fields
+ expect(userBase).not.toContain('posts');
+ expect(userBase).not.toContain('Post[]');
+ });
+
+ it('should generate relation classes with only relation fields', () => {
+ const outputPath = path.resolve(__dirname, 'generated/full-features');
+
+ // Check UserRelations class
+ const userRelationsPath = path.join(
+ outputPath,
+ 'models',
+ 'UserRelations.model.ts',
+ );
+
+ // Check if file exists
+ expect(existsSync(userRelationsPath)).toBe(true);
+
+ const userRelations = readFileSync(userRelationsPath, 'utf-8');
+
+ // Should have relation field with decorators
+ expect(userRelations).toContain('export class UserRelations');
+ expect(userRelations).toContain('posts!: Post[]');
+ expect(userRelations).toContain('@IsDefined()');
+ expect(userRelations).toContain('@ApiProperty({');
+
+ // Should NOT have scalar fields
+ expect(userRelations).not.toContain('id!: number');
+ expect(userRelations).not.toContain('email!: string');
+ });
+});