Add optional total_count for config service api's (#288)

pavan-traceable · web-flow · commit 1ca592c2c0e9 · 2025-04-21T20:17:34.000+05:30
* Add generic api's for filtering, pagination in getAllConfigs call
diff --git a/config-object-store/src/main/java/org/hypertrace/config/objectstore/ConfigsResponse.java b/config-object-store/src/main/java/org/hypertrace/config/objectstore/ConfigsResponse.java
@@ -0,0 +1,9 @@
+package org.hypertrace.config.objectstore;
+
+import java.util.List;
+
+public interface ConfigsResponse<T> {
+  List<T> getContextualConfigObjects();
+
+  long totalCount();
+}
diff --git a/config-object-store/src/main/java/org/hypertrace/config/objectstore/ConfigsResponseImpl.java b/config-object-store/src/main/java/org/hypertrace/config/objectstore/ConfigsResponseImpl.java
@@ -0,0 +1,23 @@
+package org.hypertrace.config.objectstore;
+
+import java.util.List;
+
+class ConfigsResponseImpl<T> implements ConfigsResponse<T> {
+  private final List<T> contextualConfigObjects;
+  private final long totalCount;
+
+  public ConfigsResponseImpl(List<T> contextualConfigObjects, long totalCount) {
+    this.contextualConfigObjects = contextualConfigObjects;
+    this.totalCount = totalCount;
+  }
+
+  @Override
+  public List<T> getContextualConfigObjects() {
+    return contextualConfigObjects;
+  }
+
+  @Override
+  public long totalCount() {
+    return totalCount;
+  }
+}
diff --git a/config-object-store/src/main/java/org/hypertrace/config/objectstore/IdentifiedFilterPushedDownObjectStore.java b/config-object-store/src/main/java/org/hypertrace/config/objectstore/IdentifiedFilterPushedDownObjectStore.java
@@ -9,6 +9,7 @@
 import org.hypertrace.config.service.v1.ConfigServiceGrpc;
 import org.hypertrace.config.service.v1.Filter;
 import org.hypertrace.config.service.v1.GetAllConfigsRequest;
+import org.hypertrace.config.service.v1.GetAllConfigsResponse;
 import org.hypertrace.config.service.v1.Pagination;
 import org.hypertrace.config.service.v1.SortBy;
 import org.hypertrace.core.grpcutils.context.RequestContext;
@@ -39,9 +40,9 @@ protected IdentifiedFilterPushedDownObjectStore(
 
   public List<ContextualConfigObject<T>> getMatchingObjects(
       RequestContext context, F filterInput, List<S> sortInput, @Nullable Pagination pagination) {
-    Filter filter = buildFilter(filterInput);
-    List<SortBy> sortByList = sortInput.stream().map(this::buildSort).collect(Collectors.toList());
-    return getMatchingObjects(context, filter, sortByList, pagination);
+    ConfigsResponse<ContextualConfigObject<T>> configsResponse =
+        getMatchingObjects(context, filterInput, sortInput, pagination, false);
+    return configsResponse.getContextualConfigObjects();
   }
 
   public List<ContextualConfigObject<T>> getMatchingObjects(
@@ -71,28 +72,54 @@ public Optional<T> getMatchingData(RequestContext context, F filterInput, List<S
     return getMatchingObject(context, filterInput, sortInput).map(ConfigObject::getData);
   }
 
-  List<ContextualConfigObject<T>> getMatchingObjects(
-      RequestContext context, Filter filter, List<SortBy> sortByList, Pagination pagination) {
-    return context
-        .call(
+  public ConfigsResponse<ContextualConfigObject<T>> getMatchingObjectsWithTotalCount(
+      RequestContext context, F filterInput, List<S> sortInput, @Nullable Pagination pagination) {
+    return getMatchingObjects(context, filterInput, sortInput, pagination, true);
+  }
+
+  public ConfigsResponse<T> getMatchingDataWithTotalCount(
+      RequestContext context, F filterInput, List<S> sortInput, @Nullable Pagination pagination) {
+    ConfigsResponse<ContextualConfigObject<T>> configsResponse =
+        getMatchingObjects(context, filterInput, sortInput, pagination, true);
+    return new ConfigsResponseImpl<>(
+        configsResponse.getContextualConfigObjects().stream()
+            .map(ConfigObject::getData)
+            .collect(Collectors.toUnmodifiableList()),
+        configsResponse.totalCount());
+  }
+
+  ConfigsResponse<ContextualConfigObject<T>> getMatchingObjects(
+      RequestContext context,
+      F filterInput,
+      List<S> sortInput,
+      Pagination pagination,
+      boolean totalIncluded) {
+    Filter filter = buildFilter(filterInput);
+    List<SortBy> sortByList = sortInput.stream().map(this::buildSort).collect(Collectors.toList());
+    GetAllConfigsResponse getConfigsResponse =
+        context.call(
             () ->
                 this.configServiceBlockingStub
                     .withDeadline(getDeadline())
-                    .getAllConfigs(buildGetAllConfigsRequest(filter, sortByList, pagination)))
-        .getContextSpecificConfigsList()
-        .stream()
-        .map(
-            contextSpecificConfig ->
-                ContextualConfigObjectImpl.tryBuild(
-                    contextSpecificConfig, this::buildDataFromValue))
-        .flatMap(Optional::stream)
-        .collect(
-            Collectors.collectingAndThen(
-                Collectors.toUnmodifiableList(), this::orderFetchedObjects));
+                    .getAllConfigs(
+                        buildGetAllConfigsRequest(filter, sortByList, pagination, totalIncluded)));
+
+    List<ContextualConfigObject<T>> contextualConfigObjectsList =
+        getConfigsResponse.getContextSpecificConfigsList().stream()
+            .map(
+                contextSpecificConfig ->
+                    ContextualConfigObjectImpl.tryBuild(
+                        contextSpecificConfig, this::buildDataFromValue))
+            .flatMap(Optional::stream)
+            .collect(
+                Collectors.collectingAndThen(
+                    Collectors.toUnmodifiableList(), this::orderFetchedObjects));
+    return new ConfigsResponseImpl<>(
+        contextualConfigObjectsList, getConfigsResponse.getTotalCount());
   }
 
   private GetAllConfigsRequest buildGetAllConfigsRequest(
-      Filter filter, List<SortBy> sortByList, Pagination pagination) {
+      Filter filter, List<SortBy> sortByList, Pagination pagination, boolean totalIncluded) {
     GetAllConfigsRequest.Builder getAllConfigsRequest =
         GetAllConfigsRequest.newBuilder()
             .setResourceName(this.resourceName)
@@ -102,6 +129,9 @@ private GetAllConfigsRequest buildGetAllConfigsRequest(
     if (pagination != null) {
       getAllConfigsRequest.setPagination(pagination);
     }
+    if (totalIncluded) {
+      getAllConfigsRequest.setIncludeTotal(true);
+    }
     return getAllConfigsRequest.build();
   }
 
diff --git a/config-service-api/src/main/proto/org/hypertrace/config/service/v1/config_service.proto b/config-service-api/src/main/proto/org/hypertrace/config/service/v1/config_service.proto
@@ -100,11 +100,17 @@ message GetAllConfigsRequest {
   // optional - pagination parameters to limit and offset the result set.
   // Useful for retrieving configs in pages when total count is large.
   Pagination pagination = 5;
+
+  // optional - include total count in the response
+  bool include_total = 6;
 }
 
 message GetAllConfigsResponse {
   // list of config values along with the associated contexts, sorted in the descending order of their creation time
   repeated ContextSpecificConfig context_specific_configs = 1;
+
+  // Total number of records matching the filter before pagination if include_total is true.
+  optional int64 total_count = 2;
 }
 
 message ContextSpecificConfig {
diff --git a/config-service-impl/src/main/java/org/hypertrace/config/service/ConfigServiceGrpcImpl.java b/config-service-impl/src/main/java/org/hypertrace/config/service/ConfigServiceGrpcImpl.java
@@ -5,6 +5,7 @@
 import static org.hypertrace.config.service.ConfigServiceUtils.merge;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.google.protobuf.Value;
 import io.grpc.Status;
 import io.grpc.stub.StreamObserver;
@@ -13,6 +14,9 @@
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 import org.hypertrace.config.service.store.ConfigStore;
@@ -42,6 +46,10 @@ public class ConfigServiceGrpcImpl extends ConfigServiceGrpc.ConfigServiceImplBa
   private static String DEFAULT_USER_ID = "Unknown";
   private static String DEFAULT_USER_EMAIL = "Unknown";
   private final ConfigStore configStore;
+  private static final Executor configExecutor =
+      Executors.newFixedThreadPool(
+          4,
+          new ThreadFactoryBuilder().setNameFormat("config-executor-%d").setDaemon(true).build());
 
   public ConfigServiceGrpcImpl(ConfigStore configStore) {
     this.configStore = configStore;
@@ -115,21 +123,51 @@ public void getConfig(
   public void getAllConfigs(
       GetAllConfigsRequest request, StreamObserver<GetAllConfigsResponse> responseObserver) {
     try {
-      List<ContextSpecificConfig> contextSpecificConfigList =
+      ConfigResource configResource =
+          new ConfigResource(
+              request.getResourceName(), request.getResourceNamespace(), getTenantId());
+
+      CompletableFuture<Long> totalCountFuture = null;
+
+      // Start count query first (in parallel)
+      if (request.getIncludeTotal()) {
+        totalCountFuture =
+            CompletableFuture.supplyAsync(
+                () -> {
+                  try {
+                    return configStore.getMatchingConfigsCount(configResource, request.getFilter());
+                  } catch (Exception e) {
+                    throw Status.INTERNAL
+                        .withCause(e)
+                        .withDescription("Failed to fetch total count")
+                        .asRuntimeException();
+                  }
+                },
+                configExecutor);
+      }
+
+      // Then run getAllConfigs on the current thread
+      List<ContextSpecificConfig> configList =
           configStore.getAllConfigs(
-              new ConfigResource(
-                  request.getResourceName(), request.getResourceNamespace(), getTenantId()),
+              configResource,
               request.getFilter(),
               request.getPagination(),
               request.getSortByList());
-      responseObserver.onNext(
-          GetAllConfigsResponse.newBuilder()
-              .addAllContextSpecificConfigs(contextSpecificConfigList)
-              .build());
+
+      // Build the response
+      GetAllConfigsResponse.Builder responseBuilder =
+          GetAllConfigsResponse.newBuilder().addAllContextSpecificConfigs(configList);
+
+      if (totalCountFuture != null) {
+        long totalCount = totalCountFuture.join(); // Wait if not finished
+        responseBuilder.setTotalCount(totalCount);
+      }
+
+      responseObserver.onNext(responseBuilder.build());
       responseObserver.onCompleted();
     } catch (Exception e) {
-      log.error("Get all configs failed for request:{}", request, e);
-      responseObserver.onError(e);
+      log.error("Get all configs failed for request: {}", request, e);
+      responseObserver.onError(Status.fromThrowable(e).asRuntimeException());
     }
   }
 
diff --git a/config-service-impl/src/main/java/org/hypertrace/config/service/store/ConfigStore.java b/config-service-impl/src/main/java/org/hypertrace/config/service/store/ConfigStore.java
@@ -70,6 +70,8 @@ List<ContextSpecificConfig> getAllConfigs(
       ConfigResource configResource, Filter filter, Pagination pagination, List<SortBy> sortByList)
       throws IOException;
 
+  long getMatchingConfigsCount(ConfigResource configResource, Filter filter);
+
   /**
    * Write each of the provided config value associated with the specified config resource to the
    * store.
diff --git a/config-service-impl/src/main/java/org/hypertrace/config/service/store/DocumentConfigStore.java b/config-service-impl/src/main/java/org/hypertrace/config/service/store/DocumentConfigStore.java
@@ -2,6 +2,7 @@
 
 import static com.google.common.collect.Streams.zip;
 import static org.hypertrace.config.service.store.ConfigDocument.CONTEXT_FIELD_NAME;
+import static org.hypertrace.config.service.store.ConfigDocument.CREATION_TIMESTAMP_FIELD_NAME;
 import static org.hypertrace.config.service.store.ConfigDocument.RESOURCE_FIELD_NAME;
 import static org.hypertrace.config.service.store.ConfigDocument.RESOURCE_NAMESPACE_FIELD_NAME;
 import static org.hypertrace.config.service.store.ConfigDocument.TENANT_ID_FIELD_NAME;
@@ -15,7 +16,6 @@
 import java.time.Clock;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -230,12 +230,21 @@ public List<ContextSpecificConfig> getAllConfigs(
         processDocument(documentIterator.next(), seenContexts, configList);
       }
     }
-
-    configList.sort(
-        Comparator.comparingLong(ContextSpecificConfig::getCreationTimestamp).reversed());
     return configList;
   }
 
+  @Override
+  public long getMatchingConfigsCount(
+      ConfigResource configResource, org.hypertrace.config.service.v1.Filter filter) {
+    Query query =
+        buildQuery(
+            configResource,
+            filter,
+            org.hypertrace.config.service.v1.Pagination.getDefaultInstance(),
+            Collections.emptyList());
+    return collection.count(query);
+  }
+
   private Query buildQuery(
       ConfigResource configResource,
       @NonNull org.hypertrace.config.service.v1.Filter filter,
@@ -248,15 +257,15 @@ private Query buildQuery(
       queryBuilder.setPagination(
           Pagination.builder().offset(pagination.getOffset()).limit(pagination.getLimit()).build());
     }
-
+    queryBuilder.addSort(IdentifierExpression.of(VERSION_FIELD_NAME), SortOrder.DESC);
     if (!sortByList.isEmpty()) {
       sortByList.forEach(
           sortBy ->
               queryBuilder.addSort(
                   IdentifierExpression.of(sortBy.getSelection().getConfigJsonPath()),
                   convertSortOrder(sortBy)));
     } else {
-      queryBuilder.addSort(IdentifierExpression.of(VERSION_FIELD_NAME), SortOrder.DESC);
+      queryBuilder.addSort(IdentifierExpression.of(CREATION_TIMESTAMP_FIELD_NAME), SortOrder.DESC);
     }
     return queryBuilder.build();
   }
diff --git a/config-service-impl/src/test/java/org/hypertrace/config/service/ConfigServiceGrpcImplTest.java b/config-service-impl/src/test/java/org/hypertrace/config/service/ConfigServiceGrpcImplTest.java
@@ -206,7 +206,7 @@ void getAllConfigs() throws IOException {
   }
 
   @Test
-  void getAllConfigs_withFilterPaginationAndSorting() throws IOException {
+  void getAllConfigs_withFilterPaginationSortingAndTotalCount() throws IOException {
     ConfigStore configStore = mock(ConfigStore.class);
 
     // Define filter
@@ -236,6 +236,8 @@ void getAllConfigs_withFilterPaginationAndSorting() throws IOException {
             ContextSpecificConfig.newBuilder().setConfig(config1).build(),
             ContextSpecificConfig.newBuilder().setContext(CONTEXT1).setConfig(config2).build());
 
+    long mockedTotalCount = 42L;
+
     // Mock behavior
     when(configStore.getAllConfigs(
             argThat(
@@ -248,18 +250,22 @@ void getAllConfigs_withFilterPaginationAndSorting() throws IOException {
             eq(pagination),
             eq(List.of(sortBy))))
         .thenReturn(contextSpecificConfigList);
+    when(configStore.getMatchingConfigsCount(
+            eq(new ConfigResource(RESOURCE_NAME, RESOURCE_NAMESPACE, TENANT_ID)), eq(filter)))
+        .thenReturn(mockedTotalCount);
 
     ConfigServiceGrpcImpl configServiceGrpc = new ConfigServiceGrpcImpl(configStore);
     StreamObserver<GetAllConfigsResponse> responseObserver = mock(StreamObserver.class);
 
-    // Build request with all fields
+    // Build request with all fields and includeTotal = true
     GetAllConfigsRequest request =
         GetAllConfigsRequest.newBuilder()
             .setResourceName(RESOURCE_NAME)
             .setResourceNamespace(RESOURCE_NAMESPACE)
             .setFilter(filter)
             .setPagination(pagination)
             .addSortBy(sortBy)
+            .setIncludeTotal(true)
             .build();
 
     // Execute in tenant context
@@ -276,6 +282,7 @@ void getAllConfigs_withFilterPaginationAndSorting() throws IOException {
 
     GetAllConfigsResponse actualResponse = responseCaptor.getValue();
     assertEquals(contextSpecificConfigList, actualResponse.getContextSpecificConfigsList());
+    assertEquals(mockedTotalCount, actualResponse.getTotalCount());
   }
 
   @Test
diff --git a/config-service-impl/src/test/java/org/hypertrace/config/service/store/DocumentConfigStoreTest.java b/config-service-impl/src/test/java/org/hypertrace/config/service/store/DocumentConfigStoreTest.java
@@ -274,6 +274,7 @@ void getAllConfigs() throws IOException {
             Filter.getDefaultInstance(), // for empty filter
             Pagination.getDefaultInstance(),
             Collections.emptyList());
+
     assertEquals(2, contextSpecificConfigList.size());
     assertEquals(
         ContextSpecificConfig.newBuilder()
