feat: Optimize typeorm to initialize loading entities (#1150)

midwayjs · Jul 14, 2021 · f0faf2f · f0faf2f
1 parent 7b9dfd4
commit f0faf2f
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 28 deletions.
diff --git a/packages/orm/src/configuration.ts b/packages/orm/src/configuration.ts
@@ -5,6 +5,7 @@ import {
   getClassMetadata,
   listModule,
 } from '@midwayjs/decorator';
+import { join } from 'path';
 import {
   Connection,
   ConnectionOptions,
@@ -18,8 +19,7 @@ import {
   EVENT_SUBSCRIBER_KEY,
   ORM_MODEL_KEY,
 } from '.';
-import { ORM_HOOK_KEY, OrmConnectionHook } from './hook';
-import { join } from 'path';
+import { OrmConnectionHook, ORM_HOOK_KEY } from './hook';
 
 @Configuration({
   importConfigs: [join(__dirname, './config')],
@@ -43,12 +43,32 @@ export class OrmConfiguration implements ILifeCycle {
     const entities = listModule(ENTITY_MODEL_KEY);
     const eventSubs = listModule(EVENT_SUBSCRIBER_KEY);
 
+    const connectionNameMap = { ALL: [] };
+    for (const entity of entities) {
+      const _connectionName = getClassMetadata(
+        ENTITY_MODEL_KEY,
+        entity
+      ).connectionName;
+      if (!connectionNameMap[_connectionName]) {
+        connectionNameMap[_connectionName] = [];
+      }
+      connectionNameMap[_connectionName].push(entity);
+    }
+
     const opts = this.formatConfig();
 
     for (const connectionOption of opts) {
-      connectionOption.entities = entities || [];
-      connectionOption.subscribers = eventSubs || [];
       const name = connectionOption.name || 'default';
+      const connectionEntities = [
+        ...connectionNameMap['ALL'],
+        ...(connectionNameMap[name] || []),
+      ];
+
+      connectionOption.entities = connectionOption.entities
+        ? connectionOption.entities
+        : connectionEntities || [];
+
+      connectionOption.subscribers = eventSubs || [];
       this.connectionNames.push(name);
       let isConnected = false;
       try {

diff --git a/packages/orm/src/index.ts b/packages/orm/src/index.ts
@@ -1,22 +1,27 @@
+import { attachClassMetadata, saveModule } from '@midwayjs/core';
+import { saveClassMetadata } from '@midwayjs/decorator';
 import {
-  EntityOptions,
+  Connection,
+  EntityOptions as BaseEntityOptions,
+  EntitySchema,
   getMetadataArgsStorage,
+  getRepository,
+  MongoRepository,
   ObjectType,
-  EntitySchema,
   Repository,
   TreeRepository,
-  MongoRepository,
-  Connection,
-  getRepository,
 } from 'typeorm';
 import { ViewEntityOptions } from 'typeorm/decorator/options/ViewEntityOptions';
-import { saveModule, attachClassMetadata } from '@midwayjs/core';
-import { saveClassMetadata } from '@midwayjs/decorator';
 
 export const CONNECTION_KEY = 'orm:getConnection';
 export const ENTITY_MODEL_KEY = 'entity_model_key';
 export const EVENT_SUBSCRIBER_KEY = 'event_subscriber_key';
 export const ORM_MODEL_KEY = '__orm_model_key__';
+
+export interface EntityOptions extends BaseEntityOptions {
+  connectionName?: string;
+}
+
 /**
  * Entity - typeorm
  * @param options EntityOptions
@@ -25,7 +30,7 @@ export function EntityModel(options?: EntityOptions): ClassDecorator;
 /**
  * Entity - typeorm
  * @param name string
- * @param options EntityOptions
+ * @param options EntityOptions & connectionName
  */
 export function EntityModel(
   name?: string,
@@ -43,14 +48,20 @@ export function EntityModel(
   const options =
     (typeof nameOrOptions === 'object'
       ? (nameOrOptions as EntityOptions)
-      : maybeOptions) || {};
+      : maybeOptions) || ({} as any);
   const name = typeof nameOrOptions === 'string' ? nameOrOptions : options.name;
-
+  const connectionName = options?.connectionName || 'ALL';
   return function (target) {
     if (typeof target === 'function') {
       saveModule(ENTITY_MODEL_KEY, target);
+      saveClassMetadata(ENTITY_MODEL_KEY, { connectionName }, target);
     } else {
       saveModule(ENTITY_MODEL_KEY, (target as any).constructor);
+      saveClassMetadata(
+        ENTITY_MODEL_KEY,
+        { connectionName },
+        (target as any).constructor
+      );
     }
 
     getMetadataArgsStorage().tables.push({
@@ -170,8 +181,9 @@ export type getMongoRepository = <Entity>(
  */
 export type getCustomRepository = <T>(customRepository: ObjectType<T>) => T;
 
-export * from './repository';
+export { OrmConfiguration as Configuration } from './configuration';
 export * from './hook';
+export * from './repository';
 
 export type GetConnection = (instanceName?: string) => Connection;
 
@@ -186,5 +198,3 @@ export function useEntityModel<Entity>(
 ): Repository<Entity> {
   return getRepository<Entity>(clz, connectionName);
 }
-
-export { OrmConfiguration as Configuration } from './configuration';
diff --git a/packages/orm/test/fixtures/base-fn-multiple-db/src/configuration.ts b/packages/orm/test/fixtures/base-fn-multiple-db/src/configuration.ts
@@ -1,19 +1,16 @@
+import { IMidwayApplication } from '@midwayjs/core';
 import { App, Configuration, Inject } from '@midwayjs/decorator';
-import * as orm from '../../../../src';
+import * as assert from 'assert';
 import { join } from 'path';
+import { getConnection, Repository } from 'typeorm';
+import * as orm from '../../../../src';
 import { getRepository, InjectEntityModel } from '../../../../src';
+import { OnlyTestLoadLog } from './model/onlyTestLoadLog';
 import { User } from './model/user';
-import { Repository } from 'typeorm';
-import * as assert from 'assert';
-import { IMidwayApplication } from '@midwayjs/core';
 
 @Configuration({
-  imports: [
-    orm
-  ],
-  importConfigs: [
-    join(__dirname, './config')
-  ]
+  imports: [orm],
+  importConfigs: [join(__dirname, './config')],
 })
 export class ContainerConfiguration {
   @Inject('orm:getRepository')
@@ -41,14 +38,19 @@ export class ContainerConfiguration {
     const users = await this.defaultUserModel.findAndCount(user);
     assert(users[0][0]['name'] === 'oneuser1');
 
-
     const result = await this.testUserModel.findOne(user);
     assert(!result);
 
     const aa = await container.getAsync('baseFnMultipleHook');
     assert.equal(aa.bcreate, 1);
     assert.equal(aa.acreate, 1);
 
+    const defaultConn = getConnection('default');
+    const testConn = getConnection('test');
+
+    assert(defaultConn.options.entities.includes(OnlyTestLoadLog) === false);
+    assert(testConn.options.entities.includes(OnlyTestLoadLog) === true);
+
     this.app.setAttr('result', 'hello world' + JSON.stringify(users));
   }
 }
diff --git a/packages/orm/test/fixtures/base-fn-multiple-db/src/model/onlyTestLoadLog.ts b/packages/orm/test/fixtures/base-fn-multiple-db/src/model/onlyTestLoadLog.ts
@@ -0,0 +1,13 @@
+import { Column, PrimaryGeneratedColumn } from 'typeorm';
+import { EntityModel } from '../../../../../src';
+
+@EntityModel({
+  connectionName: 'test',
+})
+export class OnlyTestLoadLog {
+  @PrimaryGeneratedColumn({ name: 'id' })
+  id: number;
+
+  @Column({ name: 'content' })
+  content: string;
+}