From 02da447b64c1e560dada87bc2ec4168a39f6f4a7 Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Thu, 13 Jun 2024 19:54:01 +0530
Subject: [PATCH 1/2] fix | fix the total query converter

---
 .../query/service/EntityQueryConverter.java   | 178 +----------------
 .../query/service/EntityQueryServiceImpl.java |  17 +-
 .../service/EntityQueryConverterTest.java     | 184 ------------------
 .../resources/configs/common/application.conf |   5 +
 4 files changed, 17 insertions(+), 367 deletions(-)
 delete mode 100644 entity-service-impl/src/test/java/org/hypertrace/entity/query/service/EntityQueryConverterTest.java

diff --git a/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryConverter.java b/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryConverter.java
index 5e3eac97..26510c47 100644
--- a/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryConverter.java
+++ b/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryConverter.java
@@ -2,119 +2,17 @@
 
 import static java.util.stream.Collectors.toMap;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
 import java.util.Map.Entry;
 import java.util.stream.Collectors;
-import lombok.Value;
-import org.hypertrace.core.grpcutils.context.RequestContext;
-import org.hypertrace.entity.attribute.translator.EntityAttributeMapping;
-import org.hypertrace.entity.data.service.v1.AttributeFilter;
+import lombok.NoArgsConstructor;
 import org.hypertrace.entity.data.service.v1.AttributeValue;
 import org.hypertrace.entity.data.service.v1.AttributeValueList;
 import org.hypertrace.entity.data.service.v1.AttributeValueMap;
-import org.hypertrace.entity.data.service.v1.Operator;
-import org.hypertrace.entity.data.service.v1.Query;
-import org.hypertrace.entity.query.service.v1.EntityQueryRequest;
-import org.hypertrace.entity.query.service.v1.Expression;
-import org.hypertrace.entity.query.service.v1.Filter;
 import org.hypertrace.entity.query.service.v1.LiteralConstant;
-import org.hypertrace.entity.query.service.v1.OrderByExpression;
-import org.hypertrace.entity.query.service.v1.SortOrder;
-import org.hypertrace.entity.query.service.v1.ValueType;
-import org.hypertrace.entity.service.constants.EntityServiceConstants;
 
+@NoArgsConstructor
 class EntityQueryConverter {
 
-  private final EntityAttributeMapping attributeMapping;
-
-  EntityQueryConverter(EntityAttributeMapping attributeMapping) {
-    this.attributeMapping = attributeMapping;
-  }
-
-  @Deprecated(forRemoval = true)
-  public Query convertToEDSQuery(RequestContext requestContext, EntityQueryRequest queryRequest) {
-    Query.Builder queryBuilder = Query.newBuilder().setEntityType(queryRequest.getEntityType());
-
-    AttributeFilter attributeFilter =
-        convertToAttributeFilter(requestContext, queryBuilder, queryRequest.getFilter());
-    if (attributeFilter != null) {
-      queryBuilder.setFilter(attributeFilter);
-    }
-
-    queryBuilder.addAllOrderBy(convertOrderBy(requestContext, queryRequest.getOrderByList()));
-    queryBuilder.setLimit(queryRequest.getLimit());
-    queryBuilder.setOffset(queryRequest.getOffset());
-
-    return queryBuilder.build();
-  }
-
-  private AttributeFilter convertToAttributeFilter(
-      RequestContext requestContext, Query.Builder queryBuilder, Filter filter) {
-    if (filter == null || filter.equals(Filter.getDefaultInstance())) {
-      return null;
-    }
-    AttributeFilter.Builder builder = null;
-    if (filter.getChildFilterCount() == 0) {
-      // Copy the lhs and rhs from the filter.
-      AttributeColumnInformation attributeData =
-          getAttributeColumnInfo(requestContext, filter.getLhs());
-      String edsColumnName = attributeData.getColumnName();
-      org.hypertrace.entity.query.service.v1.Value rhsValue =
-          filter.getRhs().getLiteral().getValue();
-      if (edsColumnName.equals(EntityServiceConstants.ENTITY_ID)) {
-        if (rhsValue.getValueType() == ValueType.STRING) {
-          queryBuilder.addEntityId(rhsValue.getString());
-        } else if (rhsValue.getValueType() == ValueType.STRING_ARRAY) {
-          queryBuilder.addAllEntityId(rhsValue.getStringArrayList());
-        }
-      } else if (edsColumnName.equals(EntityServiceConstants.ENTITY_NAME)) {
-        queryBuilder.setEntityName(rhsValue.getString());
-      } else {
-        builder = AttributeFilter.newBuilder();
-        builder.setOperator(convertOperator(filter.getOperator()));
-        builder.setName(edsColumnName);
-        builder.setIsMultiValued(attributeData.isMultiValued());
-        builder.setAttributeValue(convertToAttributeValue(filter.getRhs()));
-      }
-    } else {
-      builder = AttributeFilter.newBuilder();
-      builder.setOperator(convertOperator(filter.getOperator()));
-      for (Filter child : filter.getChildFilterList()) {
-        AttributeFilter attributeFilter =
-            convertToAttributeFilter(requestContext, queryBuilder, child);
-        if (null != attributeFilter) {
-          builder.addChildFilter(attributeFilter);
-        }
-      }
-    }
-    return (null != builder) ? builder.build() : null;
-  }
-
-  private static org.hypertrace.entity.data.service.v1.Operator convertOperator(
-      org.hypertrace.entity.query.service.v1.Operator operator) {
-    // Might not handle all cases but this is a starter.
-    return Operator.valueOf(operator.name());
-  }
-
-  private static AttributeValue.Builder convertToAttributeValue(Expression expression) {
-    switch (expression.getValueCase()) {
-      case LITERAL:
-        return convertToAttributeValue(expression.getLiteral());
-      case COLUMNIDENTIFIER:
-        throw new IllegalArgumentException("Column Identifier should be LHS");
-      case FUNCTION:
-        throw new UnsupportedOperationException(
-            "Filtering on functional expressions not supported in EDS");
-      case ORDERBY:
-        throw new IllegalArgumentException("Filter cant have a Order By clause");
-      case VALUE_NOT_SET:
-      default:
-        return AttributeValue.newBuilder();
-    }
-  }
-
   public static AttributeValue.Builder convertToAttributeValue(LiteralConstant literal) {
     AttributeValue.Builder builder = AttributeValue.newBuilder();
     org.hypertrace.entity.query.service.v1.Value value = literal.getValue();
@@ -237,76 +135,4 @@ public static AttributeValue.Builder convertToAttributeValue(LiteralConstant lit
     }
     return builder;
   }
-
-  private AttributeColumnInformation getAttributeColumnInfo(
-      RequestContext requestContext, Expression expression) {
-    switch (expression.getValueCase()) {
-      case LITERAL:
-        throw new IllegalArgumentException("LHS should be a Attribute key");
-      case COLUMNIDENTIFIER:
-        String attributeId = expression.getColumnIdentifier().getColumnName();
-        String attributeName =
-            this.attributeMapping
-                .getDocStorePathByAttributeId(requestContext, attributeId)
-                .orElseThrow(
-                    () ->
-                        new IllegalArgumentException(
-                            String.format(
-                                "Unrecognized attribute: %s does not match any known entity attribute",
-                                attributeId)));
-        final boolean isArray = this.attributeMapping.isArray(requestContext, attributeId);
-        return new AttributeColumnInformation(attributeName, isArray);
-      case FUNCTION:
-        throw new UnsupportedOperationException(
-            "Filtering on functional expressions not supported in EDS");
-      case ORDERBY:
-        throw new IllegalArgumentException("Filter cant have a Order By clause");
-      case VALUE_NOT_SET:
-      default:
-        throw new IllegalArgumentException(
-            String.format("Attribute Key for expression:%s not found", expression));
-    }
-  }
-
-  private List<org.hypertrace.entity.data.service.v1.OrderByExpression> convertOrderBy(
-      RequestContext requestContext, List<OrderByExpression> orderByExpressions) {
-    if (orderByExpressions.isEmpty()) {
-      return Collections.emptyList();
-    }
-
-    List<org.hypertrace.entity.data.service.v1.OrderByExpression> result = new ArrayList<>();
-    for (OrderByExpression orderByExpression : orderByExpressions)
-      if (orderByExpression.hasExpression()) {
-        if (orderByExpression.getExpression().hasColumnIdentifier()) {
-          String edsColumnName =
-              getAttributeColumnInfo(requestContext, orderByExpression.getExpression())
-                  .getColumnName();
-          org.hypertrace.entity.data.service.v1.OrderByExpression convertedExpression =
-              org.hypertrace.entity.data.service.v1.OrderByExpression.newBuilder()
-                  .setName(edsColumnName)
-                  .setOrder(convertSortOrder(orderByExpression.getOrder()))
-                  .build();
-          result.add(convertedExpression);
-        } else {
-          // entity data service and doc store only support field order by. There's no
-          // aggregate order by yet
-          throw new UnsupportedOperationException(
-              "OrderByExpression only support Column Identifier Expression");
-        }
-      }
-    return result;
-  }
-
-  private static org.hypertrace.entity.data.service.v1.SortOrder convertSortOrder(
-      SortOrder sortOrder) {
-    return SortOrder.DESC == sortOrder
-        ? org.hypertrace.entity.data.service.v1.SortOrder.DESC
-        : org.hypertrace.entity.data.service.v1.SortOrder.ASC;
-  }
-
-  @Value
-  private static class AttributeColumnInformation {
-    String columnName;
-    boolean multiValued;
-  }
 }
diff --git a/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryServiceImpl.java b/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryServiceImpl.java
index 4d68d0d9..295dd9df 100644
--- a/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryServiceImpl.java
+++ b/entity-service-impl/src/main/java/org/hypertrace/entity/query/service/EntityQueryServiceImpl.java
@@ -62,7 +62,6 @@
 import org.hypertrace.entity.data.service.v1.AttributeValue;
 import org.hypertrace.entity.data.service.v1.AttributeValueList;
 import org.hypertrace.entity.data.service.v1.Entity;
-import org.hypertrace.entity.data.service.v1.Query;
 import org.hypertrace.entity.fetcher.EntityFetcher;
 import org.hypertrace.entity.metric.EntityCounterMetricSender;
 import org.hypertrace.entity.query.service.converter.AliasProvider;
@@ -131,7 +130,6 @@ public class EntityQueryServiceImpl extends EntityQueryServiceImplBase {
           .collect(joining("."));
 
   private final Collection entitiesCollection;
-  private final EntityQueryConverter entityQueryConverter;
   private final EntityAttributeMapping entityAttributeMapping;
   private final int CHUNK_SIZE;
   private final Injector injector;
@@ -208,7 +206,6 @@ public EntityQueryServiceImpl(
       int maxStringLengthForUpdate) {
     this.entitiesCollection = entitiesCollection;
     this.entityAttributeMapping = entityAttributeMapping;
-    this.entityQueryConverter = new EntityQueryConverter(entityAttributeMapping);
     this.CHUNK_SIZE = chunkSize;
     this.injector = Guice.createInjector(new ConverterModule(entityAttributeMapping));
     this.maxEntitiesToDelete = maxEntitiesToDelete;
@@ -667,11 +664,17 @@ public void total(
             .build();
 
     // converting entity query request to entity data service query
-    Query query = entityQueryConverter.convertToEDSQuery(requestContext, entityQueryRequest);
+    org.hypertrace.core.documentstore.query.Query query;
+    try {
+      final Converter<EntityQueryRequest, org.hypertrace.core.documentstore.query.Query>
+          queryConverter = getQueryConverter();
+      query = queryConverter.convert(entityQueryRequest, requestContext);
+    } catch (Exception ex) {
+      responseObserver.onError(new ServiceException("Unable to parse query."));
+      return;
+    }
 
-    // TODO: Replace to use the new org.hypertrace.core.documentstore.query.Query DTO
-    long total =
-        entitiesCollection.total(DocStoreConverter.transform(tenantId.get(), query, emptyList()));
+    long total = entitiesCollection.count(query);
     responseObserver.onNext(TotalEntitiesResponse.newBuilder().setTotal(total).build());
     responseObserver.onCompleted();
   }
diff --git a/entity-service-impl/src/test/java/org/hypertrace/entity/query/service/EntityQueryConverterTest.java b/entity-service-impl/src/test/java/org/hypertrace/entity/query/service/EntityQueryConverterTest.java
deleted file mode 100644
index 1d8ce668..00000000
--- a/entity-service-impl/src/test/java/org/hypertrace/entity/query/service/EntityQueryConverterTest.java
+++ /dev/null
@@ -1,184 +0,0 @@
-package org.hypertrace.entity.query.service;
-
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.mockito.Mockito.when;
-
-import java.time.Instant;
-import java.util.Optional;
-import org.hypertrace.core.grpcutils.context.RequestContext;
-import org.hypertrace.entity.attribute.translator.EntityAttributeMapping;
-import org.hypertrace.entity.data.service.v1.Query;
-import org.hypertrace.entity.query.service.v1.ColumnIdentifier;
-import org.hypertrace.entity.query.service.v1.EntityQueryRequest;
-import org.hypertrace.entity.query.service.v1.Expression;
-import org.hypertrace.entity.query.service.v1.Filter;
-import org.hypertrace.entity.query.service.v1.Function;
-import org.hypertrace.entity.query.service.v1.LiteralConstant;
-import org.hypertrace.entity.query.service.v1.Operator;
-import org.hypertrace.entity.query.service.v1.OrderByExpression;
-import org.hypertrace.entity.query.service.v1.SortOrder;
-import org.hypertrace.entity.query.service.v1.ValueType;
-import org.hypertrace.entity.v1.entitytype.EntityType;
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.ExtendWith;
-import org.mockito.Mock;
-import org.mockito.junit.jupiter.MockitoExtension;
-
-@ExtendWith(MockitoExtension.class)
-public class EntityQueryConverterTest {
-  private static final String ATTRIBUTE_ID_NAME1 = "eqsColumn1";
-  private static final String EDS_COLUMN_NAME1 = "edsColumn1";
-  private static final String ATTRIBUTE_ID_NAME2 = "eqsColumn2";
-  private static final String EDS_COLUMN_NAME2 = "edsColumn2";
-
-  @Mock EntityAttributeMapping mockAttributeMapping;
-  @Mock RequestContext mockRequestContext;
-  private EntityQueryConverter queryConverter;
-
-  @BeforeEach
-  public void setup() {
-    this.queryConverter = new EntityQueryConverter(mockAttributeMapping);
-  }
-
-  @Test
-  public void test_convertToEDSQuery_limitAndOffset() {
-    // no offset and limit specified
-    EntityQueryRequest request = EntityQueryRequest.newBuilder().build();
-    Query convertedQuery = this.queryConverter.convertToEDSQuery(mockRequestContext, request);
-    assertEquals(0, convertedQuery.getOffset());
-    assertEquals(0, convertedQuery.getLimit());
-
-    int limit = 3;
-    int offset = 1;
-    request = EntityQueryRequest.newBuilder().setLimit(limit).setOffset(offset).build();
-    convertedQuery = this.queryConverter.convertToEDSQuery(mockRequestContext, request);
-    assertEquals(limit, convertedQuery.getLimit());
-    assertEquals(offset, convertedQuery.getOffset());
-  }
-
-  @Test
-  public void test_convertToEdsQuery_orderByExpression() {
-    mockAttribute1();
-    mockAttribute2();
-    EntityQueryRequest request =
-        EntityQueryRequest.newBuilder()
-            .addOrderBy(
-                OrderByExpression.newBuilder()
-                    .setExpression(
-                        Expression.newBuilder()
-                            .setColumnIdentifier(
-                                ColumnIdentifier.newBuilder()
-                                    .setColumnName(ATTRIBUTE_ID_NAME1)
-                                    .build()))
-                    .setOrder(SortOrder.ASC)
-                    .build())
-            .addOrderBy(
-                OrderByExpression.newBuilder()
-                    .setExpression(
-                        Expression.newBuilder()
-                            .setColumnIdentifier(
-                                ColumnIdentifier.newBuilder()
-                                    .setColumnName(ATTRIBUTE_ID_NAME2)
-                                    .build()))
-                    .setOrder(SortOrder.DESC)
-                    .build())
-            .build();
-    Query convertedQuery = queryConverter.convertToEDSQuery(mockRequestContext, request);
-    assertEquals(2, convertedQuery.getOrderByCount());
-    // ensure that the order of OrderByExpression is maintained
-    assertEquals(EDS_COLUMN_NAME1, convertedQuery.getOrderByList().get(0).getName());
-    assertEquals(
-        org.hypertrace.entity.data.service.v1.SortOrder.ASC,
-        convertedQuery.getOrderBy(0).getOrder());
-    assertEquals(EDS_COLUMN_NAME2, convertedQuery.getOrderByList().get(1).getName());
-    assertEquals(
-        org.hypertrace.entity.data.service.v1.SortOrder.DESC,
-        convertedQuery.getOrderBy(1).getOrder());
-  }
-
-  @Test
-  public void test_convertToEdsQuery_missingSortOrderByExpression_assignWithAscByDefault() {
-    mockAttribute1();
-    EntityQueryRequest request =
-        EntityQueryRequest.newBuilder()
-            .addOrderBy(
-                OrderByExpression.newBuilder()
-                    .setExpression(
-                        Expression.newBuilder()
-                            .setColumnIdentifier(
-                                ColumnIdentifier.newBuilder()
-                                    .setColumnName(ATTRIBUTE_ID_NAME1)
-                                    .build())))
-            .build();
-    Query convertedQuery = queryConverter.convertToEDSQuery(mockRequestContext, request);
-    assertEquals(1, convertedQuery.getOrderByCount());
-    // ensure that the order of OrderByExpression is maintained
-    assertEquals(EDS_COLUMN_NAME1, convertedQuery.getOrderByList().get(0).getName());
-    assertEquals(
-        org.hypertrace.entity.data.service.v1.SortOrder.ASC,
-        convertedQuery.getOrderBy(0).getOrder());
-  }
-
-  @Test
-  public void test_filter() {
-    mockAttribute2();
-    EntityQueryRequest queryRequest =
-        EntityQueryRequest.newBuilder()
-            .setEntityType(EntityType.SERVICE.name())
-            .setFilter(
-                Filter.newBuilder()
-                    .setOperatorValue(Operator.LT.getNumber())
-                    .setLhs(
-                        Expression.newBuilder()
-                            .setColumnIdentifier(
-                                ColumnIdentifier.newBuilder()
-                                    .setColumnName(ATTRIBUTE_ID_NAME2)
-                                    .build())
-                            .build())
-                    .setRhs(
-                        Expression.newBuilder()
-                            .setLiteral(
-                                LiteralConstant.newBuilder()
-                                    .setValue(
-                                        org.hypertrace.entity.query.service.v1.Value.newBuilder()
-                                            .setLong(Instant.now().toEpochMilli())
-                                            .setValueType(ValueType.LONG)
-                                            .build())
-                                    .build())
-                            .build())
-                    .build())
-            .build();
-    Query query = queryConverter.convertToEDSQuery(mockRequestContext, queryRequest);
-    Assertions.assertEquals(EntityType.SERVICE.name(), query.getEntityType());
-    Assertions.assertNotNull(query.getFilter());
-  }
-
-  @Test
-  public void test_convertToEdsQuery_functionExpressionOrderByExpression_throwsException() {
-    EntityQueryRequest request =
-        EntityQueryRequest.newBuilder()
-            .addOrderBy(
-                OrderByExpression.newBuilder()
-                    .setExpression(
-                        Expression.newBuilder().setFunction(Function.newBuilder().build()).build())
-                    .setOrder(SortOrder.ASC)
-                    .build())
-            .build();
-    assertThrows(
-        UnsupportedOperationException.class,
-        () -> queryConverter.convertToEDSQuery(mockRequestContext, request));
-  }
-
-  private void mockAttribute1() {
-    when(mockAttributeMapping.getDocStorePathByAttributeId(mockRequestContext, ATTRIBUTE_ID_NAME1))
-        .thenReturn(Optional.of(EDS_COLUMN_NAME1));
-  }
-
-  private void mockAttribute2() {
-    when(mockAttributeMapping.getDocStorePathByAttributeId(mockRequestContext, ATTRIBUTE_ID_NAME2))
-        .thenReturn(Optional.of(EDS_COLUMN_NAME2));
-  }
-}
diff --git a/entity-service/src/main/resources/configs/common/application.conf b/entity-service/src/main/resources/configs/common/application.conf
index a10974a7..1fe6a8b1 100644
--- a/entity-service/src/main/resources/configs/common/application.conf
+++ b/entity-service/src/main/resources/configs/common/application.conf
@@ -51,6 +51,11 @@ entity.service.idAttributeMap = [
 ]
 
 entity.service.attributeMap = [
+  {
+    "scope": "BACKEND",
+    "name": "BACKEND.id",
+    "subDocPath": "entityId"
+  },
   {
     "scope": "API",
     "name": "API.id",

From 4d3144eda8ca6a142f6734513dd6a91aa3da3f3f Mon Sep 17 00:00:00 2001
From: aman-bansal <bansalaman2905@gmail.com>
Date: Thu, 13 Jun 2024 22:05:13 +0530
Subject: [PATCH 2/2] fixing tests

---
 .../service/EntityQueryServiceImplTest.java   | 38 ++++++++++---------
 1 file changed, 21 insertions(+), 17 deletions(-)

diff --git a/entity-service-impl/src/test/java/org/hypertrace/entity/query/service/EntityQueryServiceImplTest.java b/entity-service-impl/src/test/java/org/hypertrace/entity/query/service/EntityQueryServiceImplTest.java
index 73745b6c..66b903d5 100644
--- a/entity-service-impl/src/test/java/org/hypertrace/entity/query/service/EntityQueryServiceImplTest.java
+++ b/entity-service-impl/src/test/java/org/hypertrace/entity/query/service/EntityQueryServiceImplTest.java
@@ -37,15 +37,16 @@
 import org.hypertrace.core.documentstore.Document;
 import org.hypertrace.core.documentstore.Filter;
 import org.hypertrace.core.documentstore.JSONDocument;
-import org.hypertrace.core.documentstore.Query;
 import org.hypertrace.core.documentstore.SingleValueKey;
 import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
 import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.LogicalExpression;
 import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
 import org.hypertrace.core.documentstore.expression.operators.RelationalOperator;
 import org.hypertrace.core.documentstore.model.options.UpdateOptions;
 import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
 import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue;
+import org.hypertrace.core.documentstore.query.Query;
 import org.hypertrace.core.grpcutils.context.RequestContext;
 import org.hypertrace.entity.attribute.translator.EntityAttributeChangeEvaluator;
 import org.hypertrace.entity.attribute.translator.EntityAttributeMapping;
@@ -1147,8 +1148,6 @@ public void test_buildTotalQuery() throws Exception {
               .setFilter(org.hypertrace.entity.query.service.v1.Filter.getDefaultInstance())
               .build();
 
-      ArgumentCaptor<Query> docStoreQueryCaptor = ArgumentCaptor.forClass(Query.class);
-
       EntityQueryServiceImpl eqs =
           new EntityQueryServiceImpl(
               entitiesCollection,
@@ -1172,19 +1171,24 @@ public void test_buildTotalQuery() throws Exception {
                 return null;
               });
 
-      verify(entitiesCollection, times(1)).total(docStoreQueryCaptor.capture());
-      Query query = docStoreQueryCaptor.getValue();
-      assertEquals(Filter.Op.AND, query.getFilter().getOp());
-      assertEquals(2, query.getFilter().getChildFilters().length);
-      // tenant id filter
-      assertEquals(Filter.Op.EQ, query.getFilter().getChildFilters()[0].getOp());
-      assertEquals("tenantId", query.getFilter().getChildFilters()[0].getFieldName());
-      assertEquals(TENANT_ID, query.getFilter().getChildFilters()[0].getValue());
-
-      // entity type filter
-      assertEquals(Filter.Op.EQ, query.getFilter().getChildFilters()[1].getOp());
-      assertEquals("entityType", query.getFilter().getChildFilters()[1].getFieldName());
-      assertEquals(TEST_ENTITY_TYPE, query.getFilter().getChildFilters()[1].getValue());
+      verify(entitiesCollection, times(1))
+          .count(
+              Query.builder()
+                  .setFilter(
+                      org.hypertrace.core.documentstore.query.Filter.builder()
+                          .expression(
+                              LogicalExpression.and(
+                                  List.of(
+                                      RelationalExpression.of(
+                                          IdentifierExpression.of("tenantId"),
+                                          RelationalOperator.EQ,
+                                          ConstantExpression.of(TENANT_ID)),
+                                      RelationalExpression.of(
+                                          IdentifierExpression.of("entityType"),
+                                          RelationalOperator.EQ,
+                                          ConstantExpression.of(TEST_ENTITY_TYPE)))))
+                          .build())
+                  .build());
     }
 
     @DisplayName("should send correct total response")
@@ -1211,7 +1215,7 @@ public void test_sendCorrectTotalResponse() throws Exception {
               5000);
       StreamObserver<TotalEntitiesResponse> mockResponseObserver = mock(StreamObserver.class);
 
-      when(entitiesCollection.total(any())).thenReturn(123L);
+      when(entitiesCollection.count(any())).thenReturn(123L);
       Context.current()
           .withValue(RequestContext.CURRENT, mockRequestContextWithTenantId())
           .call(
