Fixes #24996 Add filter function to EntityRelationsGraph to individually filter entities from graph

Signed-off-by: Jordan Slott <jordan.slott@twosigma.com>

.changeset/forty-experts-live.md

@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-catalog-graph': patch
+---
+
+Add function to `EntityRelationsGraph` filter that excludes entities from graph

plugins/catalog-graph/api-report.md

@@ -115,6 +115,7 @@ export type EntityRelationsGraphProps = {
   mergeRelations?: boolean;
   kinds?: string[];
   relations?: string[];
+  entityFilter?: (entity: Entity) => boolean;
   direction?: Direction;
   onNodeClick?: (value: EntityNode, event: MouseEvent_2<unknown>) => void;
   relationPairs?: RelationPairs;

plugins/catalog-graph/src/components/CatalogGraphCard/CatalogGraphCard.tsx

@@ -71,6 +71,7 @@ export const CatalogGraphCard = (
     direction = Direction.LEFT_RIGHT,
     kinds,
     relations,
+    entityFilter,
     height,
     className,
     rootEntityNames,
@@ -140,6 +141,7 @@ export const CatalogGraphCard = (
         mergeRelations={mergeRelations}
         direction={direction}
         relationPairs={relationPairs}
+        entityFilter={entityFilter}
         zoom={zoom}
       />
     </InfoCard>

plugins/catalog-graph/src/components/CatalogGraphPage/CatalogGraphPage.tsx

@@ -121,7 +121,11 @@ export const CatalogGraphPage = (
     };
   } & Partial<EntityRelationsGraphProps>,
 ) => {
-  const { relationPairs = ALL_RELATION_PAIRS, initialState } = props;
+  const {
+    relationPairs = ALL_RELATION_PAIRS,
+    initialState,
+    entityFilter,
+  } = props;
 
   const navigate = useNavigate();
   const classes = useStyles();
@@ -256,6 +260,7 @@ export const CatalogGraphPage = (
                 onNodeClick={onNodeClick}
                 direction={direction}
                 relationPairs={relationPairs}
+                entityFilter={entityFilter}
                 className={classes.graph}
                 zoom="enabled"
                 curve={curve}

plugins/catalog-graph/src/components/EntityRelationsGraph/EntityRelationsGraph.tsx

@@ -16,6 +16,7 @@
 
 import {
   CompoundEntityRef,
+  Entity,
   stringifyEntityRef,
 } from '@backstage/catalog-model';
 import {
@@ -77,6 +78,7 @@ export type EntityRelationsGraphProps = {
   mergeRelations?: boolean;
   kinds?: string[];
   relations?: string[];
+  entityFilter?: (entity: Entity) => boolean;
   direction?: Direction;
   onNodeClick?: (value: EntityNode, event: MouseEvent<unknown>) => void;
   relationPairs?: RelationPairs;
@@ -101,6 +103,7 @@ export const EntityRelationsGraph = (props: EntityRelationsGraphProps) => {
     mergeRelations = true,
     kinds,
     relations,
+    entityFilter,
     direction = Direction.LEFT_RIGHT,
     onNodeClick,
     relationPairs = ALL_RELATION_PAIRS,
@@ -130,6 +133,7 @@ export const EntityRelationsGraph = (props: EntityRelationsGraphProps) => {
     mergeRelations,
     kinds,
     relations,
+    entityFilter,
     onNodeClick,
     relationPairs,
   });

plugins/catalog-graph/src/components/EntityRelationsGraph/useEntityRelationGraph.test.ts

@@ -353,6 +353,28 @@ describe('useEntityRelationGraph', () => {
     });
   });
 
+  test('should filter by func', async () => {
+    const { result, rerender } = renderHook(() =>
+      useEntityRelationGraph({
+        rootEntityRefs: ['b:d/c'],
+        filter: {
+          entityFilter: e => e.metadata.name !== 'c2',
+        },
+      }),
+    );
+
+    // Simulate rerendering as this is triggered automatically due to the mock
+    for (let i = 0; i < 5; ++i) {
+      rerender();
+    }
+
+    expect(result.current.entities).toEqual({
+      'b:d/c': expect.anything(),
+      'b:d/c1': expect.anything(),
+      'k:d/a1': expect.anything(),
+    });
+  });
+
   test('should support multiple roots by kind', async () => {
     const { result, rerender } = renderHook(() =>
       useEntityRelationGraph({

plugins/catalog-graph/src/components/EntityRelationsGraph/useEntityRelationGraph.ts

@@ -16,6 +16,7 @@
 import { Entity } from '@backstage/catalog-model';
 import { useEffect } from 'react';
 import { useEntityStore } from './useEntityStore';
+import { pickBy } from 'lodash';
 
 /**
  * Discover the graph of entities connected by relations, starting from a set of
@@ -24,13 +25,19 @@ import { useEntityStore } from './useEntityStore';
  */
 export function useEntityRelationGraph({
   rootEntityRefs,
-  filter: { maxDepth = Number.POSITIVE_INFINITY, relations, kinds } = {},
+  filter: {
+    maxDepth = Number.POSITIVE_INFINITY,
+    relations,
+    kinds,
+    entityFilter,
+  } = {},
 }: {
   rootEntityRefs: string[];
   filter?: {
     maxDepth?: number;
     relations?: string[];
     kinds?: string[];
+    entityFilter?: (entity: Entity) => boolean;
   };
 }): {
   entities?: { [ref: string]: Entity };
@@ -60,6 +67,11 @@ export function useEntityRelationGraph({
         processedEntityRefs.add(entityRef);
 
         if (entity && entity.relations) {
+          // If the entity is filtered out then no need to check any
+          // of its outgoing relationships to other entities
+          if (entityFilter && !entityFilter(entity)) {
+            continue;
+          }
           for (const rel of entity.relations) {
             if (
               (!relations || relations.includes(rel.type)) &&
@@ -81,12 +93,23 @@ export function useEntityRelationGraph({
 
       ++depth;
     }
-
     requestEntities([...expectedEntities]);
-  }, [entities, rootEntityRefs, maxDepth, relations, kinds, requestEntities]);
+  }, [
+    entities,
+    rootEntityRefs,
+    maxDepth,
+    relations,
+    kinds,
+    entityFilter,
+    requestEntities,
+  ]);
+
+  const filteredEntities = entityFilter
+    ? pickBy(entities, (value, _key) => entityFilter(value))
+    : entities;
 
   return {
-    entities,
+    entities: filteredEntities,
     loading,
     error,
   };

plugins/catalog-graph/src/components/EntityRelationsGraph/useEntityRelationNodesAndEdges.ts

@@ -18,7 +18,7 @@ import useDebounce from 'react-use/esm/useDebounce';
 import { RelationPairs, ALL_RELATION_PAIRS } from './relations';
 import { EntityEdge, EntityNode } from './types';
 import { useEntityRelationGraph } from './useEntityRelationGraph';
-import { DEFAULT_NAMESPACE } from '@backstage/catalog-model';
+import { Entity, DEFAULT_NAMESPACE } from '@backstage/catalog-model';
 
 /**
  * Generate nodes and edges to render the entity graph.
@@ -30,6 +30,7 @@ export function useEntityRelationNodesAndEdges({
   mergeRelations = true,
   kinds,
   relations,
+  entityFilter,
   onNodeClick,
   relationPairs = ALL_RELATION_PAIRS,
 }: {
@@ -39,6 +40,7 @@ export function useEntityRelationNodesAndEdges({
   mergeRelations?: boolean;
   kinds?: string[];
   relations?: string[];
+  entityFilter?: (entity: Entity) => boolean;
   onNodeClick?: (value: EntityNode, event: MouseEvent<unknown>) => void;
   relationPairs?: RelationPairs;
 }): {
@@ -57,6 +59,7 @@ export function useEntityRelationNodesAndEdges({
       maxDepth,
       kinds,
       relations,
+      entityFilter,
     },
   });