ATLAS-2514: Deleting a table with lineage should not remove propagated classification edges

repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java

@@ -917,6 +917,43 @@ public final class GraphHelper {
         return ret;
     }
 
+    public static List<AtlasEdge> getClassificationEdges(AtlasVertex entityVertex) {
+        return getClassificationEdges(entityVertex, false);
+    }
+
+    public static List<AtlasEdge> getPropagatedClassificationEdges(AtlasVertex entityVertex) {
+        return getClassificationEdges(entityVertex, true);
+    }
+
+    public static List<AtlasEdge> getAllClassificationEdges(AtlasVertex entityVertex) {
+        return getClassificationEdges(entityVertex, null);
+    }
+
+    public static List<AtlasEdge> getClassificationEdges(AtlasVertex entityVertex, Boolean propagated) {
+        List<AtlasEdge>  ret   = new ArrayList<>();
+        AtlasVertexQuery query = entityVertex.query().direction(AtlasEdgeDirection.OUT).label(CLASSIFICATION_LABEL);
+
+        if (propagated != null) {
+            query = query.has(CLASSIFICATION_EDGE_IS_PROPAGATED_PROPERTY_KEY, propagated);
+        }
+
+        Iterable edges = query.edges();
+
+        if (edges != null) {
+            Iterator<AtlasEdge> iterator = edges.iterator();
+
+            while (iterator.hasNext()) {
+                AtlasEdge edge = iterator.next();
+
+                if (edge != null) {
+                    ret.add(edge);
+                }
+            }
+        }
+
+        return ret;
+    }
+
     public static List<String> getSuperTypeNames(AtlasVertex<?,?> entityVertex) {
         ArrayList<String>  superTypes     = new ArrayList<>();
         Collection<String> propertyValues = entityVertex.getPropertyValues(Constants.SUPER_TYPES_PROPERTY_KEY, String.class);

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

@@ -52,6 +52,8 @@ import static org.apache.atlas.repository.Constants.CLASSIFICATION_LABEL;
 import static org.apache.atlas.repository.Constants.PROPAGATED_TRAIT_NAMES_PROPERTY_KEY;
 import static org.apache.atlas.repository.graph.GraphHelper.EDGE_LABEL_PREFIX;
 import static org.apache.atlas.repository.graph.GraphHelper.addListProperty;
+import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEdge;
+import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEdges;
 import static org.apache.atlas.repository.graph.GraphHelper.getPropagatedEdges;
 import static org.apache.atlas.repository.graph.GraphHelper.getTraitNames;
 import static org.apache.atlas.repository.graph.GraphHelper.getTypeName;
@@ -112,7 +114,7 @@ public abstract class DeleteHandlerV1 {
 
         // Delete traits and vertices.
         for (AtlasVertex deletionCandidateVertex : deletionCandidateVertices) {
-            deleteAllTraits(deletionCandidateVertex);
+            deleteAllClassifications(deletionCandidateVertex);
             deleteTypeVertex(deletionCandidateVertex, false);
         }
     }
@@ -336,9 +338,7 @@ public abstract class DeleteHandlerV1 {
             break;
 
             case CLASSIFICATION:
-                removeTagPropagation(instanceVertex);
-
-                deleteTypeVertex(instanceVertex, force);
+                deleteClassificationVertex(instanceVertex, force);
             break;
 
             case ENTITY:
@@ -367,7 +367,7 @@ public abstract class DeleteHandlerV1 {
                                 getTypeName(propagatedEntityVertex), GraphHelper.getGuid(propagatedEntityVertex), CLASSIFICATION_LABEL);
                     }
 
-                    removePropagatedTraitName(propagatedEntityVertex, classificationName);
+                    removeFromPropagatedTraitNames(propagatedEntityVertex, classificationName);
 
                     deleteEdge(propagatedEdge, true);
 
@@ -381,7 +381,7 @@ public abstract class DeleteHandlerV1 {
         return ret;
     }
 
-    private void removePropagatedTraitName(AtlasVertex entityVertex, String classificationName) {
+    private void removeFromPropagatedTraitNames(AtlasVertex entityVertex, String classificationName) {
         if (entityVertex != null && StringUtils.isNotEmpty(classificationName)) {
             List<String> propagatedTraitNames = getTraitNames(entityVertex, true);
 
@@ -485,22 +485,15 @@ public abstract class DeleteHandlerV1 {
     }
 
     /**
-     * Delete all traits from the specified vertex.
+     * Delete all associated classifications from the specified entity vertex.
      * @param instanceVertex
      * @throws AtlasException
      */
-    private void deleteAllTraits(AtlasVertex instanceVertex) throws AtlasBaseException {
-        String       typeName   = GraphHelper.getTypeName(instanceVertex);
-        List<String> traitNames = GraphHelper.getTraitNames(instanceVertex);
-
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("Deleting traits {} for {}", traitNames, string(instanceVertex));
-        }
+    private void deleteAllClassifications(AtlasVertex instanceVertex) throws AtlasBaseException {
+        List<AtlasEdge> classificationEdges = getClassificationEdges(instanceVertex);
 
-        for (String traitNameToBeDeleted : traitNames) {
-            String relationshipLabel = GraphHelper.getTraitLabel(typeName, traitNameToBeDeleted);
-
-            deleteEdgeReference(instanceVertex, relationshipLabel, TypeCategory.CLASSIFICATION, false);
+        for (AtlasEdge edge : classificationEdges) {
+            deleteEdgeReference(edge, TypeCategory.CLASSIFICATION, false, false, instanceVertex);
         }
     }
 
@@ -689,4 +682,12 @@ public abstract class DeleteHandlerV1 {
 
         _deleteVertex(instanceVertex, force);
     }
+
+    protected void deleteClassificationVertex(AtlasVertex classificationVertex, boolean force) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Deleting classification vertex", string(classificationVertex));
+        }
+
+        _deleteVertex(classificationVertex, force);
+    }
 }

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

@@ -49,10 +49,10 @@ public class SoftDeleteHandlerV1 extends DeleteHandlerV1 {
             graphHelper.removeVertex(instanceVertex);
         } else {
             AtlasEntity.Status state = AtlasGraphUtilsV1.getState(instanceVertex);
+
             if (state != AtlasEntity.Status.DELETED) {
                 GraphHelper.setProperty(instanceVertex, STATE_PROPERTY_KEY, AtlasEntity.Status.DELETED.name());
-                GraphHelper.setProperty(instanceVertex, MODIFICATION_TIMESTAMP_PROPERTY_KEY,
-                    RequestContextV1.get().getRequestTime());
+                GraphHelper.setProperty(instanceVertex, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContextV1.get().getRequestTime());
                 GraphHelper.setProperty(instanceVertex, MODIFIED_BY_KEY, RequestContextV1.get().getUser());
             }
         }
@@ -64,13 +64,12 @@ public class SoftDeleteHandlerV1 extends DeleteHandlerV1 {
             graphHelper.removeEdge(edge);
         } else {
             AtlasEntity.Status state = AtlasGraphUtilsV1.getState(edge);
+
             if (state != AtlasEntity.Status.DELETED) {
                 GraphHelper.setProperty(edge, STATE_PROPERTY_KEY, AtlasEntity.Status.DELETED.name());
-                GraphHelper
-                    .setProperty(edge, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContextV1.get().getRequestTime());
+                GraphHelper.setProperty(edge, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContextV1.get().getRequestTime());
                 GraphHelper.setProperty(edge, MODIFIED_BY_KEY, RequestContextV1.get().getUser());
             }
         }
     }
 }
\ No newline at end of file
-