POLICY-172 changes in unlinking API

atlanhq · Dec 16, 2024 · e4224f9 · e4224f9
1 parent 97d8534
commit e4224f9
Showing 1 changed file with 8 additions and 2 deletions.
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
@@ -4876,8 +4876,8 @@ private AtlasVertex updateVertexPolicy(AtlasVertex vertex, String policyId) {
         Set<String> compliantPolicies = getMultiValuedSetProperty(vertex, ASSET_POLICY_GUIDS);
         Set<String> 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<String> policies, String policyId) {
+        Set<String> toRemove = policies.stream().filter(i-> i.contains(policyId)).collect(Collectors.toSet());
+        return policies.removeAll(toRemove);
+
+    }
+
     private Set<String> getMultiValuedSetProperty(AtlasVertex vertex, String propertyName) {
         return Optional.ofNullable(vertex.getMultiValuedSetProperty(propertyName, String.class))
                 .map(HashSet::new)