fix: multiple database connection with autoLoadEntities (#62)

mikro-orm · Feb 25, 2022 · 7dfc097 · 7dfc097
1 parent 0c438e2
commit 7dfc097
Show file tree

Hide file tree

Showing 5 changed files with 100 additions and 10 deletions.
diff --git a/src/mikro-orm.common.ts b/src/mikro-orm.common.ts
@@ -1,9 +1,8 @@
-import type { AnyEntity, EntityName } from '@mikro-orm/core';
+import type { EntityName } from '@mikro-orm/core';
 import { MikroORM, Utils } from '@mikro-orm/core';
 import { Inject, Logger } from '@nestjs/common';
 
 export const MIKRO_ORM_MODULE_OPTIONS = Symbol('mikro-orm-module-options');
-export const REGISTERED_ENTITIES = new Set<EntityName<AnyEntity>>();
 export const CONTEXT_NAMES: string[] = [];
 export const logger = new Logger(MikroORM.name);
 

diff --git a/src/mikro-orm.entities.storage.ts b/src/mikro-orm.entities.storage.ts
@@ -0,0 +1,30 @@
+import type { AnyEntity, EntityName } from '@mikro-orm/core';
+
+export class MikroOrmEntitiesStorage {
+
+  private static readonly storage = new Map<string, Set<EntityName<AnyEntity>>>();
+
+  static addEntity(entity: EntityName<AnyEntity>, contextName = 'default') {
+    let set = this.storage.get(contextName);
+    if (!set) {
+      set = new Set<EntityName<AnyEntity>>();
+      this.storage.set(contextName, set);
+    }
+
+    set.add(entity);
+  }
+
+  static getEntities(contextName = 'default') {
+    return this.storage.get(contextName)?.values() || [];
+  }
+
+  static clear(contextName = 'default') {
+    const set = this.storage.get(contextName);
+    if (!set) {
+      return;
+    }
+
+    set.clear();
+  }
+
+}
diff --git a/src/mikro-orm.module.ts b/src/mikro-orm.module.ts
@@ -10,8 +10,8 @@ import type {
   MikroOrmMiddlewareModuleOptions,
   MikroOrmModuleFeatureOptions,
 } from './typings';
-import { REGISTERED_ENTITIES } from './mikro-orm.common';
 import { MikroOrmMiddlewareModule } from './mikro-orm-middleware.module';
+import { MikroOrmEntitiesStorage } from './mikro-orm.entities.storage';
 
 @Module({})
 export class MikroOrmModule {
@@ -32,12 +32,12 @@ export class MikroOrmModule {
 
   static forFeature(options: EntityName<AnyEntity>[] | MikroOrmModuleFeatureOptions, contextName?: string): DynamicModule {
     const entities = Array.isArray(options) ? options : (options.entities || []);
-    const name = Array.isArray(options) ? contextName : options.contextName;
+    const name = (Array.isArray(options) || contextName) ? contextName : options.contextName;
     const providers = createMikroOrmRepositoryProviders(entities, name);
 
     for (const e of entities) {
       if (!Utils.isString(e)) {
-        REGISTERED_ENTITIES.add(e);
+        MikroOrmEntitiesStorage.addEntity(e, name);
       }
     }
 

diff --git a/src/mikro-orm.providers.ts b/src/mikro-orm.providers.ts
@@ -1,22 +1,23 @@
-import { getEntityManagerToken, getMikroORMToken, getRepositoryToken, logger, MIKRO_ORM_MODULE_OPTIONS, REGISTERED_ENTITIES } from './mikro-orm.common';
+import { getEntityManagerToken, getMikroORMToken, getRepositoryToken, logger, MIKRO_ORM_MODULE_OPTIONS } from './mikro-orm.common';
 import type { AnyEntity, EntityName } from '@mikro-orm/core';
 import { ConfigurationLoader, EntityManager, MetadataStorage, MikroORM } from '@mikro-orm/core';
 
 import type { MikroOrmModuleAsyncOptions, MikroOrmModuleOptions, MikroOrmOptionsFactory } from './typings';
 import type { Provider, Type } from '@nestjs/common';
 import { Scope } from '@nestjs/common';
+import { MikroOrmEntitiesStorage } from './mikro-orm.entities.storage';
 
 export function createMikroOrmProvider(contextName?: string): Provider {
   return {
     provide: contextName ? getMikroORMToken(contextName) : MikroORM,
     useFactory: async (options?: MikroOrmModuleOptions) => {
       if (options?.autoLoadEntities) {
-        options.entities = [...(options.entities || []), ...REGISTERED_ENTITIES.values()];
-        options.entitiesTs = [...(options.entitiesTs || []), ...REGISTERED_ENTITIES.values()];
+        options.entities = [...(options.entities || []), ...MikroOrmEntitiesStorage.getEntities(contextName)];
+        options.entitiesTs = [...(options.entitiesTs || []), ...MikroOrmEntitiesStorage.getEntities(contextName)];
         delete options.autoLoadEntities;
       }
 
-      REGISTERED_ENTITIES.clear();
+      MikroOrmEntitiesStorage.clear(contextName);
 
       if (!options || Object.keys(options).length === 0) {
         const config = await ConfigurationLoader.getConfiguration();

diff --git a/tests/mikro-orm.module.test.ts b/tests/mikro-orm.module.test.ts
@@ -223,6 +223,30 @@ describe('MikroORM Module', () => {
 
       await orm.close();
     });
+
+    it('forFeature and autoLoadEntities enabled should return repository', async () => {
+      const { entities, ...options } = testOptions;
+
+      const module = await Test.createTestingModule({
+        imports: [
+          MikroOrmModule.forRoot({
+            autoLoadEntities: true,
+            ...options,
+          }),
+          MikroOrmModule.forFeature([Foo]),
+        ],
+      }).compile();
+
+      const orm = module.get<MikroORM>(MikroORM);
+      const entityManager = module.get<EntityManager>(EntityManager);
+      const repository = module.get<EntityRepository<Foo>>(getRepositoryToken(Foo));
+
+      expect(orm).toBeDefined();
+      expect(entityManager).toBeDefined();
+      expect(repository).toBeDefined();
+
+      await orm.close();
+    });
   });
 
   describe('Multiple Databases', () => {
@@ -322,7 +346,41 @@ describe('MikroORM Module', () => {
             ...testOptions,
           }),
           MikroOrmModule.forFeature([Foo], 'database1'),
-          MikroOrmModule.forFeature([Bar], 'database2'),
+          MikroOrmModule.forFeature({ entities: [Bar], contextName: 'database2' }),
+        ],
+      }).compile();
+
+      const orm1 = module.get<MikroORM>(getMikroORMToken('database1'));
+      const orm2 = module.get<MikroORM>(getMikroORMToken('database2'));
+      const repository1 = module.get<EntityRepository<Foo>>(getRepositoryToken(Foo, 'database1'));
+      const repository2 = module.get<EntityRepository<Bar>>(getRepositoryToken(Bar, 'database2'));
+
+      expect(orm1).toBeDefined();
+      expect(repository1).toBeDefined();
+      expect(orm2).toBeDefined();
+      expect(repository2).toBeDefined();
+
+      await orm1.close();
+      await orm2.close();
+    });
+
+    it('forFeature and autoLoadEntities enabled should return repository', async () => {
+      const { entities, ...options } = testOptions;
+
+      const module = await Test.createTestingModule({
+        imports: [
+          MikroOrmModule.forRoot({
+            contextName: 'database1',
+            autoLoadEntities: true,
+            ...options,
+          }),
+          MikroOrmModule.forRoot({
+            contextName: 'database2',
+            autoLoadEntities: true,
+            ...options,
+          }),
+          MikroOrmModule.forFeature([Foo], 'database1'),
+          MikroOrmModule.forFeature({ entities: [Bar] }, 'database2'),
         ],
       }).compile();
 
@@ -331,6 +389,8 @@ describe('MikroORM Module', () => {
       const repository1 = module.get<EntityRepository<Foo>>(getRepositoryToken(Foo, 'database1'));
       const repository2 = module.get<EntityRepository<Bar>>(getRepositoryToken(Bar, 'database2'));
 
+      expect(() => module.get<EntityRepository<Foo>>(getRepositoryToken(Foo, 'database2'))).toThrow();
+      expect(() => module.get<EntityRepository<Bar>>(getRepositoryToken(Bar, 'database1'))).toThrow();
       expect(orm1).toBeDefined();
       expect(repository1).toBeDefined();
       expect(orm2).toBeDefined();