diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java index 83fd29c79d..292ce1de09 100644 --- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java +++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java @@ -193,24 +193,42 @@ public void setClassifications(Map map) throws AtlasB AtlasPerfMetrics.MetricRecorder recorder = RequestContext.get().startMetricRecord("commitChanges.notify"); Map> deleted = RequestContext.get().getDeletedClassificationAndVertices(); + Set allVertices = new HashSet<>(); if (MapUtils.isNotEmpty(deleted)) { - Map> entityClassification = getEntityClassificationsMapping(deleted); - - for (Map.Entry> atlasEntityListEntry : entityClassification.entrySet()) { - entityChangeNotifier.onClassificationDeletedFromEntity(atlasEntityListEntry.getKey(), atlasEntityListEntry.getValue()); + for (AtlasClassification deletedClassification : deleted.keySet()) { + Collection vertices = deleted.get(deletedClassification); + List 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> added = RequestContext.get().getAddedClassificationAndVertices(); if (MapUtils.isNotEmpty(added)) { - Map> entityClassification = getEntityClassificationsMapping(added); - - for (Map.Entry> atlasEntityListEntry : entityClassification.entrySet()) { - entityChangeNotifier.onClassificationAddedToEntity(atlasEntityListEntry.getKey(), atlasEntityListEntry.getValue()); + for (AtlasClassification addedClassification : added.keySet()) { + Collection vertices = added.get(addedClassification); + List 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);