From 68f9c4193ab1dacd4433cdc027914ff646a798e4 Mon Sep 17 00:00:00 2001
From: Chareice <chareice@live.com>
Date: Thu, 1 Feb 2024 22:54:46 +0800
Subject: [PATCH] chore: await load datasource at application start

---
 packages/core/data-source-manager/src/data-source.ts |  5 ++++-
 .../src/server/models/data-source.ts                 | 10 ++++++++--
 .../plugin-data-source-manager/src/server/plugin.ts  | 12 +++++++-----
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/packages/core/data-source-manager/src/data-source.ts b/packages/core/data-source-manager/src/data-source.ts
index 0add540ea6..4bec506a9f 100644
--- a/packages/core/data-source-manager/src/data-source.ts
+++ b/packages/core/data-source-manager/src/data-source.ts
@@ -4,13 +4,16 @@ import compose from 'koa-compose';
 import { ResourceManager } from './resource-manager';
 import { loadDefaultActions } from './load-default-actions';
 import { ICollectionManager } from './types';
+import EventEmitter from 'events';
 
-export abstract class DataSource {
+export abstract class DataSource extends EventEmitter {
   public collectionManager: ICollectionManager;
   public resourceManager: ResourceManager;
   public acl: ACL;
 
   constructor(protected options: any) {
+    super();
+
     this.acl = this.createACL();
 
     this.resourceManager = this.createResourceManager({
diff --git a/packages/plugins/@nocobase/plugin-data-source-manager/src/server/models/data-source.ts b/packages/plugins/@nocobase/plugin-data-source-manager/src/server/models/data-source.ts
index b4fd7e31ba..b3e3624819 100644
--- a/packages/plugins/@nocobase/plugin-data-source-manager/src/server/models/data-source.ts
+++ b/packages/plugins/@nocobase/plugin-data-source-manager/src/server/models/data-source.ts
@@ -45,8 +45,8 @@ const availableActions: {
 };
 
 export class DataSourceModel extends Model {
-  async loadIntoApplication(options: { app: Application; transaction?: Transaction }) {
-    const { app, transaction } = options;
+  async loadIntoApplication(options: { app: Application; transaction?: Transaction; loadAtAfterStart?: boolean }) {
+    const { app, loadAtAfterStart } = options;
 
     const dataSourceKey = this.get('key');
 
@@ -72,6 +72,12 @@ export class DataSourceModel extends Model {
       name: this.get('key'),
     });
 
+    if (loadAtAfterStart) {
+      dataSource.on('loadMessage', ({ message }) => {
+        app.setMaintainingMessage(`${message} in data source ${this.get('displayName')}`);
+      });
+    }
+
     const acl = dataSource.acl;
 
     for (const [actionName, actionParams] of Object.entries(availableActions)) {
diff --git a/packages/plugins/@nocobase/plugin-data-source-manager/src/server/plugin.ts b/packages/plugins/@nocobase/plugin-data-source-manager/src/server/plugin.ts
index aaced9cc94..a60f22e41b 100644
--- a/packages/plugins/@nocobase/plugin-data-source-manager/src/server/plugin.ts
+++ b/packages/plugins/@nocobase/plugin-data-source-manager/src/server/plugin.ts
@@ -227,11 +227,13 @@ export class PluginDataSourceManagerServer extends Plugin {
 
     this.app.on('afterStart', async (app: Application) => {
       const dataSourcesRecords: DataSourceModel[] = await this.app.db.getRepository('dataSources').find();
-      for (const dataSourceRecord of dataSourcesRecords) {
-        dataSourceRecord.loadIntoApplication({
-          app,
-        });
-      }
+
+      const loadPromises = dataSourcesRecords.map((dataSourceRecord) =>
+        dataSourceRecord.loadIntoApplication({ app, loadAtAfterStart: true }),
+      );
+
+      this.app.setMaintainingMessage('Loading data sources...');
+      await Promise.all(loadPromises);
     });
 
     this.app.db.on('dataSourcesRolesResources.afterSaveWithAssociations', async (model, options) => {