Revery changes from ClassificationAssociator

atlanhq · Dec 21, 2023 · 3faa2a0 · 3faa2a0
1 parent b7dcd5e
commit 3faa2a0
Showing 1 changed file with 26 additions and 8 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
@@ -193,24 +193,42 @@ public void setClassifications(Map<String, AtlasEntityHeader> map) throws AtlasB
 
             AtlasPerfMetrics.MetricRecorder recorder = RequestContext.get().startMetricRecord("commitChanges.notify");
             Map<AtlasClassification, Collection<Object>> deleted = RequestContext.get().getDeletedClassificationAndVertices();
+            Set<AtlasVertex> allVertices = new HashSet<>();
 
             if (MapUtils.isNotEmpty(deleted)) {
-                Map<AtlasEntity, List<AtlasClassification>> entityClassification = getEntityClassificationsMapping(deleted);
-
-                for (Map.Entry<AtlasEntity, List<AtlasClassification>> atlasEntityListEntry : entityClassification.entrySet()) {
-                    entityChangeNotifier.onClassificationDeletedFromEntity(atlasEntityListEntry.getKey(), atlasEntityListEntry.getValue());
+                for (AtlasClassification deletedClassification : deleted.keySet()) {
+                    Collection<Object> vertices = deleted.get(deletedClassification);
+                    List<AtlasEntity> propagatedEntities = new ArrayList<>();
+
+                    for (Object obj : vertices) {
+                        AtlasVertex vertex = (AtlasVertex) obj;
+                        AtlasEntity entity = instanceConverter.getAndCacheEntity(GraphHelper.getGuid(vertex), IGNORE_REL);
+
+                        allVertices.add(vertex);
+                        propagatedEntities.add(entity);
+                    }
+                    entityChangeNotifier.onClassificationsDeletedFromEntities(propagatedEntities, Collections.singletonList(deletedClassification));
                 }
             }
 
             Map<AtlasClassification, Collection<Object>> added = RequestContext.get().getAddedClassificationAndVertices();
             if (MapUtils.isNotEmpty(added)) {
-                Map<AtlasEntity, List<AtlasClassification>> entityClassification = getEntityClassificationsMapping(added);
-
-                for (Map.Entry<AtlasEntity, List<AtlasClassification>> atlasEntityListEntry : entityClassification.entrySet()) {
-                    entityChangeNotifier.onClassificationAddedToEntity(atlasEntityListEntry.getKey(), atlasEntityListEntry.getValue());
+                for (AtlasClassification addedClassification : added.keySet()) {
+                    Collection<Object> vertices = added.get(addedClassification);
+                    List<AtlasEntity> propagatedEntities = new ArrayList<>();
+
+                    for (Object obj : vertices) {
+                        AtlasVertex vertex = (AtlasVertex) obj;
+                        AtlasEntity entity = instanceConverter.getAndCacheEntity(GraphHelper.getGuid(vertex), IGNORE_REL);
+
+                        allVertices.add(vertex);
+                        propagatedEntities.add(entity);
+                    }
+                    entityChangeNotifier.onClassificationsAddedToEntities(propagatedEntities, Collections.singletonList(addedClassification), false);
                 }
             }
 
+            entityGraphMapper.updateClassificationText(null, allVertices);
             transactionInterceptHelper.intercept();
             RequestContext.get().endMetricRecord(recorder);
             RequestContext.get().setDelayTagNotifications(false);