Enhanced query logic (#848)

* Fetch vertex details for fragments

* Fix tests

* Update changelog

Author: kmcginnes

Changelog.md

@@ -8,7 +8,8 @@
   ([#826](https://github.com/aws/graph-explorer/pull/826),
   [#840](https://github.com/aws/graph-explorer/pull/840))
 - **Added** query editor for Gremlin connections
-  ([#853](https://github.com/aws/graph-explorer/pull/853),
+  ([#848](https://github.com/aws/graph-explorer/pull/848),
+  [#853](https://github.com/aws/graph-explorer/pull/853),
   [#850](https://github.com/aws/graph-explorer/pull/850),
   [#843](https://github.com/aws/graph-explorer/pull/843),
   [#842](https://github.com/aws/graph-explorer/pull/842),

packages/graph-explorer/src/connector/gremlin/mappers/mapResults.test.ts

@@ -91,12 +91,21 @@ describe("mapResults", () => {
 
   it("should remove duplicate vertices", () => {
     const edge = createRandomEdge(createRandomVertex(), createRandomVertex());
+    const sourceFragment = createVertex({
+      id: edge.source,
+      types: edge.sourceTypes,
+    });
+    const targetFragment = createVertex({
+      id: edge.target,
+      types: edge.targetTypes,
+    });
     const gEdge = createGEdge(edge);
     const gList = createGList([gEdge, gEdge]);
     const results = mapResults(gList);
     expect(results).toEqual(
       toMappedQueryResults({
         edges: [edge],
+        vertices: [sourceFragment, targetFragment],
       })
     );
   });
@@ -149,6 +158,10 @@ describe("mapResults", () => {
             },
           }),
         ],
+        vertices: [
+          createVertex({ id: "2", types: ["Person"] }),
+          createVertex({ id: "1", types: ["Person"] }),
+        ],
       })
     );
   });

packages/graph-explorer/src/connector/gremlin/mappers/mapResults.ts

@@ -1,4 +1,4 @@
-import { Vertex, Edge, toNodeMap, toEdgeMap } from "@/core";
+import { Vertex, Edge, toNodeMap, toEdgeMap, createVertex } from "@/core";
 import { GAnyValue } from "../types";
 import mapApiEdge from "./mapApiEdge";
 import mapApiVertex from "./mapApiVertex";
@@ -11,8 +11,26 @@ export function mapResults(data: GAnyValue) {
   const vertexMap = toNodeMap(
     values.filter(e => "vertex" in e).map(e => e.vertex)
   );
+
   const edgeMap = toEdgeMap(values.filter(e => "edge" in e).map(e => e.edge));
 
+  // Add fragment vertices from the edges if they are missing
+  for (const edge of edgeMap.values()) {
+    if (!vertexMap.has(edge.source)) {
+      vertexMap.set(
+        edge.source,
+        createVertex({ id: edge.source, types: edge.sourceTypes })
+      );
+    }
+
+    if (!vertexMap.has(edge.target)) {
+      vertexMap.set(
+        edge.target,
+        createVertex({ id: edge.target, types: edge.targetTypes })
+      );
+    }
+  }
+
   const vertices = vertexMap.values().toArray();
   const edges = edgeMap.values().toArray();
 

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

@@ -11,3 +11,4 @@ export { default as usePrevious } from "./usePrevious";
 export { default as useTranslations } from "./useTranslations";
 export * from "./useEdgeDetailsQuery";
 export * from "./useVertexDetailsQuery";
+export * from "./useDisplayVertexFromFragment";

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

@@ -21,10 +21,24 @@ import {
   VertexTypeConfig,
 } from "@/core";
 import { waitFor } from "@testing-library/react";
+import { useMaterializeVertices } from "./useMaterializeVertices";
+import { cloneDeep } from "lodash";
+
+vi.mock("./useMaterializeVertices", () => ({
+  useMaterializeVertices: vi.fn(),
+}));
+
+beforeEach(() => {
+  vi.resetAllMocks();
+});
 
 test("should add one node", async () => {
   const vertex = createRandomVertex();
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([vertex]))
+  );
+
   const { result } = renderHookWithRecoilRoot(() => {
     const callback = useAddToGraph();
     const vertices = useRecoilValue(nodesAtom);
@@ -43,11 +57,43 @@ test("should add one node", async () => {
   });
 });
 
+test("should materialize fragment vertices", async () => {
+  const vertex = createRandomVertex();
+  const clonedVertex = cloneDeep(vertex);
+  vertex.__isFragment = true;
+  vertex.attributes = {};
+
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([clonedVertex]))
+  );
+
+  const { result } = renderHookWithRecoilRoot(() => {
+    const callback = useAddToGraph();
+    const vertices = useRecoilValue(nodesAtom);
+    const edges = useRecoilValue(edgesAtom);
+
+    return { callback, vertices, edges };
+  });
+
+  act(() => {
+    result.current.callback({ vertices: [vertex] });
+  });
+
+  await waitFor(() => {
+    const actual = result.current.vertices.get(vertex.id);
+    expect(actual).toEqual(clonedVertex);
+  });
+});
+
 test("should add one edge", async () => {
   const node1 = createRandomVertex();
   const node2 = createRandomVertex();
   const edge = createRandomEdge(node1, node2);
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([node1, node2]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -74,6 +120,10 @@ test("should add one edge", async () => {
 test("should add multiple nodes and edges", async () => {
   const randomEntities = createRandomEntities();
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(randomEntities.nodes)
+  );
+
   const { result } = renderHookWithRecoilRoot(() => {
     const callback = useAddToGraph();
     const vertices = useRecoilValue(nodesAtom);
@@ -105,6 +155,10 @@ test("should update schema when adding a node", async () => {
   const expectedVertexType = extractConfigFromEntity(vertex);
   const dbState = new DbState();
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([vertex]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -136,6 +190,10 @@ test("should update schema when adding a node with no label", async () => {
   const expectedVertexType = extractConfigFromEntity(vertex);
   const dbState = new DbState();
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([vertex]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -169,6 +227,10 @@ test("should update schema when adding an edge", async () => {
   dbState.addVertexToGraph(node1);
   dbState.addVertexToGraph(node2);
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([node1, node2]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -205,6 +267,10 @@ test("should add missing attributes to the schema when adding a node", async ()
   };
   dbState.activeSchema.vertices.push(initialVtConfig);
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([vertex]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -245,6 +311,10 @@ test("should add missing attributes to the schema when adding an edge", async ()
   };
   dbState.activeSchema.edges.push(initialEtConfig);
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([node1, node2]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -270,8 +340,13 @@ test("should add missing attributes to the schema when adding an edge", async ()
 });
 
 test("should update graph storage when adding a node", async () => {
+  const vertex = createRandomVertex();
   const dbState = new DbState();
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([vertex]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -281,7 +356,6 @@ test("should update graph storage when adding a node", async () => {
     snapshot => dbState.applyTo(snapshot)
   );
 
-  const vertex = createRandomVertex();
   act(() => {
     result.current.callback({ vertices: [vertex] });
   });
@@ -305,6 +379,10 @@ test("should update graph storage when adding an edge", async () => {
   dbState.addVertexToGraph(node1);
   dbState.addVertexToGraph(node2);
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([node1, node2]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -329,8 +407,17 @@ test("should update graph storage when adding an edge", async () => {
 });
 
 test("should ignore blank nodes when updating graph storage", async () => {
+  const vertex = createRandomVertexForRdf();
+  const blankNode = createRandomVertexForRdf();
+  blankNode.__isBlank = true;
+  const edge = createRandomEdge(vertex, blankNode);
+
   const dbState = new DbState();
 
+  vi.mocked(useMaterializeVertices).mockReturnValue(() =>
+    Promise.resolve(toNodeMap([vertex, blankNode]))
+  );
+
   const { result } = renderHookWithRecoilRoot(
     () => {
       const callback = useAddToGraph();
@@ -340,11 +427,6 @@ test("should ignore blank nodes when updating graph storage", async () => {
     snapshot => dbState.applyTo(snapshot)
   );
 
-  const vertex = createRandomVertexForRdf();
-  const blankNode = createRandomVertexForRdf();
-  blankNode.__isBlank = true;
-  const edge = createRandomEdge(vertex, blankNode);
-
   act(() => {
     result.current.callback({ vertices: [vertex, blankNode], edges: [edge] });
   });