Merge pull request #2745 from atlanhq/dev/feat/removeRelationshipAttr…

…ibute

PLT-573 | Beta | Add removeSupport in bulk API

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

@@ -364,6 +364,12 @@ public Object getAppendRelationshipAttribute(String name) {
         return a != null ? a.get(name) : null;
     }
 
+    public Object getRemoveRelationshipAttribute(String name) {
+        Map<String, Object> a = this.removeRelationshipAttributes;
+
+        return a != null ? a.get(name) : null;
+    }
+
     public boolean hasRelationshipAttribute(String name) {
         Map<String, Object> r = this.relationshipAttributes;
 
@@ -404,7 +410,24 @@ public void setRemoveRelationshipAttributes(Map<String, Object> removeRelationsh
         this.removeRelationshipAttributes = removeRelationshipAttributes;
     }
 
+    public boolean hasRemoveRelationshipAttribute(String name) {
+        Map<String, Object> r = this.removeRelationshipAttributes;
+
+        return r != null ? r.containsKey(name) : false;
+    }
+
+    public void setRemoveRelationshipAttribute(String name, Object value) {
+        Map<String, Object> r = this.removeRelationshipAttributes;
 
+        if (r != null) {
+            r.put(name, value);
+        } else {
+            r = new HashMap<>();
+            r.put(name, value);
+
+            this.removeRelationshipAttributes = r;
+        }
+    }
     public Map<String, String> getCustomAttributes() {
         return customAttributes;
     }

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

@@ -70,6 +70,7 @@ public interface AtlasRelationshipStore {
 
     AtlasEdge getOrCreate(AtlasVertex end1Vertex, AtlasVertex end2Vertex, AtlasRelationship relationship) throws AtlasBaseException;
 
+    AtlasEdge getRelationship(AtlasVertex fromVertex, AtlasVertex toVertex, AtlasRelationship relationship) throws AtlasBaseException;
 
     /**
      * Retrieve a relationship if it exists or creates a new relationship instance.

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

@@ -45,6 +45,7 @@ public final class EntityGraphDiscoveryContext {
     private final Set<String>                     localGuids               = new HashSet<>();
 
     private  boolean isAppendRelationshipAttributeVisited;
+    private  boolean isRemoveRelationshipAttributeVisited;
 
     public EntityGraphDiscoveryContext(AtlasTypeRegistry typeRegistry, EntityStream entityStream) {
         this.typeRegistry = typeRegistry;
@@ -164,4 +165,12 @@ public boolean isAppendRelationshipAttributeVisited() {
     public void setAppendRelationshipAttributeVisited(boolean appendRelationshipAttributeVisited) {
         isAppendRelationshipAttributeVisited = appendRelationshipAttributeVisited;
     }
+
+    public boolean isRemoveRelationshipAttributeVisited() {
+        return isRemoveRelationshipAttributeVisited;
+    }
+
+    public void setRemoveRelationshipAttributeVisited(boolean removeRelationshipAttributeVisited) {
+        isRemoveRelationshipAttributeVisited = removeRelationshipAttributeVisited;
+    }
 }

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

@@ -376,15 +376,28 @@ private void visitRelationships(AtlasEntityType entityType, AtlasEntity entity,
                 visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());
 
                 visitedAttributes.add(attrName);
-            } else if (entity.hasAppendRelationshipAttribute(attrName)) {
-                discoveryContext.setAppendRelationshipAttributeVisited(true);
-                Object         attrVal          = entity.getAppendRelationshipAttribute(attrName);
-                String         relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
-                AtlasAttribute attribute        = entityType.getRelationshipAttribute(attrName, relationshipType);
+            } else {
+                if (entity.hasAppendRelationshipAttribute(attrName)) {
+                    discoveryContext.setAppendRelationshipAttributeVisited(true);
+                    Object attrVal = entity.getAppendRelationshipAttribute(attrName);
+                    String relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
+                    AtlasAttribute attribute = entityType.getRelationshipAttribute(attrName, relationshipType);
 
-                visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());
+                    visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());
 
-                visitedAttributes.add(attrName);
+                    visitedAttributes.add(attrName);
+                }
+
+                if (entity.hasRemoveRelationshipAttribute(attrName)) {
+                    discoveryContext.setRemoveRelationshipAttributeVisited(true);
+                    Object attrVal = entity.getRemoveRelationshipAttribute(attrName);
+                    String relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
+                    AtlasAttribute attribute = entityType.getRelationshipAttribute(attrName, relationshipType);
+
+                    visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());
+
+                    visitedAttributes.add(attrName);
+                }
             }
         }
     }

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

@@ -1708,6 +1708,10 @@ private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, Entit
                     context.setUpdatedWithRelationshipAttributes(entity);
                 }
 
+                if (discoveryContext.isRemoveRelationshipAttributeVisited() && MapUtils.isNotEmpty(entity.getRemoveRelationshipAttributes())) {
+                    context.setUpdatedWithRemoveRelationshipAttributes(entity);
+                }
+
                 if (isRestoreRequested) {
                     Status currStatus = AtlasGraphUtilsV2.getState(vertex);
                     if (currStatus == Status.DELETED) {

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

@@ -399,7 +399,8 @@ public AtlasEdge getOrCreate(AtlasVertex end1Vertex, AtlasVertex end2Vertex, Atl
         return ret;
     }
 
-    private AtlasEdge getRelationship(AtlasVertex fromVertex, AtlasVertex toVertex, AtlasRelationship relationship) throws AtlasBaseException {
+    @Override
+    public AtlasEdge getRelationship(AtlasVertex fromVertex, AtlasVertex toVertex, AtlasRelationship relationship) throws AtlasBaseException {
         String relationshipLabel = getRelationshipEdgeLabel(fromVertex, toVertex, relationship.getTypeName());
 
         return getRelationshipEdge(fromVertex, toVertex, relationshipLabel);