Refactor useAddToGraph() hook

packages/graph-explorer/src/hooks/useAddToGraph.test.ts

@@ -13,14 +13,14 @@ import { nodesAtom, toNodeMap } from "@/core/StateProvider/nodes";
 test("should add one node", async () => {
   const vertex = createRandomVertex();
   const { result } = renderHookWithRecoilRoot(() => {
-    const callback = useAddToGraph(vertex);
+    const callback = useAddToGraph();
     const [entities] = useEntities();
 
     return { callback, entities };
   });
 
   await act(async () => {
-    result.current.callback();
+    result.current.callback({ vertices: [vertex] });
   });
 
   const actual = result.current.entities.nodes.get(vertex.id);
@@ -34,7 +34,7 @@ test("should add one edge", async () => {
 
   const { result } = renderHookWithRecoilRoot(
     () => {
-      const callback = useAddToGraph(edge);
+      const callback = useAddToGraph();
       const [entities] = useEntities();
 
       return { callback, entities };
@@ -45,7 +45,7 @@ test("should add one edge", async () => {
   );
 
   await act(async () => {
-    result.current.callback();
+    result.current.callback({ edges: [edge] });
   });
 
   const actual = result.current.entities.edges.get(edge.id);
@@ -55,17 +55,17 @@ test("should add one edge", async () => {
 test("should add multiple nodes and edges", async () => {
   const randomEntities = createRandomEntities();
   const { result } = renderHookWithRecoilRoot(() => {
-    const callback = useAddToGraph(
-      ...randomEntities.nodes.values(),
-      ...randomEntities.edges.values()
-    );
+    const callback = useAddToGraph();
     const [entities] = useEntities();
 
     return { callback, entities };
   });
 
   await act(async () => {
-    result.current.callback();
+    result.current.callback({
+      vertices: [...randomEntities.nodes.values()],
+      edges: [...randomEntities.edges.values()],
+    });
   });
 
   const actualNodes = result.current.entities.nodes.values().toArray();

packages/graph-explorer/src/hooks/useAddToGraph.ts

@@ -6,20 +6,32 @@ import { useCallback } from "react";
 import { useSetRecoilState } from "recoil";
 
 /** Returns a callback that adds an array of nodes and edges to the graph. */
-export function useAddToGraph(...entitiesToAdd: (Vertex | Edge)[]) {
+export function useAddToGraph() {
   const setEntities = useSetRecoilState(entitiesSelector);
 
-  return useCallback(() => {
-    const nodes = entitiesToAdd.filter(e => e.entityType === "vertex");
-    const edges = entitiesToAdd.filter(e => e.entityType === "edge");
+  return useCallback(
+    (entities: { vertices?: Vertex[]; edges?: Edge[] }) => {
+      const vertices = entities.vertices ?? [];
+      const edges = entities.edges ?? [];
 
-    if (nodes.length === 0 && edges.length === 0) {
-      return;
-    }
+      if (vertices.length === 0 && edges.length === 0) {
+        return;
+      }
 
-    setEntities({
-      nodes: toNodeMap(nodes),
-      edges: toEdgeMap(edges),
-    });
-  }, [entitiesToAdd, setEntities]);
+      setEntities({
+        nodes: toNodeMap(vertices),
+        edges: toEdgeMap(edges),
+      });
+    },
+    [setEntities]
+  );
+}
+
+/** Returns a callback the given vertex to the graph. */
+export function useAddVertexToGraph(vertex: Vertex) {
+  const callback = useAddToGraph();
+  return useCallback(
+    () => callback({ vertices: [vertex] }),
+    [callback, vertex]
+  );
 }

packages/graph-explorer/src/hooks/useExpandNode.tsx

@@ -7,14 +7,12 @@ import {
   type NeighborsResponse,
 } from "@/connector";
 import { loggerSelector, useExplorer } from "@/core/connector";
-import useEntities from "./useEntities";
 import { useRecoilValue } from "recoil";
 import { useMutation, useQueryClient } from "@tanstack/react-query";
 import { Vertex } from "@/core";
 import { createDisplayError } from "@/utils/createDisplayError";
-import { toNodeMap } from "@/core/StateProvider/nodes";
-import { toEdgeMap } from "@/core/StateProvider/edges";
 import { useNeighborsCallback } from "@/core";
+import { useAddToGraph } from "./useAddToGraph";
 
 export type ExpandNodeFilters = Omit<
   NeighborsRequest,
@@ -33,7 +31,7 @@ export type ExpandNodeRequest = {
 export default function useExpandNode() {
   const queryClient = useQueryClient();
   const explorer = useExplorer();
-  const [_, setEntities] = useEntities();
+  const addToGraph = useAddToGraph();
   const { enqueueNotification, clearNotification } = useNotification();
   const remoteLogger = useRecoilValue(loggerSelector);
 
@@ -82,10 +80,7 @@ export default function useExpandNode() {
       updateEdgeDetailsCache(explorer, queryClient, data.edges);
 
       // Update nodes and edges in the graph
-      setEntities({
-        nodes: toNodeMap(data.vertices),
-        edges: toEdgeMap(data.edges),
-      });
+      addToGraph(data);
     },
     onError: error => {
       remoteLogger.error(`Failed to expand node: ${error.message}`);

packages/graph-explorer/src/modules/SearchSidebar/NodeSearchResult.tsx

@@ -7,7 +7,7 @@ import {
   VertexRow,
 } from "@/components";
 import {
-  useAddToGraph,
+  useAddVertexToGraph,
   useHasVertexBeenAddedToGraph,
   useRemoveNodeFromGraph,
 } from "@/hooks";
@@ -25,7 +25,7 @@ export function NodeSearchResult({ node }: { node: Vertex }) {
   const [expanded, setExpanded] = useState(false);
   const displayNode = useDisplayVertexFromVertex(node);
 
-  const addToGraph = useAddToGraph(node);
+  const addToGraph = useAddVertexToGraph(node);
   const removeFromGraph = useRemoveNodeFromGraph(node.id);
   const hasBeenAdded = useHasVertexBeenAddedToGraph(node.id);
 

packages/graph-explorer/src/modules/SearchSidebar/SearchResultsList.tsx

@@ -60,7 +60,7 @@ export function SearchResultsList({
 }
 
 function LoadedResults({ vertices, edges, scalars }: MappedQueryResults) {
-  const sendToGraph = useAddToGraph(...vertices, ...edges);
+  const sendToGraph = useAddToGraph();
   const canSendToGraph = vertices.length > 0 || edges.length > 0;
 
   const counts = [
@@ -89,7 +89,7 @@ function LoadedResults({ vertices, edges, scalars }: MappedQueryResults) {
       <PanelFooter className="sticky bottom-0 flex flex-row items-center justify-between">
         <Button
           icon={<PlusCircleIcon />}
-          onPress={sendToGraph}
+          onPress={() => sendToGraph({ vertices, edges })}
           isDisabled={!canSendToGraph}
         >
           Add all

packages/graph-explorer/src/workspaces/DataExplorer/DataExplorer.tsx

@@ -45,7 +45,7 @@ import {
   userStylingAtom,
   VertexPreferences,
 } from "@/core/StateProvider/userPreferences";
-import { useAddToGraph, useHasVertexBeenAddedToGraph } from "@/hooks";
+import { useAddVertexToGraph, useHasVertexBeenAddedToGraph } from "@/hooks";
 import usePrefixesUpdater from "@/hooks/usePrefixesUpdater";
 import useTranslations from "@/hooks/useTranslations";
 import useUpdateVertexTypeCounts from "@/hooks/useUpdateVertexTypeCounts";
@@ -264,7 +264,7 @@ function DisplayNameAndDescriptionOptions({
 }
 
 function AddToExplorerButton({ vertex }: { vertex: Vertex }) {
-  const addToGraph = useAddToGraph(vertex);
+  const addToGraph = useAddVertexToGraph(vertex);
   const isInExplorer = useHasVertexBeenAddedToGraph(vertex.id);
 
   return (