From ee659eaa038a94ee13801801e84152df4d79683d Mon Sep 17 00:00:00 2001
From: Guy Ben-Aharon <baguy3@gmail.com>
Date: Wed, 1 Jan 2025 10:48:42 +0200
Subject: [PATCH] fix: add loadDiagramFromData logic to chartdb provider (#513)

---
 .../chartdb-context/chartdb-context.tsx       |  2 +
 .../chartdb-context/chartdb-provider.tsx      | 52 +++++++++++--------
 2 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/src/context/chartdb-context/chartdb-context.tsx b/src/context/chartdb-context/chartdb-context.tsx
index b117270d..9cba9ae5 100644
--- a/src/context/chartdb-context/chartdb-context.tsx
+++ b/src/context/chartdb-context/chartdb-context.tsx
@@ -84,6 +84,7 @@ export interface ChartDBContext {
         options?: { updateHistory: boolean }
     ) => Promise<void>;
     loadDiagram: (diagramId: string) => Promise<Diagram | undefined>;
+    loadDiagramFromData: (diagram: Diagram) => void;
     updateDiagramUpdatedAt: () => Promise<void>;
     clearDiagramData: () => Promise<void>;
     deleteDiagram: () => Promise<void>;
@@ -246,6 +247,7 @@ export const chartDBContext = createContext<ChartDBContext>({
     updateDiagramName: emptyFn,
     updateDiagramUpdatedAt: emptyFn,
     loadDiagram: emptyFn,
+    loadDiagramFromData: emptyFn,
     clearDiagramData: emptyFn,
     deleteDiagram: emptyFn,
 
diff --git a/src/context/chartdb-context/chartdb-provider.tsx b/src/context/chartdb-context/chartdb-provider.tsx
index 521217ec..26691ea3 100644
--- a/src/context/chartdb-context/chartdb-provider.tsx
+++ b/src/context/chartdb-context/chartdb-provider.tsx
@@ -1336,15 +1336,9 @@ export const ChartDBProvider: React.FC<
         ]
     );
 
-    const loadDiagram: ChartDBContext['loadDiagram'] = useCallback(
-        async (diagramId: string) => {
-            const diagram = await db.getDiagram(diagramId, {
-                includeRelationships: true,
-                includeTables: true,
-                includeDependencies: true,
-            });
-
-            if (diagram) {
+    const loadDiagramFromData: ChartDBContext['loadDiagramFromData'] =
+        useCallback(
+            async (diagram) => {
                 setDiagramId(diagram.id);
                 setDiagramName(diagram.name);
                 setDatabaseType(diagram.databaseType);
@@ -1356,23 +1350,36 @@ export const ChartDBProvider: React.FC<
                 setDiagramUpdatedAt(diagram.updatedAt);
 
                 events.emit({ action: 'load_diagram', data: { diagram } });
+            },
+            [
+                setDiagramId,
+                setDiagramName,
+                setDatabaseType,
+                setDatabaseEdition,
+                setTables,
+                setRelationships,
+                setDependencies,
+                setDiagramCreatedAt,
+                setDiagramUpdatedAt,
+                events,
+            ]
+        );
+
+    const loadDiagram: ChartDBContext['loadDiagram'] = useCallback(
+        async (diagramId: string) => {
+            const diagram = await db.getDiagram(diagramId, {
+                includeRelationships: true,
+                includeTables: true,
+                includeDependencies: true,
+            });
+
+            if (diagram) {
+                loadDiagramFromData(diagram);
             }
 
             return diagram;
         },
-        [
-            db,
-            setDiagramId,
-            setDiagramName,
-            setDatabaseType,
-            setDatabaseEdition,
-            setTables,
-            setRelationships,
-            setDependencies,
-            setDiagramCreatedAt,
-            setDiagramUpdatedAt,
-            events,
-        ]
+        [db, loadDiagramFromData]
     );
 
     return (
@@ -1393,6 +1400,7 @@ export const ChartDBProvider: React.FC<
                 updateDiagramId,
                 updateDiagramName,
                 loadDiagram,
+                loadDiagramFromData,
                 updateDatabaseType,
                 updateDatabaseEdition,
                 clearDiagramData,