Merge branch 'feat/removeRelationshipAttribute' into dev/feat/removeR…

…elationshipAttribute
atlanhq · Jan 10, 2024 · cad5a54 · cad5a54
2 parents 48d8ef0 + 4295aea
commit cad5a54
Show file tree

Hide file tree

Showing 6 changed files with 227 additions and 69 deletions.
diff --git a/intg/src/main/java/org/apache/atlas/model/instance/AtlasEntity.java b/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;
     }

diff --git a/...ry/src/main/java/org/apache/atlas/repository/store/graph/EntityGraphDiscoveryContext.java b/...ry/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;
+    }
 }
diff --git a/...src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityGraphDiscoveryV2.java b/...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);
+                }
             }
         }
     }

diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java b/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) {