From d65b268e78d199d7cb7ae68cb34b591e7c603695 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Ad=C3=A1mek?= <banan23@gmail.com>
Date: Sun, 19 May 2024 18:43:17 +0200
Subject: [PATCH] feat: register driver specific MikroORM provider (#176)

Closes #123
---
 src/mikro-orm-core.module.ts | 67 ++++++++++++++++++++++++++++++++++--
 src/mikro-orm.providers.ts   |  7 ++--
 2 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/src/mikro-orm-core.module.ts b/src/mikro-orm-core.module.ts
index f12b47f..c6064a0 100644
--- a/src/mikro-orm-core.module.ts
+++ b/src/mikro-orm-core.module.ts
@@ -20,6 +20,18 @@ async function tryRequire(name: string): Promise<Dictionary | undefined> {
   }
 }
 
+// TODO: provide the package name via some platform method, prefer that over the static map when available
+const PACKAGES = {
+  MongoDriver: '@mikro-orm/mongo',
+  MySqlDriver: '@mikro-orm/mysql',
+  MsSqlDriver: '@mikro-orm/mssql',
+  MariaDbDriver: '@mikro-orm/mariadb',
+  PostgreSqlDriver: '@mikro-orm/postgresql',
+  SqliteDriver: '@mikro-orm/sqlite',
+  LibSqlDriver: '@mikro-orm/libsql',
+  BetterSqliteDriver: '@mikro-orm/better-sqlite',
+} as const;
+
 @Global()
 @Module({})
 export class MikroOrmCoreModule implements OnApplicationShutdown {
@@ -34,11 +46,36 @@ export class MikroOrmCoreModule implements OnApplicationShutdown {
     const mongo = await tryRequire('@mikro-orm/mongodb');
     const em = await this.createEntityManager(options);
 
+    if (em && !contextName) {
+      const packageName = PACKAGES[em.getDriver().constructor.name as keyof typeof PACKAGES];
+      const driverPackage = await tryRequire(packageName);
+
+      if (driverPackage) {
+        return {
+          module: MikroOrmCoreModule,
+          providers: [
+            { provide: MIKRO_ORM_MODULE_OPTIONS, useValue: options || {} },
+            createMikroOrmProvider(contextName),
+            createMikroOrmProvider(contextName, driverPackage.MikroORM),
+            createEntityManagerProvider(options?.scope, EntityManager),
+            createEntityManagerProvider(options?.scope, driverPackage.EntityManager),
+          ],
+          exports: [
+            MikroORM,
+            EntityManager,
+            driverPackage.EntityManager,
+            driverPackage.MikroORM,
+          ],
+        };
+      }
+    }
+
     return {
       module: MikroOrmCoreModule,
       providers: [
         { provide: MIKRO_ORM_MODULE_OPTIONS, useValue: options || {} },
         createMikroOrmProvider(contextName),
+        ...(mongo ? [createMikroOrmProvider(contextName, mongo.MikroORM)] : []),
         createEntityManagerProvider(options?.scope, EntityManager, contextName),
         ...(em ? [createEntityManagerProvider(options?.scope, em.constructor as Type, contextName)] : []),
         ...(knex ? [createEntityManagerProvider(options?.scope, knex.EntityManager, contextName)] : []),
@@ -49,7 +86,7 @@ export class MikroOrmCoreModule implements OnApplicationShutdown {
         contextName ? getEntityManagerToken(contextName) : EntityManager,
         ...(em && !contextName ? [em.constructor] : []),
         ...(knex && !contextName ? [knex.EntityManager] : []),
-        ...(mongo && !contextName ? [mongo.EntityManager] : []),
+        ...(mongo && !contextName ? [mongo.EntityManager, mongo.MikroORM] : []),
       ],
     };
   }
@@ -60,6 +97,32 @@ export class MikroOrmCoreModule implements OnApplicationShutdown {
     const mongo = await tryRequire('@mikro-orm/mongodb');
     const em = await this.createEntityManager(options);
 
+    if (em && !contextName) {
+      const packageName = PACKAGES[em.getDriver().constructor.name as keyof typeof PACKAGES];
+      const driverPackage = await tryRequire(packageName);
+
+      if (driverPackage) {
+        return {
+          module: MikroOrmCoreModule,
+          imports: options.imports || [],
+          providers: [
+            ...(options.providers || []),
+            ...createAsyncProviders({ ...options, contextName: options.contextName }),
+            createMikroOrmProvider(contextName),
+            createMikroOrmProvider(contextName, driverPackage.MikroORM),
+            createEntityManagerProvider(options?.scope, EntityManager),
+            createEntityManagerProvider(options?.scope, driverPackage.EntityManager),
+          ],
+          exports: [
+            MikroORM,
+            EntityManager,
+            driverPackage.EntityManager,
+            driverPackage.MikroORM,
+          ],
+        };
+      }
+    }
+
     return {
       module: MikroOrmCoreModule,
       imports: options.imports || [],
@@ -77,7 +140,7 @@ export class MikroOrmCoreModule implements OnApplicationShutdown {
         contextName ? getEntityManagerToken(contextName) : EntityManager,
         ...(em && !contextName ? [em.constructor] : []),
         ...(knex && !contextName ? [knex.EntityManager] : []),
-        ...(mongo && !contextName ? [mongo.EntityManager] : []),
+        ...(mongo && !contextName ? [mongo.EntityManager, mongo.MikroORM] : []),
       ],
     };
   }
diff --git a/src/mikro-orm.providers.ts b/src/mikro-orm.providers.ts
index af1370a..e704b75 100644
--- a/src/mikro-orm.providers.ts
+++ b/src/mikro-orm.providers.ts
@@ -7,9 +7,12 @@ import type { InjectionToken, Provider, Type } from '@nestjs/common';
 import { Scope } from '@nestjs/common';
 import { MikroOrmEntitiesStorage } from './mikro-orm.entities.storage';
 
-export function createMikroOrmProvider(contextName?: string): Provider {
+export function createMikroOrmProvider(
+  contextName?: string,
+  type: Type = MikroORM,
+): Provider {
   return {
-    provide: contextName ? getMikroORMToken(contextName) : MikroORM,
+    provide: contextName ? getMikroORMToken(contextName) : type,
     useFactory: async (options?: MikroOrmModuleOptions) => {
       options = { ...options };