ATLAS-3806: Classifications information missing in notification events during entity update/delete

intg/src/main/java/org/apache/atlas/model/instance/EntityMutationResponse.java

@@ -215,10 +215,10 @@ public class EntityMutationResponse {
 
     @JsonIgnore
     public void addEntity(EntityOperation op, AtlasEntityHeader header) {
-        // if an entity is already included in CREATE, ignore subsequent UPDATE, PARTIAL_UPDATE
+        // if an entity is already included in CREATE, update the header, to capture propagated classifications
         if (op == EntityOperation.UPDATE || op == EntityOperation.PARTIAL_UPDATE) {
             if (entityHeaderExists(getCreatedEntities(), header.getGuid())) {
-                return;
+                op = EntityOperation.CREATE;
             }
         }
 

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -277,7 +277,7 @@ public class EntityGraphMapper {
                 mapAttributes(createdEntity, entityType, vertex, CREATE, context);
                 setCustomAttributes(vertex,createdEntity);
 
-                resp.addEntity(CREATE, constructHeader(createdEntity, entityType, vertex));
+                resp.addEntity(CREATE, constructHeader(createdEntity, vertex));
                 addClassifications(context, guid, createdEntity.getClassifications());
 
                 addOrUpdateBusinessAttributes(vertex, entityType, createdEntity.getBusinessAttributes());
@@ -299,7 +299,7 @@ public class EntityGraphMapper {
                 mapAttributes(updatedEntity, entityType, vertex, updateType, context);
                 setCustomAttributes(vertex,updatedEntity);
 
-                resp.addEntity(updateType, constructHeader(updatedEntity, entityType, vertex));
+                resp.addEntity(updateType, constructHeader(updatedEntity, vertex));
 
                 if (replaceClassifications) {
                     deleteClassifications(guid);
@@ -1853,15 +1853,10 @@ public class EntityGraphMapper {
         }
     }
 
-    private AtlasEntityHeader constructHeader(AtlasEntity entity, final AtlasEntityType type, AtlasVertex vertex) {
-        AtlasEntityHeader header = new AtlasEntityHeader(entity.getTypeName());
-
-        header.setGuid(getIdFromVertex(vertex));
-        header.setStatus(entity.getStatus());
-        header.setIsIncomplete(entity.getIsIncomplete());
-
-        for (AtlasAttribute attribute : type.getUniqAttributes().values()) {
-            header.setAttribute(attribute.getName(), entity.getAttribute(attribute.getName()));
+    private AtlasEntityHeader constructHeader(AtlasEntity entity, AtlasVertex vertex) throws AtlasBaseException {
+        AtlasEntityHeader header = entityRetriever.toAtlasEntityHeaderWithClassifications(vertex);
+        if (entity.getClassifications() == null) {
+            entity.setClassifications(header.getClassifications());
         }
 
         return header;