From 3bf04d4a96b433c67df8353d0097069a9a323a93 Mon Sep 17 00:00:00 2001
From: Harry Chen <czy88840616@gmail.com>
Date: Wed, 10 Apr 2024 21:40:40 +0800
Subject: [PATCH 1/2] feat: support multi-datasource for typegoose

---
 packages/mongoose/src/manager.ts              |  4 ++
 packages/typegoose/src/configuration.ts       | 42 +++++++------------
 .../src/decorator/injectEntityModel.ts        |  3 +-
 3 files changed, 21 insertions(+), 28 deletions(-)

diff --git a/packages/mongoose/src/manager.ts b/packages/mongoose/src/manager.ts
index 763e587b34a0..c71f7314828a 100644
--- a/packages/mongoose/src/manager.ts
+++ b/packages/mongoose/src/manager.ts
@@ -87,6 +87,10 @@ export class MongooseDataSourceManager extends DataSourceManager<mongoose.Connec
   ): Promise<boolean> {
     return dataSource.readyState === mongoose.ConnectionStates.connected;
   }
+
+  public getDataSourceNameByModel(modelOrRepository: any): string | undefined {
+    return super.getDataSourceNameByModel(modelOrRepository);
+  }
 }
 
 /**
diff --git a/packages/typegoose/src/configuration.ts b/packages/typegoose/src/configuration.ts
index 050c6e75391a..3edc88eb4f22 100644
--- a/packages/typegoose/src/configuration.ts
+++ b/packages/typegoose/src/configuration.ts
@@ -1,7 +1,6 @@
 import {
   Config,
   Configuration,
-  getClassMetadata,
   Init,
   Inject,
   listModule,
@@ -9,6 +8,7 @@ import {
   IMidwayApplication,
   IMidwayContainer,
   MidwayDecoratorService,
+  MidwayCommonError,
 } from '@midwayjs/core';
 import * as mongoose from '@midwayjs/mongoose';
 import { ENTITY_MODEL_KEY } from './interface';
@@ -31,7 +31,7 @@ export class TypegooseConfiguration {
   @Inject()
   decoratorService: MidwayDecoratorService;
 
-  modelMap = new WeakMap();
+  dataSourceManager: mongoose.MongooseDataSourceManager;
 
   @Init()
   async init() {
@@ -44,39 +44,27 @@ export class TypegooseConfiguration {
           connectionName: string;
         }
       ) => {
-        return this.modelMap.get(meta.modelKey);
+        // get connection
+        const conn = this.dataSourceManager.getDataSource(
+          meta.connectionName ||
+            this.dataSourceManager.getDataSourceNameByModel(meta.modelKey) ||
+            this.dataSourceManager.getDefaultDataSourceName()
+        );
+        if (!conn) {
+          throw new MidwayCommonError(
+            `DataSource ${meta.connectionName} not found with current model ${meta.modelKey}, please check it.`
+          );
+        }
+        return getModelForClass(meta.modelKey, { existingConnection: conn });
       }
     );
   }
 
   async onReady(container: IMidwayContainer) {
-    const connectionFactory = await container.getAsync(
+    this.dataSourceManager = await container.getAsync(
       mongoose.MongooseDataSourceManager
     );
-
-    for (const dataSourceName of connectionFactory.getDataSourceNames()) {
-      const conn = connectionFactory.getDataSource(dataSourceName);
-      if (conn && (conn as any).entities) {
-        for (const Model of (conn as any).entities) {
-          const model = getModelForClass(Model, { existingConnection: conn });
-          this.modelMap.set(Model, model);
-        }
-      }
-    }
-
     const Models = listModule(ENTITY_MODEL_KEY);
-    // 兼容老代码
-    for (const Model of Models) {
-      const metadata = getClassMetadata(ENTITY_MODEL_KEY, Model) ?? {};
-      const connectionName = metadata.connectionName ?? 'default';
-      const conn = connectionFactory.getDataSource(connectionName);
-      if (conn) {
-        const model = getModelForClass(Model, { existingConnection: conn });
-        this.modelMap.set(Model, model);
-      } else {
-        throw new Error(`connection name ${metadata.connectionName} not found`);
-      }
-    }
 
     // 兼容老代码
     if (Models.length === 0 && this.oldMongooseConfig['uri']) {
diff --git a/packages/typegoose/src/decorator/injectEntityModel.ts b/packages/typegoose/src/decorator/injectEntityModel.ts
index d2d835ff1e6b..71730144a81b 100644
--- a/packages/typegoose/src/decorator/injectEntityModel.ts
+++ b/packages/typegoose/src/decorator/injectEntityModel.ts
@@ -1,8 +1,9 @@
 import { createCustomPropertyDecorator } from '@midwayjs/core';
 import { ENTITY_MODEL_KEY } from '../interface';
 
-export function InjectEntityModel(modelKey: any) {
+export function InjectEntityModel(modelKey: any, connectionName?: string) {
   return createCustomPropertyDecorator(ENTITY_MODEL_KEY, {
     modelKey,
+    connectionName,
   });
 }

From 3f3ae369cc8901b54f73f7f2bbca17e6fc500cd2 Mon Sep 17 00:00:00 2001
From: Harry Chen <czy88840616@gmail.com>
Date: Wed, 10 Apr 2024 21:54:18 +0800
Subject: [PATCH 2/2] chore: revert

---
 packages/mongoose/src/manager.ts | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/packages/mongoose/src/manager.ts b/packages/mongoose/src/manager.ts
index c71f7314828a..763e587b34a0 100644
--- a/packages/mongoose/src/manager.ts
+++ b/packages/mongoose/src/manager.ts
@@ -87,10 +87,6 @@ export class MongooseDataSourceManager extends DataSourceManager<mongoose.Connec
   ): Promise<boolean> {
     return dataSource.readyState === mongoose.ConnectionStates.connected;
   }
-
-  public getDataSourceNameByModel(modelOrRepository: any): string | undefined {
-    return super.getDataSourceNameByModel(modelOrRepository);
-  }
 }
 
 /**