feat: register once hook in datasource manager (#4024)

* chore: datasource hook

* feat: register once hook in datasource manager

* chore: api name

* chore: test

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

@@ -1,5 +1,6 @@
 import { createMockServer, mockDatabase, supertest } from '@nocobase/test';
 import { SequelizeDataSource } from '../sequelize-data-source';
+import { vi } from 'vitest';
 
 describe('example', () => {
   test.skip('case1', async () => {
@@ -153,4 +154,48 @@ describe('example', () => {
 
     await app.destroy();
   });
+
+  it('should register every datasource instance', async () => {
+    const hook = vi.fn();
+
+    const app = await createMockServer({
+      acl: false,
+      resourcer: {
+        prefix: '/api/',
+      },
+      name: 'update-filter',
+    });
+
+    app.dataSourceManager.afterAddDataSource(hook);
+    // it should be called on main datasource
+    expect(hook).toBeCalledTimes(1);
+
+    const database = mockDatabase({
+      tablePrefix: 'ds1_',
+    });
+
+    // it should be called when adding a new datasource
+    const ds1 = new SequelizeDataSource({
+      name: 'ds1',
+      resourceManager: {},
+      collectionManager: {
+        database,
+      },
+    });
+
+    ds1.collectionManager.defineCollection({
+      name: 'test1',
+      fields: [{ type: 'string', name: 'name' }],
+    });
+
+    await ds1.collectionManager.sync();
+
+    ds1.acl.allow('test1', 'update', 'public');
+
+    await app.dataSourceManager.add(ds1);
+
+    expect(hook).toBeCalledTimes(2);
+
+    await app.destroy();
+  });
 });

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

@@ -2,10 +2,14 @@ import { ToposortOptions } from '@nocobase/utils';
 import { DataSource } from './data-source';
 import { DataSourceFactory } from './data-source-factory';
 
+type DataSourceHook = (dataSource: DataSource) => void;
+
 export class DataSourceManager {
   dataSources: Map<string, DataSource>;
   factory: DataSourceFactory = new DataSourceFactory();
 
+  onceHooks: Array<DataSourceHook> = [];
+
   protected middlewares = [];
 
   constructor(public options = {}) {
@@ -16,6 +20,10 @@ export class DataSourceManager {
   async add(dataSource: DataSource, options: any = {}) {
     await dataSource.load(options);
     this.dataSources.set(dataSource.name, dataSource);
+
+    for (const hook of this.onceHooks) {
+      hook(dataSource);
+    }
   }
 
   use(fn: any, options?: ToposortOptions) {
@@ -36,4 +44,15 @@ export class DataSourceManager {
       return ds.middleware(this.middlewares)(ctx, next);
     };
   }
+
+  afterAddDataSource(hook: DataSourceHook) {
+    this.addHookAndRun(hook);
+  }
+
+  private addHookAndRun(hook: DataSourceHook) {
+    this.onceHooks.push(hook);
+    for (const dataSource of this.dataSources.values()) {
+      hook(dataSource);
+    }
+  }
 }

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

@@ -79,18 +79,23 @@ export abstract class DataSource extends EventEmitter {
   }
 
   middleware(middlewares: any = []) {
+    const dataSource = this;
+
     if (!this['_used']) {
       for (const [fn, options] of middlewares) {
         this.resourceManager.use(fn, options);
       }
       this['_used'] = true;
     }
+
     return async (ctx, next) => {
       ctx.getCurrentRepository = () => {
         const { resourceName, resourceOf } = ctx.action;
         return this.collectionManager.getRepository(resourceName, resourceOf);
       };
 
+      ctx.dataSource = dataSource;
+
       return compose([this.collectionToResourceMiddleware(), this.resourceManager.middleware()])(ctx, next);
     };
   }

packages/plugins/@nocobase/plugin-export/src/server/actions/export-xlsx.ts

@@ -6,12 +6,11 @@ import { columns2Appends } from '../utils';
 
 export async function exportXlsx(ctx: Context, next: Next) {
   const { title, filter, sort, fields, except } = ctx.action.params;
-  const { resourceName, resourceOf } = ctx.action;
   let columns = ctx.action.params.values?.columns || ctx.action.params?.columns;
   if (typeof columns === 'string') {
     columns = JSON.parse(columns);
   }
-  const repository = ctx.db.getRepository<any>(resourceName, resourceOf) as Repository;
+  const repository = ctx.getCurrentRepository() as Repository;
   const collection = repository.collection;
   columns = columns?.filter((col) => collection.hasField(col.dataIndex[0]) && col?.dataIndex?.length > 0);
   const appends = columns2Appends(columns, ctx);

packages/plugins/@nocobase/plugin-export/src/server/index.ts

@@ -5,11 +5,18 @@ export class PluginExportServer extends Plugin {
   beforeLoad() {}
 
   async load() {
-    this.app.resourcer.registerActionHandler('export', exportXlsx);
-    this.app.acl.setAvailableAction('export', {
+    this.app.dataSourceManager.afterAddDataSource((dataSource) => {
+      // @ts-ignore
+      if (!dataSource.collectionManager?.db) {
+        return;
+      }
+
+      dataSource.resourceManager.registerActionHandler('export', exportXlsx);
+      dataSource.acl.setAvailableAction('export', {
         displayName: '{{t("Export")}}',
         allowConfigureFields: true,
       });
+    });
   }
 
   async install(options: InstallOptions) {}

packages/plugins/@nocobase/plugin-export/src/server/utils/columns2Appends.ts

@@ -2,13 +2,13 @@ export function columns2Appends(columns, ctx) {
   const { resourceName } = ctx.action;
   const appends = new Set([]);
   for (const column of columns) {
-    let collection = ctx.db.getCollection(resourceName);
+    let collection = ctx.dataSource.collectionManager.getCollection(resourceName);
     const appendColumns = [];
     for (let i = 0, iLen = column.dataIndex.length; i < iLen; i++) {
       const field = collection.getField(column.dataIndex[i]);
       if (field?.target) {
         appendColumns.push(column.dataIndex[i]);
-        collection = ctx.db.getCollection(field.target);
+        collection = ctx.dataSource.collectionManager.getCollection(field.target);
       }
     }
     if (appendColumns.length > 0) {

packages/plugins/@nocobase/plugin-import/src/server/index.ts

@@ -11,18 +11,25 @@ export class PluginImportServer extends Plugin {
   }
 
   async load() {
-    this.app.resourcer.use(importMiddleware);
-    this.app.resourcer.registerActionHandler('downloadXlsxTemplate', downloadXlsxTemplate);
-    this.app.resourcer.registerActionHandler('importXlsx', importXlsx);
+    this.app.dataSourceManager.afterAddDataSource((dataSource) => {
+      // @ts-ignore
+      if (!dataSource.collectionManager?.db) {
+        return;
+      }
 
-    this.app.acl.setAvailableAction('importXlsx', {
+      dataSource.resourceManager.use(importMiddleware);
+      dataSource.resourceManager.registerActionHandler('downloadXlsxTemplate', downloadXlsxTemplate);
+      dataSource.resourceManager.registerActionHandler('importXlsx', importXlsx);
+
+      dataSource.acl.setAvailableAction('importXlsx', {
         displayName: '{{t("Import")}}',
         allowConfigureFields: true,
         type: 'new-data',
         onNewRecord: true,
       });
 
-    this.app.acl.allow('*', 'downloadXlsxTemplate', 'loggedIn');
+      dataSource.acl.allow('*', 'downloadXlsxTemplate', 'loggedIn');
+    });
   }
 
   async install(options: InstallOptions) {