SCCommon/wise-cost-connect/src/connect/mongoDBConnect.ts

import { IMidwayApplication } from '@midwayjs/core';
import { Inject } from '@midwayjs/decorator';
import mongoose, { Connection, ConnectOptions, Mongoose } from 'mongoose';

export interface IMongooseApp extends IMidwayApplication {
  mongoose: Mongoose;
  mongooseDB: MongoDBConnect;
}

export class MongoDBConnect {
  @Inject()
  logger;

  private clients: Map<string, Connection>;

  private newConnection(
    name = 'default',
    config: { url: string; options: ConnectOptions }
  ) {
    const { url, options } = config;
    if (!url) {
      return this.logger.error('url 不能为空！');
    }
    const conn = mongoose.createConnection(url, options);
    this.clients[name] = conn;
  }

  private createClients(config: any) {
    const { client, clients } = config;
    if (client) {
      this.newConnection(client.name, client);
    }

    if (clients) {
      Object.keys(clients).forEach(dbName =>
        this.newConnection(dbName, clients[dbName])
      );
    }
  }

  constructor(config: any, app: IMongooseApp) {
    this.clients = new Map();
    this.createClients(config);
    app.mongoose = mongoose;
    app.mongooseDB = this;
  }

  get(dbName: string) {
    return this.clients[dbName];
  }

  async closeDB() {
    const connects: Connection[] = Object.values(this.clients);
    for (const connect of connects) {
      await connect.close();
    }
  }
}