From 76a225e3542e70d7f9d068e15bd1bc1bade52f69 Mon Sep 17 00:00:00 2001
From: ChengLei Shao <chareice@live.com>
Date: Fri, 17 Nov 2023 11:22:46 +0800
Subject: [PATCH] feat: register collection sync logic (#3055)

---
 .../src/__tests__/collection-factory.test.ts  |  27 +++-
 .../core/database/src/collection-factory.ts   |  32 +++--
 packages/core/database/src/database.ts        |  18 ++-
 packages/core/database/src/model.ts           | 117 ++++++++++--------
 4 files changed, 121 insertions(+), 73 deletions(-)

diff --git a/packages/core/database/src/__tests__/collection-factory.test.ts b/packages/core/database/src/__tests__/collection-factory.test.ts
index e5022b294..ffbef3297 100644
--- a/packages/core/database/src/__tests__/collection-factory.test.ts
+++ b/packages/core/database/src/__tests__/collection-factory.test.ts
@@ -20,8 +20,8 @@ describe('collection factory', function () {
       static type = 'child';
     }
 
-    db.collectionFactory.registerCollectionType(ChildCollection, (options) => {
-      return options.child == true;
+    db.collectionFactory.registerCollectionType(ChildCollection, {
+      condition: (options) => options.child,
     });
 
     const collection = db.collectionFactory.createCollection({
@@ -37,4 +37,27 @@ describe('collection factory', function () {
 
     expect(collection2).toBeInstanceOf(Collection);
   });
+
+  it('should register collection type with sync logic', async () => {
+    class ChildCollection extends Collection {
+      static type = 'child';
+    }
+
+    const fn = jest.fn();
+
+    db.collectionFactory.registerCollectionType(ChildCollection, {
+      condition: (options) => options.child,
+      onSync(model, options) {
+        fn();
+      },
+    });
+
+    const collection = db.collectionFactory.createCollection({
+      name: 'child',
+      child: true,
+    });
+
+    await collection.sync();
+    expect(fn).toHaveBeenCalledTimes(1);
+  });
 });
diff --git a/packages/core/database/src/collection-factory.ts b/packages/core/database/src/collection-factory.ts
index 56a31d942..5d3324e39 100644
--- a/packages/core/database/src/collection-factory.ts
+++ b/packages/core/database/src/collection-factory.ts
@@ -1,28 +1,38 @@
 import { Collection, CollectionOptions } from './collection';
 import Database from './database';
+import { Model } from './model';
+
+type CollectionTypeOptions = {
+  condition: (options: CollectionOptions) => boolean;
+  onSync?: (model: typeof Model, options: any) => Promise<void>;
+};
 
 export class CollectionFactory {
-  private collectionTypes: Array<{
-    ctor: typeof Collection;
-    condition: (options: CollectionOptions) => boolean;
-  }> = [];
+  // Using a Map with the collection subclass as the key and options as the value
+  public collectionTypes: Map<typeof Collection, CollectionTypeOptions> = new Map();
 
   constructor(private database: Database) {}
 
-  registerCollectionType(collectionClass: typeof Collection, condition: (options: CollectionOptions) => boolean) {
-    this.collectionTypes.push({ ctor: collectionClass, condition });
+  registerCollectionType(
+    collectionClass: typeof Collection, // Using the collection class as the key
+    options: CollectionTypeOptions,
+  ) {
+    // Storing the options associated with the collection class
+    this.collectionTypes.set(collectionClass, options);
   }
 
-  createCollection<T extends Collection>(options: CollectionOptions): T {
-    let klass = Collection;
-    for (const { ctor, condition } of this.collectionTypes) {
-      if (condition(options)) {
+  createCollection<T extends Collection>(collectionOptions: CollectionOptions): T {
+    let klass: typeof Collection = Collection;
+
+    // Iterating over the map to find the right class based on the condition
+    for (const [ctor, options] of this.collectionTypes) {
+      if (options.condition(collectionOptions)) {
         klass = ctor;
         break;
       }
     }
 
-    return new klass(options, {
+    return new klass(collectionOptions, {
       database: this.database,
     }) as T;
   }
diff --git a/packages/core/database/src/database.ts b/packages/core/database/src/database.ts
index 613f5386a..9809254f9 100644
--- a/packages/core/database/src/database.ts
+++ b/packages/core/database/src/database.ts
@@ -313,16 +313,22 @@ export class Database extends EventEmitter implements AsyncEmitter {
   }
 
   registerCollectionType() {
-    this.collectionFactory.registerCollectionType(InheritedCollection, (options) => {
-      return options.inherits && lodash.castArray(options.inherits).length > 0;
+    this.collectionFactory.registerCollectionType(InheritedCollection, {
+      condition: (options) => {
+        return options.inherits && lodash.castArray(options.inherits).length > 0;
+      },
     });
 
-    this.collectionFactory.registerCollectionType(ViewCollection, (options) => {
-      return options.viewName || options.view;
+    this.collectionFactory.registerCollectionType(ViewCollection, {
+      condition: (options) => {
+        return options.viewName || options.view;
+      },
     });
 
-    this.collectionFactory.registerCollectionType(SqlCollection, (options) => {
-      return options.sql;
+    this.collectionFactory.registerCollectionType(SqlCollection, {
+      condition: (options) => {
+        return options.sql;
+      },
     });
   }
 
diff --git a/packages/core/database/src/model.ts b/packages/core/database/src/model.ts
index 4cec90e6b..ef1f4eee3 100644
--- a/packages/core/database/src/model.ts
+++ b/packages/core/database/src/model.ts
@@ -31,6 +31,69 @@ export class Model<TModelAttributes extends {} = any, TCreationAttributes extend
   protected _changedWithAssociations = new Set();
   protected _previousDataValuesWithAssociations = {};
 
+  static async sync(options) {
+    if (this.collection.isView()) {
+      return;
+    }
+
+    if (this.collection.options.sync === false) {
+      return;
+    }
+
+    // @ts-ignore
+    const collectionSyncOptions = this.database.collectionFactory.collectionTypes.get(this.collection.constructor)
+      ?.onSync;
+
+    if (collectionSyncOptions) {
+      await collectionSyncOptions(this, options);
+      return;
+    }
+
+    const model = this as any;
+
+    const _schema = model._schema;
+
+    if (_schema && _schema != 'public') {
+      await this.sequelize.query(`CREATE SCHEMA IF NOT EXISTS "${_schema}";`, {
+        raw: true,
+        transaction: options?.transaction,
+      });
+    }
+
+    // fix sequelize sync with model that not have any column
+    if (Object.keys(model.tableAttributes).length === 0) {
+      if (this.database.inDialect('sqlite', 'mysql')) {
+        console.error(`Zero-column tables aren't supported in ${this.database.sequelize.getDialect()}`);
+        return;
+      }
+
+      // @ts-ignore
+      const queryInterface = this.sequelize.queryInterface;
+
+      if (!queryInterface.patched) {
+        const oldDescribeTable = queryInterface.describeTable;
+        queryInterface.describeTable = async function (...args) {
+          try {
+            return await oldDescribeTable.call(this, ...args);
+          } catch (err) {
+            if (err.message.includes('No description found for')) {
+              return [];
+            } else {
+              throw err;
+            }
+          }
+        };
+        queryInterface.patched = true;
+      }
+    }
+
+    if (this.collection.isInherited()) {
+      return SyncRunner.syncInheritModel(model, options);
+    }
+
+    return SequelizeModel.sync.call(this, options);
+  }
+
   // TODO
   public toChangedWithAssociations() {
     // @ts-ignore
@@ -149,58 +212,4 @@ export class Model<TModelAttributes extends {} = any, TCreationAttributes extend
 
     return lodash.orderBy(data, orderItems, orderDirections);
   }
-
-  static async sync(options) {
-    if (this.collection.isView()) {
-      return;
-    }
-
-    if (this.collection.options.sync === false) {
-      return;
-    }
-
-    const model = this as any;
-
-    const _schema = model._schema;
-
-    if (_schema && _schema != 'public') {
-      await this.sequelize.query(`CREATE SCHEMA IF NOT EXISTS "${_schema}";`, {
-        raw: true,
-        transaction: options?.transaction,
-      });
-    }
-
-    // fix sequelize sync with model that not have any column
-    if (Object.keys(model.tableAttributes).length === 0) {
-      if (this.database.inDialect('sqlite', 'mysql')) {
-        console.error(`Zero-column tables aren't supported in ${this.database.sequelize.getDialect()}`);
-        return;
-      }
-
-      // @ts-ignore
-      const queryInterface = this.sequelize.queryInterface;
-
-      if (!queryInterface.patched) {
-        const oldDescribeTable = queryInterface.describeTable;
-        queryInterface.describeTable = async function (...args) {
-          try {
-            return await oldDescribeTable.call(this, ...args);
-          } catch (err) {
-            if (err.message.includes('No description found for')) {
-              return [];
-            } else {
-              throw err;
-            }
-          }
-        };
-        queryInterface.patched = true;
-      }
-    }
-
-    if (this.collection.isInherited()) {
-      return SyncRunner.syncInheritModel(model, options);
-    }
-
-    return SequelizeModel.sync.call(this, options);
-  }
 }