chore | remove the trace entities graph from structured trace graph (#45)

Co-authored-by: kotharironak <[email protected]>

Author: aman bansal

data-model/src/main/java/org/hypertrace/core/datamodel/shared/StructuredTraceGraph.java

@@ -4,7 +4,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import org.hypertrace.core.datamodel.Entity;
 import org.hypertrace.core.datamodel.Event;
 import org.hypertrace.core.datamodel.StructuredTrace;
 
@@ -15,49 +14,27 @@
 public class StructuredTraceGraph {
 
   private TraceEventsGraph traceEventsGraph;
-  private TraceEntitiesGraph traceEntitiesGraph;
 
   public StructuredTraceGraph(StructuredTrace trace) {
     this.traceEventsGraph = new TraceEventsGraph(trace);
-    this.traceEntitiesGraph = new TraceEntitiesGraph(trace);
   }
 
   public void reCreateTraceEventsGraph(StructuredTrace trace) {
     this.traceEventsGraph = new TraceEventsGraph(trace);
   }
 
-  public void reCreateTraceEntitiesGraph(StructuredTrace trace) {
-    this.traceEntitiesGraph = new TraceEntitiesGraph(trace);
-  }
-
   public Set<Event> getRootEvents() {
     return traceEventsGraph.getRootEvents();
   }
 
-  public Set<Entity> getRootEntities() {
-    return traceEntitiesGraph.getRootEntities();
-  }
-
   public Event getParentEvent(Event event) {
     return traceEventsGraph.getParentEvent(event);
   }
 
-  public List<Entity> getParentEntities(Entity entity) {
-    return traceEntitiesGraph.getParentEntities(entity);
-  }
-
   public List<Event> getChildrenEvents(Event event) {
     return traceEventsGraph.getChildrenEvents(event);
   }
 
-  public List<Entity> getChildrenEntities(Entity entity) {
-    return traceEntitiesGraph.getChildrenEntities(entity);
-  }
-
-  public Map<String, Entity> getEntityMap() {
-    return traceEntitiesGraph.getEntityMap();
-  }
-
   public Map<ByteBuffer, Event> getEventMap() {
     return traceEventsGraph.getEventMap();
   }
@@ -73,8 +50,4 @@ public Map<ByteBuffer, List<ByteBuffer>> getParentToChildEventIds() {
   public TraceEventsGraph getTraceEventsGraph() {
     return traceEventsGraph;
   }
-
-  public TraceEntitiesGraph getTraceEntitiesGraph() {
-    return traceEntitiesGraph;
-  }
 }

data-model/src/test/java/org/hypertrace/core/datamodel/shared/StructuredTraceGraphTest.java

@@ -7,7 +7,6 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
@@ -86,14 +85,8 @@ private void createGraph_shouldCreateCorrectGraph() {
         Sets.newHashSet(entities.get(rootIndex1), entities.get(rootIndex2));
 
     graph = new StructuredTraceGraph(trace);
-    assertEquals(expectedRootEntities, graph.getRootEntities());
     assertEquals(expectedRootEvents, graph.getRootEvents());
     assertEquals(events.get(sourceIdx1), graph.getParentEvent(events.get(targetIdx1)));
-    assertEquals(
-        Lists.newArrayList(entities.get(sourceIdx1)),
-        graph.getParentEntities(entities.get(targetIdx1)));
-    List<Entity> root1Children = graph.getChildrenEntities(entities.get(rootIndex1));
-    assertEquals(Lists.newArrayList(entities.get(targetIdx1)), root1Children);
     assertEquals(expectedEventMap, graph.getEventMap());
     assertEquals(childToParentIds, graph.getChildIdsToParentIds());
   }
@@ -194,7 +187,6 @@ void testOnMockedEvents() {
 
     StructuredTraceGraph graph = new StructuredTraceGraph(trace);
     assertEquals(3, graph.getEventMap().size());
-    assertEquals(1, graph.getRootEntities().size());
     assertEquals(1, graph.getRootEvents().size());
     Map<ByteBuffer, ByteBuffer> childIdToParentIds = graph.getChildIdsToParentIds();
     Map<ByteBuffer, List<ByteBuffer>> parentToChildEventIds = graph.getParentToChildEventIds();
@@ -204,8 +196,6 @@ void testOnMockedEvents() {
     assertTrue(
         parentToChildEventIds.containsKey(e1.getEventId())
             && parentToChildEventIds.containsKey(e2.getEventId()));
-    assertTrue(graph.getParentEntities(entity2).contains(entity1));
-    assertTrue(graph.getParentEntities(entity3).contains(entity2));
     assertEquals(e1, graph.getParentEvent(e2));
     assertEquals(e2, graph.getParentEvent(e3));
   }
@@ -296,11 +286,9 @@ void test_recreate_partialGraph() {
 
     StructuredTraceGraph graph = new StructuredTraceGraph(trace);
     assertEquals(2, graph.getEventMap().size());
-    assertEquals(1, graph.getRootEntities().size());
     assertEquals(1, graph.getRootEvents().size());
     assertTrue(graph.getChildIdsToParentIds().containsKey(e2.getEventId()));
     assertTrue(graph.getParentToChildEventIds().containsKey(e1.getEventId()));
-    assertTrue(graph.getParentEntities(entity2).contains(entity1));
     assertEquals(e1, graph.getParentEvent(e2));
 
     // add a new event and rebuild the event graph
@@ -328,34 +316,24 @@ void test_recreate_partialGraph() {
         StructuredTraceBuilder.buildStructuredTraceFromRawSpans(
             List.of(rawSpan1, rawSpan2, rawSpan3), traceId, CUSTOMER_ID);
     TraceEventsGraph traceEventsGraph = graph.getTraceEventsGraph();
-    TraceEntitiesGraph traceEntitiesGraph = graph.getTraceEntitiesGraph();
 
     graph.reCreateTraceEventsGraph(trace);
-    assertSame(traceEntitiesGraph, graph.getTraceEntitiesGraph());
     assertNotSame(traceEventsGraph, graph.getTraceEventsGraph());
     assertEquals(3, graph.getEventMap().size());
-    assertEquals(1, graph.getRootEntities().size());
     assertEquals(1, graph.getRootEvents().size());
     assertTrue(graph.getChildIdsToParentIds().containsKey(e2.getEventId()));
     assertTrue(graph.getChildIdsToParentIds().containsKey(e3.getEventId()));
     assertTrue(graph.getParentToChildEventIds().containsKey(e1.getEventId()));
-    assertTrue(graph.getParentEntities(entity2).contains(entity1));
-    assertEquals(1, graph.getChildrenEntities(entity1).size());
     assertEquals(e1, graph.getParentEvent(e2));
     assertEquals(e1, graph.getParentEvent(e3));
 
     traceEventsGraph = graph.getTraceEventsGraph();
-    graph.reCreateTraceEntitiesGraph(trace);
-    assertNotSame(traceEntitiesGraph, graph.getTraceEntitiesGraph());
     assertSame(traceEventsGraph, graph.getTraceEventsGraph());
     assertEquals(3, graph.getEventMap().size());
-    assertEquals(1, graph.getRootEntities().size());
     assertEquals(1, graph.getRootEvents().size());
     assertTrue(graph.getChildIdsToParentIds().containsKey(e2.getEventId()));
     assertTrue(graph.getChildIdsToParentIds().containsKey(e3.getEventId()));
     assertTrue(graph.getParentToChildEventIds().containsKey(e1.getEventId()));
-    assertTrue(graph.getParentEntities(entity2).contains(entity1));
-    assertEquals(2, graph.getChildrenEntities(entity1).size());
     assertEquals(e1, graph.getParentEvent(e2));
     assertEquals(e1, graph.getParentEvent(e3));
   }