Merge pull request #2747 from atlanhq/feat/removeRelationshipAttribute

PLT-324 | Staging | Add append and remove support for bulk endpoint

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

@@ -92,6 +92,8 @@ public enum Status { ACTIVE, DELETED, PURGED }
     private Boolean starred       = null;
 
     private Map<String, Object>              relationshipAttributes;
+    private Map<String, Object>              appendRelationshipAttributes;
+    private Map<String, Object>              removeRelationshipAttributes;
     private List<AtlasClassification>        classifications;
     private List<AtlasTermAssignmentHeader>  meanings;
     private Map<String, String>              customAttributes;
@@ -230,6 +232,8 @@ public AtlasEntity(AtlasEntity other) {
             setBusinessAttributes(other.getBusinessAttributes());
             setLabels(other.getLabels());
             setPendingTasks(other.getPendingTasks());
+            setAppendRelationshipAttributes(other.getAppendRelationshipAttributes());
+            setRemoveRelationshipAttributes(other.getRemoveRelationshipAttributes());
         }
     }
 
@@ -354,12 +358,76 @@ public Object getRelationshipAttribute(String name) {
         return a != null ? a.get(name) : null;
     }
 
+    public Object getAppendRelationshipAttribute(String name) {
+        Map<String, Object> a = this.appendRelationshipAttributes;
+
+        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;
 
         return r != null ? r.containsKey(name) : false;
     }
 
+    public boolean hasAppendRelationshipAttribute(String name) {
+        Map<String, Object> r = this.appendRelationshipAttributes;
+
+        return r != null ? r.containsKey(name) : false;
+    }
+
+    public Map<String, Object> getAppendRelationshipAttributes() {
+        return appendRelationshipAttributes;
+    }
+
+    public void setAppendRelationshipAttributes(Map<String, Object> appendRelationshipAttributes) {
+        this.appendRelationshipAttributes = appendRelationshipAttributes;
+    }
+
+    public void setAppendRelationshipAttribute(String name, Object value) {
+        Map<String, Object> r = this.appendRelationshipAttributes;
+
+        if (r != null) {
+            r.put(name, value);
+        } else {
+            r = new HashMap<>();
+            r.put(name, value);
+
+            this.appendRelationshipAttributes = r;
+        }
+    }
+    public Map<String, Object> getRemoveRelationshipAttributes() {
+        return removeRelationshipAttributes;
+    }
+
+    public void setRemoveRelationshipAttributes(Map<String, Object> removeRelationshipAttributes) {
+        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

@@ -44,6 +44,9 @@ public final class EntityGraphDiscoveryContext {
     private final Map<AtlasObjectId, AtlasVertex> resolvedIdsByUniqAttribs = new HashMap<>();
     private final Set<String>                     localGuids               = new HashSet<>();
 
+    private  boolean isAppendRelationshipAttributeVisited;
+    private  boolean isRemoveRelationshipAttributeVisited;
+
     public EntityGraphDiscoveryContext(AtlasTypeRegistry typeRegistry, EntityStream entityStream) {
         this.typeRegistry = typeRegistry;
         this.entityStream = entityStream;
@@ -154,4 +157,20 @@ public void cleanUp() {
         resolvedIdsByUniqAttribs.clear();
         localGuids.clear();
     }
+
+    public boolean isAppendRelationshipAttributeVisited() {
+        return 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

@@ -358,6 +358,8 @@ private void visitRelationships(AtlasEntityType entityType, AtlasEntity entity,
         for (String attrName : entityType.getRelationshipAttributes().keySet()) {
 
             // if attribute is not in 'relationshipAttributes', try 'attributes'
+            // appendRelationshipAttribute will be ignored if same attribute is present
+            // in relationshipAttribute
             if (entity.hasRelationshipAttribute(attrName)) {
                 Object         attrVal          = entity.getRelationshipAttribute(attrName);
                 String         relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
@@ -374,10 +376,31 @@ 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);
+
+                    visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());
+
+                    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);
+                }
             }
         }
     }
-
     void visitStruct(AtlasStructType structType, AtlasStruct struct, String referringEntityGuid) throws AtlasBaseException {
         for (AtlasAttribute attribute : structType.getAllAttributes().values()) {
             AtlasType attrType = attribute.getAttributeType();

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

@@ -1438,7 +1438,6 @@ private EntityMutationResponse createOrUpdate(EntityStream entityStream, boolean
                     }
                 }
             }
-
             // for existing entities, skip update if incoming entity doesn't have any change
             if (CollectionUtils.isNotEmpty(context.getUpdatedEntities())) {
                 MetricRecorder checkForUnchangedEntities = RequestContext.get().startMetricRecord("checkForUnchangedEntities");
@@ -1644,6 +1643,14 @@ private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, Entit
                 String entityActiveKey = Status.ACTIVE.toString();
                 boolean isRestoreRequested = ((StringUtils.isNotEmpty(entityStateValue) && entityStateValue.equals(entityActiveKey)) || (StringUtils.isNotEmpty(entityStatusValue) && entityStatusValue.equals(entityActiveKey)));
 
+                if (discoveryContext.isAppendRelationshipAttributeVisited() && MapUtils.isNotEmpty(entity.getAppendRelationshipAttributes())) {
+                    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

@@ -397,7 +397,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);