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

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;
@@ -342,7 +343,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 +535,7 @@ private JSONDocument convertToJsonDocument(LiteralConstant literalConstant) {
   }
 
   private List<Document> getProjectedDocuments(
+      final String tenantId,
       final Iterable<String> entityIds,
       final List<Expression> selectionList,
       final RequestContext requestContext)
@@ -550,6 +555,11 @@ private List<Document> getProjectedDocuments(
                     IdentifierExpression.of(EntityServiceConstants.ENTITY_ID),
                     IN,
                     ConstantExpression.ofStrings(entityIdList)))
+            .expression(
+                RelationalExpression.of(
+                    IdentifierExpression.of(EntityServiceConstants.TENANT_ID),
+                    EQ,
+                    ConstantExpression.of(tenantId)))
             .build();
 
     final org.hypertrace.core.documentstore.query.Query query =