enhance logic

Author: aarshi0301

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java

@@ -1927,29 +1927,63 @@ private void addRemovedProcessRelationshipToDiffEntity(AtlasEntity diffEntity, A
             // Get vertices from the edge
             AtlasVertex vertexA = removedEdge.getOutVertex();
             AtlasVertex vertexB = removedEdge.getInVertex();
+            String edgeLabel = removedEdge.getLabel();
+
+            if (edgeLabel == null) {
+                LOG.warn("Cannot add removed process relationship: edgeLabel is null");
+                return;
+            }
+
+            // Determine which end the diffEntity corresponds to
+            // vertexA (OutVertex) corresponds to EndDef1, vertexB (InVertex) corresponds to EndDef2
+            AtlasRelationshipDef relationshipDef = typeRegistry.getRelationshipDefByLabel(edgeLabel);
+            if (relationshipDef == null) {
+                LOG.warn("Cannot add removed process relationship: relationshipDef not found for edgeLabel: {}", edgeLabel);
+                return;
+            }
+
             AtlasVertex otherVertex;
+            AtlasVertex diffVertex;
+            String attributeName;
             if (diffEntity.getGuid().equals(GraphHelper.getGuid(vertexA))) {
                 otherVertex = vertexB;
+                diffVertex = vertexA;
+                // diffEntity is EndDef1 (OutVertex)
+                attributeName = relationshipDef.getEndDef1().getName();
             } else {
                 otherVertex = vertexA;
+                diffVertex = vertexB;
+                // diffEntity is EndDef2 (InVertex)
+                attributeName = relationshipDef.getEndDef2().getName();
             }
 
             if (otherVertex == null) {
-                LOG.warn("Cannot add removed process relationship: processVertex is null");
+                LOG.warn("Cannot add removed process relationship: otherVertex is null");
                 return;
             }
 
-
-            String edgeLabel = removedEdge.getLabel();
-
-            if (edgeLabel == null) {
-                LOG.warn("Cannot add removed process relationship: edgeLabel is null");
+            if (diffVertex == null) {
+                LOG.warn("Cannot add removed process relationship: diffVertex is null");
                 return;
             }
 
-            String attributeName = typeRegistry.getRelationshipDefByLabel(edgeLabel).getEndDef1().getName();
-            // Add to removedRelationshipAttributes
-            diffEntity.setRemovedRelationshipAttribute(attributeName, new AtlasObjectId(getGuid(otherVertex), getTypeName(otherVertex)));
+
+
+            AtlasEntityType entityType = typeRegistry.getEntityTypeByName(getTypeName(diffVertex));
+            AtlasAttribute relationshipAttribute = entityType != null ? 
+                entityType.getRelationshipAttribute(attributeName, relationshipDef.getName()) : null;
+            
+            TypeCategory typeCategory = relationshipAttribute != null && relationshipAttribute.getAttributeType() != null ?
+                relationshipAttribute.getAttributeType().getTypeCategory() : null;
+            
+            AtlasObjectId objectId = new AtlasObjectId(getGuid(otherVertex), getTypeName(otherVertex));
+            
+            // If attribute is an array type, pass as a list; otherwise pass as a single object
+            if (typeCategory == TypeCategory.ARRAY) {
+                diffEntity.setRemovedRelationshipAttribute(attributeName, Collections.singletonList(objectId));
+            } else {
+                diffEntity.setRemovedRelationshipAttribute(attributeName, objectId);
+            }
         } catch (Exception e) {
             LOG.warn("Failed to add removed process relationship to diff entity for notification: {}", e.getMessage(), e);
         }