Merge pull request #293 from liam-hq/resolving-the-issue-of-remaining…

…-highlights Resolving the issue of remaining highlights
liam-hq · Dec 18, 2024 · f0f3780 · f0f3780
2 parents 0557af3 + fd422ef
commit f0f3780
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 29 deletions.
diff --git a/frontend/.changeset/forty-berries-hammer.md b/frontend/.changeset/forty-berries-hammer.md
@@ -0,0 +1,6 @@
+---
+"@liam-hq/erd-core": patch
+"@liam-hq/cli": patch
+---
+
+Resolving the issue of remaining highlights
diff --git a/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/ERDContent.tsx b/frontend/packages/erd-core/src/components/ERDRenderer/ERDContent/ERDContent.tsx
@@ -68,6 +68,22 @@ export const isRelatedToTable = (
   )
 }
 
+const getHighlightedHandles = (
+  edges: Edge[],
+  targetId: string,
+  nodeId: string,
+) => {
+  const highlightedTargetHandles = edges
+    .filter((edge) => edge.source === targetId && edge.target === nodeId)
+    .map((edge) => edge.targetHandle)
+
+  const highlightedSourceHandles = edges
+    .filter((edge) => edge.target === targetId && edge.source === nodeId)
+    .map((edge) => edge.sourceHandle)
+
+  return highlightedTargetHandles.concat(highlightedSourceHandles) || []
+}
+
 export const ERDContentInner: FC<Props> = ({
   nodes: _nodes,
   edges: _edges,
@@ -134,44 +150,79 @@ export const ERDContentInner: FC<Props> = ({
         (e) => e.source === id || e.target === id,
       )
 
-      const updatedEdges = edges.map((e) =>
-        relatedEdges.includes(e) ? highlightEdge(e) : e,
+      const relatedToActiveNodeEdges = edges.filter(
+        (e) => e.source === activeNodeId || e.target === activeNodeId,
       )
 
+      const updatedEdges = edges.map((e) => {
+        if (relatedEdges.includes(e)) {
+          return highlightEdge(e)
+        }
+        if (relatedToActiveNodeEdges.includes(e)) {
+          return highlightEdge(e)
+        }
+        return unhighlightEdge(e)
+      })
+
       const updatedNodes = nodes.map((node) => {
         if (node.id === id) {
           return { ...node, data: { ...node.data, isHighlighted: true } }
         }
 
         const isRelated = isRelatedToTable(relationships, node.id, id)
+        const isRelatedToActiveNode =
+          activeNodeId && isRelatedToTable(relationships, node.id, activeNodeId)
 
         if (isRelated) {
-          const highlightedTargetHandles = relatedEdges
-            .filter((edge) => edge.source === id && edge.target === node.id)
-            .map((edge) => edge.targetHandle)
-
-          const highlightedSourceHandles = relatedEdges
-            .filter((edge) => edge.target === id && edge.source === node.id)
-            .map((edge) => edge.sourceHandle)
+          let highlightedHandles = getHighlightedHandles(edges, id, node.id)
+          if (isRelatedToActiveNode) {
+            highlightedHandles = highlightedHandles.concat(
+              getHighlightedHandles(edges, activeNodeId, node.id),
+            )
+          }
 
           return {
             ...node,
             data: {
               ...node.data,
               isHighlighted: true,
-              highlightedHandles:
-                highlightedTargetHandles.concat(highlightedSourceHandles) || [],
+              highlightedHandles: highlightedHandles,
+            },
+          }
+        }
+        if (isRelatedToActiveNode) {
+          const highlightedHandles = getHighlightedHandles(
+            edges,
+            activeNodeId,
+            node.id,
+          ).concat(getHighlightedHandles(edges, id, node.id))
+
+          const isHighlighted = node.id === activeNodeId
+
+          return {
+            ...node,
+            data: {
+              ...node.data,
+              isHighlighted: isHighlighted,
+              highlightedHandles: highlightedHandles,
             },
           }
         }
 
-        return node
+        return {
+          ...node,
+          data: {
+            ...node.data,
+            isHighlighted: false,
+            highlightedHandles: [],
+          },
+        }
       })
 
       setEdges(updatedEdges)
       setNodes(updatedNodes)
     },
-    [edges, nodes, setNodes, setEdges, relationships],
+    [edges, nodes, setNodes, setEdges, relationships, activeNodeId],
   )
 
   const handleMouseLeaveNode: NodeMouseHandler<Node> = useCallback(
@@ -194,19 +245,11 @@ export const ERDContentInner: FC<Props> = ({
           )
 
           if (node.id === activeNodeId || isRelated) {
-            const highlightedTargetHandles = relatedEdges
-              .filter(
-                (edge) =>
-                  edge.source === activeNodeId && edge.target === node.id,
-              )
-              .map((edge) => edge.targetHandle)
-
-            const highlightedSourceHandles = relatedEdges
-              .filter(
-                (edge) =>
-                  edge.target === activeNodeId && edge.source === node.id,
-              )
-              .map((edge) => edge.sourceHandle)
+            const highlightedHandles = getHighlightedHandles(
+              edges,
+              activeNodeId,
+              node.id,
+            )
 
             const isHighlighted = node.id === activeNodeId
 
@@ -215,9 +258,7 @@ export const ERDContentInner: FC<Props> = ({
               data: {
                 ...node.data,
                 isHighlighted,
-                highlightedHandles:
-                  highlightedTargetHandles.concat(highlightedSourceHandles) ||
-                  [],
+                highlightedHandles: highlightedHandles,
               },
             }
           }