feat: export singleton orm client

core/orm-decorator/src/decorator/Model.ts

@@ -3,6 +3,7 @@ import { ModelInfoUtil } from '../util/ModelInfoUtil';
 
 export interface ModelParams {
   tableName?: string;
+  dataSource?: string;
 }
 
 export const MODEL_PROTO_IMPL_TYPE = 'MODEL_PROTO';
@@ -18,6 +19,9 @@ export function Model(param?: ModelParams) {
     if (param?.tableName) {
       ModelInfoUtil.setTableName(param.tableName, clazz);
     }
+    if (param?.dataSource) {
+      ModelInfoUtil.setDataSource(param.dataSource, clazz);
+    }
     func(clazz);
   };
 }

plugin/orm/app.ts

@@ -6,13 +6,15 @@ import { ModelProtoHook } from './lib/ModelProtoHook';
 import { MODEL_PROTO_IMPL_TYPE } from '@eggjs/tegg-orm-decorator';
 import ContextModelProto from './lib/ContextModelProto';
 import { ContextModeObject } from './lib/ContextModeObject';
+import { ORMLoadUnitHook } from './lib/ORMLoadUnitHook';
 
 export default class OrmAppBootHook {
   private readonly app: Application;
   private readonly dataSourceManager: DataSourceManager;
   private readonly leoricRegister: LeoricRegister;
   private readonly modelProtoManager: ModelProtoManager;
   private readonly modelProtoHook: ModelProtoHook;
+  private readonly ormLoadUnitHook: ORMLoadUnitHook;
 
   constructor(app) {
     this.app = app;
@@ -22,11 +24,13 @@ export default class OrmAppBootHook {
     this.modelProtoHook = new ModelProtoHook(this.modelProtoManager);
     this.app.eggPrototypeCreatorFactory.registerPrototypeCreator(MODEL_PROTO_IMPL_TYPE, ContextModelProto.createProto);
     this.app.leoricRegister = this.leoricRegister;
+    this.ormLoadUnitHook = new ORMLoadUnitHook();
   }
 
   configWillLoad() {
     this.app.eggPrototypeLifecycleUtil.registerLifecycle(this.modelProtoHook);
     this.app.eggObjectFactory.registerEggObjectCreateMethod(ContextModelProto, ContextModeObject.createObject);
+    this.app.loadUnitLifecycleUtil.registerLifecycle(this.ormLoadUnitHook);
   }
 
   configDidLoad() {

plugin/orm/lib/LeoricRegister.ts

@@ -17,22 +17,35 @@ export class LeoricRegister extends Base {
     this.realmMap = new Map();
   }
 
-  getOrCreateRealm(datasource: string | undefined): any {
+  getConfig(datasource?: string) {
     let config: OrmConfig | undefined;
     if (!datasource) {
       config = this.dataSourceManager.getDefaultConfig();
     } else {
       config = this.dataSourceManager.getConfig(datasource);
     }
-    if (!config) {
-      throw new Error(`not found datasource for ${datasource}`);
+    return config;
+  }
+
+  getRealm(config: OrmConfig | undefined): Realm | undefined {
+    if (!config?.database) {
+      return undefined;
     }
-    let realm = this.realmMap.get(config.database);
-    if (realm) {
-      return realm;
+    const realm = this.realmMap.get(config.database);
+    return realm;
+  }
+
+  getOrCreateRealm(datasource: string | undefined): any {
+    const config = this.getConfig(datasource);
+    let realm: Realm | undefined;
+    if (config) {
+      realm = this.getRealm(config);
+      if (realm) {
+        return realm;
+      }
     }
     realm = new (Realm as any)({ ...config });
-    this.realmMap.set(config.database, realm);
+    this.realmMap.set(config!.database, realm);
     return realm;
   }
 

plugin/orm/lib/ORMLoadUnitHook.ts

@@ -0,0 +1,24 @@
+import { LifecycleHook } from '@eggjs/tegg';
+import {
+  EggLoadUnitType,
+  EggPrototypeCreatorFactory,
+  EggPrototypeFactory,
+  LoadUnit,
+  LoadUnitLifecycleContext,
+} from '@eggjs/tegg-metadata';
+import { Orm } from './SingletonORM';
+
+const REGISTER_CLAZZ = [
+  Orm,
+];
+
+export class ORMLoadUnitHook implements LifecycleHook<LoadUnitLifecycleContext, LoadUnit> {
+  async postCreate(_ctx: LoadUnitLifecycleContext, loadUnit: LoadUnit): Promise<void> {
+    if (loadUnit.type === EggLoadUnitType.APP) {
+      for (const clazz of REGISTER_CLAZZ) {
+        const proto = await EggPrototypeCreatorFactory.createProto(clazz, loadUnit);
+        EggPrototypeFactory.instance.registerPrototype(proto, loadUnit);
+      }
+    }
+  }
+}

plugin/orm/lib/SingletonORM.ts

@@ -0,0 +1,30 @@
+import {
+  AccessLevel,
+  Inject,
+  SingletonProto,
+} from '@eggjs/tegg';
+import { LeoricRegister } from './LeoricRegister';
+import Realm from 'leoric';
+
+@SingletonProto({
+  accessLevel: AccessLevel.PUBLIC,
+})
+export class Orm {
+  @Inject()
+  private leoricRegister: LeoricRegister;
+
+  // default dataSource
+  get client(): Realm {
+    const defaultConfig = this.leoricRegister.getConfig();
+    return this.leoricRegister.getRealm(defaultConfig)!;
+  }
+
+  getClient(datasource: string): Realm {
+    const config = this.leoricRegister.getConfig(datasource);
+    if (!config) {
+      throw new Error(`not found ${datasource} datasource`);
+    }
+    return this.leoricRegister.getOrCreateRealm(config.database)!;
+  }
+
+}

plugin/orm/test/fixtures/apps/orm-app/config/config.default.js

@@ -1,6 +1,6 @@
 'use strict';
 
-module.exports = function() {
+module.exports = function () {
   const config = {
     keys: 'test key',
     security: {
@@ -9,19 +9,53 @@ module.exports = function() {
       },
     },
     orm: {
-      client: 'mysql',
-      database: 'test',
-      host: '127.0.0.1',
-      port: 3306,
-      user: 'root',
+      datasources: [
+        {
+          client: 'mysql',
+          database: 'test',
+          host: '127.0.0.1',
+          port: 3306,
+          user: 'root',
 
-      delegate: 'model',
-      baseDir: 'model',
-      migrations: 'database',
+          delegate: 'model',
+          baseDir: 'model',
+          migrations: 'database',
 
-      define: {
-        underscored: true,
-      },
+          define: {
+            underscored: true,
+          },
+        },
+        {
+          client: 'mysql',
+          database: 'apple',
+          host: '127.0.0.1',
+          port: 3306,
+          user: 'root',
+
+          delegate: 'model',
+          baseDir: 'model',
+          migrations: 'database',
+
+          define: {
+            underscored: true,
+          },
+        },
+        {
+          client: 'mysql',
+          database: 'banana',
+          host: '127.0.0.1',
+          port: 3306,
+          user: 'root',
+
+          delegate: 'model',
+          baseDir: 'model',
+          migrations: 'database',
+
+          define: {
+            underscored: true,
+          },
+        },
+      ],
     },
   };
   return config;

plugin/orm/test/fixtures/apps/orm-app/modules/orm-module/AppService.ts

@@ -1,11 +1,15 @@
 import { ContextProto, Inject } from '@eggjs/tegg';
+import { Orm } from '@eggjs/tegg-orm-plugin/lib/SingletonORM';
 import { App } from './model/App';
 
 @ContextProto()
 export class AppService {
   @Inject()
   App: typeof App;
 
+  @Inject()
+  private readonly orm: Orm;
+
   async createApp(data: {
     name: string;
     desc: string;
@@ -18,4 +22,17 @@ export class AppService {
     const app = await this.App.findOne({ name });
     return app as App;
   }
+
+  async rawQuery(dataSource: string, sql: string, values?: any[]) {
+    return await this.orm.getClient(dataSource).query(sql, values);
+  }
+
+  async getClient(name: string) {
+    return this.orm.getClient(name);
+  }
+
+  async getDefaultClient() {
+    return this.orm.client;
+  }
+
 }

plugin/orm/test/fixtures/apps/orm-app/modules/orm-module/model/App.ts

@@ -1,7 +1,9 @@
 import { Attribute, Model } from '@eggjs/tegg-orm-decorator';
 import { DataTypes, Bone } from 'leoric';
 
-@Model()
+@Model({
+  dataSource: 'test',
+})
 export class App extends Bone {
   @Attribute(DataTypes.STRING)
   name: string;

plugin/orm/test/fixtures/apps/orm-app/modules/orm-module/model/Pkg.ts

@@ -1,7 +1,9 @@
 import { Attribute, Model } from '@eggjs/tegg-orm-decorator';
 import { DataTypes, Bone } from 'leoric';
 
-@Model()
+@Model({
+  dataSource: 'test',
+})
 export class Pkg extends Bone {
   @Attribute(DataTypes.STRING)
   name: string;

plugin/orm/test/fixtures/prepare.js

@@ -33,6 +33,10 @@ async function query(sql) {
 async function init() {
   await query('DROP DATABASE IF EXISTS `test`;');
   await query('CREATE DATABASE test;');
+  await query('DROP DATABASE IF EXISTS `apple`;');
+  await query('CREATE DATABASE apple;');
+  await query('DROP DATABASE IF EXISTS `banana`;');
+  await query('CREATE DATABASE banana;');
   await query('use test;');
   await query('CREATE TABLE `apps` (\n' +
     '  `id` bigint unsigned NOT NULL AUTO_INCREMENT,\n' +

plugin/orm/test/index.test.ts

@@ -14,6 +14,7 @@ describe('test/orm.test.ts', () => {
   }
 
   let app: MockApplication;
+  let appService: AppService;
   let ctx;
 
   afterEach(async () => {
@@ -82,4 +83,57 @@ describe('test/orm.test.ts', () => {
     await appService.findApp('egg');
     app.expectLog(/sql: SELECT \* FROM `apps` WHERE `name` = 'egg' LIMIT 1 path: \//);
   });
+
+  it('singleton ORM client', async () => {
+    ctx = await app.mockModuleContext();
+    appService = await ctx.getEggObject(AppService);
+
+    describe('raw query', () => {
+      before(async () => {
+        const appModel = await appService.createApp({
+          name: 'egg',
+          desc: 'the framework',
+        });
+        assert(appModel);
+        assert(appModel.name === 'egg');
+        assert(appModel.desc === 'the framework');
+      });
+
+      it('query success', async () => {
+        const res = await appService.rawQuery('test', 'select * from apps where name = "egg"');
+        assert(res.rows.length === 1);
+        assert(res.rows[0].name === 'egg');
+      });
+
+      it('query success for args', async () => {
+        const res = await appService.rawQuery('test', 'select * from apps where name = ?', [ 'egg' ]);
+        assert(res.rows.length === 1);
+        assert(res.rows[0].name === 'egg');
+      });
+    });
+
+    describe('multi db', () => {
+
+      it('should work for multi database', async () => {
+        const appleClient = await appService.getClient('apple');
+        const bananaClient = await appService.getClient('banana');
+        assert(appleClient.options.database === 'apple');
+        assert(appleClient.options.database === 'apple');
+        assert(bananaClient.options.database === 'banana');
+        assert(bananaClient.options.database === 'banana');
+      });
+
+      it('should throw when invalid database', async () => {
+        await assert.rejects(async () => {
+          await appService.getClient('orange');
+        }, /not found orange datasource/);
+      });
+
+      it('should return undefined when get default client', async () => {
+        const defaultClient = await appService.getDefaultClient();
+        assert(defaultClient === undefined);
+      });
+    });
+
+  });
 });

plugin/orm/typings/index.d.ts

@@ -4,6 +4,7 @@ import '@eggjs/tegg-plugin';
 import { DataType } from 'leoric';
 import { AttributeOptions } from '@eggjs/tegg-orm-decorator';
 import { LeoricRegister } from '../lib/LeoricRegister';
+import { Orm } from '../lib/SingletonORM';
 
 declare module '@eggjs/tegg-orm-decorator' {
   export declare function Attribute(dataType: DataType, options?: AttributeOptions): (target: any, propertyKey: PropertyKey) => void;
@@ -12,6 +13,7 @@ declare module '@eggjs/tegg-orm-decorator' {
 declare module 'egg' {
   export interface TeggOrmApplication {
     leoricRegister: LeoricRegister;
+    orm: Orm;
   }
 
   interface Application extends TeggOrmApplication {