Merge pull request #2732 from atlanhq/dev/feat/classificationEvents

Update Classification add/delete behaviour
atlanhq · Jan 6, 2024 · 539af91 · 539af91
2 parents da8cd17 + a6b6b6e
commit 539af91
Show file tree

Hide file tree

Showing 6 changed files with 145 additions and 116 deletions.
diff --git a/intg/src/main/java/org/apache/atlas/model/audit/EntityAuditEventV2.java b/intg/src/main/java/org/apache/atlas/model/audit/EntityAuditEventV2.java
@@ -136,6 +136,9 @@ public static EntityAuditActionV2 fromString(String strValue) {
 
     private AtlasEntityHeader   entityDetail;
     private Map<String, String> headers;
+    private List<Map<String,Object>> classificationDetails;
+    @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
+    private String classificationDetail;
 
     public EntityAuditEventV2() { }
 
@@ -291,12 +294,13 @@ public boolean equals(Object o) {
                Objects.equals(created, that.created) &&
                Objects.equals(typeName, that.typeName) &&
                Objects.equals(entityQualifiedName, that.entityQualifiedName) &&
-               Objects.equals(headers, that.headers);
+                Objects.equals(headers, that.headers) &&
+                Objects.equals(classificationDetails, that.classificationDetails);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(entityId, timestamp, user, action, details, eventKey, entity, type, detail, created, entityQualifiedName, typeName, headers);
+        return Objects.hash(entityId, timestamp, user, action, details, eventKey, entity, type, detail, created, entityQualifiedName, typeName, headers, classificationDetails);
     }
 
     @Override
@@ -316,6 +320,7 @@ public String toString() {
         sb.append(", detail=").append(detail);
         sb.append(", created=").append(created);
         sb.append(", headers=").append(headers);
+        sb.append(", classificationDetails").append(classificationDetails);
         sb.append('}');
 
         return sb.toString();
@@ -347,6 +352,7 @@ public void clear() {
         detail = null;
         created = 0L;
         headers = null;
+        classificationDetails = null;
     }
 
     private String getJsonPartFromDetails() {
@@ -416,4 +422,20 @@ public static void sortEvents(List<EntityAuditEventV2> events, String sortByColu
 
         events.sort(sortOrderDesc ? comparator.reversed() : comparator);
     }
+
+    public List<Map<String, Object>> getClassificationDetails() {
+        return classificationDetails;
+    }
+
+    public void setClassificationDetails(List<Map<String, Object>> classificationDetails) {
+        this.classificationDetails = classificationDetails;
+    }
+
+    public String getClassificationDetail() {
+        return classificationDetail;
+    }
+
+    public void setClassificationDetail(String classificationDetail) {
+        this.classificationDetail = classificationDetail;
+    }
 }
diff --git a/intg/src/main/java/org/apache/atlas/type/AtlasTypeRegistry.java b/intg/src/main/java/org/apache/atlas/type/AtlasTypeRegistry.java
@@ -85,6 +85,10 @@ public AtlasType getType(String typeName) throws AtlasBaseException {
             LOG.debug("==> AtlasTypeRegistry.getType({})", typeName);
         }
 
+        if (typeName == null) {
+            return null;
+        }
+
         AtlasType ret = registryData.allTypes.getTypeByName(typeName);
 
         if (ret == null) {

diff --git a/repository/src/main/java/org/apache/atlas/repository/audit/ESBasedAuditRepository.java b/repository/src/main/java/org/apache/atlas/repository/audit/ESBasedAuditRepository.java
@@ -84,6 +84,7 @@ public class ESBasedAuditRepository extends AbstractStorageBasedAuditRepository
     private static final String USER = "user";
     private static final String DETAIL = "detail";
     private static final String ENTITY = "entity";
+    private static final String CLASSIFICATION_DETAIL= "classificationDetail";
     private static final String bulkMetadata = String.format("{ \"index\" : { \"_index\" : \"%s\" } }%n", INDEX_NAME);
 
     /*
@@ -134,7 +135,8 @@ public void putEventsV2(List<EntityAuditEventV2> events) throws AtlasBaseExcepti
                             event.getEntityQualifiedName(),
                             event.getEntity().getTypeName(),
                             created,
-                            "" + event.getEntity().getUpdateTime().getTime());
+                            "" + event.getEntity().getUpdateTime().getTime(),
+                            event.getClassificationDetail());
 
                     bulkRequestBody.append(bulkMetadata);
                     bulkRequestBody.append(bulkItem);
@@ -174,7 +176,7 @@ private String getQueryTemplate(Map<String, String> requestContextHeaders) {
         StringBuilder template = new StringBuilder();
 
         template.append("'{'\"entityId\":\"{0}\",\"action\":\"{1}\",\"detail\":{2},\"user\":\"{3}\", \"eventKey\":\"{4}\", " +
-                        "\"entityQualifiedName\": {5}, \"typeName\": \"{6}\",\"created\":{7}, \"timestamp\":{8}");
+                        "\"entityQualifiedName\": {5}, \"typeName\": \"{6}\",\"created\":{7}, \"timestamp\":{8}, \"classificationDetail\":{9}");
 
         if (MapUtils.isNotEmpty(requestContextHeaders)) {
             template.append(",")
@@ -226,7 +228,14 @@ private EntityAuditSearchResult getResultFromResponse(String responseString) thr
             EntityAuditEventV2 event = new EntityAuditEventV2();
             event.setEntityId(entityGuid);
             event.setAction(EntityAuditEventV2.EntityAuditActionV2.fromString((String) source.get(ACTION)));
-            event.setDetail((Map<String, Object>) source.get(DETAIL));
+            if (source.get(DETAIL) != null) {
+                if (source.get(DETAIL) instanceof Map) {
+                    event.setDetail((Map<String, Object>) source.get(DETAIL));
+                }
+            }
+            if (source.get(CLASSIFICATION_DETAIL) instanceof List) {
+                event.setClassificationDetails((List<Map<String, Object>>) source.get(CLASSIFICATION_DETAIL));
+            }
             event.setUser((String) source.get(USER));
             event.setCreated((long) source.get(CREATED));
             if (source.get(TIMESTAMP) != null) {