feat: update datasource loading progress (#5518)

* chore: tmp commit

* chore: datasource loading progress

* chore: merge

* fix: loading status

* chore: test

* chore: test

packages/core/data-source-manager/src/data-source.ts

@@ -17,25 +17,24 @@ import { Logger } from '@nocobase/logger';
 
 export type DataSourceOptions = any;
 
+export type LoadingProgress = {
+  total: number;
+  loaded: number;
+};
+
 export abstract class DataSource extends EventEmitter {
   public collectionManager: ICollectionManager;
   public resourceManager: ResourceManager;
   public acl: ACL;
+
   logger: Logger;
-  _sqlLogger: Logger;
 
   constructor(protected options: DataSourceOptions) {
     super();
     this.init(options);
   }
 
-  setLogger(logger: Logger) {
-    this.logger = logger;
-  }
-
-  setSqlLogger(logger: Logger) {
-    this._sqlLogger = logger;
-  }
+  _sqlLogger: Logger;
 
   get sqlLogger() {
     return this._sqlLogger || this.logger;
@@ -49,6 +48,14 @@ export abstract class DataSource extends EventEmitter {
     return Promise.resolve(true);
   }
 
+  setLogger(logger: Logger) {
+    this.logger = logger;
+  }
+
+  setSqlLogger(logger: Logger) {
+    this._sqlLogger = logger;
+  }
+
   init(options: DataSourceOptions = {}) {
     this.acl = this.createACL();
 
@@ -100,6 +107,10 @@ export abstract class DataSource extends EventEmitter {
     return null;
   }
 
+  emitLoadingProgress(progress: LoadingProgress) {
+    this.emit('loadingProgress', progress);
+  }
+
   async load(options: any = {}) {}
   async close() {}
 

packages/plugins/@nocobase/plugin-data-source-manager/src/server/__tests__/data-sources.test.ts

@@ -23,6 +23,86 @@ describe('data source', async () => {
     await app.destroy();
   });
 
+  it('should return datasource status when datasource is loading or reloading', async () => {
+    class MockDataSource extends DataSource {
+      static testConnection(options?: any): Promise<boolean> {
+        return Promise.resolve(true);
+      }
+
+      async load(): Promise<void> {
+        await waitSecond(1000);
+      }
+
+      createCollectionManager(options?: any): any {
+        return undefined;
+      }
+    }
+
+    app.dataSourceManager.factory.register('mock', MockDataSource);
+
+    app.dataSourceManager.beforeAddDataSource(async (dataSource: DataSource) => {
+      const total = 1000;
+      for (let i = 0; i < total; i++) {
+        dataSource.emitLoadingProgress({
+          total,
+          loaded: i,
+        });
+      }
+    });
+
+    await app.db.getRepository('dataSources').create({
+      values: {
+        key: 'mockInstance1',
+        type: 'mock',
+        displayName: 'Mock',
+        options: {},
+      },
+    });
+
+    await waitSecond(200);
+
+    // get data source status
+    const plugin: any = app.pm.get('data-source-manager');
+    expect(plugin.dataSourceStatus['mockInstance1']).toBe('loading');
+
+    const loadingStatus = plugin.dataSourceLoadingProgress['mockInstance1'];
+    expect(loadingStatus).toBeDefined();
+  });
+
+  it('should get error when datasource loading failed', async () => {
+    class MockDataSource extends DataSource {
+      static testConnection(options?: any): Promise<boolean> {
+        return Promise.resolve(true);
+      }
+
+      async load(): Promise<void> {
+        throw new Error(`load failed`);
+      }
+
+      createCollectionManager(options?: any): any {
+        return undefined;
+      }
+    }
+
+    app.dataSourceManager.factory.register('mock', MockDataSource);
+
+    await app.db.getRepository('dataSources').create({
+      values: {
+        key: 'mockInstance1',
+        type: 'mock',
+        displayName: 'Mock',
+        options: {},
+      },
+    });
+
+    await waitSecond(2000);
+    // get data source status
+    const plugin: any = app.pm.get('data-source-manager');
+    expect(plugin.dataSourceStatus['mockInstance1']).toBe('loading-failed');
+
+    expect(plugin.dataSourceErrors['mockInstance1'].message).toBe('load failed');
+  });
+
   it('should list main datasource in api', async () => {
     const listResp = await app.agent().resource('dataSources').list();
     expect(listResp.status).toBe(200);

packages/plugins/@nocobase/plugin-data-source-manager/src/server/models/data-source.ts

@@ -100,6 +100,10 @@ export class DataSourceModel extends Model {
         sqlLogger: app.sqlLogger.child({ dataSourceKey }),
       });
 
+      dataSource.on('loadingProgress', (progress) => {
+        pluginDataSourceManagerServer.dataSourceLoadingProgress[dataSourceKey] = progress;
+      });
+
       if (loadAtAfterStart) {
         dataSource.on('loadMessage', ({ message }) => {
           app.setMaintainingMessage(`${message} in data source ${this.get('displayName')}`);

packages/plugins/@nocobase/plugin-data-source-manager/src/server/plugin.ts

@@ -22,6 +22,7 @@ import { DataSourcesRolesResourcesModel } from './models/connections-roles-resou
 import { DataSourcesRolesResourcesActionModel } from './models/connections-roles-resources-action';
 import { DataSourceModel } from './models/data-source';
 import { DataSourcesRolesModel } from './models/data-sources-roles-model';
+import { LoadingProgress } from '@nocobase/data-source-manager';
 
 type DataSourceState = 'loading' | 'loaded' | 'loading-failed' | 'reloading' | 'reloading-failed';
 
@@ -36,6 +37,10 @@ export class PluginDataSourceManagerServer extends Plugin {
     [dataSourceKey: string]: DataSourceState;
   } = {};
 
+  public dataSourceLoadingProgress: {
+    [dataSourceKey: string]: LoadingProgress;
+  } = {};
+
   async beforeLoad() {
     this.app.db.registerModels({
       DataSourcesCollectionModel,
@@ -87,6 +92,10 @@ export class PluginDataSourceManagerServer extends Plugin {
 
         const klass = this.app.dataSourceManager.factory.getClass(type);
 
+        if (!klass) {
+          throw new Error(`Data source type "${type}" is not registered`);
+        }
+
         try {
           await klass.testConnection(dataSourceOptions);
         } catch (error) {

packages/plugins/@nocobase/plugin-data-source-manager/src/server/resourcers/data-sources-collections.ts

@@ -17,6 +17,29 @@ export default {
 
       const { associatedIndex: dataSourceKey } = params;
       const dataSource = ctx.app.dataSourceManager.dataSources.get(dataSourceKey);
+      const plugin: any = ctx.app.pm.get('data-source-manager');
+
+      const dataSourceStatus = plugin.dataSourceStatus[dataSourceKey];
+
+      if (dataSourceStatus === 'loading-failed') {
+        const error = plugin.dataSourceErrors[dataSourceKey];
+        if (error) {
+          throw new Error(`dataSource ${dataSourceKey} loading failed: ${error.message}`);
+        }
+
+        throw new Error(`dataSource ${dataSourceKey} loading failed`);
+      }
+
+      if (['loading', 'reloading'].includes(dataSourceStatus)) {
+        const progress = plugin.dataSourceLoadingProgress[dataSourceKey];
+
+        if (progress) {
+          throw new Error(`dataSource ${dataSourceKey} is ${dataSourceStatus} (${progress.loaded}/${progress.total})`);
+        }
+
+        throw new Error(`dataSource ${dataSourceKey} is ${dataSourceStatus}`);
+      }
+
       if (!dataSource) {
         throw new Error(`dataSource ${dataSourceKey} not found`);
       }