Merge pull request #3889 from atlanhq/DG-1893

DG-1893
atlanhq · Dec 13, 2024 · 9602288 · 9602288
2 parents 554b20c + 45c1c3d
commit 9602288
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/...ry/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java b/...ry/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java
@@ -51,6 +51,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -245,13 +246,16 @@ private Map<String, List<AtlasClassification>> validateAndTransfer(AtlasEntityHe
             Map<String, List<AtlasClassification>> operationListMap = new HashMap<>();
 
             Set<String> requiredClassificationKeys = Stream.concat(
-                            incomingEntityHeader.getRemoveClassifications().stream(),
-                            incomingEntityHeader.getUpdateClassifications().stream()
-                    ).map(this::generateClassificationComparisonKey)
+                            Optional.ofNullable(incomingEntityHeader.getRemoveClassifications()).orElse(Collections.emptyList()).stream(),
+                            Optional.ofNullable(incomingEntityHeader.getUpdateClassifications()).orElse(Collections.emptyList()).stream()
+                    ).filter(classification -> classification.getEntityGuid().equals(entityToBeChanged.getGuid()))
+                    .map(this::generateClassificationComparisonKey)
                     .collect(Collectors.toSet());
 
-            Set<String> preExistingClassificationKeys = entityToBeChanged.getClassifications()
+            Set<String> preExistingClassificationKeys = Optional.ofNullable(entityToBeChanged.getClassifications())
+                    .orElse(Collections.emptyList())
                     .stream()
+                    .filter(classification -> classification.getEntityGuid().equals(entityToBeChanged.getGuid()))
                     .map(this::generateClassificationComparisonKey)
                     .collect(Collectors.toSet());
 
@@ -264,8 +268,10 @@ private Map<String, List<AtlasClassification>> validateAndTransfer(AtlasEntityHe
                 throw new AtlasBaseException(AtlasErrorCode.CLASSIFICATION_NOT_ASSOCIATED_WITH_ENTITY, firstTypeName);
             }
 
-            List<AtlasClassification> filteredClassifications = incomingEntityHeader.getAppendClassifications()
+            List<AtlasClassification> filteredClassifications = Optional.ofNullable(incomingEntityHeader.getAppendClassifications())
+                    .orElse(Collections.emptyList())
                     .stream()
+                    .filter(classification -> classification.getEntityGuid().equals(entityToBeChanged.getGuid()))
                     .filter(appendClassification -> !preExistingClassificationKeys.contains(generateClassificationComparisonKey(appendClassification)))
                     .collect(Collectors.toList());