fix: fetching total from EDS (#64)

* fix: set total in data fetcher node

* fix: fetching total from eds

* fix tests

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcher.java

@@ -183,6 +183,16 @@ public EntityFetcherResponse getTimeAggregatedMetrics(
 
   @Override
   public int getTotalEntities(EntitiesRequestContext requestContext, EntitiesRequest entitiesRequest) {
-    throw new UnsupportedOperationException("Fetching total entities not supported by EDS");
+    EntityFetcherResponse entityFetcherResponse = getEntities(
+        requestContext,
+        EntitiesRequest.newBuilder(entitiesRequest)
+            .clearSelection()
+            .clearTimeAggregation()
+            .clearOrderBy()
+            .clearLimit()
+            .setOffset(0)
+            .build()
+    );
+    return entityFetcherResponse.size();
   }
 }

gateway-service-impl/src/main/java/org/hypertrace/gateway/service/entity/query/ExecutionTreeBuilder.java

@@ -182,6 +182,13 @@ private QueryNode buildExecutionTreeForEdsFilterAndSelection() {
 
     QueryNode rootNode = new DataFetcherNode(EDS.name(), filter, selectionLimit, selectionOffset, orderBys);
     executionContext.setSortAndPaginationNodeAdded(true);
+
+    if (ExecutionTreeUtils.hasEntityIdEqualsFilter(executionContext)) {
+      executionContext.setTotal(1);
+    } else {
+      rootNode = new TotalFetcherNode(rootNode, EDS.name());
+    }
+
     return rootNode;
   }
 

gateway-service-impl/src/test/java/org/hypertrace/gateway/service/common/datafetcher/EntityDataServiceEntityFetcherTests.java

@@ -176,17 +176,6 @@ public void test_getEntities_WithoutPagination() {
         entityDataServiceEntityFetcher.getEntities(entitiesRequestContext, entitiesRequest).size());
   }
 
-  @Test
-  public void test_getTotalEntities() {
-    assertThrows(
-        UnsupportedOperationException.class,
-        () -> {
-          entityDataServiceEntityFetcher.getTotalEntities(
-              new EntitiesRequestContext(TENANT_ID, 0, 1, "API", "API.startTime", Map.of()),
-              EntitiesRequest.newBuilder().build());
-        });
-  }
-
   @Test
   public void test_getAggregatedMetrics() {
     assertThrows(

gateway-service-impl/src/test/java/org/hypertrace/gateway/service/entity/query/ExecutionTreeBuilderTest.java

@@ -435,13 +435,16 @@ public void testExecutionTreeBuilderWithSelectFilterOrderPagination() {
       ExecutionTreeBuilder executionTreeBuilder = new ExecutionTreeBuilder(executionContext);
       QueryNode executionTree = executionTreeBuilder.build();
       assertNotNull(executionTree);
+      assertTrue(executionTree instanceof TotalFetcherNode);
+      assertEquals("EDS", ((TotalFetcherNode) executionTree).getSource());
 
-      assertTrue(executionTree instanceof DataFetcherNode);
-      assertEquals("EDS", ((DataFetcherNode) executionTree).getSource());
-      assertEquals(20, ((DataFetcherNode) executionTree).getOffset());
-      assertEquals(10, ((DataFetcherNode) executionTree).getLimit());
+      QueryNode childNode = ((TotalFetcherNode) executionTree).getChildNode();
+      assertTrue(childNode instanceof DataFetcherNode);
+      assertEquals("EDS", ((DataFetcherNode) childNode).getSource());
+      assertEquals(20, ((DataFetcherNode) childNode).getOffset());
+      assertEquals(10, ((DataFetcherNode) childNode).getLimit());
       assertEquals(
-          List.of(orderByExpression), ((DataFetcherNode) executionTree).getOrderByExpressionList());
+          List.of(orderByExpression), ((DataFetcherNode) childNode).getOrderByExpressionList());
     }
 
     {
@@ -497,10 +500,14 @@ public void testExecutionTreeBuilderWithSelectPagination() {
     ExecutionTreeBuilder executionTreeBuilder = new ExecutionTreeBuilder(executionContext);
     QueryNode executionTree = executionTreeBuilder.build();
     assertNotNull(executionTree);
-    assertTrue(executionTree instanceof DataFetcherNode);
-    assertEquals(10, ((DataFetcherNode)executionTree).getLimit());
-    assertEquals(20, ((DataFetcherNode)executionTree).getOffset());
-    assertEquals(List.of(), ((DataFetcherNode)executionTree).getOrderByExpressionList());
+    assertTrue(executionTree instanceof TotalFetcherNode);
+    assertEquals("EDS", ((TotalFetcherNode) executionTree).getSource());
+
+    QueryNode childNode = ((TotalFetcherNode) executionTree).getChildNode();
+    assertTrue(childNode instanceof DataFetcherNode);
+    assertEquals(10, ((DataFetcherNode)childNode).getLimit());
+    assertEquals(20, ((DataFetcherNode)childNode).getOffset());
+    assertEquals(List.of(), ((DataFetcherNode)childNode).getOrderByExpressionList());
   }
 
   @Test