From 3988b8d990a5125f1be532212e8d4289a899f2e7 Mon Sep 17 00:00:00 2001
From: 1ilit <lilitmousheghyan@gmail.com>
Date: Fri, 31 May 2024 20:17:21 +0300
Subject: [PATCH] Recover relationships after undoing table delete

---
 src/components/EditorHeader/ControlPanel.jsx | 5 +++--
 src/context/TablesContext.jsx                | 8 +++++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/components/EditorHeader/ControlPanel.jsx b/src/components/EditorHeader/ControlPanel.jsx
index dadfbea5..d3e97c9a 100644
--- a/src/components/EditorHeader/ControlPanel.jsx
+++ b/src/components/EditorHeader/ControlPanel.jsx
@@ -146,7 +146,8 @@ export default function ControlPanel({
       }
     } else if (a.action === Action.DELETE) {
       if (a.element === ObjectType.TABLE) {
-        addTable(a.data, false);
+        a.data.relationship.forEach((x) => addRelationship(x, false));
+        addTable(a.data.table, false);
       } else if (a.element === ObjectType.RELATIONSHIP) {
         addRelationship(a.data, false);
       } else if (a.element === ObjectType.NOTE) {
@@ -337,7 +338,7 @@ export default function ControlPanel({
       }
     } else if (a.action === Action.DELETE) {
       if (a.element === ObjectType.TABLE) {
-        deleteTable(a.data.id, false);
+        deleteTable(a.data.table.id, false);
       } else if (a.element === ObjectType.RELATIONSHIP) {
         deleteRelationship(a.data.id, false);
       } else if (a.element === ObjectType.NOTE) {
diff --git a/src/context/TablesContext.jsx b/src/context/TablesContext.jsx
index 218e5089..564745e4 100644
--- a/src/context/TablesContext.jsx
+++ b/src/context/TablesContext.jsx
@@ -68,12 +68,18 @@ export default function TablesContextProvider({ children }) {
   const deleteTable = (id, addToHistory = true) => {
     if (addToHistory) {
       Toast.success(t("table_deleted"));
+      const rels = relationships.reduce((acc, r) => {
+        if (r.startTableId === id || r.endTableId === id) {
+          acc.push(r);
+        }
+        return acc;
+      }, []);
       setUndoStack((prev) => [
         ...prev,
         {
           action: Action.DELETE,
           element: ObjectType.TABLE,
-          data: tables[id],
+          data: { table: tables[id], relationship: rels },
           message: t("delete_table", { tableName: tables[id] }),
         },
       ]);