diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java index ad804726f1..ca7e60a171 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java @@ -4876,8 +4876,8 @@ private AtlasVertex updateVertexPolicy(AtlasVertex vertex, String policyId) { Set compliantPolicies = getMultiValuedSetProperty(vertex, ASSET_POLICY_GUIDS); Set nonCompliantPolicies = getMultiValuedSetProperty(vertex, NON_COMPLIANT_ASSET_POLICY_GUIDS); - boolean removed = compliantPolicies.remove(policyId); - removed |= nonCompliantPolicies.remove(policyId); + boolean removed = removePolicyAndRule(compliantPolicies, policyId); + removed |= removePolicyAndRule(nonCompliantPolicies,policyId); if (removed) { vertex.removePropertyValue(ASSET_POLICY_GUIDS, policyId); @@ -4893,6 +4893,12 @@ private AtlasVertex updateVertexPolicy(AtlasVertex vertex, String policyId) { return vertex; } + private boolean removePolicyAndRule(Set policies, String policyId) { + Set toRemove = policies.stream().filter(i-> i.contains(policyId)).collect(Collectors.toSet()); + return policies.removeAll(toRemove); + + } + private Set getMultiValuedSetProperty(AtlasVertex vertex, String propertyName) { return Optional.ofNullable(vertex.getMultiValuedSetProperty(propertyName, String.class)) .map(HashSet::new)