Merge pull request #2249 from atlanhq/fix/gov-1255

[Staging] Fix for refresh propagation propagating tag to source vertex
atlanhq · Aug 10, 2023 · 79794df · 79794df
2 parents 8446626 + 9001140
commit 79794df
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java
@@ -3523,7 +3523,7 @@ public void classificationRefreshPropagation(String classificationId) throws Atl
         LOG.info("{} entity vertices have classification with id {} attached", propagatedVerticesIds.size(), classificationId);
 
         List<String> verticesIdsToAddClassification =  new ArrayList<>();
-        List<String> propagatedVerticesIdWithoutEdge = entityRetriever.getImpactedVerticesIds(sourceEntityVertex , classificationId,
+        List<String> propagatedVerticesIdWithoutEdge = entityRetriever.getImpactedVerticesIdsClassificationAttached(sourceEntityVertex , classificationId,
                 CLASSIFICATION_PROPAGATION_EXCLUSION_MAP.get(propagationMode), verticesIdsToAddClassification);
 
         LOG.info("To add classification with id {} to {} vertices", classificationId, verticesIdsToAddClassification.size());

diff --git a/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java b/...sitory/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java
@@ -616,10 +616,15 @@ public List<String> getImpactedVerticesIds(AtlasVertex entityVertex, String rela
         return ret;
     }
 
-    public List<String> getImpactedVerticesIds(AtlasVertex entityVertex, String classificationId, List<String> edgeLabelsToExclude, List<String> verticesWithoutClassification) {
+    public List<String> getImpactedVerticesIdsClassificationAttached(AtlasVertex entityVertex, String classificationId, List<String> edgeLabelsToExclude, List<String> verticesWithoutClassification) {
         List<String> ret = new ArrayList<>();
 
-        traverseImpactedVerticesByLevel(entityVertex, null, classificationId, ret, edgeLabelsToExclude, verticesWithoutClassification);
+        GraphHelper.getClassificationEdges(entityVertex).forEach(classificationEdge -> {
+            AtlasVertex classificationVertex = classificationEdge.getInVertex();
+            if (classificationVertex != null && classificationId.equals(classificationVertex.getIdForDisplay())) {
+                traverseImpactedVerticesByLevel(entityVertex, null, classificationId, ret, edgeLabelsToExclude, verticesWithoutClassification);
+            }
+        });
 
         return ret;
     }

diff --git a/repository/src/main/java/org/apache/atlas/tasks/AtlasTaskService.java b/repository/src/main/java/org/apache/atlas/tasks/AtlasTaskService.java
@@ -151,7 +151,7 @@ public List<AtlasTask> createAtlasTasks(List<AtlasTask> tasks) throws AtlasBaseE
                 if (isClassificationTaskType(taskType)) {
                     String classificationName = task.getClassificationName();
                     String entityGuid = task.getEntityGuid();
-                    String classificationId = resolveAndReturnClassificationId(classificationName, entityGuid);
+                    String classificationId = StringUtils.isEmpty(task.getClassificationId()) ? resolveAndReturnClassificationId(classificationName, entityGuid) : task.getClassificationId();
                     if (StringUtils.isEmpty(classificationId)) {
                         throw new AtlasBaseException(AtlasErrorCode.TASK_INVALID_PARAMETERS, task.toString());
                     }