feat: add discriminator

nestjs · Nov 8, 2020 · 2d8e33c · 2d8e33c
1 parent ae066d8
commit 2d8e33c
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 9 deletions.
diff --git a/lib/interfaces/async-model-factory.interface.ts b/lib/interfaces/async-model-factory.interface.ts
@@ -3,7 +3,7 @@ import { ModelDefinition } from './model-definition.interface';
 
 export interface AsyncModelFactory
   extends Pick<ModuleMetadata, 'imports'>,
-    Pick<ModelDefinition, 'name' | 'collection'> {
+    Pick<ModelDefinition, 'name' | 'collection' | 'discriminators'> {
   useFactory: (
     ...args: any[]
   ) => ModelDefinition['schema'] | Promise<ModelDefinition['schema']>;

diff --git a/lib/interfaces/model-definition.interface.ts b/lib/interfaces/model-definition.interface.ts
@@ -1,5 +1,13 @@
+import { Schema } from 'mongoose';
+
+export type DiscriminatorOptions = {
+  name: string;
+  schema: Schema;
+};
+
 export type ModelDefinition = {
   name: string;
   schema: any;
   collection?: string;
+  discriminators?: DiscriminatorOptions[];
 };
diff --git a/lib/mongoose.providers.ts b/lib/mongoose.providers.ts
@@ -1,16 +1,27 @@
 import { flatten } from '@nestjs/common';
-import { Connection, Schema } from 'mongoose';
+import { Connection } from 'mongoose';
 import { getConnectionToken, getModelToken } from './common/mongoose.utils';
-import { AsyncModelFactory } from './interfaces';
+import { AsyncModelFactory, ModelDefinition } from './interfaces';
 
 export function createMongooseProviders(
   connectionName?: string,
-  models: { name: string; schema: Schema; collection?: string }[] = [],
+  options: ModelDefinition[] = [],
 ) {
-  const providers = (models || []).map(model => ({
-    provide: getModelToken(model.name),
-    useFactory: (connection: Connection) =>
-      connection.model(model.name, model.schema, model.collection),
+  const providers = (options || []).map((option) => ({
+    provide: getModelToken(option.name),
+    useFactory: (connection: Connection) => {
+      const model = connection.model(
+        option.name,
+        option.schema,
+        option.collection,
+      );
+      if (option.discriminators) {
+        for (const { name, schema } of option.discriminators) {
+          model.discriminator(name, schema);
+        }
+      }
+      return model;
+    },
     inject: [getConnectionToken(connectionName)],
   }));
   return providers;
@@ -20,7 +31,7 @@ export function createMongooseAsyncProviders(
   connectionName?: string,
   modelFactories: AsyncModelFactory[] = [],
 ) {
-  const providers = (modelFactories || []).map(model => [
+  const providers = (modelFactories || []).map((model) => [
     {
       provide: getModelToken(model.name),
       useFactory: async (connection: Connection, ...args: unknown[]) => {