From 8b670ebb6455c785a95e1da5c7a7967edb6e3fb9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Ad=C3=A1mek?= <banan23@gmail.com>
Date: Sat, 23 Dec 2023 18:47:07 +0100
Subject: [PATCH] fix: defer automatic `clear` call on `autoLoadEntities` cache
 done on shutdown

When the `addEntity` is called next, the storage will be cleared automatically before it.
We want to keep the cache, as it's populated on require time, but sometimes (tests) the contexts could be cleared.
This resolves both cases by deferring the `clear` call to the first `addEntity` call.
---
 src/mikro-orm-core.module.ts      |  2 +-
 src/mikro-orm.entities.storage.ts | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/mikro-orm-core.module.ts b/src/mikro-orm-core.module.ts
index e7d6fe3..f12b47f 100644
--- a/src/mikro-orm-core.module.ts
+++ b/src/mikro-orm-core.module.ts
@@ -112,7 +112,7 @@ export class MikroOrmCoreModule implements OnApplicationShutdown {
 
     if (orm) {
       await orm.close();
-      MikroOrmEntitiesStorage.clear(orm.config.get('contextName'));
+      MikroOrmEntitiesStorage.clearLater();
     }
 
     CONTEXT_NAMES.length = 0;
diff --git a/src/mikro-orm.entities.storage.ts b/src/mikro-orm.entities.storage.ts
index 9ee994d..d435ef6 100644
--- a/src/mikro-orm.entities.storage.ts
+++ b/src/mikro-orm.entities.storage.ts
@@ -4,9 +4,16 @@ import type { EntityName } from './typings';
 export class MikroOrmEntitiesStorage {
 
   private static readonly storage = new Map<string, Set<EntityName<AnyEntity>>>();
+  private static shouldClear = false;
 
   static addEntity(entity: EntityName<AnyEntity>, contextName = 'default') {
+    if (this.shouldClear) {
+      this.clear(contextName);
+      this.shouldClear = false;
+    }
+
     let set = this.storage.get(contextName);
+
     if (!set) {
       set = new Set<EntityName<AnyEntity>>();
       this.storage.set(contextName, set);
@@ -23,4 +30,13 @@ export class MikroOrmEntitiesStorage {
     this.storage.get(contextName)?.clear();
   }
 
+  /**
+   * When the `addEntity` is called next, the storage will be cleared automatically before it.
+   * We want to keep the cache, as it's populated on require time, but sometimes (tests) the contexts could be cleared.
+   * This resolves both cases by deferring the `clear` call to the first `addEntity` call.
+   */
+  static clearLater() {
+    this.shouldClear = true;
+  }
+
 }