ATLAS-3492: updated object-id attributes in audit logs to replace…

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

@@ -597,6 +597,8 @@ public class EntityGraphMapper {
                                                       true, ctx.getAttribute().getRelationshipEdgeDirection(), ctx.getReferringVertex());
                 }
 
+                setAssignedGuid(ctx.getValue(), context.getGuidAssignments());
+
                 return newEdge;
             }
 
@@ -626,6 +628,8 @@ public class EntityGraphMapper {
             }
         }
 
+        setAssignedGuid(ctx.getValue(), context.getGuidAssignments());
+
         return ret;
     }
 
@@ -886,12 +890,18 @@ public class EntityGraphMapper {
         AtlasVertex entityVertex = context.getDiscoveryContext().getResolvedEntityVertex(guid);
 
         if (entityVertex == null) {
+            if (AtlasTypeUtil.isAssignedGuid(guid)) {
+                entityVertex = context.getVertex(guid);
+            }
+
+            if (entityVertex == null) {
                 AtlasObjectId objId = getObjectId(ctx.getValue());
 
                 if (objId != null) {
                     entityVertex = context.getDiscoveryContext().getResolvedEntityVertex(objId);
                 }
             }
+        }
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, (ctx.getValue() == null ? null : ctx.getValue().toString()));
@@ -921,15 +931,22 @@ public class EntityGraphMapper {
             LOG.debug("==> mapObjectIdValueUsingRelationship({})", ctx);
         }
 
-        AtlasVertex attributeVertex = context.getDiscoveryContext().getResolvedEntityVertex(getGuid(ctx.getValue()));
+        String      guid            = getGuid(ctx.getValue());
+        AtlasVertex attributeVertex = context.getDiscoveryContext().getResolvedEntityVertex(guid);
         AtlasVertex entityVertex    = ctx.getReferringVertex();
         AtlasEdge   ret;
 
         if (attributeVertex == null) {
+            if (AtlasTypeUtil.isAssignedGuid(guid)) {
+                attributeVertex = context.getVertex(guid);
+            }
+
+            if (attributeVertex == null) {
                 AtlasObjectId objectId = getObjectId(ctx.getValue());
 
                 attributeVertex = (objectId != null) ? context.getDiscoveryContext().getResolvedEntityVertex(objectId) : null;
             }
+        }
 
         if (attributeVertex == null) {
             if(RequestContext.get().isImportInProgress()) {
@@ -1287,6 +1304,39 @@ public class EntityGraphMapper {
         return null;
     }
 
+    private static void setAssignedGuid(Object val, Map<String, String> guidAssignements) {
+        if (val != null && MapUtils.isNotEmpty(guidAssignements)) {
+            if (val instanceof AtlasObjectId) {
+                AtlasObjectId objId = (AtlasObjectId) val;
+                String        guid  = objId.getGuid();
+
+                if (StringUtils.isNotEmpty(guid) && !AtlasTypeUtil.isAssignedGuid(guid)) {
+                    String assignedGuid = guidAssignements.get(guid);
+
+                    if (StringUtils.isNotEmpty(assignedGuid)) {
+                        RequestContext.get().recordEntityGuidUpdate(objId, guid);
+
+                        objId.setGuid(assignedGuid);
+                    }
+                }
+            } else if (val instanceof Map) {
+                Map    objId   = (Map) val;
+                Object guidVal = objId.get(AtlasObjectId.KEY_GUID);
+                String guid    = objId != null ? guidVal.toString() : null;
+
+                if (StringUtils.isNotEmpty(guid) && !AtlasTypeUtil.isAssignedGuid(guid)) {
+                    String assignedGuid = guidAssignements.get(guid);
+
+                    if (StringUtils.isNotEmpty(assignedGuid)) {
+                        RequestContext.get().recordEntityGuidUpdate(objId, guid);
+
+                        objId.put(AtlasObjectId.KEY_GUID, assignedGuid);
+                    }
+                }
+            }
+        }
+    }
+
     private static Map<String, Object> getRelationshipAttributes(Object val) throws AtlasBaseException {
         if (val instanceof AtlasRelatedObjectId) {
             AtlasStruct relationshipStruct = ((AtlasRelatedObjectId) val).getRelationshipAttributes();

server-api/src/main/java/org/apache/atlas/RequestContext.java

@@ -22,6 +22,7 @@ import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
+import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.store.DeleteType;
 import org.apache.atlas.utils.AtlasPerfMetrics;
 import org.apache.atlas.utils.AtlasPerfMetrics.MetricRecorder;
@@ -37,6 +38,8 @@ import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.HashMap;
 
+import static org.apache.atlas.model.instance.AtlasObjectId.KEY_GUID;
+
 
 public class RequestContext {
     private static final Logger METRICS = LoggerFactory.getLogger("METRICS");
@@ -337,11 +340,23 @@ public class RequestContext {
     }
 
     public void recordEntityGuidUpdate(AtlasEntity entity, String guidInRequest) {
+        recordEntityGuidUpdate(new EntityGuidPair(entity, guidInRequest));
+    }
+
+    public void recordEntityGuidUpdate(AtlasObjectId entity, String guidInRequest) {
+        recordEntityGuidUpdate(new EntityGuidPair(entity, guidInRequest));
+    }
+
+    public void recordEntityGuidUpdate(Map entity, String guidInRequest) {
+        recordEntityGuidUpdate(new EntityGuidPair(entity, guidInRequest));
+    }
+
+    public void recordEntityGuidUpdate(EntityGuidPair record) {
         if (entityGuidInRequest == null) {
             entityGuidInRequest = new ArrayList<>();
         }
 
-        entityGuidInRequest.add(new EntityGuidPair(entity, guidInRequest));
+        entityGuidInRequest.add(record);
     }
 
     public void resetEntityGuidUpdates() {
@@ -353,7 +368,7 @@ public class RequestContext {
     }
 
     public class EntityGuidPair {
-        private final AtlasEntity entity;
+        private final Object entity;
         private final String guid;
 
         public EntityGuidPair(AtlasEntity entity, String guid) {
@@ -361,8 +376,24 @@ public class RequestContext {
             this.guid   = guid;
         }
 
+        public EntityGuidPair(AtlasObjectId entity, String guid) {
+            this.entity = entity;
+            this.guid   = guid;
+        }
+
+        public EntityGuidPair(Map entity, String guid) {
+            this.entity = entity;
+            this.guid   = guid;
+        }
+
         public void resetEntityGuid() {
-            entity.setGuid(guid);
+            if (entity instanceof AtlasEntity) {
+                ((AtlasEntity) entity).setGuid(guid);
+            } else if (entity instanceof AtlasObjectId) {
+                ((AtlasObjectId) entity).setGuid(guid);
+            } else if (entity instanceof Map) {
+                ((Map) entity).put(KEY_GUID, guid);
+            }
         }
     }