fix | fix the query for fetching the updated document (#284)

* fix | fix the query for fetching the updated document

Author: aman-bansal

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

@@ -6,6 +6,7 @@
 import static java.util.Objects.isNull;
 import static java.util.stream.Collectors.joining;
 import static java.util.stream.Collectors.toUnmodifiableList;
+import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
 import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.IN;
 import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.NONE;
 import static org.hypertrace.entity.data.service.v1.AttributeValue.VALUE_LIST_FIELD_NUMBER;
@@ -45,6 +46,7 @@
 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.model.options.UpdateOptions;
 import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
@@ -342,7 +344,10 @@ public void update(EntityUpdateRequest request, StreamObserver<ResultSetChunk> r
       // Finally, return the selections
       List<Document> documents =
           getProjectedDocuments(
-              request.getEntityIdsList(), request.getSelectionList(), requestContext);
+              maybeTenantId.get(),
+              request.getEntityIdsList(),
+              request.getSelectionList(),
+              requestContext);
 
       responseObserver.onNext(
           convertDocumentsToResultSetChunk(documents, request.getSelectionList()));
@@ -531,6 +536,7 @@ private JSONDocument convertToJsonDocument(LiteralConstant literalConstant) {
   }
 
   private List<Document> getProjectedDocuments(
+      final String tenantId,
       final Iterable<String> entityIds,
       final List<Expression> selectionList,
       final RequestContext requestContext)
@@ -546,10 +552,16 @@ private List<Document> getProjectedDocuments(
     final Filter filter =
         Filter.builder()
             .expression(
-                RelationalExpression.of(
-                    IdentifierExpression.of(EntityServiceConstants.ENTITY_ID),
-                    IN,
-                    ConstantExpression.ofStrings(entityIdList)))
+                LogicalExpression.and(
+                    List.of(
+                        RelationalExpression.of(
+                            IdentifierExpression.of(EntityServiceConstants.ENTITY_ID),
+                            IN,
+                            ConstantExpression.ofStrings(entityIdList)),
+                        RelationalExpression.of(
+                            IdentifierExpression.of(EntityServiceConstants.TENANT_ID),
+                            EQ,
+                            ConstantExpression.of(tenantId)))))
             .build();
 
     final org.hypertrace.core.documentstore.query.Query query =