POLICY-166: Linking and Unlinking API improvements

intg/src/main/java/org/apache/atlas/model/instance/BusinessPolicyRequest.java

@@ -0,0 +1,133 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas.model.instance;
+
+import com.fasterxml.jackson.annotation.JsonAutoDetect;
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+
+import javax.xml.bind.annotation.XmlAccessType;
+import javax.xml.bind.annotation.XmlAccessorType;
+import javax.xml.bind.annotation.XmlRootElement;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.NONE;
+import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.PUBLIC_ONLY;
+
+/**
+ * Request to link/unlink policies from an asset.
+ */
+@JsonAutoDetect(getterVisibility = PUBLIC_ONLY, setterVisibility = PUBLIC_ONLY, fieldVisibility = NONE)
+@JsonSerialize(include = JsonSerialize.Inclusion.NON_NULL)
+@JsonIgnoreProperties(ignoreUnknown = true)
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@XmlRootElement
+@XmlAccessorType(XmlAccessType.PROPERTY)
+public class BusinessPolicyRequest implements Serializable {
+    private static final long serialVersionUID = 1L;
+
+    private List<AssetComplianceInfo> data = new ArrayList<>();
+
+    public List<AssetComplianceInfo> getData() {
+        return data;
+    }
+
+    public void setData(List<AssetComplianceInfo> data) {
+        this.data = data;
+    }
+
+    @Override
+    public String toString() {
+        return "LinkBusinessPolicyRequest{" +
+                "data=" + data +
+                '}';
+    }
+
+    public static class AssetComplianceInfo implements Serializable {
+        private static final long serialVersionUID = 1L;
+
+        private String assetId;
+
+        private Set<String> addCompliantGUIDs = new HashSet<>();
+
+        private Set<String> removeCompliantGUIDs = new HashSet<>();
+
+
+        private Set<String> addNonCompliantGUIDs = new HashSet<>();
+
+        private Set<String> removeNonCompliantGUIDs = new HashSet<>();
+
+
+        public String getAssetId() {
+            return assetId;
+        }
+
+        public void setAssetId(String assetId) {
+            this.assetId = assetId;
+        }
+
+        public Set<String> getAddCompliantGUIDs() {
+            return addCompliantGUIDs;
+        }
+
+        public void setAddCompliantGUIDs(Set<String> addCompliantGUIDs) {
+            this.addCompliantGUIDs = addCompliantGUIDs;
+        }
+
+        public Set<String> getRemoveCompliantGUIDs() {
+            return removeCompliantGUIDs;
+        }
+
+        public void setRemoveCompliantGUIDs(Set<String> removeCompliantGUIDs) {
+            this.removeCompliantGUIDs = removeCompliantGUIDs;
+        }
+
+        public Set<String> getAddNonCompliantGUIDs() {
+            return addNonCompliantGUIDs;
+        }
+
+        public void setAddNonCompliantGUIDs(Set<String> addNonCompliantGUIDs) {
+            this.addNonCompliantGUIDs = addNonCompliantGUIDs;
+        }
+
+        public Set<String> getRemoveNonCompliantGUIDs() {
+            return removeNonCompliantGUIDs;
+        }
+
+        public void setRemoveNonCompliantGUIDs(Set<String> removeNonCompliantGUIDs) {
+            this.removeNonCompliantGUIDs = removeNonCompliantGUIDs;
+        }
+
+        @Override
+        public String toString() {
+            final StringBuilder sb = new StringBuilder("AssetComplianceInfo{");
+            sb.append("assetId='").append(assetId).append('\'');
+            sb.append(", addCompliantGUIDs=").append(addCompliantGUIDs);
+            sb.append(", removeCompliantGUIDs=").append(removeCompliantGUIDs);
+            sb.append(", addNonCompliantGUIDs=").append(addNonCompliantGUIDs);
+            sb.append(", removeNonCompliantGUIDs=").append(removeNonCompliantGUIDs);
+            sb.append('}');
+            return sb.toString();
+        }
+    }
+}

intg/src/main/java/org/apache/atlas/model/instance/LinkBusinessPolicyRequest.java

@@ -47,6 +47,8 @@ public class LinkBusinessPolicyRequest implements Serializable {
     private Set<String> linkGuids;
     private Set<String> unlinkGuids;
 
+    private Set<String> businessPolicyGuids;
+
     public Set<String> getLinkGuids() {
         return linkGuids;
     }
@@ -63,11 +65,20 @@ public void setUnlinkGuids(Set<String> unlinkGuids) {
         this.unlinkGuids = unlinkGuids;
     }
 
+    public Set<String> getBusinessPolicyGuids() {
+        return businessPolicyGuids;
+    }
+
+    public void setBusinessPolicyGuids(Set<String> businessPolicyGuids) {
+        this.businessPolicyGuids = businessPolicyGuids;
+    }
+
     @Override
     public String toString() {
         final StringBuilder sb = new StringBuilder("LinkBusinessPolicyRequest{");
         sb.append("linkGuids=").append(linkGuids);
         sb.append(", unlinkGuids=").append(unlinkGuids);
+        sb.append(", businessPolicyGuids=").append(businessPolicyGuids);
         sb.append('}');
         return sb.toString();
     }

repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java

@@ -369,7 +369,7 @@ EntityMutationResponse deleteByUniqueAttributes(List<AtlasObjectId> objectIds)
 
     void unlinkMeshEntityFromAssets(String meshEntityId, Set<String> unlinkGuids) throws AtlasBaseException;
 
-    void linkBusinessPolicy(String policyId, Set<String> linkGuids) throws AtlasBaseException;
+    void linkBusinessPolicy(List<BusinessPolicyRequest.AssetComplianceInfo> data) throws AtlasBaseException;
 
     void unlinkBusinessPolicy(String policyId, Set<String> unlinkGuids) throws AtlasBaseException;
 

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java

@@ -2781,18 +2781,17 @@ public void repairAccesscontrolAlias(String guid) throws AtlasBaseException {
 
     @Override
     @GraphTransaction
-    public void linkBusinessPolicy(String policyGuid, Set<String> linkGuids) throws AtlasBaseException {
+    public void linkBusinessPolicy(List<BusinessPolicyRequest.AssetComplianceInfo> data) throws AtlasBaseException {
         AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("linkBusinessPolicy.GraphTransaction");
-
+        List<AtlasVertex> atlasVertices = new ArrayList<>();
         try {
-            List<AtlasVertex> vertices = this.entityGraphMapper.linkBusinessPolicy(policyGuid, linkGuids);
-            if (CollectionUtils.isEmpty(vertices)) {
-                return;
+            for (BusinessPolicyRequest.AssetComplianceInfo ad : data) {
+                AtlasVertex av = this.entityGraphMapper.linkBusinessPolicy(ad);
+                atlasVertices.add(av);
             }
-
-            handleEntityMutation(vertices);
+            handleEntityMutation(atlasVertices);
         } catch (Exception e) {
-            LOG.error("Error during linkBusinessPolicy for policyGuid: {}", policyGuid, e);
+            LOG.error("Error during linkBusinessPolicy for policyGuid: ", e);
             throw e;
         } finally {
             RequestContext.get().endMetricRecord(metric);

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -84,7 +84,6 @@
 
 import static org.apache.atlas.AtlasConfiguration.LABEL_MAX_LENGTH;
 import static org.apache.atlas.AtlasConfiguration.STORE_DIFFERENTIAL_AUDITS;
-import static org.apache.atlas.AtlasErrorCode.CLASSIFICATION_NOT_FOUND;
 import static org.apache.atlas.model.TypeCategory.ARRAY;
 import static org.apache.atlas.model.TypeCategory.CLASSIFICATION;
 import static org.apache.atlas.model.instance.AtlasEntity.Status.ACTIVE;
@@ -4771,20 +4770,96 @@ public void addHasLineage(Set<AtlasEdge> inputOutputEdges, boolean isRestoreEnti
     }
 
 
-    public List<AtlasVertex> linkBusinessPolicy(String policyId, Set<String> linkGuids) {
-        return linkGuids.stream().map(guid -> findByGuid(graph, guid)).filter(Objects::nonNull).filter(ev -> {
-            Set<String> existingValues = ev.getMultiValuedSetProperty(ASSET_POLICY_GUIDS, String.class);
-            return !existingValues.contains(policyId);
-        }).peek(ev -> {
-            Set<String> existingValues = ev.getMultiValuedSetProperty(ASSET_POLICY_GUIDS, String.class);
-            existingValues.add(policyId);
-            ev.setProperty(ASSET_POLICY_GUIDS, policyId);
-            ev.setProperty(ASSET_POLICIES_COUNT, existingValues.size());
+    public AtlasVertex linkBusinessPolicy(final BusinessPolicyRequest.AssetComplianceInfo data) {
+        String assetGuid = data.getAssetId();
+        AtlasVertex vertex = findByGuid(graph, assetGuid);
 
-            updateModificationMetadata(ev);
+        // Retrieve existing policies
+        Set<String> existingCompliant = getVertexPolicies(vertex, ASSET_POLICY_GUIDS);
+        Set<String> existingNonCompliant = getVertexPolicies(vertex, NON_COMPLIANT_ASSET_POLICY_GUIDS);
 
-            cacheDifferentialEntity(ev, existingValues, ev.getMultiValuedSetProperty(NON_COMPLIANT_ASSET_POLICY_GUIDS, String.class));
-        }).collect(Collectors.toList());
+        // Retrieve new policies
+        Set<String> addCompliantGUIDs = getOrCreateEmptySet(data.getAddCompliantGUIDs());
+        Set<String> addNonCompliantGUIDs = getOrCreateEmptySet(data.getAddNonCompliantGUIDs());
+        Set<String> removeCompliantGUIDs = getOrCreateEmptySet(data.getRemoveCompliantGUIDs());
+        Set<String> removeNonCompliantGUIDs = getOrCreateEmptySet(data.getRemoveNonCompliantGUIDs());
+
+
+        // Update vertex properties
+        addToAttribute(vertex, ASSET_POLICY_GUIDS, addCompliantGUIDs);
+        removeFromAttribute(vertex, ASSET_POLICY_GUIDS, removeCompliantGUIDs);
+
+
+        addToAttribute(vertex, NON_COMPLIANT_ASSET_POLICY_GUIDS, addNonCompliantGUIDs);
+        removeFromAttribute(vertex, NON_COMPLIANT_ASSET_POLICY_GUIDS, removeNonCompliantGUIDs);
+
+        // Count and set policies
+        Set<String> effectiveCompliantGUIDs = getVertexPolicies(vertex, ASSET_POLICY_GUIDS);
+        Set<String> effectiveNonCompliantGUIDs = getVertexPolicies(vertex, NON_COMPLIANT_ASSET_POLICY_GUIDS);
+
+        int compliantPolicyCount = countPoliciesExcluding(effectiveCompliantGUIDs, "rule");
+        int nonCompliantPolicyCount = countPoliciesExcluding(effectiveNonCompliantGUIDs, "rule");
+
+        int totalPolicyCount = compliantPolicyCount + nonCompliantPolicyCount;
+
+        vertex.setProperty(ASSET_POLICIES_COUNT, totalPolicyCount);
+        updateModificationMetadata(vertex);
+
+        // Create and cache differential entity
+        AtlasEntity diffEntity = createDifferentialEntity(
+                vertex, effectiveCompliantGUIDs, effectiveNonCompliantGUIDs, existingCompliant, existingNonCompliant, totalPolicyCount);
+
+        RequestContext.get().cacheDifferentialEntity(diffEntity);
+        return vertex;
+    }
+
+    private static Set<String> getOrCreateEmptySet(Set<String> input) {
+        return input == null ? new HashSet<>() : input;
+    }
+
+    private void addToAttribute(AtlasVertex vertex, String propertyKey, Set<String> policies) {
+        String targetProperty = determineTargetProperty(propertyKey);
+        policies.stream()
+                .filter(StringUtils::isNotEmpty)
+                .forEach(policyGuid -> vertex.setProperty(targetProperty, policyGuid));
+    }
+
+    private void removeFromAttribute(AtlasVertex vertex, String propertyKey, Set<String> policies) {
+        String targetProperty = determineTargetProperty(propertyKey);
+        policies.stream()
+                .filter(StringUtils::isNotEmpty)
+                .forEach(policyGuid -> vertex.removePropertyValue(targetProperty, policyGuid));
+    }
+
+    private String determineTargetProperty(String propertyKey) {
+        return ASSET_POLICY_GUIDS.equals(propertyKey)
+                ? ASSET_POLICY_GUIDS
+                : NON_COMPLIANT_ASSET_POLICY_GUIDS;
+    }
+
+    private int countPoliciesExcluding(Set<String> policies, String substring) {
+        return (int) policies.stream().filter(policy -> !policy.contains(substring)).count();
+    }
+
+    private AtlasEntity createDifferentialEntity(AtlasVertex vertex, Set<String> effectiveCompliant, Set<String> effectiveNonCompliant,
+                                                 Set<String> existingCompliant, Set<String> existingNonCompliant, int totalPolicyCount) {
+        AtlasEntity diffEntity = new AtlasEntity(vertex.getProperty(TYPE_NAME_PROPERTY_KEY, String.class));
+        setEntityCommonAttributes(vertex, diffEntity);
+        diffEntity.setAttribute(ASSET_POLICIES_COUNT, totalPolicyCount);
+
+        if (!existingCompliant.equals(effectiveCompliant)) {
+            diffEntity.setAttribute(ASSET_POLICY_GUIDS, effectiveCompliant);
+        }
+        if (!existingNonCompliant.equals(effectiveNonCompliant)) {
+            diffEntity.setAttribute(NON_COMPLIANT_ASSET_POLICY_GUIDS, effectiveNonCompliant);
+        }
+
+        return diffEntity;
+    }
+
+    private Set<String> getVertexPolicies(AtlasVertex vertex, String propertyKey) {
+        return Optional.ofNullable(vertex.getMultiValuedSetProperty(propertyKey, String.class))
+                .orElse(Collections.emptySet());
     }
 
 
@@ -4811,15 +4886,17 @@ 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);
             vertex.removePropertyValue(NON_COMPLIANT_ASSET_POLICY_GUIDS, policyId);
 
-            int totalPolicies = compliantPolicies.size() + nonCompliantPolicies.size();
-            vertex.setProperty(ASSET_POLICIES_COUNT, totalPolicies);
+            int compliantPolicyCount = countPoliciesExcluding(compliantPolicies, "rule");
+            int nonCompliantPolicyCount = countPoliciesExcluding(nonCompliantPolicies, "rule");
+            int totalPolicyCount = compliantPolicyCount + nonCompliantPolicyCount;
+            vertex.setProperty(ASSET_POLICIES_COUNT, totalPolicyCount);
 
             updateModificationMetadata(vertex);
             cacheDifferentialEntity(vertex, compliantPolicies, nonCompliantPolicies);
@@ -4828,6 +4905,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)
@@ -4927,7 +5010,7 @@ public AtlasVertex moveBusinessPolicies(Set<String> policyIds, String assetId, S
 
         // Check if the asset already has the given policy IDs
         if (policyIds.isEmpty()) {
-            throw new AtlasBaseException(AtlasErrorCode.BAD_REQUEST, "Asset already has the given policy id");
+            return assetVertex;
         }
 
         // Move policies to the appropriate set