From c5fcc62bed04f18c2dfb88a94a37999fc7ea787b Mon Sep 17 00:00:00 2001
From: ChengLei Shao <chareice@live.com>
Date: Wed, 16 Nov 2022 20:10:26 +0800
Subject: [PATCH] fix: remove node after collection removed (#1095)

---
 .../inhertits/collection-inherits.test.ts     | 22 +++++++++++++++++++
 packages/core/database/src/database.ts        |  4 ++++
 packages/core/database/src/inherited-map.ts   | 16 ++++++++++++++
 3 files changed, 42 insertions(+)

diff --git a/packages/core/database/src/__tests__/inhertits/collection-inherits.test.ts b/packages/core/database/src/__tests__/inhertits/collection-inherits.test.ts
index 9c3e50fd1..8c0cef029 100644
--- a/packages/core/database/src/__tests__/inhertits/collection-inherits.test.ts
+++ b/packages/core/database/src/__tests__/inhertits/collection-inherits.test.ts
@@ -15,6 +15,28 @@ pgOnly()('collection inherits', () => {
     await db.close();
   });
 
+  it('should remove Node after collection destroy', async () => {
+    const table1 = db.collection({
+      name: 'table1',
+      fields: [{ type: 'string', name: 'name' }],
+    });
+
+    db.collection({
+      name: 'table2',
+      fields: [{ type: 'string', name: 'integer' }],
+    });
+
+    const collection3 = db.collection({
+      name: 'table3',
+      inherits: ['table1', 'table2'],
+    });
+
+    await db.removeCollection(collection3.name);
+
+    expect(db.inheritanceMap.getNode('table3')).toBeUndefined();
+    expect(table1.isParent()).toBeFalsy();
+  });
+
   it('can update relation with child table', async () => {
     const A = db.collection({
       name: 'a',
diff --git a/packages/core/database/src/database.ts b/packages/core/database/src/database.ts
index 72e5c22d9..6cbd47195 100644
--- a/packages/core/database/src/database.ts
+++ b/packages/core/database/src/database.ts
@@ -262,6 +262,10 @@ export class Database extends EventEmitter implements AsyncEmitter {
         transaction: options.transaction,
       });
     });
+
+    this.on('afterRemoveCollection', (collection) => {
+      this.inheritanceMap.removeNode(collection.name);
+    });
   }
 
   addMigration(item: MigrationItem) {
diff --git a/packages/core/database/src/inherited-map.ts b/packages/core/database/src/inherited-map.ts
index 2ece2af4d..1963fd16c 100644
--- a/packages/core/database/src/inherited-map.ts
+++ b/packages/core/database/src/inherited-map.ts
@@ -1,4 +1,5 @@
 import lodash from 'lodash';
+import Database from './database';
 
 class TableNode {
   name: string;
@@ -14,6 +15,21 @@ class TableNode {
 export default class InheritanceMap {
   nodes: Map<string, TableNode> = new Map<string, TableNode>();
 
+  removeNode(name: string) {
+    const node = this.nodes.get(name);
+    if (!node) return;
+
+    for (const parent of node.parents) {
+      parent.children.delete(node);
+    }
+
+    for (const child of node.children) {
+      child.parents.delete(node);
+    }
+
+    this.nodes.delete(name);
+  }
+
   getOrCreateNode(name: string) {
     if (!this.nodes.has(name)) {
       this.nodes.set(name, new TableNode(name));