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 d682dd04..4d68d0d9 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 @@ -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 r // Finally, return the selections List 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 getProjectedDocuments( + final String tenantId, final Iterable entityIds, final List selectionList, final RequestContext requestContext) @@ -546,10 +552,16 @@ private List 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 =