Added support for discriminators

lib/factories/tenancy.factory.ts

@@ -1,5 +1,5 @@
 import { Provider } from '@nestjs/common';
-import { Connection } from 'mongoose';
+import { Connection, Model } from 'mongoose';
 import { ModelDefinition } from '../interfaces';
 import {
   CONNECTION_MAP,
@@ -25,13 +25,13 @@ export const createTenancyProviders = (
         connectionMap: ConnectionMap,
       ) => {
         const exists = modelDefinitionMap.has(name);
-        if (!exists) {
-          modelDefinitionMap.set(name, { ...definition });
+        if (exists) return;
 
-          connectionMap.forEach((connection: Connection) => {
-            connection.model(name, schema, collection);
-          });
-        }
+        modelDefinitionMap.set(name, { ...definition });
+
+        connectionMap.forEach((connection: Connection) => {
+          connection.model(name, schema, collection);
+        });
       },
       inject: [MODEL_DEFINITION_MAP, CONNECTION_MAP],
     });
@@ -47,6 +47,33 @@ export const createTenancyProviders = (
       },
       inject: [TENANT_CONNECTION],
     });
+
+    // Create descriminators
+    const discriminators = definition.discriminators || [];
+    providers.push(
+      ...discriminators.map((discriminator) => ({
+        provide: getTenantModelToken(discriminator.name),
+        useFactory: (
+          baseModel: Model<Document>,
+          tenantConnection: Connection,
+        ) => {
+          const modelOnConnection = tenantConnection.models[discriminator.name];
+          if (modelOnConnection) return modelOnConnection;
+
+          const discriminatorModel = baseModel.discriminator(
+            discriminator.name,
+            discriminator.schema,
+            discriminator.value,
+          );
+          return tenantConnection.model(
+            discriminator.name,
+            discriminatorModel.schema,
+            collection,
+          );
+        },
+        inject: [getTenantModelToken(name), TENANT_CONNECTION],
+      })),
+    );
   }
 
   // Return the list of providers mapping

lib/interfaces/model-definition.interface.ts

@@ -1,7 +1,14 @@
-import { Schema } from 'mongoose';
+import { DiscriminatorOptions, ObjectId, Schema } from 'mongoose';
+
+export type Discriminator = {
+  name: string;
+  schema: Schema;
+  value?: string | number | ObjectId | DiscriminatorOptions;
+};
 
 export interface ModelDefinition {
   name: string;
   schema: Schema;
   collection?: string;
+  discriminators?: Discriminator[];
 }

tests/e2e/animal-tenancy.spec.ts

@@ -0,0 +1,134 @@
+import { INestApplication } from '@nestjs/common';
+import { Test } from '@nestjs/testing';
+import { Server } from 'http';
+import * as request from 'supertest';
+import { AnimalType } from '../src/animals/dto/animal-type.enum';
+import { CreateMaineCoonDto } from '../src/animals/dto/create-maine-coon.dto';
+import { CreateBeagleDto } from '../src/animals/dto/create-beagle.dto';
+import { AppModule } from '../src/app.module';
+
+describe('AnimalTenancy', () => {
+  let server: Server;
+  let app: INestApplication;
+  const dogs = generateDogs();
+  const cats = generateCats();
+  const TENANT_HEADER = 'X-TENANT-ID';
+  const DOG_CLUB = 'dog-club';
+  const CAT_CLUB = 'cat-club';
+
+  beforeEach(async () => {
+    const module = await Test.createTestingModule({
+      imports: [AppModule],
+    }).compile();
+
+    app = module.createNestApplication();
+    server = app.getHttpServer();
+    await app.init();
+
+    await cleanupTenantDb(CAT_CLUB);
+    await cleanupTenantDb(DOG_CLUB);
+
+    for (const dog of dogs) {
+      await request(server)
+        .post('/animals')
+        .set(TENANT_HEADER, DOG_CLUB)
+        .send(dog)
+        .expect(201);
+    }
+
+    for (const cat of cats) {
+      await request(server)
+        .post('/animals')
+        .set(TENANT_HEADER, CAT_CLUB)
+        .send(cat)
+        .expect(201);
+    }
+  });
+
+  it(`each animal type must have all relevant fields`, (done) => {
+    request(server)
+      .get(`/animals`)
+      .set(TENANT_HEADER, DOG_CLUB)
+      .expect(200, (_err, { body }) => {
+        body.forEach(({ animalType, isMajestic, isGood }) => {
+          expect(animalType).toBe(AnimalType.BEAGLE);
+          expect(isGood).toBe(true);
+          expect(isMajestic).toBe(undefined);
+        });
+
+        done();
+      });
+    request(server)
+      .get(`/animals`)
+      .set(TENANT_HEADER, CAT_CLUB)
+      .expect(200, (_err, { body }) => {
+        body.forEach(({ animalType, isMajestic, isGood }) => {
+          expect(animalType).toBe(AnimalType.MAINE_COON);
+          expect(isMajestic).toBe(true);
+          expect(isGood).toBe(undefined);
+        });
+
+        done();
+      });
+  });
+
+  it(`dog-club tenant should see only dogs`, (done) => {
+    request(server)
+      .get(`/animals`)
+      .set(TENANT_HEADER, DOG_CLUB)
+      .expect(200, (_err, { body }) => {
+        expect(body).toHaveLength(3);
+        body.forEach(({ animalType }) => {
+          expect(animalType).toBe(AnimalType.BEAGLE);
+        });
+
+        done();
+      });
+  });
+
+  it(`cat-club tenant should see only cats`, (done) => {
+    request(server)
+      .get(`/animals`)
+      .set(TENANT_HEADER, CAT_CLUB)
+      .expect(200, (_err, { body }) => {
+        expect(body).toHaveLength(3);
+        body.forEach(({ animalType }) => {
+          expect(animalType).toBe(AnimalType.MAINE_COON);
+        });
+
+        done();
+      });
+  });
+
+  afterEach(async () => {
+    await app.close();
+  });
+
+  async function cleanupTenantDb(tenantId: string) {
+    return request(server).delete(`/animals`).set(TENANT_HEADER, tenantId);
+  }
+});
+
+function generateCats(count = 3): CreateMaineCoonDto[] {
+  return Array(count)
+    .fill(0)
+    .map((_, i) => ({
+      animalType: AnimalType.MAINE_COON,
+      age: i,
+      breed: 'Maine coon',
+      name: `Furrball ${i}`,
+      isMajestic: true,
+    }));
+}
+
+function generateDogs(count = 3): CreateBeagleDto[] {
+  return Array(count)
+    .fill(0)
+    .map((_, i) => ({
+      animalType: AnimalType.BEAGLE,
+      age: i,
+      breed: 'Beagle',
+      name: `Fluffball ${i}`,
+      isGood: true,
+    }));
+}

tests/src/animals/animals.controller.ts

@@ -0,0 +1,49 @@
+import { Body, Controller, Delete, Get, Post } from '@nestjs/common';
+import { BeaglesService } from './beagles.service';
+import { AnimalType } from './dto/animal-type.enum';
+import { AnimalsDto, CreateAnimalDto } from './dto/composite-dtos';
+import { CreateBeagleDto } from './dto/create-beagle.dto';
+import { CreateMaineCoonDto } from './dto/create-maine-coon.dto';
+import { MaineCoonsService } from './maine-coons.service';
+
+@Controller('animals')
+export class AnimalsController {
+  constructor(
+    private readonly catsService: MaineCoonsService,
+    private readonly dogsService: BeaglesService,
+  ) {}
+
+  @Post()
+  async create(@Body() createAnimalDto: CreateAnimalDto): Promise<AnimalsDto> {
+    switch (createAnimalDto.animalType) {
+      case AnimalType.MAINE_COON:
+        return this.catsService.create(createAnimalDto as CreateMaineCoonDto);
+      case AnimalType.BEAGLE:
+        return this.dogsService.create(createAnimalDto as CreateBeagleDto);
+      default:
+        throw new Error('No such animal');
+    }
+  }
+
+  @Get()
+  async findAll(): Promise<AnimalsDto[]> {
+    const promises = Promise.all([
+      this.catsService.findAll(),
+      this.dogsService.findAll(),
+    ]);
+    const results = await promises;
+    return results.reduce((acc, curr) => {
+      acc.push(...curr);
+      return acc;
+    }, [] as AnimalsDto[]);
+  }
+
+  // For e2e tests cleanup
+  @Delete()
+  async removeAll(): Promise<void> {
+    await Promise.all([
+      this.catsService.removeAll(),
+      this.dogsService.removeAll(),
+    ]);
+  }
+}

tests/src/animals/animals.module.ts

@@ -0,0 +1,26 @@
+import { Module } from '@nestjs/common';
+import { TenancyModule } from '../../../lib';
+import { AnimalsController } from './animals.controller';
+import { BeaglesService } from './beagles.service';
+import { MaineCoonsService } from './maine-coons.service';
+import { Animal, AnimalSchema } from './schemas/animal.schema';
+import { Beagle, BeagleSchema } from './schemas/beagle.schema';
+import { MaineCoon, MaineCoonSchema } from './schemas/maine-coon.schema';
+
+@Module({
+  imports: [
+    TenancyModule.forFeature([
+      {
+        name: Animal.name,
+        schema: AnimalSchema,
+        discriminators: [
+          { name: MaineCoon.name, schema: MaineCoonSchema },
+          { name: Beagle.name, schema: BeagleSchema },
+        ],
+      },
+    ]),
+  ],
+  controllers: [AnimalsController],
+  providers: [BeaglesService, MaineCoonsService],
+})
+export class AnimalsModule {}

tests/src/animals/beagles.service.ts

@@ -0,0 +1,26 @@
+import { Injectable } from '@nestjs/common';
+import { InjectModel } from '@nestjs/mongoose';
+import { Model } from 'mongoose';
+import { BeagleDto } from './dto/composite-dtos';
+import { CreateBeagleDto } from './dto/create-beagle.dto';
+import { Beagle } from './schemas/beagle.schema';
+
+@Injectable()
+export class BeaglesService {
+  constructor(
+    @InjectModel(Beagle.name) private readonly beagleModel: Model<Beagle>,
+  ) {}
+
+  async create(createDto: CreateBeagleDto): Promise<BeagleDto> {
+    const created = new this.beagleModel(createDto);
+    return created.save();
+  }
+
+  async findAll(): Promise<BeagleDto[]> {
+    return this.beagleModel.find().lean().exec();
+  }
+
+  async removeAll(): Promise<void> {
+    await this.beagleModel.deleteMany();
+  }
+}

tests/src/animals/dto/animal-type.enum.ts

@@ -0,0 +1,4 @@
+export enum AnimalType {
+  BEAGLE = 'Beagle',
+  MAINE_COON = 'MaineCoon',
+}

tests/src/animals/dto/animal.dto.ts

@@ -0,0 +1,8 @@
+import { AnimalType } from './animal-type.enum';
+
+export class AnimalDto {
+  readonly animalType: AnimalType;
+  readonly name: string;
+  readonly age: number;
+  readonly breed: string;
+}

tests/src/animals/dto/composite-dtos.ts

@@ -0,0 +1,16 @@
+import { AnimalDto } from './animal.dto';
+import { CreateMaineCoonDto } from './create-maine-coon.dto';
+import { CreateBeagleDto } from './create-beagle.dto';
+
+export type CreateAnimalDto = CreateBeagleDto | CreateMaineCoonDto;
+
+export interface MaineCoonDto extends AnimalDto {
+  isMajestic: boolean;
+  isGood?: never;
+}
+export interface BeagleDto extends AnimalDto {
+  isGood: boolean;
+  isMajestic?: never;
+}
+
+export type AnimalsDto = BeagleDto | MaineCoonDto;

tests/src/animals/dto/create-beagle.dto.ts

@@ -0,0 +1,7 @@
+import { AnimalType } from './animal-type.enum';
+import { AnimalDto } from './animal.dto';
+
+export class CreateBeagleDto extends AnimalDto {
+  readonly animalType: AnimalType = AnimalType.BEAGLE;
+  readonly isGood: true;
+}

tests/src/animals/dto/create-maine-coon.dto.ts

@@ -0,0 +1,7 @@
+import { AnimalType } from './animal-type.enum';
+import { AnimalDto } from './animal.dto';
+
+export class CreateMaineCoonDto extends AnimalDto {
+  readonly animalType: AnimalType = AnimalType.MAINE_COON;
+  readonly isMajestic: true;
+}

tests/src/animals/maine-coons.service.ts

@@ -0,0 +1,27 @@
+import { Injectable } from '@nestjs/common';
+import { InjectModel } from '@nestjs/mongoose';
+import { Model } from 'mongoose';
+import { MaineCoonDto } from './dto/composite-dtos';
+import { CreateMaineCoonDto } from './dto/create-maine-coon.dto';
+import { MaineCoon } from './schemas/maine-coon.schema';
+
+@Injectable()
+export class MaineCoonsService {
+  constructor(
+    @InjectModel(MaineCoon.name)
+    private readonly maineCoonModel: Model<MaineCoon>,
+  ) {}
+
+  async create(createDto: CreateMaineCoonDto): Promise<MaineCoonDto> {
+    const created = new this.maineCoonModel(createDto);
+    return created.save();
+  }
+
+  async findAll(): Promise<MaineCoonDto[]> {
+    return this.maineCoonModel.find().lean().exec();
+  }
+
+  async removeAll(): Promise<void> {
+    await this.maineCoonModel.deleteMany();
+  }
+}