ATLAS-1544: implementation of REST endpoints for entity create/update/bulk-get

client/src/main/java/org/apache/atlas/AtlasEntitiesClientV2.java

@@ -17,7 +17,6 @@
  */
 package org.apache.atlas;
 
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -25,11 +24,10 @@ import javax.ws.rs.HttpMethod;
 import javax.ws.rs.core.MultivaluedMap;
 import javax.ws.rs.core.Response;
 
-import org.apache.atlas.model.SearchFilter;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasClassification.AtlasClassifications;
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasEntity.AtlasEntities;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.commons.collections.MapUtils;
@@ -37,34 +35,33 @@ import org.apache.commons.configuration.Configuration;
 import org.apache.hadoop.security.UserGroupInformation;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.sun.jersey.api.client.GenericType;
 import com.sun.jersey.api.client.WebResource;
 import com.sun.jersey.core.util.MultivaluedMapImpl;
 
 public class AtlasEntitiesClientV2 extends AtlasBaseClient {
-
-    private static final GenericType<List<AtlasEntity>> ENTITY_LIST_TYPE = new GenericType<List<AtlasEntity>>(){};
     public static final String ENTITY_API = BASE_URI + "v2/entity/";
-    public static final String ENTITIES_API = BASE_URI + "v2/entities/";
 
-    private static final APIInfo GET_ENTITY_BY_GUID = new APIInfo(ENTITY_API + "guid/", HttpMethod.GET, Response.Status.OK);
-    private static final APIInfo CREATE_ENTITY = new APIInfo(ENTITIES_API, HttpMethod.POST, Response.Status.OK);
-    private static final APIInfo UPDATE_ENTITY = CREATE_ENTITY;
+    private static final APIInfo GET_ENTITY_BY_GUID         = new APIInfo(ENTITY_API + "guid/%s", HttpMethod.GET, Response.Status.OK);
     private static final APIInfo GET_ENTITY_BY_ATTRIBUTE    = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s", HttpMethod.GET, Response.Status.OK);
-    private static final APIInfo UPDATE_ENTITY_BY_ATTRIBUTE = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s/attribute/%s", HttpMethod.PUT, Response.Status.OK);
-    private static final APIInfo DELETE_ENTITY_BY_ATTRIBUTE = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s/attribute/%s", HttpMethod.DELETE, Response.Status.OK);
-    private static final APIInfo DELETE_ENTITY_BY_GUID = new APIInfo(ENTITY_API + "guid/", HttpMethod.DELETE, Response.Status.OK);
-    private static final APIInfo DELETE_ENTITY_BY_GUIDS = new APIInfo(ENTITIES_API + "guids/", HttpMethod.DELETE, Response.Status.OK);
+    private static final APIInfo CREATE_ENTITY              = new APIInfo(ENTITY_API, HttpMethod.POST, Response.Status.OK);
+    private static final APIInfo UPDATE_ENTITY              = CREATE_ENTITY;
+    private static final APIInfo UPDATE_ENTITY_BY_ATTRIBUTE = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s", HttpMethod.PUT, Response.Status.OK);
+    private static final APIInfo DELETE_ENTITY_BY_GUID      = new APIInfo(ENTITY_API + "guid/%s", HttpMethod.DELETE, Response.Status.OK);
+    private static final APIInfo DELETE_ENTITY_BY_ATTRIBUTE = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s", HttpMethod.DELETE, Response.Status.OK);
+
+    private static final APIInfo GET_ENTITIES_BY_GUIDS    = new APIInfo(ENTITY_API + "bulk/", HttpMethod.GET, Response.Status.OK);
+    private static final APIInfo CREATE_ENTITIES          = new APIInfo(ENTITY_API + "bulk/", HttpMethod.POST, Response.Status.OK);
+    private static final APIInfo UPDATE_ENTITIES          = CREATE_ENTITIES;
+    private static final APIInfo DELETE_ENTITIES_BY_GUIDS = new APIInfo(ENTITY_API + "bulk/", HttpMethod.DELETE, Response.Status.OK);
+
     private static final APIInfo GET_CLASSIFICATIONS    = new APIInfo(ENTITY_API + "guid/%s/classifications", HttpMethod.GET, Response.Status.OK);
     private static final APIInfo ADD_CLASSIFICATIONS    = new APIInfo(ENTITY_API + "guid/%s/classifications", HttpMethod.POST, Response.Status.NO_CONTENT);
-
     private static final APIInfo UPDATE_CLASSIFICATIONS = new APIInfo(ENTITY_API + "guid/%s/classifications", HttpMethod.PUT, Response.Status.OK);
     private static final APIInfo DELETE_CLASSIFICATION  = new APIInfo(ENTITY_API + "guid/%s/classification/%s", HttpMethod.DELETE, Response.Status.NO_CONTENT);
-    private static final APIInfo GET_ENTITIES = new APIInfo(ENTITIES_API + "guids/", HttpMethod.GET, Response.Status.OK);
-    private static final APIInfo CREATE_ENTITIES = new APIInfo(ENTITIES_API, HttpMethod.POST, Response.Status.OK);
-    private static final APIInfo UPDATE_ENTITIES = CREATE_ENTITIES;
-    private static final APIInfo DELETE_ENTITIES = new APIInfo(ENTITIES_API + "guids/", HttpMethod.GET, Response.Status.NO_CONTENT);
-    private static final APIInfo SEARCH_ENTITIES = new APIInfo(ENTITIES_API, HttpMethod.GET, Response.Status.OK);
+
+
+    public static final String PREFIX_ATTR = "attr:";
+
 
     public AtlasEntitiesClientV2(String[] baseUrl, String[] basicAuthUserNamePassword) {
         super(baseUrl, basicAuthUserNamePassword);
@@ -88,50 +85,61 @@ public class AtlasEntitiesClientV2 extends AtlasBaseClient {
     }
 
     public AtlasEntityWithExtInfo getEntityByGuid(String guid) throws AtlasServiceException {
-
-        return callAPI(GET_ENTITY_BY_GUID, null, AtlasEntityWithExtInfo.class, guid);
+        return callAPI(formatPathForPathParams(GET_ENTITY_BY_GUID, guid), AtlasEntityWithExtInfo.class, null);
     }
 
-    public AtlasEntities getEntityByGuids(List<String> guids) throws AtlasServiceException {
-        return callAPI(GET_ENTITY_BY_GUID, AtlasEntities.class, "guid", guids);
+    public AtlasEntityWithExtInfo getEntityByAttribute(String type, Map<String, String> attributes) throws AtlasServiceException {
+        MultivaluedMap<String, String> queryParams = attributesToQueryParams(attributes);
+
+        return callAPI(formatPathForPathParams(GET_ENTITY_BY_ATTRIBUTE, type), AtlasEntityWithExtInfo.class, queryParams);
     }
 
-    public static final String PREFIX_ATTR = "attr:";
+    public EntityMutationResponse updateEntityByAttribute(String type, Map<String, String> attributes, AtlasEntity entity) throws AtlasServiceException {
+        MultivaluedMap<String, String> queryParams = attributesToQueryParams(attributes);
 
-    public AtlasEntityWithExtInfo getEntityByAttribute(String type, Map<String, String> attributes) throws AtlasServiceException {
+        return callAPI(formatPathForPathParams(UPDATE_ENTITY_BY_ATTRIBUTE, type), entity, EntityMutationResponse.class, queryParams);
+    }
+
+    public EntityMutationResponse deleteEntityByAttribute(String type, Map<String, String> attributes) throws AtlasServiceException {
         MultivaluedMap<String, String> queryParams = attributesToQueryParams(attributes);
 
-        return callAPI(formatPathForPathParams(GET_ENTITY_BY_ATTRIBUTE, type), AtlasEntityWithExtInfo.class, queryParams);
+        return callAPI(formatPathForPathParams(DELETE_ENTITY_BY_ATTRIBUTE, type), null, EntityMutationResponse.class, queryParams);
     }
 
-    public EntityMutationResponse updateEntityByAttribute(String type, String attribute, String value, AtlasEntity entity) throws AtlasServiceException {
-        MultivaluedMap<String, String> queryParams = new MultivaluedMapImpl();
-        queryParams.add("value", value);
-        return callAPI(formatPathForPathParams(UPDATE_ENTITY_BY_ATTRIBUTE, type, attribute), entity, EntityMutationResponse.class, queryParams);
+    public EntityMutationResponse createEntity(AtlasEntityWithExtInfo entity) throws AtlasServiceException {
+        return callAPI(CREATE_ENTITY, entity, EntityMutationResponse.class);
     }
 
-    public EntityMutationResponse deleteEntityByAttribute(String type, String attribute, String value) throws AtlasServiceException {
-        MultivaluedMap<String, String> queryParams = new MultivaluedMapImpl();
-        queryParams.add("value", value);
-        return callAPI(formatPathForPathParams(DELETE_ENTITY_BY_ATTRIBUTE, type, attribute), null, EntityMutationResponse.class, queryParams);
+    public EntityMutationResponse updateEntity(AtlasEntityWithExtInfo entity) throws AtlasServiceException {
+        return callAPI(UPDATE_ENTITY, entity, EntityMutationResponse.class);
+    }
+
+    public EntityMutationResponse deleteEntityByGuid(String guid) throws AtlasServiceException {
+        return callAPI(formatPathForPathParams(DELETE_ENTITY_BY_GUID, guid), EntityMutationResponse.class, null);
     }
 
-    public EntityMutationResponse createEntity(final AtlasEntity atlasEntity) throws AtlasServiceException {
-        return callAPI(CREATE_ENTITY, new HashMap<String, AtlasEntity>(1) {{ put(atlasEntity.getGuid(), atlasEntity); }}, EntityMutationResponse.class);
+
+    public AtlasEntitiesWithExtInfo getEntitiesByGuids(List<String> guids) throws AtlasServiceException {
+        MultivaluedMap<String, String> queryParams = new MultivaluedMapImpl();
+
+        queryParams.put("guid", guids);
+
+        return callAPI(GET_ENTITIES_BY_GUIDS, AtlasEntitiesWithExtInfo.class, queryParams);
     }
 
-    public EntityMutationResponse updateEntity(final AtlasEntity atlasEntity) throws AtlasServiceException {
-        return callAPI(UPDATE_ENTITY, new HashMap<String, AtlasEntity>(1) {{ put(atlasEntity.getGuid(), atlasEntity); }}, EntityMutationResponse.class);
+    public EntityMutationResponse createEntities(AtlasEntitiesWithExtInfo atlasEntities) throws AtlasServiceException {
+        return callAPI(CREATE_ENTITIES, atlasEntities, EntityMutationResponse.class);
     }
 
-    public AtlasEntity deleteEntityByGuid(String guid) throws AtlasServiceException {
-        return callAPI(DELETE_ENTITY_BY_GUID, null, AtlasEntity.class, guid);
+    public EntityMutationResponse updateEntities(AtlasEntitiesWithExtInfo atlasEntities) throws AtlasServiceException {
+        return callAPI(UPDATE_ENTITIES, atlasEntities, EntityMutationResponse.class);
     }
 
-    public EntityMutationResponse deleteEntityByGuid(List<String> guids) throws AtlasServiceException {
-        return callAPI(DELETE_ENTITY_BY_GUIDS, EntityMutationResponse.class, "guid", guids);
+    public EntityMutationResponse deleteEntitiesByGuids(List<String> guids) throws AtlasServiceException {
+        return callAPI(DELETE_ENTITIES_BY_GUIDS, EntityMutationResponse.class, "guid", guids);
     }
 
+
     public AtlasClassifications getClassifications(String guid) throws AtlasServiceException {
         return callAPI(formatPathForPathParams(GET_CLASSIFICATIONS, guid), null, AtlasClassifications.class);
     }
@@ -158,26 +166,6 @@ public class AtlasEntitiesClientV2 extends AtlasBaseClient {
         return null;
     }
 
-    public EntityMutationResponse createEntities(List<AtlasEntity> atlasEntities) throws AtlasServiceException {
-
-        return callAPI(CREATE_ENTITIES, entityListToMap(atlasEntities), EntityMutationResponse.class);
-    }
-
-    private Map<String, AtlasEntity> entityListToMap(List<AtlasEntity> atlasEntities) {
-        Map<String,AtlasEntity> toSend = new HashMap<String, AtlasEntity>(atlasEntities.size());
-        for(AtlasEntity entity : atlasEntities) {
-            toSend.put(entity.getGuid(), entity);
-        }
-        return toSend;
-    }
-
-    public EntityMutationResponse updateEntities(List<AtlasEntity> atlasEntities) throws AtlasServiceException {
-        return callAPI(UPDATE_ENTITIES, entityListToMap(atlasEntities), EntityMutationResponse.class);
-    }
-
-    public AtlasEntity.AtlasEntities searchEntities(SearchFilter searchFilter) throws AtlasServiceException {
-        return callAPI(GET_ENTITIES, AtlasEntity.AtlasEntities.class, searchFilter.getParams());
-    }
 
     private MultivaluedMap<String, String> attributesToQueryParams(Map<String, String> attributes) {
         return attributesToQueryParams(attributes, null);

intg/src/main/java/org/apache/atlas/AtlasErrorCode.java

@@ -64,6 +64,7 @@ public enum AtlasErrorCode {
     UNRESOLVED_REFERENCES_FOUND(400, "ATLAS40033E", "Unresolved references: byId={0}; byUniqueAttributes={1}"),
     UNKNOWN_ATTRIBUTE(400, "ATLAS40034E", "Attribute {0} not found for type {1}"),
     SYSTEM_TYPE(400, "ATLAS40035E", "{0} is a System-type"),
+    INVALID_STRUCT_VALUE(400, "ATLAS40036E", "not a valid struct value {0}"),
 
     // All Not found enums go here
     TYPE_NAME_NOT_FOUND(404, "ATLAS4041E", "Given typename {0} was invalid"),
@@ -74,8 +75,8 @@ public enum AtlasErrorCode {
     INSTANCE_LINEAGE_QUERY_FAILED(404, "ATLAS4047E", "Instance lineage query failed {0}"),
     DISCOVERY_QUERY_FAILED(404, "ATLAS4048E", "Discovery query failed {0}"),
     INSTANCE_CRUD_INVALID_PARAMS(404, "ATLAS4049E", "Invalid instance creation/updation parameters passed : {0}"),
-    INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND(404, "ATLAS40410E", "Instance {0} with unique attribute {1}={2} does not exist"),
-
+    INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND(404, "ATLAS40410E", "Instance {0} with unique attribute {1} does not exist"),
+    REFERENCED_ENTITY_NOT_FOUND(404, "ATLAS40411E", "Referenced entity {0} is not found"),
 
     // All data conflict errors go here
     TYPE_ALREADY_EXISTS(409, "ATLAS4091E", "Given type {0} already exists"),

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

@@ -34,6 +34,7 @@ import javax.xml.bind.annotation.XmlRootElement;
 import javax.xml.bind.annotation.XmlSeeAlso;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
@@ -87,16 +88,16 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
         this(entityDef != null ? entityDef.getName() : null, null);
     }
 
+    public AtlasEntity(String typeName, String attrName, Object attrValue) {
+        super(typeName, attrName, attrValue);
+
+        init();
+    }
+
     public AtlasEntity(String typeName, Map<String, Object> attributes) {
         super(typeName, attributes);
 
-        setGuid(nextInternalId());
-        setStatus(null);
-        setCreatedBy(null);
-        setUpdatedBy(null);
-        setCreateTime(null);
-        setUpdateTime(null);
-        setClassifications(null);
+        init();
     }
 
     public AtlasEntity(AtlasEntity other) {
@@ -186,13 +187,15 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
 
     @JsonIgnore
     public static boolean isAssigned(String guid) {
+        if (guid != null) {
             try {
                 UUID.fromString(guid);
+                return true;
             } catch (IllegalArgumentException e) {
-            return false;
+                // ignore
             }
-
-        return true;
+        }
+        return false;
     }
 
     @JsonIgnore
@@ -200,13 +203,23 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
         return guid != null && guid.length() > 0 && guid.charAt(0) == '-';
     }
 
+    private void init() {
+        setGuid(nextInternalId());
+        setStatus(null);
+        setCreatedBy(null);
+        setUpdatedBy(null);
+        setCreateTime(null);
+        setUpdateTime(null);
+        setClassifications(null);
+    }
+
     private static String nextInternalId() {
         return "-" + Long.toString(s_nextId.getAndIncrement());
     }
 
     @JsonIgnore
     public AtlasObjectId getAtlasObjectId() {
-        return new AtlasObjectId(getTypeName(), getGuid());
+        return new AtlasObjectId(getGuid(), getTypeName());
     }
 
     @Override
@@ -279,6 +292,10 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
             setReferredEntities(null);
         }
 
+        public AtlasEntityExtInfo(AtlasEntity referredEntity) {
+            addReferredEntity(referredEntity);
+        }
+
         public AtlasEntityExtInfo(Map<String, AtlasEntity> referredEntities) {
             setReferredEntities(referredEntities);
         }
@@ -294,6 +311,11 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
         public void setReferredEntities(Map<String, AtlasEntity> referredEntities) { this.referredEntities = referredEntities; }
 
         @JsonIgnore
+        public final void addReferredEntity(AtlasEntity entity) {
+            addReferredEntity(entity.getGuid(), entity);
+        }
+
+        @JsonIgnore
         public final void addReferredEntity(String guid, AtlasEntity entity) {
             Map<String, AtlasEntity> r = this.referredEntities;
 
@@ -494,10 +516,17 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
             this(null, null);
         }
 
+        public AtlasEntitiesWithExtInfo(AtlasEntity entity) { this(Arrays.asList(entity), null);
+        }
+
         public AtlasEntitiesWithExtInfo(List<AtlasEntity> entities) {
             this(entities, null);
         }
 
+        public AtlasEntitiesWithExtInfo(AtlasEntityWithExtInfo entity) {
+            this(Arrays.asList(entity.getEntity()), entity);
+        }
+
         public AtlasEntitiesWithExtInfo(List<AtlasEntity> entities, AtlasEntityExtInfo extInfo) {
             super(extInfo);
 

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

@@ -20,6 +20,7 @@ package org.apache.atlas.model.instance;
 import java.io.Serializable;
 import java.util.List;
 import java.util.Map;
+import java.util.HashMap;
 import java.util.Objects;
 
 import javax.xml.bind.annotation.XmlAccessType;
@@ -51,88 +52,83 @@ import org.codehaus.jackson.map.annotate.JsonSerialize;
 public class AtlasObjectId  implements Serializable {
     private static final long serialVersionUID = 1L;
 
-    public static final String KEY_TYPENAME          = "typeName";
     public static final String KEY_GUID              = "guid";
+    public static final String KEY_TYPENAME          = "typeName";
     public static final String KEY_UNIQUE_ATTRIBUTES = "uniqueAttributes";
 
-    private String              typeName;
     private String              guid;
+    private String              typeName;
     private Map<String, Object> uniqueAttributes;
 
-    @JsonIgnore
-    private boolean isAssignedGuid = false;
-
-    @JsonIgnore
-    private boolean isUnAssignedGuid = false;
 
     public AtlasObjectId() {
-        this(null, null, null);
+        this(null, null, (Map<String, Object>)null);
     }
 
-    public AtlasObjectId(String typeName) {
-        this(typeName, null, null);
+    public AtlasObjectId(String guid) {
+        this(guid, null, (Map<String, Object>)null);
     }
 
-    public AtlasObjectId(String typeName, String guid) {
-        this(typeName, guid, null);
+    public AtlasObjectId(String guid, String typeName) {
+        this(guid, typeName, (Map<String, Object>)null);
     }
 
     public AtlasObjectId(String typeName, Map<String, Object> uniqueAttributes) {
-        this(typeName, null, uniqueAttributes);
+        this(null, typeName, uniqueAttributes);
     }
 
-    public AtlasObjectId(String typeName, String guid, Map<String, Object> uniqueAttributes) {
-        setTypeName(typeName);
+    public AtlasObjectId(String typeName, final String attrName, final Object attrValue) {
+        this(null, typeName, new HashMap<String, Object>() {{ put(attrName, attrValue); }});
+    }
+
+    public AtlasObjectId(String guid, String typeName, Map<String, Object> uniqueAttributes) {
         setGuid(guid);
+        setTypeName(typeName);
         setUniqueAttributes(uniqueAttributes);
     }
 
     public AtlasObjectId(AtlasObjectId other) {
         if (other != null) {
-            setTypeName(other.getTypeName());
             setGuid(other.getGuid());
+            setTypeName(other.getTypeName());
             setUniqueAttributes(other.getUniqueAttributes());
         }
     }
 
     public AtlasObjectId(Map objIdMap) {
         if (objIdMap != null) {
-            Object t = objIdMap.get(KEY_TYPENAME);
             Object g = objIdMap.get(KEY_GUID);
+            Object t = objIdMap.get(KEY_TYPENAME);
             Object u = objIdMap.get(KEY_UNIQUE_ATTRIBUTES);
 
-            if (t != null) {
-                setTypeName(t.toString());
-            }
-
             if (g != null) {
                 setGuid(g.toString());
             }
 
+            if (t != null) {
+                setTypeName(t.toString());
+            }
+
             if (u != null && u instanceof Map) {
                 setUniqueAttributes((Map)u);
             }
         }
     }
 
-    public String getTypeName() {
-        return typeName;
-    }
-
-    public void setTypeName(String typeName) {
-        this.typeName = typeName;
-    }
-
     public String getGuid() {
         return guid;
     }
 
     public void setGuid(String guid) {
         this.guid = guid;
-        if ( guid != null) {
-            this.isAssignedGuid = AtlasEntity.isAssigned(guid);
-            this.isUnAssignedGuid = AtlasEntity.isUnAssigned(guid);
     }
+
+    public String getTypeName() {
+        return typeName;
+    }
+
+    public void setTypeName(String typeName) {
+        this.typeName = typeName;
     }
 
     public Map<String, Object> getUniqueAttributes() {
@@ -144,28 +140,29 @@ public class AtlasObjectId  implements Serializable {
     }
 
     @JsonIgnore
+    public boolean isValidGuid() {
+        return isAssignedGuid() || isUnAssignedGuid();
+    }
+
+    @JsonIgnore
     public boolean isAssignedGuid() {
-        return isAssignedGuid;
+        return AtlasEntity.isAssigned(guid);
     }
 
     @JsonIgnore
     public boolean isUnAssignedGuid() {
-        return isUnAssignedGuid;
+        return AtlasEntity.isUnAssigned(guid);
     }
 
     @JsonIgnore
     public boolean isValid() {
-        if (StringUtils.isEmpty(typeName)) {
-            return false;
-        } else if (StringUtils.isNotEmpty(guid)) {
-            if (!isAssignedGuid && !isUnAssignedGuid) {
-                return false;
-            }
-        } else if (MapUtils.isEmpty(uniqueAttributes)) {
-            return false;
+        if (isAssignedGuid() || isUnAssignedGuid()) {
+            return true;
+        } else if (StringUtils.isNotEmpty(typeName) && MapUtils.isNotEmpty(uniqueAttributes)) {
+            return true;
         }
 
-        return true;
+        return false;
     }
 
     public StringBuilder toString(StringBuilder sb) {
@@ -174,8 +171,8 @@ public class AtlasObjectId  implements Serializable {
         }
 
         sb.append("AtlasObjectId{");
-        sb.append("typeName='").append(typeName).append('\'');
-        sb.append(", guid='").append(guid).append('\'');
+        sb.append("guid='").append(guid).append('\'');
+        sb.append(", typeName='").append(typeName).append('\'');
         sb.append(", uniqueAttributes={");
         AtlasBaseTypeDef.dumpObjects(uniqueAttributes, sb);
         sb.append('}');
@@ -186,17 +183,27 @@ public class AtlasObjectId  implements Serializable {
 
     @Override
     public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
+        if (this == o) {
+            return true;
+        }
+
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
         AtlasObjectId that = (AtlasObjectId) o;
+
+        if (isValidGuid() && Objects.equals(guid, that.guid)) {
+            return true;
+        }
+
         return Objects.equals(typeName, that.typeName) &&
-                Objects.equals(guid, that.guid) &&
                Objects.equals(uniqueAttributes, that.uniqueAttributes);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(typeName, guid, uniqueAttributes);
+        return isValidGuid() ? Objects.hash(guid) : Objects.hash(typeName, uniqueAttributes);
     }
 
     @Override

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

@@ -81,7 +81,7 @@ public class AtlasStruct implements Serializable {
     public AtlasStruct(AtlasStruct other) {
         if (other != null) {
             setTypeName(other.getTypeName());
-            setAttributes(other.getAttributes());
+            setAttributes(new HashMap<>(other.getAttributes()));
         }
     }
 

intg/src/main/java/org/apache/atlas/type/AtlasBuiltInTypes.java

@@ -519,7 +519,7 @@ public class AtlasBuiltInTypes {
 
         @Override
         public AtlasObjectId createDefaultValue() {
-            return new AtlasObjectId(AtlasBaseTypeDef.ATLAS_TYPE_ASSET, "test");
+            return new AtlasObjectId("-1", AtlasBaseTypeDef.ATLAS_TYPE_ASSET);
         }
 
         @Override
@@ -551,13 +551,16 @@ public class AtlasBuiltInTypes {
         }
 
         private boolean isValidMap(Map map) {
-            if (map.containsKey(AtlasObjectId.KEY_TYPENAME)) {
-                if (map.containsKey(AtlasObjectId.KEY_GUID)) {
+            Object guid = map.get(AtlasObjectId.KEY_GUID);
+
+            if (guid != null && StringUtils.isNotEmpty(guid.toString())) {
                 return true;
             } else {
+                Object typeName = map.get(AtlasObjectId.KEY_TYPENAME);
+                if (typeName != null && StringUtils.isNotEmpty(typeName.toString())) {
                     Object uniqueAttributes = map.get(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES);
 
-                    if (uniqueAttributes instanceof Map && MapUtils.isNotEmpty((Map)uniqueAttributes)) {
+                    if (uniqueAttributes instanceof Map && MapUtils.isNotEmpty((Map) uniqueAttributes)) {
                         return true;
                     }
                 }

intg/src/main/java/org/apache/atlas/type/AtlasEntityType.java

@@ -190,6 +190,9 @@ public class AtlasEntityType extends AtlasStructType {
             if (obj instanceof AtlasObjectId) {
                 AtlasObjectId objId = (AtlasObjectId ) obj;
                 return isAssignableFrom(objId);
+            } else if (obj instanceof Map) {
+                AtlasObjectId objId = new AtlasObjectId((Map)obj);
+                return isAssignableFrom(objId);
             }
 
             for (AtlasEntityType superType : superTypes) {

intg/src/main/java/org/apache/atlas/type/AtlasStructType.java

@@ -254,7 +254,7 @@ public class AtlasStructType extends AtlasType {
                     }
                 }
             } else if (obj instanceof Map) {
-                Map map = (Map) obj;
+                Map attributes = AtlasTypeUtil.toStructAttributes((Map)obj);
 
                 for (AtlasAttributeDef attributeDef : structDef.getAttributeDefs()) {
                     String    attrName = attributeDef.getName();
@@ -262,7 +262,7 @@ public class AtlasStructType extends AtlasType {
 
                     if (attribute != null) {
                         AtlasType dataType = attribute.getAttributeType();
-                        Object value     = map.get(attrName);
+                        Object value     = attributes.get(attrName);
                         String fieldName = objName + "." + attrName;
 
                         if (value != null) {
@@ -408,6 +408,7 @@ public class AtlasStructType extends AtlasType {
         private final AtlasType         attributeType;
         private final AtlasAttributeDef attributeDef;
         private final String            qualifiedName;
+        private final String            vertexPropertyName;
         private final boolean           isOwnedRef;
         private final String            inverseRefAttribute;
 
@@ -416,6 +417,7 @@ public class AtlasStructType extends AtlasType {
             this.attributeDef       = attrDef;
             this.attributeType      = attributeType;
             this.qualifiedName      = getQualifiedAttributeName(definedInType.getStructDef(), attributeDef.getName());
+            this.vertexPropertyName = encodePropertyKey(this.qualifiedName);
 
             boolean isOwnedRef          = false;
             String  inverseRefAttribute = null;
@@ -458,17 +460,47 @@ public class AtlasStructType extends AtlasType {
 
         public String getQualifiedName() { return qualifiedName; }
 
-        public String getQualifiedAttributeName() {
-            return qualifiedName;
-        }
+        public String getVertexPropertyName() { return vertexPropertyName; }
 
         public boolean isOwnedRef() { return isOwnedRef; }
 
         public String getInverseRefAttribute() { return inverseRefAttribute; }
 
-        public static String getQualifiedAttributeName(AtlasStructDef structDef, String attrName) {
+        private static String getQualifiedAttributeName(AtlasStructDef structDef, String attrName) {
             final String typeName = structDef.getName();
             return attrName.contains(".") ? attrName : String.format("%s.%s", typeName, attrName);
         }
+
+        private static String encodePropertyKey(String key) {
+            if (StringUtils.isBlank(key)) {
+                return key;
+            }
+
+            for (String[] strMap : RESERVED_CHAR_ENCODE_MAP) {
+                key = key.replace(strMap[0], strMap[1]);
+            }
+
+            return key;
+        }
+
+        private static String decodePropertyKey(String key) {
+            if (StringUtils.isBlank(key)) {
+                return key;
+            }
+
+            for (String[] strMap : RESERVED_CHAR_ENCODE_MAP) {
+                key = key.replace(strMap[1], strMap[0]);
+            }
+
+            return key;
+        }
+
+        private static String[][] RESERVED_CHAR_ENCODE_MAP = new String[][] {
+                new String[] { "{",  "_o" },
+                new String[] { "}",  "_c" },
+                new String[] { "\"", "_q" },
+                new String[] { "$",  "_d" },
+                new String[] { "%", "_p"  },
+        };
     }
 }

intg/src/main/java/org/apache/atlas/type/AtlasTypeUtil.java

@@ -21,6 +21,8 @@ import com.google.common.collect.ImmutableSet;
 
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.model.typedef.AtlasClassificationDef;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
@@ -35,14 +37,16 @@ import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang.StringUtils;
 
+import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
-
 import java.util.Map;
 import java.util.Set;
-import java.util.Arrays;
+
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -298,4 +302,26 @@ public class AtlasTypeUtil {
 
         return headerList;
     }
+
+    public static Collection<AtlasObjectId> toObjectIds(Collection<AtlasEntity> entities) {
+        List<AtlasObjectId> ret = new ArrayList<>();
+
+        if (CollectionUtils.isNotEmpty(entities)) {
+            for (AtlasEntity entity : entities) {
+                if (entity != null) {
+                    ret.add(entity.getAtlasObjectId());
+                }
+            }
+        }
+
+        return ret;
+    }
+
+    public static Map toStructAttributes(Map map) {
+        if (map != null && map.containsKey("typeName") && map.containsKey("attributes") && map.get("attributes") instanceof Map) {
+            return (Map)map.get("attributes");
+        }
+
+        return map;
+    }
 }
\ No newline at end of file

intg/src/test/java/org/apache/atlas/TestUtilsV2.java

@@ -21,6 +21,8 @@ package org.apache.atlas;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasStruct;
 import org.apache.atlas.model.typedef.AtlasClassificationDef;
@@ -275,8 +277,98 @@ public final class TestUtilsV2 {
     }
 
     public static final String DEPARTMENT_TYPE = "Department";
-    public static final String PERSON_TYPE = "Person";
     public static final String EMPLOYEE_TYPE   = "Employee";
+    public static final String MANAGER_TYPE    = "Manager";
+    public static final String ADDRESS_TYPE    = "Address";
+
+    public static AtlasEntitiesWithExtInfo createDeptEg2() {
+        AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo();
+
+        /******* Department - HR *******/
+        AtlasEntity   hrDept   = new AtlasEntity(DEPARTMENT_TYPE, "name", "hr");
+        AtlasObjectId hrDeptId = hrDept.getAtlasObjectId();
+
+        /******* Address Entities *******/
+        AtlasStruct janeAddr = new AtlasStruct(ADDRESS_TYPE);
+            janeAddr.setAttribute("street", "Great America Parkway");
+            janeAddr.setAttribute("city", "Santa Clara");
+
+        AtlasStruct juliusAddr = new AtlasStruct(ADDRESS_TYPE);
+            juliusAddr.setAttribute("street", "Madison Ave");
+            juliusAddr.setAttribute("city", "Newtonville");
+
+        AtlasStruct maxAddr = new AtlasStruct(ADDRESS_TYPE);
+            maxAddr.setAttribute("street", "Ripley St");
+            maxAddr.setAttribute("city", "Newton");
+
+        AtlasStruct johnAddr = new AtlasStruct(ADDRESS_TYPE);
+            johnAddr.setAttribute("street", "Stewart Drive");
+            johnAddr.setAttribute("city", "Sunnyvale");
+
+        /******* Manager - Jane (John and Max subordinates) *******/
+        AtlasEntity   jane   = new AtlasEntity(MANAGER_TYPE);
+        AtlasObjectId janeId = jane.getAtlasObjectId();
+            jane.setAttribute("name", "Jane");
+            jane.setAttribute("department", hrDeptId);
+            jane.setAttribute("address", janeAddr);
+
+        /******* Manager - Julius (no subordinates) *******/
+        AtlasEntity   julius   = new AtlasEntity(MANAGER_TYPE);
+        AtlasObjectId juliusId = julius.getAtlasObjectId();
+            julius.setAttribute("name", "Julius");
+            julius.setAttribute("department", hrDeptId);
+            julius.setAttribute("address", juliusAddr);
+            julius.setAttribute("subordinates", ImmutableList.of());
+
+        /******* Employee - Max (Manager: Jane, Mentor: Julius) *******/
+        AtlasEntity   max   = new AtlasEntity(EMPLOYEE_TYPE);
+        AtlasObjectId maxId = max.getAtlasObjectId();
+            max.setAttribute("name", "Max");
+            max.setAttribute("department", hrDeptId);
+            max.setAttribute("address", maxAddr);
+            max.setAttribute("manager", janeId);
+            max.setAttribute("mentor", juliusId);
+            max.setAttribute("birthday",new Date(1979, 3, 15));
+            max.setAttribute("hasPets", true);
+            max.setAttribute("age", 36);
+            max.setAttribute("numberOfCars", 2);
+            max.setAttribute("houseNumber", 17);
+            max.setAttribute("carMileage", 13);
+            max.setAttribute("shares", Long.MAX_VALUE);
+            max.setAttribute("salary", Double.MAX_VALUE);
+            max.setAttribute("numberOfStarsEstimate", new BigInteger("1000000000000000000000000000000"));
+            max.setAttribute("approximationOfPi", new BigDecimal("3.1415926535897932"));
+
+        /******* Employee - John (Manager: Jane, Mentor: Max) *******/
+        AtlasEntity   john   = new AtlasEntity(EMPLOYEE_TYPE);
+        AtlasObjectId johnId = john.getAtlasObjectId();
+            john.setAttribute("name", "John");
+            john.setAttribute("department", hrDeptId);
+            john.setAttribute("address", johnAddr);
+            john.setAttribute("manager", janeId);
+            john.setAttribute("mentor", maxId);
+            john.setAttribute("birthday",new Date(1950, 5, 15));
+            john.setAttribute("hasPets", true);
+            john.setAttribute("numberOfCars", 1);
+            john.setAttribute("houseNumber", 153);
+            john.setAttribute("carMileage", 13364);
+            john.setAttribute("shares", 15000);
+            john.setAttribute("salary", 123345.678);
+            john.setAttribute("age", 50);
+            john.setAttribute("numberOfStarsEstimate", new BigInteger("1000000000000000000000"));
+            john.setAttribute("approximationOfPi", new BigDecimal("3.141592653589793238462643383279502884197169399375105820974944592307816406286"));
+
+        jane.setAttribute("subordinates", ImmutableList.of(johnId, maxId));
+        hrDept.setAttribute("employees", ImmutableList.of(janeId, juliusId, maxId, johnId));
+
+        entitiesWithExtInfo.addEntity(hrDept);
+        entitiesWithExtInfo.addEntity(jane);
+        entitiesWithExtInfo.addEntity(julius);
+        entitiesWithExtInfo.addEntity(max);
+        entitiesWithExtInfo.addEntity(john);
+
+        return entitiesWithExtInfo;
+    }
 
     public static Map<String, AtlasEntity> createDeptEg1() {
         Map<String, AtlasEntity> deptEmpEntities = new HashMap<>();
@@ -292,7 +384,7 @@ public final class TestUtilsV2 {
         AtlasEntity max = new AtlasEntity(EMPLOYEE_TYPE);
         AtlasEntity maxAddr = new AtlasEntity("Address");
 
-        AtlasObjectId deptId = new AtlasObjectId(hrDept.getTypeName(), hrDept.getGuid());
+        AtlasObjectId deptId = new AtlasObjectId(hrDept.getGuid(), hrDept.getTypeName());
         hrDept.setAttribute("name", "hr");
         john.setAttribute("name", "John");
         john.setAttribute("department", deptId);
@@ -329,8 +421,8 @@ public final class TestUtilsV2 {
         AtlasObjectId johnId = john.getAtlasObjectId();
 
         //TODO - Change to MANAGER_TYPE for JULIUS
-        AtlasObjectId maxId = new AtlasObjectId(EMPLOYEE_TYPE, max.getGuid());
-        AtlasObjectId juliusId = new AtlasObjectId(EMPLOYEE_TYPE, julius.getGuid());
+        AtlasObjectId maxId = new AtlasObjectId(max.getGuid(), EMPLOYEE_TYPE);
+        AtlasObjectId juliusId = new AtlasObjectId(julius.getGuid(), EMPLOYEE_TYPE);
 
         max.setAttribute("name", "Max");
         max.setAttribute("department", deptId);
@@ -356,10 +448,6 @@ public final class TestUtilsV2 {
 
         jane.setAttribute("subordinates", ImmutableList.of(johnId, maxId));
 
-//        Map<String, Integer> secClearanceLevelMap = new HashMap<>();
-//        secClearanceLevelMap.put("level", 1);
-//        jane.setAttribute("SecurityClearance", secClearanceLevelMap);
-
         deptEmpEntities.put(jane.getGuid(), jane);
         deptEmpEntities.put(john.getGuid(), john);
         deptEmpEntities.put(julius.getGuid(), julius);
@@ -659,27 +747,32 @@ public final class TestUtilsV2 {
         return RandomStringUtils.randomAlphanumeric(10);
     }
 
-    public static Map<String, AtlasEntity> createDBEntity() {
-        Map<String, AtlasEntity> ret = new HashMap<>();
+    public static AtlasEntity createDBEntity() {
         AtlasEntity entity = new AtlasEntity(DATABASE_TYPE);
         String dbName = RandomStringUtils.randomAlphanumeric(10);
         entity.setAttribute(NAME, dbName);
         entity.setAttribute("description", "us db");
 
-        ret.put(entity.getGuid(), entity);
-        return ret;
+        return entity;
     }
 
-    public static Map<String, AtlasEntity> createTableEntity(String dbId) {
-        Map<String, AtlasEntity> ret = new HashMap<>();
+    public static AtlasEntityWithExtInfo createDBEntityV2() {
+        AtlasEntity dbEntity = new AtlasEntity(DATABASE_TYPE);
+
+        dbEntity.setAttribute(NAME, RandomStringUtils.randomAlphanumeric(10));
+        dbEntity.setAttribute("description", "us db");
 
+        return new AtlasEntityWithExtInfo(dbEntity);
+    }
+
+    public static AtlasEntity createTableEntity(AtlasEntity dbEntity) {
         AtlasEntity entity = new AtlasEntity(TABLE_TYPE);
         String tableName = RandomStringUtils.randomAlphanumeric(10);
         entity.setAttribute(NAME, tableName);
         entity.setAttribute("description", "random table");
         entity.setAttribute("type", "type");
         entity.setAttribute("tableType", "MANAGED");
-        entity.setAttribute("database", new AtlasObjectId(DATABASE_TYPE, dbId));
+        entity.setAttribute("database", dbEntity.getAtlasObjectId());
         entity.setAttribute("created", new Date());
 
         Map<String, Object> partAttributes = new HashMap<String, Object>() {{
@@ -692,16 +785,39 @@ public final class TestUtilsV2 {
             put("key1", "value1");
         }});
 
-        ret.put(entity.getGuid(), entity);
+        return entity;
+    }
+
+    public static AtlasEntityWithExtInfo createTableEntityV2(AtlasEntity dbEntity) {
+        AtlasEntity tblEntity = new AtlasEntity(TABLE_TYPE);
+
+        tblEntity.setAttribute(NAME, RandomStringUtils.randomAlphanumeric(10));
+        tblEntity.setAttribute("description", "random table");
+        tblEntity.setAttribute("type", "type");
+        tblEntity.setAttribute("tableType", "MANAGED");
+        tblEntity.setAttribute("database", dbEntity.getAtlasObjectId());
+        tblEntity.setAttribute("created", new Date());
+
+        final AtlasStruct partitionStruct = new AtlasStruct("partition_struct_type", "name", "part0");
+
+        tblEntity.setAttribute("partitions", new ArrayList<AtlasStruct>() {{ add(partitionStruct); }});
+        tblEntity.setAttribute("parametersMap",
+                new java.util.HashMap<String, String>() {{ put("key1", "value1"); }});
+
+
+        AtlasEntityWithExtInfo ret = new AtlasEntityWithExtInfo(tblEntity);
+
+        ret.addReferredEntity(dbEntity);
+
         return ret;
     }
 
-    public static AtlasEntity createColumnEntity(String tableId) {
+    public static AtlasEntity createColumnEntity(AtlasEntity tableEntity) {
 
         AtlasEntity entity = new AtlasEntity(COLUMN_TYPE);
         entity.setAttribute(NAME, RandomStringUtils.randomAlphanumeric(10));
         entity.setAttribute("type", "VARCHAR(32)");
-        entity.setAttribute("table", new AtlasObjectId(TABLE_TYPE, tableId));
+        entity.setAttribute("table", tableEntity.getAtlasObjectId());
         return entity;
     }
 

intg/src/test/java/org/apache/atlas/type/TestAtlasObjectIdType.java

@@ -45,37 +45,41 @@ public class TestAtlasObjectIdType {
         Map<Object, Object> invalidObj2 = new HashMap<>();
         Map<Object, Object> invalidObj3 = new HashMap<>();
         Map<Object, Object> invalidObj4 = new HashMap<>();
-        Map<Object, Object> invalidObj5 = new HashMap<>();
-
-        validObj1.put(AtlasObjectId.KEY_TYPENAME, "testType");
-        validObj1.put(AtlasObjectId.KEY_GUID, "guid-1234");
-        validObj2.put(AtlasObjectId.KEY_TYPENAME, "testType");
-        validObj2.put(AtlasObjectId.KEY_GUID, 1234);
 
         Map<String, Object> uniqAttribs = new HashMap<String, Object>();
         uniqAttribs.put("name", "testTypeInstance-1");
 
+        // guid
+        validObj1.put(AtlasObjectId.KEY_GUID, "guid-1234");
+
+        // typeName & unique-attributes
+        validObj2.put(AtlasObjectId.KEY_TYPENAME, "testType");
+        validObj2.put(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES, uniqAttribs);
+
+        // guid, typeName & unique-attributes
+        validObj3.put(AtlasObjectId.KEY_GUID, "guid-1234");
         validObj3.put(AtlasObjectId.KEY_TYPENAME, "testType");
         validObj3.put(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES, uniqAttribs);
 
-        invalidObj1.put(AtlasObjectId.KEY_TYPENAME, "testType"); // no guid
-
-        invalidObj2.put(AtlasObjectId.KEY_GUID, "guid-1234");    // no typeName or uniqueAttribute
-        invalidObj2.put(AtlasObjectId.KEY_TYPENAME + "-invalid", "testType");
+        // no guid or typeName/unique-attributes
+        invalidObj1.put(AtlasObjectId.KEY_GUID + "-invalid", "guid-1234"); // no guid or typename or uniqueAttribute
 
-        invalidObj3.put(AtlasObjectId.KEY_GUID + "-invalid", "guid-1234"); // no guid or typename or uniqueAttribute
+        // no unique-attributes
+        invalidObj2.put(AtlasObjectId.KEY_TYPENAME, "testType"); // no guid
 
-        invalidObj4.put(AtlasObjectId.KEY_TYPENAME, "testType");    // empty uniqueAttribute
-        invalidObj4.put(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES, new HashMap<String, Object>());
+        // empty uniqueAttribute
+        invalidObj3.put(AtlasObjectId.KEY_TYPENAME, "testType");
+        invalidObj3.put(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES, new HashMap<String, Object>());
 
-        invalidObj5.put(AtlasObjectId.KEY_TYPENAME, "testType");    // non-map uniqueAttribute
-        invalidObj5.put(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES, new ArrayList<String>());
+        // non-map uniqueAttribute
+        invalidObj4.put(AtlasObjectId.KEY_TYPENAME, "testType");
+        invalidObj4.put(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES, new ArrayList<String>());
 
         validValues = new Object[] {
-            null, validObj1, validObj2, validObj3, new AtlasObjectId(), new AtlasObjectId("testType", "guid-1234"), };
+            null, validObj1, validObj2, validObj3, new AtlasObjectId(), new AtlasObjectId("guid-1234", "testType"), };
 
         invalidValues = new Object[] {
-            invalidObj1, invalidObj2, invalidObj3, invalidObj4, invalidObj5,
+            invalidObj1, invalidObj2, invalidObj3, invalidObj4,
             Byte.valueOf((byte)1), Short.valueOf((short)1), Integer.valueOf(1),
             Long.valueOf(1L), Float.valueOf(1), Double.valueOf(1), BigInteger.valueOf(1), BigDecimal.valueOf(1), "1",
             "", "12ab", "abcd", "-12ab",
@@ -116,12 +120,12 @@ public class TestAtlasObjectIdType {
             if (value instanceof AtlasObjectId) {
                 assertEquals(normalizedValue, value, "value=" + value);
             } else if (value instanceof Map) {
-                assertEquals(normalizedValue.getTypeName(), ((Map)value).get(AtlasObjectId.KEY_TYPENAME).toString(),
+                assertEquals(normalizedValue.getTypeName(), ((Map)value).get(AtlasObjectId.KEY_TYPENAME),
                              "value=" + value);
                 if (((Map)value).get(AtlasObjectId.KEY_GUID) == null) {
                     assertEquals(normalizedValue.getGuid(), ((Map)value).get(AtlasObjectId.KEY_GUID),  "value=" + value);
                 } else {
-                    assertEquals(normalizedValue.getGuid().toString(), ((Map) value).get(AtlasObjectId.KEY_GUID).toString(), "value=" + value);
+                    assertEquals(normalizedValue.getGuid(), ((Map) value).get(AtlasObjectId.KEY_GUID), "value=" + value);
                 }
 
                 assertEquals(normalizedValue.getUniqueAttributes(), ((Map)value).get(AtlasObjectId.KEY_UNIQUE_ATTRIBUTES),

repository/src/main/java/org/apache/atlas/RepositoryMetadataModule.java

@@ -42,17 +42,10 @@ import org.apache.atlas.repository.graph.DeleteHandler;
 import org.apache.atlas.repository.graph.GraphBackedMetadataRepository;
 import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
-import org.apache.atlas.repository.store.graph.EntityGraphDiscovery;
-import org.apache.atlas.repository.store.graph.EntityResolver;
-import org.apache.atlas.repository.store.graph.v1.ArrayVertexMapper;
-import org.apache.atlas.repository.store.graph.v1.AtlasEntityGraphDiscoveryV1;
 import org.apache.atlas.repository.store.graph.v1.AtlasEntityStoreV1;
 import org.apache.atlas.repository.store.graph.v1.AtlasTypeDefGraphStoreV1;
 import org.apache.atlas.repository.store.graph.v1.DeleteHandlerV1;
 import org.apache.atlas.repository.store.graph.v1.EntityGraphMapper;
-import org.apache.atlas.repository.store.graph.v1.IDBasedEntityResolver;
-import org.apache.atlas.repository.store.graph.v1.MapVertexMapper;
-import org.apache.atlas.repository.store.graph.v1.UniqAttrBasedEntityResolver;
 import org.apache.atlas.repository.typestore.GraphBackedTypeStore;
 import org.apache.atlas.repository.typestore.ITypeStore;
 import org.apache.atlas.service.Service;
@@ -121,15 +114,6 @@ public class RepositoryMetadataModule extends com.google.inject.AbstractModule {
 
         bind(EntityGraphMapper.class);
 
-        bind(MapVertexMapper.class).asEagerSingleton();
-
-        bind(ArrayVertexMapper.class).asEagerSingleton();
-
-        Multibinder<EntityResolver> entityRefResolver =
-            Multibinder.newSetBinder(binder(), EntityResolver.class);
-        entityRefResolver.addBinding().to(IDBasedEntityResolver.class);
-        entityRefResolver.addBinding().to(UniqAttrBasedEntityResolver.class);
-
         //Add EntityAuditListener as EntityChangeListener
         Multibinder<EntityChangeListener> entityChangeListenerBinder =
                 Multibinder.newSetBinder(binder(), EntityChangeListener.class);
@@ -138,8 +122,6 @@ public class RepositoryMetadataModule extends com.google.inject.AbstractModule {
         MethodInterceptor interceptor = new GraphTransactionInterceptor();
         requestInjection(interceptor);
         bindInterceptor(Matchers.any(), Matchers.annotatedWith(GraphTransaction.class), interceptor);
-
-        bind(EntityGraphDiscovery.class).to(AtlasEntityGraphDiscoveryV1.class);
     }
 
     protected Configuration getConfiguration() {

repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java

@@ -244,7 +244,7 @@ public final class GraphHelper {
         /**
         Iterable<AtlasEdge> edges = outVertex.getEdges(Direction.OUT, edgeLabel);
         for (AtlasEdge edge : edges) {
-            if (edge.getId().toString().equals(edgeId)) {
+            if (edge.getObjectId().toString().equals(edgeId)) {
                 return edge;
             }
         }

repository/src/main/java/org/apache/atlas/repository/store/bootstrap/AtlasTypeDefStoreInitializer.java

@@ -86,6 +86,26 @@ public class AtlasTypeDefStoreInitializer {
                     continue;
                 }
 
+                AtlasTypesDef typesToCreate = getTypesToCreate(typesDef, typeRegistry);
+
+                if (typesToCreate.isEmpty()) {
+                    LOG.info("No new type in file {}", typeDefFile.getAbsolutePath());
+
+                    continue;
+                }
+
+                LOG.info("Loading types defined in file {}", typeDefFile.getAbsolutePath());
+
+                typeDefStore.createTypesDef(typesToCreate);
+            } catch (Throwable t) {
+                LOG.error("error while registering types in file {}", typeDefFile.getAbsolutePath(), t);
+            }
+        }
+
+        applyTypePatches(typeDefStore, typeRegistry, typesDirName);
+    }
+
+    public static AtlasTypesDef getTypesToCreate(AtlasTypesDef typesDef, AtlasTypeRegistry typeRegistry) {
         AtlasTypesDef typesToCreate = new AtlasTypesDef();
 
         if (CollectionUtils.isNotEmpty(typesDef.getEnumDefs())) {
@@ -120,21 +140,7 @@ public class AtlasTypeDefStoreInitializer {
             }
         }
 
-                if (typesToCreate.isEmpty()) {
-                    LOG.info("No new type in file {}", typeDefFile.getAbsolutePath());
-
-                    continue;
-                }
-
-                LOG.info("Loading types defined in file {}", typeDefFile.getAbsolutePath());
-
-                typeDefStore.createTypesDef(typesToCreate);
-            } catch (Throwable t) {
-                LOG.error("error while registering types in file {}", typeDefFile.getAbsolutePath(), t);
-            }
-        }
-
-        applyTypePatches(typeDefStore, typeRegistry, typesDirName);
+        return typesToCreate;
     }
 
     private void applyTypePatches(AtlasTypeDefStore typeDefStore, AtlasTypeRegistry typeRegistry, String typesDirName) {

repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasEntityStore.java

@@ -24,8 +24,8 @@ import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.EntityMutationResponse;
+import org.apache.atlas.repository.store.graph.v1.EntityStream;
 import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasTypeRegistry;
 
 import java.util.List;
 import java.util.Map;
@@ -34,13 +34,6 @@ import java.util.Map;
  * Persistence/Retrieval API for AtlasEntity
  */
 public interface AtlasEntityStore {
-
-    /**
-     * Initialization
-     */
-    void init(AtlasTypeRegistry typeRegistry) throws AtlasBaseException;
-
-
     /**
      *
      * Get entity definition by its guid
@@ -50,22 +43,6 @@ public interface AtlasEntityStore {
     AtlasEntityWithExtInfo getById(String guid) throws AtlasBaseException;
 
     /**
-     * Delete an entity by its guid
-     * @param guid
-     * @return
-     */
-    EntityMutationResponse deleteById(String guid);
-
-    /**
-     * Create or update  entities
-     * @param entities Map of the entity Id(guid or transient Id) to AtlasEntity objects that need to be created
-     * @return EntityMutationResponse Entity mutations operations with the correspomding set of entities on which these operations were performed
-     * @throws AtlasBaseException
-     */
-
-    EntityMutationResponse createOrUpdate(Map<String, AtlasEntity> entities) throws AtlasBaseException;
-
-    /**
      * Batch GET to retrieve entities by their ID
      * @param guid
      * @return
@@ -73,43 +50,57 @@ public interface AtlasEntityStore {
      */
     AtlasEntitiesWithExtInfo getByIds(List<String> guid) throws AtlasBaseException;
 
-    /*
-     * Return list of deleted entity guids
-     */
-    EntityMutationResponse deleteByIds(List<String> guid) throws AtlasBaseException;
-
     /**
      *
      * Get an eneity by its unique attribute
-     * @param entityType
-     * @param uniqAttributes
-     * @return AtlasEntity
+     * @param entityType     type of the entity
+     * @param uniqAttributes Attributes that uniquely identify the entity
+     * @return EntityMutationResponse details of the updates performed by this call
      */
-    AtlasEntityWithExtInfo getByUniqueAttribute(AtlasEntityType entityType, Map<String, Object> uniqAttributes)
+    AtlasEntityWithExtInfo getByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes)
             throws AtlasBaseException;
 
     /**
+     * Create or update  entities in the stream
+     * @param entityStream AtlasEntityStream
+     * @return EntityMutationResponse Entity mutations operations with the corresponding set of entities on which these operations were performed
+     * @throws AtlasBaseException
+     */
+    EntityMutationResponse createOrUpdate(EntityStream entityStream) throws AtlasBaseException;
+
+    /**
      * @deprecated
-     * Create or update a single entity
-     * @param typeName The entity's type
-     * @param attributeName Attribute that uniquely identifies the entity
-     * @param attributeValue The unqiue attribute's value
-     * @return EntityMutationResponse Entity mutations operations with the correspomding set of entities on which these operations were performed
+     * Update a single entity
+     * @param entityType     type of the entity
+     * @param uniqAttributes Attributes that uniquely identify the entity
+     * @return EntityMutationResponse details of the updates performed by this call
      * @throws AtlasBaseException
      *
      */
+    EntityMutationResponse updateByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes,
+                                                    AtlasEntity entity) throws AtlasBaseException;
 
-    EntityMutationResponse updateByUniqueAttribute(String typeName, String attributeName, String attributeValue, AtlasEntity entity) throws AtlasBaseException;
+    /**
+     * Delete an entity by its guid
+     * @param guid
+     * @return
+     */
+    EntityMutationResponse deleteById(String guid) throws AtlasBaseException;
 
     /**
      * @deprecated
-     * @param typeName
-     * @param attributeName
-     * @param attributeValue
-     * @return
+     * @param entityType      type of the entity
+     * @param uniqAttributes Attributes that uniquely identify the entity
+     * @return EntityMutationResponse details of the updates performed by this call
      * @throws AtlasBaseException
      */
-    EntityMutationResponse deleteByUniqueAttribute(String typeName, String attributeName, String attributeValue) throws AtlasBaseException;
+    EntityMutationResponse deleteByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes)
+                                                                                             throws AtlasBaseException;
+
+    /*
+     * Return list of deleted entity guids
+     */
+    EntityMutationResponse deleteByIds(List<String> guid) throws AtlasBaseException;
 
     /**
      * Add classification(s)

repository/src/main/java/org/apache/atlas/repository/store/graph/EntityGraphDiscovery.java

@@ -20,8 +20,10 @@ package org.apache.atlas.repository.store.graph;
 
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.repository.store.graph.v1.AtlasEntityStream;
 
 import java.util.List;
+import java.util.Map;
 
 public interface EntityGraphDiscovery {
 
@@ -33,7 +35,7 @@ public interface EntityGraphDiscovery {
      * Resolved references already exist in the ATLAS repository and have an assigned unique GUID
      * Unresolved attribute references result in an error if they are not composite (managed by a parent entity)
      */
-    EntityGraphDiscoveryContext discoverEntities(List<AtlasEntity> entities) throws AtlasBaseException;
+    EntityGraphDiscoveryContext discoverEntities() throws AtlasBaseException;
 
     void cleanUp() throws AtlasBaseException;
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/EntityGraphDiscoveryContext.java

@@ -17,179 +17,117 @@
  */
 package org.apache.atlas.repository.store.graph;
 
-import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.repository.store.graph.v1.EntityStream;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
-import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 
 public final class EntityGraphDiscoveryContext {
     private static final Logger LOG = LoggerFactory.getLogger(EntityGraphDiscoveryContext.class);
 
     private final AtlasTypeRegistry               typeRegistry;
-    private List<AtlasEntity>               rootEntities               = new ArrayList<>();
-    private Map<AtlasObjectId, AtlasVertex> resolvedIds                = new LinkedHashMap<>();
-    private Set<AtlasObjectId>              unresolvedIds              = new HashSet<>();
-    private List<AtlasObjectId>             unresolvedIdsByUniqAttribs = new ArrayList<>();
-
-    public EntityGraphDiscoveryContext(AtlasTypeRegistry typeRegistry) {
+    private final EntityStream                    entityStream;
+    private final List<String>                    referencedGuids          = new ArrayList<>();
+    private final Set<AtlasObjectId>              referencedByUniqAttribs  = new HashSet<>();
+    private final Map<String, AtlasVertex>        resolvedGuids            = new HashMap<>();
+    private final Map<AtlasObjectId, AtlasVertex> resolvedIdsByUniqAttribs = new HashMap<>();
+    private final Set<String>                     localGuids               = new HashSet<>();
+
+    public EntityGraphDiscoveryContext(AtlasTypeRegistry typeRegistry, EntityStream entityStream) {
         this.typeRegistry = typeRegistry;
+        this.entityStream = entityStream;
     }
 
-
-    public Collection<AtlasEntity> getRootEntities() {
-        return rootEntities;
+    public EntityStream getEntityStream() {
+        return entityStream;
     }
 
-    public Map<AtlasObjectId, AtlasVertex> getResolvedIds() {
-        return resolvedIds;
-    }
+    public List<String> getReferencedGuids() { return referencedGuids; }
 
-    public Set<AtlasObjectId> getUnresolvedIds() {
-        return unresolvedIds;
-    }
+    public Set<AtlasObjectId> getReferencedByUniqAttribs() { return referencedByUniqAttribs; }
 
-    public List<AtlasObjectId> getUnresolvedIdsByUniqAttribs() {
-        return unresolvedIdsByUniqAttribs;
+    public Map<String, AtlasVertex> getResolvedGuids() {
+        return resolvedGuids;
     }
 
-
-    public void addRootEntity(AtlasEntity rootEntity) {
-        this.rootEntities.add(rootEntity);
+    public Map<AtlasObjectId, AtlasVertex> getResolvedIdsByUniqAttribs() {
+        return resolvedIdsByUniqAttribs;
     }
 
+    public Set<String> getLocalGuids() { return localGuids; }
 
-    public void addResolvedId(AtlasObjectId objId, AtlasVertex vertex) {
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("addResolvedId({})", objId);
-        }
 
-        resolvedIds.put(objId, vertex);
+    public void addReferencedGuid(String guid) {
+        if (! referencedGuids.contains(guid)) {
+            referencedGuids.add(guid);
         }
-
-    public boolean removeUnResolvedId(AtlasObjectId objId) {
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("removeUnResolvedId({})", objId);
     }
 
-        return unresolvedIds.remove(objId);
-    }
+    public void addReferencedByUniqAttribs(AtlasObjectId objId) { referencedByUniqAttribs.add(objId); }
 
 
-    public void addUnResolvedId(AtlasObjectId objId) {
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("addUnResolvedId({})", objId);
-        }
-
-        this.unresolvedIds.add(objId);
-    }
-
-    public boolean removeUnResolvedIds(List<AtlasObjectId> objIds) {
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("removeUnResolvedIds({})", objIds);
-        }
-
-        return unresolvedIds.removeAll(objIds);
-    }
+    public void addResolvedGuid(String guid, AtlasVertex vertex) { resolvedGuids.put(guid, vertex); }
 
+    public void addResolvedIdByUniqAttribs(AtlasObjectId objId, AtlasVertex vertex) { resolvedIdsByUniqAttribs.put(objId, vertex); }
 
-    public void addUnresolvedIdByUniqAttribs(AtlasObjectId objId) {
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("addUnresolvedIdByUniqAttribs({})", objId);
-        }
+    public void addLocalGuidReference(String guid) { localGuids.add(guid); }
 
-        this.unresolvedIdsByUniqAttribs.add(objId);
-    }
+    public boolean isResolvedGuid(String guid) { return resolvedGuids.containsKey(guid); }
 
-    public boolean removeUnresolvedIdsByUniqAttribs(List<AtlasObjectId> objIds) {
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("removeUnresolvedIdsByUniqAttribs({})", objIds);
-        }
+    public boolean isResolvedIdByUniqAttrib(AtlasObjectId objId) { return resolvedIdsByUniqAttribs.containsKey(objId); }
 
-        return unresolvedIdsByUniqAttribs.removeAll(objIds);
-    }
 
-    public boolean hasUnresolvedReferences() {
-        return unresolvedIdsByUniqAttribs.size() > 0 || unresolvedIds.size() > 0;
-    }
+    public AtlasVertex getResolvedEntityVertex(String guid) throws AtlasBaseException {
+        AtlasVertex ret = resolvedGuids.get(guid);
 
-    public boolean isResolvedId(AtlasObjectId id) {
-        return resolvedIds.containsKey(id);
+        return ret;
     }
 
-    public AtlasVertex getResolvedEntityVertex(AtlasObjectId ref) throws AtlasBaseException {
-        AtlasVertex vertex = resolvedIds.get(ref);
+    public AtlasVertex getResolvedEntityVertex(AtlasObjectId objId) {
+        AtlasVertex vertex = resolvedIdsByUniqAttribs.get(objId);
 
         // check also for sub-types; ref={typeName=Asset; guid=abcd} should match {typeName=hive_table; guid=abcd}
         if (vertex == null) {
-            final AtlasEntityType entityType  = typeRegistry.getEntityTypeByName(ref.getTypeName());
+            final AtlasEntityType entityType  = typeRegistry.getEntityTypeByName(objId.getTypeName());
             final Set<String>     allSubTypes = entityType.getAllSubTypes();
 
             for (String subType : allSubTypes) {
-                AtlasObjectId subTypeObjId = new AtlasObjectId(subType, ref.getGuid(), ref.getUniqueAttributes());
+                AtlasObjectId subTypeObjId = new AtlasObjectId(objId.getGuid(), subType, objId.getUniqueAttributes());
 
-                vertex = resolvedIds.get(subTypeObjId);
+                vertex = resolvedIdsByUniqAttribs.get(subTypeObjId);
 
                 if (vertex != null) {
-                    resolvedIds.put(ref, vertex);
+                    resolvedIdsByUniqAttribs.put(objId, vertex);
                     break;
                 }
             }
         }
 
-        if (vertex == null) {
-            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS,
-                                         " : Could not find an entity with " + ref.toString());
-        }
-
         return vertex;
     }
 
-    @Override
-    public boolean equals(Object obj) {
-        if (obj == null) {
-            return false;
-        } else if (obj == this) {
-            return true;
-        } else if (obj.getClass() != getClass()) {
-            return false;
-        } else {
-            EntityGraphDiscoveryContext ctx = (EntityGraphDiscoveryContext) obj;
-            return Objects.equals(rootEntities, ctx.getRootEntities()) &&
-                Objects.equals(resolvedIds, ctx.getResolvedIds()) &&
-                Objects.equals(unresolvedIdsByUniqAttribs, ctx.getUnresolvedIdsByUniqAttribs()) &&
-                Objects.equals(unresolvedIds, ctx.getUnresolvedIds());
-        }
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(rootEntities, resolvedIds, unresolvedIdsByUniqAttribs, unresolvedIds);
-    }
-
     public StringBuilder toString(StringBuilder sb) {
         if (sb == null) {
             sb = new StringBuilder();
         }
 
         sb.append("EntityGraphDiscoveryCtx{");
-        sb.append("rootEntities='").append(rootEntities).append('\'');
-        sb.append(", resolvedIds=").append(resolvedIds);
-        sb.append(", unresolvedIdsByUniqAttribs='").append(unresolvedIdsByUniqAttribs).append('\'');
-        sb.append(", unresolvedIds='").append(unresolvedIds).append('\'');
+        sb.append("referencedGuids=").append(referencedGuids);
+        sb.append(", referencedByUniqAttribs=").append(referencedByUniqAttribs);
+        sb.append(", resolvedGuids='").append(resolvedGuids);
+        sb.append(", resolvedIdsByUniqAttribs='").append(resolvedIdsByUniqAttribs);
+        sb.append(", localGuids='").append(localGuids);
         sb.append('}');
 
         return sb;
@@ -201,9 +139,10 @@ public final class EntityGraphDiscoveryContext {
     }
 
     public void cleanUp() {
-        rootEntities.clear();
-        unresolvedIdsByUniqAttribs.clear();
-        resolvedIds.clear();
-        unresolvedIds.clear();
+        referencedGuids.clear();
+        referencedByUniqAttribs.clear();
+        resolvedGuids.clear();
+        resolvedIdsByUniqAttribs.clear();
+        localGuids.clear();
     }
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/EntityResolver.java

@@ -21,10 +21,5 @@ import org.apache.atlas.exception.AtlasBaseException;
 
 
 public interface EntityResolver {
-
-    void init(EntityGraphDiscoveryContext entities) throws AtlasBaseException;
-
-    EntityGraphDiscoveryContext resolveEntityReferences() throws AtlasBaseException;
-
-    void cleanUp() throws AtlasBaseException;
+    EntityGraphDiscoveryContext resolveEntityReferences(EntityGraphDiscoveryContext entities) throws AtlasBaseException;
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/ArrayVertexMapper.java

-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.atlas.repository.store.graph.v1;
-
-import com.google.common.base.Optional;
-import com.google.inject.Inject;
-import com.google.inject.Singleton;
-import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
-import org.apache.atlas.repository.graph.GraphHelper;
-import org.apache.atlas.repository.graphdb.AtlasEdge;
-import org.apache.atlas.repository.graphdb.AtlasVertex;
-import org.apache.atlas.type.AtlasArrayType;
-import org.apache.atlas.type.AtlasStructType;
-import org.apache.atlas.type.AtlasType;
-import org.apache.commons.collections.CollectionUtils;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
-
-import static org.apache.atlas.repository.graph.GraphHelper.string;
-
-@Singleton
-public class ArrayVertexMapper implements InstanceGraphMapper<List> {
-
-    private static final Logger LOG = LoggerFactory.getLogger(ArrayVertexMapper.class);
-
-    protected final DeleteHandlerV1 deleteHandler;
-
-    protected StructVertexMapper structVertexMapper;
-
-    @Inject
-    public ArrayVertexMapper(DeleteHandlerV1 deleteHandler) {
-        this.deleteHandler = deleteHandler;
-    }
-
-    void init(StructVertexMapper structVertexMapper) {
-        this.structVertexMapper = structVertexMapper;
-    }
-
-    @Override
-    public List toGraph(GraphMutationContext ctx) throws AtlasBaseException {
-
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("Mapping instance to vertex {} for array attribute {}", string(ctx.getReferringVertex()), ctx.getAttrType().getTypeName());
-        }
-
-        List newElements = (List) ctx.getValue();
-        boolean newAttributeEmpty = (newElements == null || newElements.isEmpty());
-
-        AtlasArrayType arrType = (AtlasArrayType) ctx.getAttrType();
-        AtlasType elementType = arrType.getElementType();
-        List<Object> currentElements = getArrayElementsProperty(elementType, ctx.getReferringVertex(), ctx.getVertexPropertyKey());
-
-        List<Object> newElementsCreated = new ArrayList<>();
-
-        if (!newAttributeEmpty) {
-            for (int index = 0; index < newElements.size(); index++) {
-
-                LOG.debug("Adding/updating element at position {}, current element {}, new element {}", index,
-                    (currentElements != null && index < currentElements.size()) ? currentElements.get(index) : null, newElements.get(index));
-
-                Optional<AtlasEdge> existingEdge = getEdgeAt(currentElements, index, arrType.getElementType());
-
-                GraphMutationContext arrCtx = new GraphMutationContext.Builder(ctx.getOp(), ctx.getAttribute(),
-                    arrType.getElementType(), newElements.get(index))
-                    .referringVertex(ctx.getReferringVertex())
-                    .edge(existingEdge)
-                    .vertexProperty(ctx.getVertexPropertyKey()).build();
-
-                Object newEntry = structVertexMapper.mapCollectionElementsToVertex(arrCtx);
-                newElementsCreated.add(newEntry);
-            }
-        }
-
-        if (AtlasGraphUtilsV1.isReference(elementType)) {
-            List<AtlasEdge> additionalEdges = removeUnusedArrayEntries(ctx.getParentType(), ctx.getAttributeDef(), (List) currentElements, (List) newElementsCreated, elementType);
-            newElementsCreated.addAll(additionalEdges);
-        }
-
-        // for dereference on way out
-        setArrayElementsProperty(elementType, ctx.getReferringVertex(), ctx.getVertexPropertyKey(), newElementsCreated);
-        return newElementsCreated;
-    }
-
-    @Override
-    public void cleanUp() throws AtlasBaseException {
-
-    }
-
-    //Removes unused edges from the old collection, compared to the new collection
-    private List<AtlasEdge> removeUnusedArrayEntries(
-        AtlasStructType entityType,
-        AtlasAttributeDef attributeDef,
-        List<AtlasEdge> currentEntries,
-        List<AtlasEdge> newEntries,
-        AtlasType entryType) throws AtlasBaseException {
-        if (currentEntries != null && !currentEntries.isEmpty()) {
-            LOG.debug("Removing unused entries from the old collection");
-            if (AtlasGraphUtilsV1.isReference(entryType)) {
-
-                Collection<AtlasEdge> edgesToRemove = CollectionUtils.subtract(currentEntries, newEntries);
-
-                LOG.debug("Removing unused entries from the old collection - {}", edgesToRemove);
-
-                if (!edgesToRemove.isEmpty()) {
-                    //Remove the edges for (current edges - new edges)
-                    List<AtlasEdge> additionalElements = new ArrayList<>();
-
-                    for (AtlasEdge edge : edgesToRemove) {
-                        boolean deleteChildReferences = StructVertexMapper.shouldManageChildReferences(entityType, attributeDef.getName());
-                        boolean deleted = deleteHandler.deleteEdgeReference(edge, entryType.getTypeCategory(),
-                            deleteChildReferences, true);
-                        if (!deleted) {
-                            additionalElements.add(edge);
-                        }
-                    }
-
-                    return additionalElements;
-                }
-            }
-        }
-        return Collections.emptyList();
-    }
-
-    public static List<Object> getArrayElementsProperty(AtlasType elementType, AtlasVertex instanceVertex, String propertyName) {
-        String actualPropertyName = GraphHelper.encodePropertyKey(propertyName);
-        if (AtlasGraphUtilsV1.isReference(elementType)) {
-            return (List)instanceVertex.getListProperty(actualPropertyName, AtlasEdge.class);
-        }
-        else {
-            return (List)instanceVertex.getListProperty(actualPropertyName);
-        }
-    }
-
-    private Optional<AtlasEdge> getEdgeAt(List<Object> currentElements, int index, AtlasType elemType) {
-        Optional<AtlasEdge> existingEdge = Optional.absent();
-        if ( AtlasGraphUtilsV1.isReference(elemType) ) {
-            Object currentElement = (currentElements != null && index < currentElements.size()) ?
-                currentElements.get(index) : null;
-
-            if ( currentElement != null) {
-                existingEdge = Optional.of((AtlasEdge) currentElement);
-            }
-        }
-
-        return existingEdge;
-    }
-
-    private void setArrayElementsProperty(AtlasType elementType, AtlasVertex instanceVertex, String propertyName, List<Object> values) {
-        String actualPropertyName = GraphHelper.encodePropertyKey(propertyName);
-        if (AtlasGraphUtilsV1.isReference(elementType)) {
-            GraphHelper.setListPropertyFromElementIds(instanceVertex, actualPropertyName, (List) values);
-        }
-        else {
-            GraphHelper.setProperty(instanceVertex, actualPropertyName, values);
-        }
-    }
-}

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityGraphDiscoveryV1.java

@@ -17,10 +17,10 @@
  */
 package org.apache.atlas.repository.store.graph.v1;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.Iterator;
-import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -31,7 +31,6 @@ import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasStruct;
-import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscovery;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.repository.store.graph.EntityResolver;
@@ -41,38 +40,21 @@ import org.apache.atlas.type.AtlasMapType;
 import org.apache.atlas.type.AtlasStructType;
 import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.type.AtlasTypeRegistry;
-import org.apache.commons.lang3.StringUtils;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.inject.Inject;
-import com.google.inject.Provider;
+import org.apache.atlas.type.AtlasTypeUtil;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 
 public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
+    private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityGraphDiscoveryV1.class);
 
     private final AtlasTypeRegistry           typeRegistry;
-    private final EntityGraphDiscoveryContext discoveredEntities;
-    private final Set<String>                 processedIds    = new HashSet<>();
-    private final Collection<EntityResolver>  entityResolvers = new LinkedHashSet<>();
+    private final EntityGraphDiscoveryContext discoveryContext;
 
-    @Inject
-    public AtlasEntityGraphDiscoveryV1(AtlasTypeRegistry typeRegistry, Collection<Provider<EntityResolver>> entityResolverProviders) {
+    public AtlasEntityGraphDiscoveryV1(AtlasTypeRegistry typeRegistry, EntityStream entityStream) {
         this.typeRegistry     = typeRegistry;
-        this.discoveredEntities = new EntityGraphDiscoveryContext(typeRegistry);
-
-        for (Provider<EntityResolver> entityResolverProvider : entityResolverProviders) {
-             entityResolvers.add(entityResolverProvider.get());
-        }
-    }
-
-    @VisibleForTesting
-    public AtlasEntityGraphDiscoveryV1(AtlasTypeRegistry typeRegistry, List<EntityResolver> entityResolvers) {
-        this.typeRegistry       = typeRegistry;
-        this.discoveredEntities = new EntityGraphDiscoveryContext(typeRegistry);
-
-        for (EntityResolver entityResolver : entityResolvers) {
-            this.entityResolvers.add(entityResolver);
-        }
+        this.discoveryContext = new EntityGraphDiscoveryContext(typeRegistry, entityStream);
     }
 
     @Override
@@ -81,125 +63,148 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
     }
 
     @Override
-    public EntityGraphDiscoveryContext discoverEntities(final List<AtlasEntity> entities) throws AtlasBaseException {
+    public EntityGraphDiscoveryContext discoverEntities() throws AtlasBaseException {
 
-        //walk the graph and discover entity references
-        discover(entities);
+        // walk through entities in stream and validate them; record entity references
+        discoverAndValidate();
 
-        //resolve root and referred entities
+        // resolve entity references discovered in previous step
         resolveReferences();
 
-        return discoveredEntities;
+        return discoveryContext;
     }
 
     @Override
     public void cleanUp() throws AtlasBaseException {
-        processedIds.clear();
-        discoveredEntities.cleanUp();
+        discoveryContext.cleanUp();
+    }
 
-        for (EntityResolver resolver : entityResolvers) {
-            resolver.cleanUp();
+
+    protected void discoverAndValidate() throws AtlasBaseException {
+        EntityStream entityStream = discoveryContext.getEntityStream();
+
+        Set<String> walkedEntities = new HashSet<>();
+
+        // walk through top-level entities and find entity references
+        while (entityStream.hasNext()) {
+            AtlasEntity entity = entityStream.next();
+
+            if (entity != null) {
+                walkEntityGraph(entity);
+
+                walkedEntities.add(entity.getGuid());
             }
         }
 
+        // walk through entities referenced by other entities
+        // referencedGuids will be updated within this for() loop; avoid use of iterators
+        List<String> referencedGuids = discoveryContext.getReferencedGuids();
+        for (int i = 0; i < referencedGuids.size(); i++) {
+            String guid = referencedGuids.get(i);
 
-    protected void discover(List<AtlasEntity> entities) throws AtlasBaseException {
-        for (AtlasEntity entity : entities) {
-            AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());
-
-            if (type == null) {
-                throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
+            if (walkedEntities.contains(guid)) {
+                continue;
             }
 
-            discoveredEntities.addRootEntity(entity);
+            AtlasEntity entity = entityStream.getByGuid(guid);
 
-            walkEntityGraph(type, entity);
+            if (entity != null) {
+                walkEntityGraph(entity);
+
+                walkedEntities.add(entity.getGuid());
+            }
         }
     }
 
     protected void resolveReferences() throws AtlasBaseException {
-        for (EntityResolver resolver : entityResolvers) {
-            resolver.init(discoveredEntities);
+        EntityResolver[] entityResolvers = new EntityResolver[] { new IDBasedEntityResolver(),
+                                                                  new UniqAttrBasedEntityResolver(typeRegistry)
+                                                                };
 
-            resolver.resolveEntityReferences();
+        for (EntityResolver resolver : entityResolvers) {
+            resolver.resolveEntityReferences(discoveryContext);
         }
-
-        if (discoveredEntities.hasUnresolvedReferences()) {
-            throw new AtlasBaseException(AtlasErrorCode.UNRESOLVED_REFERENCES_FOUND,
-                                                         discoveredEntities.getUnresolvedIds().toString(),
-                                                         discoveredEntities.getUnresolvedIdsByUniqAttribs().toString());
     }
+
+    private void visitReference(AtlasEntityType type, Object val) throws AtlasBaseException {
+        if (type == null || val == null) {
+            return;
         }
 
-    private void visitReference(AtlasEntityType type, Object entity) throws AtlasBaseException {
-        if (entity != null) {
-            if (entity instanceof AtlasObjectId) {
-                AtlasObjectId objId = (AtlasObjectId)entity;
+        if (val instanceof AtlasObjectId) {
+            AtlasObjectId objId = (AtlasObjectId)val;
 
             if (!objId.isValid()) {
-                    throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, "Invalid object id " + objId);
+                throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, objId.toString());
             }
 
-                if (!StringUtils.isEmpty(objId.getGuid()) && (objId.isAssignedGuid() || objId.isUnAssignedGuid())) {
-                    discoveredEntities.addUnResolvedId(objId);
-                } else {
-                    discoveredEntities.addUnresolvedIdByUniqAttribs(objId);
+            recordObjectReference(objId);
+        } else if (val instanceof Map) {
+            AtlasObjectId objId = new AtlasObjectId((Map)val);
+
+            if (!objId.isValid()) {
+                throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, objId.toString());
             }
-            } else if (entity instanceof AtlasEntity) {
-                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, "Use AtlasObjectId to refer to another instance instead of AtlasEntity " + type.getTypeName());
+
+            recordObjectReference(objId);
+        } else if (val instanceof AtlasEntity) {
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, "found AtlasEntity");
         } else {
-                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, "Invalid object type " + entity.getClass());
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, val.toString());
         }
     }
+
+    void visitAttribute(AtlasType attrType, Object val) throws AtlasBaseException {
+        if (attrType == null || val == null) {
+            return;
         }
 
-    void visitAttribute(AtlasStructType parentType, AtlasType attrType, AtlasAttributeDef attrDef, Object val) throws AtlasBaseException {
-        if (val != null) {
-            if ( isPrimitive(attrType.getTypeCategory()) ) {
+        if (isPrimitive(attrType.getTypeCategory()) ) {
             return;
         }
         if (attrType.getTypeCategory() == TypeCategory.ARRAY) {
             AtlasArrayType arrayType = (AtlasArrayType) attrType;
             AtlasType      elemType  = arrayType.getElementType();
 
-                visitCollectionReferences(parentType, attrType, attrDef, elemType, val);
+            visitCollectionReferences(elemType, val);
         } else if (attrType.getTypeCategory() == TypeCategory.MAP) {
             AtlasType keyType   = ((AtlasMapType) attrType).getKeyType();
             AtlasType valueType = ((AtlasMapType) attrType).getValueType();
 
-                visitMapReferences(parentType, attrType, attrDef, keyType, valueType, val);
+            visitMapReferences(keyType, valueType, val);
         } else if (attrType.getTypeCategory() == TypeCategory.STRUCT) {
             visitStruct((AtlasStructType)attrType, val);
         } else if (attrType.getTypeCategory() == TypeCategory.ENTITY) {
             visitReference((AtlasEntityType) attrType,  val);
         }
     }
+
+    void visitMapReferences(AtlasType keyType, AtlasType valueType, Object val) throws AtlasBaseException {
+        if (keyType == null || valueType == null || val == null) {
+            return;
         }
 
-    void visitMapReferences(AtlasStructType parentType, final AtlasType attrType, AtlasAttributeDef attrDef, AtlasType keyType, AtlasType valueType, Object val) throws AtlasBaseException {
         if (isPrimitive(keyType.getTypeCategory()) && isPrimitive(valueType.getTypeCategory())) {
             return;
         }
 
-        if (val != null) {
         if (Map.class.isAssignableFrom(val.getClass())) {
             Iterator<Map.Entry> it = ((Map) val).entrySet().iterator();
             while (it.hasNext()) {
                 Map.Entry e = it.next();
-                    visitAttribute(parentType, keyType, attrDef, e.getKey());
-                    visitAttribute(parentType, valueType, attrDef, e.getValue());
-                }
+                visitAttribute(keyType, e.getKey());
+                visitAttribute(valueType, e.getValue());
             }
         }
     }
 
-    void visitCollectionReferences(final AtlasStructType parentType, final AtlasType attrType, final AtlasAttributeDef attrDef, AtlasType elemType, Object val) throws AtlasBaseException {
-        if (isPrimitive(elemType.getTypeCategory())) {
+    void visitCollectionReferences(AtlasType elemType, Object val) throws AtlasBaseException {
+        if (elemType == null || val == null || isPrimitive(elemType.getTypeCategory())) {
             return;
         }
 
-        if (val != null) {
         Iterator it = null;
+
         if (val instanceof Collection) {
             it = ((Collection) val).iterator();
         } else if (val instanceof Iterable) {
@@ -207,35 +212,90 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
         } else if (val instanceof Iterator) {
             it = (Iterator) val;
         }
+
         if (it != null) {
             while (it.hasNext()) {
                 Object elem = it.next();
-                    visitAttribute(parentType, elemType, attrDef, elem);
-                }
+                visitAttribute(elemType, elem);
             }
         }
     }
 
     void visitStruct(AtlasStructType structType, Object val) throws AtlasBaseException {
-        if (structType == null) {
+        if (structType == null || val == null) {
             return;
         }
 
-        for (AtlasStructType.AtlasAttribute attribute : structType.getAllAttributes().values()) {
+        AtlasStruct struct;
+
+        if (val instanceof AtlasStruct) {
+            struct = (AtlasStruct) val;
+        } else if (val instanceof Map) {
+            Map attributes = AtlasTypeUtil.toStructAttributes((Map) val);
+
+            struct = new AtlasStruct(structType.getTypeName(), attributes);
+        } else {
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_STRUCT_VALUE, val.toString());
+        }
+
+        for (AtlasAttribute attribute : structType.getAllAttributes().values()) {
             AtlasType attrType = attribute.getAttributeType();
-            Object    attrVal  = ((AtlasStruct) val).getAttribute(attribute.getName());
+            Object    attrVal  = struct.getAttribute(attribute.getName());
 
-            visitAttribute(structType, attrType, attribute.getAttributeDef(), attrVal);
+            visitAttribute(attrType, attrVal);
         }
     }
 
 
-    void walkEntityGraph(AtlasEntityType entityType, AtlasEntity entity) throws AtlasBaseException {
-        visitStruct(entityType, entity);
+    void walkEntityGraph(AtlasEntity entity) throws AtlasBaseException {
+        if (entity == null) {
+            return;
+        }
+
+        validateAndNormalize(entity);
+        AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());
+
+        recordObjectReference(entity.getGuid());
+
+        visitStruct(type, entity);
     }
 
 
     boolean isPrimitive(TypeCategory typeCategory) {
         return typeCategory == TypeCategory.PRIMITIVE || typeCategory == TypeCategory.ENUM;
     }
+
+    private void validateAndNormalize(AtlasEntity entity) throws AtlasBaseException {
+        List<String> messages = new ArrayList<>();
+
+        if (!AtlasEntity.isAssigned(entity.getGuid()) && !AtlasEntity.isUnAssigned(entity.getGuid())) {
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, "invalid guid " + entity.getGuid());
+        }
+
+        AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());
+
+        if (type == null) {
+            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
+        }
+
+        type.validateValue(entity, entity.getTypeName(), messages);
+
+        if (!messages.isEmpty()) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, messages);
+        }
+
+        type.getNormalizedValue(entity);
+    }
+
+    private void recordObjectReference(String guid) {
+        discoveryContext.addReferencedGuid(guid);
+    }
+
+    private void recordObjectReference(AtlasObjectId objId) {
+        if (objId.isValidGuid()) {
+            discoveryContext.addReferencedGuid(objId.getGuid());
+        } else {
+            discoveryContext.addReferencedByUniqAttribs(objId);
+        }
+    }
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java

@@ -19,234 +19,217 @@ package org.apache.atlas.repository.store.graph.v1;
 
 
 import com.google.inject.Inject;
+import com.google.inject.Singleton;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.GraphTransaction;
 import org.apache.atlas.RequestContextV1;
 import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.model.TypeCategory;
 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.AtlasEntity.AtlasEntitiesWithExtInfo;
-import org.apache.atlas.model.instance.AtlasEntity.Status;
-import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.EntityMutationResponse;
-import org.apache.atlas.repository.graph.AtlasGraphProvider;
-import org.apache.atlas.repository.graph.GraphHelper;
-import org.apache.atlas.repository.Constants;
-import org.apache.atlas.repository.graphdb.AtlasGraph;
-import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscovery;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
-import org.apache.atlas.repository.store.graph.EntityResolver;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 
 
+@Singleton
 public class AtlasEntityStoreV1 implements AtlasEntityStore {
-
-    protected AtlasTypeRegistry typeRegistry;
-
-    private final EntityGraphMapper graphMapper;
-    private final AtlasGraph        graph;
-
     private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityStoreV1.class);
 
-    @Inject
-    public AtlasEntityStoreV1(EntityGraphMapper vertexMapper) {
-        this.graphMapper  = vertexMapper;
-        this.graph        = AtlasGraphProvider.getGraphInstance();
-    }
+
+    private final DeleteHandlerV1   deleteHandler;
+    private final AtlasTypeRegistry typeRegistry;
 
     @Inject
-    public void init(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
+    public AtlasEntityStoreV1(DeleteHandlerV1 deleteHandler, AtlasTypeRegistry typeRegistry) {
+        this.deleteHandler = deleteHandler;
         this.typeRegistry  = typeRegistry;
     }
 
     @Override
-    public AtlasEntityWithExtInfo getById(final String guid) throws AtlasBaseException {
+    @GraphTransaction
+    public AtlasEntityWithExtInfo getById(String guid) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
-            LOG.debug("Retrieving entity with guid={}", guid);
+            LOG.debug("==> getById({})", guid);
         }
 
         EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry);
 
-        return entityRetriever.toAtlasEntityWithExtInfo(guid);
+        AtlasEntityWithExtInfo ret = entityRetriever.toAtlasEntityWithExtInfo(guid);
+
+        if (ret == null) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
+        }
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== getById({}): {}", guid, ret);
+        }
+
+        return ret;
     }
 
     @Override
-    public AtlasEntityWithExtInfo getByUniqueAttribute(AtlasEntityType entityType, Map<String, Object> uniqAttributes) throws AtlasBaseException {
-        String entityTypeName = entityType.getTypeName();
+    @GraphTransaction
+    public AtlasEntitiesWithExtInfo getByIds(List<String> guids) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> getByIds({})", guids);
+        }
+
+        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry);
+
+        AtlasEntitiesWithExtInfo ret = entityRetriever.toAtlasEntitiesWithExtInfo(guids);
 
         if (LOG.isDebugEnabled()) {
-            LOG.debug("Retrieving entity with type={} and attributes={}: values={}", entityTypeName, uniqAttributes);
+            LOG.debug("<== getByIds({}): {}", guids, ret);
+        }
+
+        return ret;
+    }
+
+    @Override
+    @GraphTransaction
+    public AtlasEntityWithExtInfo getByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes)
+                                                                                            throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> getByUniqueAttribute({}, {})", entityType.getTypeName(), uniqAttributes);
         }
 
         AtlasVertex entityVertex = AtlasGraphUtilsV1.getVertexByUniqueAttributes(entityType, uniqAttributes);
 
         EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry);
 
-        return entityRetriever.toAtlasEntityWithExtInfo(entityVertex);
+        AtlasEntityWithExtInfo ret = entityRetriever.toAtlasEntityWithExtInfo(entityVertex);
+
+        if (ret == null) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, entityType.getTypeName(),
+                                         uniqAttributes.toString());
         }
 
-    @Override
-    public EntityMutationResponse deleteById(final String guid) {
-        return null;
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== getByUniqueAttribute({}, {}): {}", entityType.getTypeName(), uniqAttributes, ret);
+        }
+
+        return ret;
     }
 
     @Override
     @GraphTransaction
-    public EntityMutationResponse createOrUpdate(final Map<String, AtlasEntity> entities) throws AtlasBaseException {
-
+    public EntityMutationResponse createOrUpdate(EntityStream entityStream) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
-            LOG.debug("==> AtlasEntityStoreV1.createOrUpdate({}, {})", entities);
+            LOG.debug("==> createOrUpdate()");
         }
 
-        //Validate
-        List<AtlasEntity> normalizedEntities = validateAndNormalize(entities);
+        if (entityStream == null || !entityStream.hasNext()) {
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "no entities to create/update.");
+        }
+
+        EntityGraphMapper entityGraphMapper = new EntityGraphMapper(deleteHandler, typeRegistry);
+
+        // Create/Update entities
+        EntityMutationContext context = preCreateOrUpdate(entityStream, entityGraphMapper);
 
-        //Discover entities, create vertices
-        EntityMutationContext ctx = preCreateOrUpdate(normalizedEntities);
+        EntityMutationResponse ret = entityGraphMapper.mapAttributes(context);
+
+        ret.setGuidAssignments(context.getGuidAssignments());
 
         if (LOG.isDebugEnabled()) {
-            LOG.debug("<== AtlasStructDefStoreV1.createOrUpdate({}, {}): {}", entities);
+            LOG.debug("<== createOrUpdate()");
         }
 
-        return graphMapper.mapAttributes(ctx);
+        return ret;
     }
 
     @Override
-    public AtlasEntitiesWithExtInfo getByIds(final List<String> guids) throws AtlasBaseException {
-        return null;
+    @GraphTransaction
+    public EntityMutationResponse updateByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes,
+                                                          AtlasEntity entity) throws AtlasBaseException {
+        throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "updateByUniqueAttributes() not implemented yet");
     }
 
     @Override
-    public EntityMutationResponse deleteByIds(final List<String> guid) throws AtlasBaseException {
-        return null;
+    @GraphTransaction
+    public EntityMutationResponse deleteById(String guid) throws AtlasBaseException {
+        throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "deleteById() not implemented yet");
     }
 
     @Override
-    public EntityMutationResponse updateByUniqueAttribute(final String typeName, final String attributeName, final String attributeValue, final AtlasEntity entity) throws AtlasBaseException {
-        return null;
+    @GraphTransaction
+    public EntityMutationResponse deleteByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes)
+            throws AtlasBaseException {
+        throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "deleteByUniqueAttributes() not implemented yet");
     }
 
     @Override
-    public EntityMutationResponse deleteByUniqueAttribute(final String typeName, final String attributeName, final String attributeValue) throws AtlasBaseException {
-        return null;
+    @GraphTransaction
+    public EntityMutationResponse deleteByIds(List<String> guids) throws AtlasBaseException {
+        throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "deleteByIds() not implemented yet");
     }
 
     @Override
-    public void addClassifications(final String guid, final List<AtlasClassification> classification) throws AtlasBaseException {
-
+    @GraphTransaction
+    public void addClassifications(String guid, List<AtlasClassification> classification) throws AtlasBaseException {
+        throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "addClassifications() not implemented yet");
     }
 
     @Override
-    public void updateClassifications(final String guid, final List<AtlasClassification> classification) throws AtlasBaseException {
-
+    @GraphTransaction
+    public void updateClassifications(String guid, List<AtlasClassification> classification) throws AtlasBaseException {
+        throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "updateClassifications() not implemented yet");
     }
 
     @Override
-    public void deleteClassifications(final String guid, final List<String> classificationNames) throws AtlasBaseException {
-
+    @GraphTransaction
+    public void deleteClassifications(String guid, List<String> classificationNames) throws AtlasBaseException {
+        throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "deleteClassifications() not implemented yet");
     }
 
-    private EntityMutationContext preCreateOrUpdate(final List<AtlasEntity> atlasEntities) throws AtlasBaseException {
-        List<EntityResolver> entityResolvers = new ArrayList<>();
-
-        entityResolvers.add(new IDBasedEntityResolver());
-        entityResolvers.add(new UniqAttrBasedEntityResolver(typeRegistry));
 
-        EntityGraphDiscovery        graphDiscoverer    = new AtlasEntityGraphDiscoveryV1(typeRegistry, entityResolvers);
-        EntityGraphDiscoveryContext discoveredEntities = graphDiscoverer.discoverEntities(atlasEntities);
-        EntityMutationContext       context            = new EntityMutationContext(discoveredEntities);
+    private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, EntityGraphMapper entityGraphMapper) throws AtlasBaseException {
+        EntityGraphDiscovery        graphDiscoverer  = new AtlasEntityGraphDiscoveryV1(typeRegistry, entityStream);
+        EntityGraphDiscoveryContext discoveryContext = graphDiscoverer.discoverEntities();
+        EntityMutationContext       context          = new EntityMutationContext(discoveryContext);
 
-        for (AtlasEntity entity : discoveredEntities.getRootEntities()) {
-            if (LOG.isDebugEnabled()) {
-                LOG.debug("==> AtlasEntityStoreV1.preCreateOrUpdate({}): {}", entity);
-            }
+        for (String guid : discoveryContext.getReferencedGuids()) {
+            AtlasVertex vertex = discoveryContext.getResolvedEntityVertex(guid);
+            AtlasEntity entity = entityStream.getByGuid(guid);
 
+            if (vertex != null) {
+                // entity would be null if guid is not in the stream but referenced by an entity in the stream
+                if (entity != null) {
                     AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
 
-            if (entityType == null) {
-                throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
-            }
-
-            final AtlasVertex vertex;
-            AtlasObjectId     objId = entity.getAtlasObjectId();
-
-            if (discoveredEntities.isResolvedId(objId) ) {
-                vertex = discoveredEntities.getResolvedEntityVertex(objId);
-
                     context.addUpdated(entity, entityType, vertex);
 
-                String guid = AtlasGraphUtilsV1.getIdFromVertex(vertex);
-
-                RequestContextV1.get().recordEntityUpdate(new AtlasObjectId(entityType.getTypeName(), guid));
+                    RequestContextV1.get().recordEntityUpdate(entity.getAtlasObjectId());
+                }
             } else {
+                AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
+
                 //Create vertices which do not exist in the repository
-                vertex = graphMapper.createVertexTemplate(entity, entityType);
+                vertex = entityGraphMapper.createVertex(entity);
 
-                context.addCreated(entity, entityType, vertex);
+                discoveryContext.addResolvedGuid(guid, vertex);
 
-                discoveredEntities.addResolvedId(objId, vertex);
-                discoveredEntities.removeUnResolvedId(objId);
+                String generatedGuid = AtlasGraphUtilsV1.getIdFromVertex(vertex);
 
-                String guid = AtlasGraphUtilsV1.getIdFromVertex(vertex);
+                entity.setGuid(generatedGuid);
 
-                RequestContextV1.get().recordEntityCreate(new AtlasObjectId(entityType.getTypeName(), guid));
-            }
+                context.addCreated(guid, entity, entityType, vertex);
 
-            if (LOG.isDebugEnabled()) {
-                LOG.debug("<== AtlasEntityStoreV1.preCreateOrUpdate({}): {}", entity, vertex);
+                RequestContextV1.get().recordEntityCreate(entity.getAtlasObjectId());
             }
         }
 
         return context;
     }
-
-    private List<AtlasEntity> validateAndNormalize(final Map<String, AtlasEntity> entities) throws AtlasBaseException {
-        List<AtlasEntity> normalizedEntities = new ArrayList<>();
-        List<String>      messages           = new ArrayList<>();
-
-        for (String entityId : entities.keySet()) {
-            if ( !AtlasEntity.isAssigned(entityId) && !AtlasEntity.isUnAssigned(entityId)) {
-                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, ": Guid in map key is invalid " + entityId);
-            }
-
-            AtlasEntity entity = entities.get(entityId);
-
-            if ( entity == null) {
-                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, ": Entity is null for guid " + entityId);
-            }
-
-            AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());
-            if (type == null) {
-                throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
-            }
-
-            type.validateValue(entity, entity.getTypeName(), messages);
-
-            if ( !messages.isEmpty()) {
-                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, messages);
-            }
-
-            AtlasEntity normalizedEntity = (AtlasEntity) type.getNormalizedValue(entity);
-
-            normalizedEntities.add(normalizedEntity);
-        }
-
-        return normalizedEntities;
-    }
-
-    public void cleanUp() throws AtlasBaseException {
-    }
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStream.java

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas.repository.store.graph.v1;
+
+import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
+
+import java.util.Iterator;
+
+public class AtlasEntityStream implements EntityStream {
+    private AtlasEntitiesWithExtInfo entitiesWithExtInfo = new AtlasEntitiesWithExtInfo();
+    private Iterator<AtlasEntity>    iterator;
+
+    public AtlasEntityStream() {
+    }
+
+    public AtlasEntityStream(AtlasEntity entity) {
+        this(new AtlasEntitiesWithExtInfo(entity));
+    }
+
+    public AtlasEntityStream(AtlasEntityWithExtInfo entityWithExtInfo) {
+        this(new AtlasEntitiesWithExtInfo(entityWithExtInfo));
+    }
+
+    public AtlasEntityStream(AtlasEntitiesWithExtInfo entitiesWithExtInfo) {
+        this.entitiesWithExtInfo = entitiesWithExtInfo;
+        this.iterator            = this.entitiesWithExtInfo.getEntities().iterator();
+    }
+
+    @Override
+    public boolean hasNext() {
+        return iterator.hasNext();
+    }
+
+    @Override
+    public AtlasEntity next() {
+        return iterator.hasNext() ? iterator.next() : null;
+    }
+
+    @Override
+    public void reset() {
+        this.iterator = entitiesWithExtInfo.getEntities().iterator();
+    }
+
+    @Override
+    public AtlasEntity getByGuid(String guid) {
+        return entitiesWithExtInfo.getEntity(guid);
+    }
+
+    @Override
+    public String toString() {
+        StringBuffer sb = new StringBuffer("AtlasEntityStream{");
+
+        sb.append("entitiesWithExtInfo=").append(entitiesWithExtInfo);
+        sb.append(", iterator=").append(iterator);
+        sb.append('}');
+
+        return sb.toString();
+    }
+}

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasGraphUtilsV1.java

@@ -217,7 +217,7 @@ public class AtlasGraphUtilsV1 {
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, entityType.getTypeName(),
-                                         uniqAttributes.keySet().toString(), uniqAttributes.values().toString());
+                                         uniqAttributes.toString());
         }
 
         return entityVertex;

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/GraphMutationContext.java → repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AttributeMutationContext.java

@@ -18,26 +18,25 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 
-import com.google.common.base.Optional;
-import org.apache.atlas.model.instance.EntityMutations;
+import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
 import org.apache.atlas.model.typedef.AtlasStructDef;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasStructType;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.atlas.type.AtlasType;
 
-import java.util.Objects;
-
-public class GraphMutationContext {
 
+import java.util.Objects;
 
-    private EntityMutations.EntityOperation op;
+public class AttributeMutationContext {
+    private EntityOperation op;
     /**
      * Atlas Attribute
      */
 
-    private AtlasStructType.AtlasAttribute attribute;
+    private AtlasAttribute attribute;
 
     /**
      * Overriding type for which elements are being mapped
@@ -49,6 +48,8 @@ public class GraphMutationContext {
      */
     private Object value;
 
+    private String vertexProperty;
+
     /**
      *
      * The vertex which corresponds to the entity/struct for which we are mapping a complex attributes like struct, traits
@@ -56,34 +57,28 @@ public class GraphMutationContext {
     AtlasVertex referringVertex;
 
     /**
-     * the vertex property that we are updating
-     */
-
-    String vertexPropertyKey;
-
-    /**
      * The current edge(in case of updates) from the parent entity/struct to the complex attribute like struct, trait
      */
-    Optional<AtlasEdge> existingEdge;
-
+    AtlasEdge existingEdge;
 
-    private GraphMutationContext(final Builder builder) {
-        this.op = builder.op;
-        this.attribute = builder.attribute;
-        this.currentElementType = builder.elementType;
-        this.existingEdge = builder.currentEdge;
-        this.value = builder.currentValue;
-        this.referringVertex = builder.referringVertex;
-        this.vertexPropertyKey = builder.vertexPropertyKey;
+    public AttributeMutationContext(EntityOperation op, AtlasVertex referringVertex, AtlasAttribute attribute, Object value) {
+        this(op, referringVertex, attribute, value, attribute.getVertexPropertyName(), null, null);
     }
 
-    public String getVertexPropertyKey() {
-        return vertexPropertyKey;
+    public AttributeMutationContext(EntityOperation op, AtlasVertex referringVertex, AtlasAttribute attribute, Object value,
+                                    String vertexProperty, AtlasType currentElementType, AtlasEdge currentEdge) {
+        this.op                 = op;
+        this.referringVertex    = referringVertex;
+        this.attribute          = attribute;
+        this.value              = value;
+        this.vertexProperty     = vertexProperty;
+        this.currentElementType = currentElementType;
+        this.existingEdge       = currentEdge;
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(op, attribute, value, referringVertex, vertexPropertyKey, existingEdge);
+        return Objects.hash(op, referringVertex, attribute, value, vertexProperty, currentElementType, existingEdge);
     }
 
     @Override
@@ -95,69 +90,17 @@ public class GraphMutationContext {
         } else if (obj.getClass() != getClass()) {
             return false;
         } else {
-            GraphMutationContext rhs = (GraphMutationContext) obj;
-            return Objects.equals(attribute, rhs.getAttribute())
-                 && Objects.equals(value, rhs.getValue())
-                 && Objects.equals(referringVertex, rhs.getReferringVertex())
-                 && Objects.equals(vertexPropertyKey, rhs.getReferringVertex())
-                 && Objects.equals(existingEdge, rhs.getCurrentEdge())
-                 && Objects.equals(op, rhs.getOp());
-        }
-    }
-
-
-    public static final class Builder {
-
-        private final EntityMutations.EntityOperation op;
-
-        private final AtlasStructType.AtlasAttribute attribute;
-
-        private final AtlasType elementType;
-
-        private final Object currentValue;
-
-        private AtlasVertex referringVertex;
-
-        private Optional<AtlasEdge> currentEdge = Optional.absent();
-
-        private  String vertexPropertyKey;
-
-
-        public Builder(EntityMutations.EntityOperation op, AtlasStructType.AtlasAttribute attribute, AtlasType currentElementType, Object currentValue) {
-            this.op = op;
-            this.attribute = attribute;
-            this.elementType = currentElementType;
-            this.currentValue = currentValue;
+            AttributeMutationContext rhs = (AttributeMutationContext) obj;
+            return Objects.equals(op, rhs.op)
+                    && Objects.equals(referringVertex, rhs.referringVertex)
+                    && Objects.equals(attribute, rhs.attribute)
+                    && Objects.equals(value, rhs.value)
+                    && Objects.equals(vertexProperty, rhs.vertexProperty)
+                    && Objects.equals(currentElementType, rhs.currentElementType)
+                    && Objects.equals(existingEdge, rhs.existingEdge);
         }
-
-        public Builder(EntityMutations.EntityOperation op, AtlasStructType.AtlasAttribute attribute, Object currentValue) {
-           this(op, attribute, null, currentValue);
     }
 
-        Builder referringVertex(AtlasVertex referringVertex) {
-            this.referringVertex = referringVertex;
-            return this;
-        }
-
-        Builder edge(AtlasEdge edge) {
-            this.currentEdge = Optional.of(edge);
-            return this;
-        }
-
-        Builder edge(Optional<AtlasEdge> edge) {
-            this.currentEdge = edge;
-            return this;
-        }
-
-        Builder vertexProperty(String propertyKey) {
-            this.vertexPropertyKey = propertyKey;
-            return this;
-        }
-
-        GraphMutationContext build() {
-            return new GraphMutationContext(this);
-        }
-    }
 
     public AtlasStructType getParentType() {
         return attribute.getDefinedInType();
@@ -183,11 +126,11 @@ public class GraphMutationContext {
         return value;
     }
 
-    public AtlasVertex getReferringVertex() {
-        return referringVertex;
-    }
+    public String getVertexProperty() { return vertexProperty; }
 
-    public Optional<AtlasEdge> getCurrentEdge() {
+    public AtlasVertex getReferringVertex() { return referringVertex; }
+
+    public AtlasEdge getCurrentEdge() {
         return existingEdge;
     }
 
@@ -195,15 +138,13 @@ public class GraphMutationContext {
         this.currentElementType = attrType;
     }
 
-    public AtlasStructType.AtlasAttribute getAttribute() {
+    public AtlasAttribute getAttribute() {
         return attribute;
     }
 
-    public EntityMutations.EntityOperation getOp() {
+    public EntityOperation getOp() {
         return op;
     }
 
-    public void setExistingEdge(final Optional<AtlasEdge> existingEdge) {
-        this.existingEdge = existingEdge;
-    }
+    public void setExistingEdge(AtlasEdge existingEdge) { this.existingEdge = existingEdge; }
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java

@@ -98,7 +98,7 @@ public abstract class DeleteHandlerV1 {
             // Record all deletion candidate GUIDs in RequestContext
             // and gather deletion candidate vertices.
             for (GraphHelper.VertexInfo vertexInfo : compositeVertices) {
-                requestContext.recordEntityDelete(new AtlasObjectId(vertexInfo.getTypeName(), vertexInfo.getGuid()));
+                requestContext.recordEntityDelete(new AtlasObjectId(vertexInfo.getGuid(), vertexInfo.getTypeName()));
                 deletionCandidateVertices.add(vertexInfo.getVertex());
             }
         }
@@ -324,7 +324,7 @@ public abstract class DeleteHandlerV1 {
                     String propertyName = AtlasGraphUtilsV1.getQualifiedAttributePropertyKey(structType, attributeInfo.getName());
 
                     if (AtlasGraphUtilsV1.isReference(valueTypeCategory)) {
-                        List<Object> keys = ArrayVertexMapper.getArrayElementsProperty(keyType, instanceVertex, propertyName);
+                        List<Object> keys = EntityGraphMapper.getArrayElementsProperty(keyType, instanceVertex, propertyName);
                         if (keys != null) {
                             for (Object key : keys) {
                                 String mapEdgeLabel = GraphHelper.getQualifiedNameForMapKey(edgeLabel, (String) key);
@@ -513,7 +513,7 @@ public abstract class DeleteHandlerV1 {
             GraphHelper.setProperty(outVertex, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY,
                 requestContext.getRequestTime());
             GraphHelper.setProperty(outVertex, Constants.MODIFIED_BY_KEY, requestContext.getUser());
-            requestContext.recordEntityUpdate(new AtlasObjectId(typeName, outId));
+            requestContext.recordEntityUpdate(new AtlasObjectId(outId, typeName));
         }
     }
 

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

@@ -20,6 +20,7 @@ package org.apache.atlas.repository.store.graph.v1;
 
 import com.google.inject.Inject;
 import org.apache.atlas.AtlasErrorCode;
+import org.apache.atlas.AtlasException;
 import org.apache.atlas.RequestContextV1;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
@@ -27,85 +28,616 @@ import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasStruct;
-import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
+import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
+import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.RepositoryException;
+import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.type.AtlasArrayType;
 import org.apache.atlas.type.AtlasEntityType;
+import org.apache.atlas.type.AtlasMapType;
 import org.apache.atlas.type.AtlasStructType;
 import org.apache.atlas.type.AtlasType;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.atlas.type.AtlasTypeUtil;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.collections.MapUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 
-public class EntityGraphMapper implements InstanceGraphMapper<AtlasEdge> {
+import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.CREATE;
+import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.UPDATE;
+import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.DELETE;
 
-    private static final Logger LOG = LoggerFactory.getLogger(EntityGraphMapper.class);
 
-    protected final GraphHelper graphHelper = GraphHelper.getInstance();
+public class EntityGraphMapper {
+    private static final Logger LOG = LoggerFactory.getLogger(EntityGraphMapper.class);
 
-    protected EntityMutationContext context;
+    private final AtlasGraph        graph       = AtlasGraphProvider.getGraphInstance();
+    private final GraphHelper       graphHelper = GraphHelper.getInstance();
+    private final DeleteHandlerV1   deleteHandler;
+    private final AtlasTypeRegistry typeRegistry;
 
-    protected final StructVertexMapper structVertexMapper;
 
     @Inject
-    public EntityGraphMapper(ArrayVertexMapper arrayVertexMapper, MapVertexMapper mapVertexMapper, DeleteHandlerV1 deleteHandler) {
-        this.structVertexMapper = new StructVertexMapper(arrayVertexMapper, mapVertexMapper, deleteHandler);
-        arrayVertexMapper.init(structVertexMapper);
-        mapVertexMapper.init(structVertexMapper);
+    public EntityGraphMapper(DeleteHandlerV1 deleteHandler, AtlasTypeRegistry typeRegistry) {
+        this.deleteHandler = deleteHandler;
+        this.typeRegistry  = typeRegistry;
     }
 
-    public AtlasVertex createVertexTemplate(final AtlasStruct instance, final AtlasStructType structType) {
-        AtlasVertex vertex = structVertexMapper.createVertexTemplate(instance, structType);
+    public AtlasVertex createVertex(AtlasEntity entity) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> createVertex({})", entity.getTypeName());
+        }
 
-        AtlasEntityType entityType = (AtlasEntityType) structType;
-        AtlasEntity entity = (AtlasEntity) instance;
+        AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
+
+        AtlasVertex ret = createStructVertex(entity);
 
-        // add super types
         for (String superTypeName : entityType.getAllSuperTypes()) {
-            AtlasGraphUtilsV1.addProperty(vertex, Constants.SUPER_TYPES_PROPERTY_KEY, superTypeName);
+            AtlasGraphUtilsV1.addProperty(ret, Constants.SUPER_TYPES_PROPERTY_KEY, superTypeName);
         }
 
         final String guid = UUID.randomUUID().toString();
 
-        // add identity
-        AtlasGraphUtilsV1.setProperty(vertex, Constants.GUID_PROPERTY_KEY, guid);
+        AtlasGraphUtilsV1.setProperty(ret, Constants.GUID_PROPERTY_KEY, guid);
+        AtlasGraphUtilsV1.setProperty(ret, Constants.VERSION_PROPERTY_KEY, getEntityVersion(entity));
+
+        return ret;
+    }
+
+    public EntityMutationResponse mapAttributes(EntityMutationContext context) throws AtlasBaseException {
+        EntityMutationResponse resp = new EntityMutationResponse();
+
+        Collection<AtlasEntity> createdEntities = context.getCreatedEntities();
+        Collection<AtlasEntity> updatedEntities = context.getUpdatedEntities();
+
+        if (CollectionUtils.isNotEmpty(createdEntities)) {
+            for (AtlasEntity createdEntity : createdEntities) {
+                String          guid       = createdEntity.getGuid();
+                AtlasVertex     vertex     = context.getVertex(guid);
+                AtlasEntityType entityType = context.getType(guid);
+
+                mapAttributes(createdEntity, vertex, CREATE, context);
+
+                resp.addEntity(CREATE, constructHeader(createdEntity, entityType, vertex));
+            }
+        }
+
+        if (CollectionUtils.isNotEmpty(updatedEntities)) {
+            for (AtlasEntity updatedEntity : updatedEntities) {
+                String          guid       = updatedEntity.getGuid();
+                AtlasVertex     vertex     = context.getVertex(guid);
+                AtlasEntityType entityType = context.getType(guid);
+
+                mapAttributes(updatedEntity, vertex, UPDATE, context);
+
+                resp.addEntity(UPDATE, constructHeader(updatedEntity, entityType, vertex));
+            }
+        }
+
+        RequestContextV1 req = RequestContextV1.get();
+
+        for (AtlasObjectId id : req.getDeletedEntityIds()) {
+            resp.addEntity(DELETE, constructHeader(id));
+        }
+
+        return resp;
+    }
+
+    private AtlasVertex createStructVertex(AtlasStruct struct) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> createStructVertex({})", struct.getTypeName());
+        }
+
+        final AtlasVertex ret = graph.addVertex();
+
+        AtlasGraphUtilsV1.setProperty(ret, Constants.ENTITY_TYPE_PROPERTY_KEY, struct.getTypeName());
+        AtlasGraphUtilsV1.setProperty(ret, Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
+        AtlasGraphUtilsV1.setProperty(ret, Constants.TIMESTAMP_PROPERTY_KEY, RequestContextV1.get().getRequestTime());
+        AtlasGraphUtilsV1.setProperty(ret, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContextV1.get().getRequestTime());
+        AtlasGraphUtilsV1.setProperty(ret, Constants.CREATED_BY_KEY, RequestContextV1.get().getUser());
+        GraphHelper.setProperty(ret, Constants.MODIFIED_BY_KEY, RequestContextV1.get().getUser());
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== createStructVertex({})", struct.getTypeName());
+        }
+
+        return ret;
+    }
+
+    private void mapAttributes(AtlasStruct struct, AtlasVertex vertex, EntityOperation op, EntityMutationContext context) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> mapAttributes({}, {})", op, struct.getTypeName());
+        }
+
+        if (MapUtils.isNotEmpty(struct.getAttributes())) {
+            AtlasStructType structType = getStructType(struct.getTypeName());
+
+            if (op.equals(CREATE)) {
+                for (AtlasAttribute attribute : structType.getAllAttributes().values()) {
+                    Object attrValue = struct.getAttribute(attribute.getName());
+
+                    mapAttribute(attribute, attrValue, vertex, op, context);
+                }
+            } else if (op.equals(UPDATE)) {
+                for (String attrName : struct.getAttributes().keySet()) {
+                    AtlasAttribute attribute = structType.getAttribute(attrName);
 
-        // add version information
-        AtlasGraphUtilsV1.setProperty(vertex, Constants.VERSION_PROPERTY_KEY, Integer.valueOf(entity.getVersion().intValue()));
+                    if (attribute != null) {
+                        Object attrValue = struct.getAttribute(attrName);
 
-        return vertex;
+                        mapAttribute(attribute, attrValue, vertex, op, context);
+                    } else {
+                        LOG.warn("mapAttributes(): invalid attribute {}.{}. Ignored..", struct.getTypeName(), attrName);
+                    }
+                }
             }
 
+            updateModificationMetadata(vertex);
+        }
 
-    @Override
-    public AtlasEdge toGraph(GraphMutationContext ctx) throws AtlasBaseException {
-        AtlasEdge result = null;
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== mapAttributes({}, {})", op, struct.getTypeName());
+        }
+    }
+
+    private void mapAttribute(AtlasAttribute attribute, Object attrValue, AtlasVertex vertex, EntityOperation op, EntityMutationContext context) throws AtlasBaseException {
+        if (attrValue == null) {
+            AtlasType attrType = attribute.getAttributeType();
+
+            if (attrType.getTypeCategory() == TypeCategory.PRIMITIVE) {
+                if (attribute.getAttributeDef().getIsOptional()) {
+                    attrValue = attrType.createOptionalDefaultValue();
+                } else {
+                    attrValue = attrType.createDefaultValue();
+                }
+            }
+        }
+
+        AttributeMutationContext ctx = new AttributeMutationContext(op, vertex, attribute, attrValue);
+
+        mapToVertexByTypeCategory(ctx, context);
+    }
+
+    private Object mapToVertexByTypeCategory(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
+        if (ctx.getOp() == EntityMutations.EntityOperation.CREATE && ctx.getValue() == null) {
+            return null;
+        }
+
+        switch (ctx.getAttrType().getTypeCategory()) {
+            case PRIMITIVE:
+            case ENUM:
+                return mapPrimitiveValue(ctx);
+
+            case STRUCT: {
+                String    edgeLabel   = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexProperty());
+                AtlasEdge currentEdge = graphHelper.getEdgeForLabel(ctx.getReferringVertex(), edgeLabel);
+                AtlasEdge edge        = currentEdge != null ? currentEdge : null;
+
+                ctx.setExistingEdge(edge);
+
+                AtlasEdge newEdge = mapStructValue(ctx, context);
+
+                if (currentEdge != null && !currentEdge.equals(newEdge)) {
+                    deleteHandler.deleteEdgeReference(currentEdge, ctx.getAttrType().getTypeCategory(), false, true);
+                }
+
+                return newEdge;
+            }
+
+            case ENTITY: {
+                String          edgeLabel    = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexProperty());
+                AtlasEdge       currentEdge  = graphHelper.getEdgeForLabel(ctx.getReferringVertex(), edgeLabel);
+                AtlasEntityType instanceType = getInstanceType(ctx.getValue());
+                AtlasEdge       edge         = currentEdge != null ? currentEdge : null;
+
+                ctx.setElementType(instanceType);
+                ctx.setExistingEdge(edge);
+
+                AtlasEdge newEdge = mapEntityValue(ctx, context);
+
+                if (currentEdge != null && !currentEdge.equals(newEdge)) {
+                    deleteHandler.deleteEdgeReference(currentEdge, ctx.getAttrType().getTypeCategory(), ctx.getAttribute().isOwnedRef(), true);
+                }
+
+                return newEdge;
+            }
+
+            case MAP:
+                return mapMapValue(ctx, context);
+
+            case ARRAY:
+                return mapArrayValue(ctx, context);
+
+            default:
+                throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, ctx.getAttrType().getTypeCategory().name());
+        }
+    }
+
+    private Object mapPrimitiveValue(AttributeMutationContext ctx) {
+        AtlasGraphUtilsV1.setProperty(ctx.getReferringVertex(), ctx.getVertexProperty(), ctx.getValue());
+
+        return ctx.getValue();
+    }
+
+    private AtlasEdge mapStructValue(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> mapStructValue({})", ctx);
+        }
+
+        AtlasEdge ret = null;
+
+        if (ctx.getCurrentEdge() != null) {
+            updateVertex((AtlasStruct) ctx.getValue(), ctx.getCurrentEdge().getInVertex(), context);
+
+            ret = ctx.getCurrentEdge();
+        } else if (ctx.getValue() != null) {
+            String edgeLabel = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexProperty());
+
+            if (ctx.getValue() instanceof AtlasStruct) {
+                ret = createVertex((AtlasStruct) ctx.getValue(), ctx.getReferringVertex(), edgeLabel, context);
+            } else if (ctx.getValue() instanceof Map) {
+                AtlasStruct stuct = new AtlasStruct(ctx.getAttrType().getTypeName(), (Map) AtlasTypeUtil.toStructAttributes((Map)ctx.getValue()));
+
+                ret = createVertex(stuct, ctx.getReferringVertex(), edgeLabel, context);
+            }
+        }
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== mapStructValue({})", ctx);
+        }
+
+        return ret;
+    }
+
+    private AtlasEdge mapEntityValue(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> mapEntityValue({})", ctx);
+        }
+
+        AtlasEdge ret = null;
+
+        String guid = getGuid(ctx.getValue());
 
-        AtlasObjectId guid = getId(ctx.getValue());
         AtlasVertex entityVertex = context.getDiscoveryContext().getResolvedEntityVertex(guid);
-        if ( ctx.getCurrentEdge().isPresent() ) {
-            result = updateEdge(ctx.getAttributeDef(), ctx.getValue(), ctx.getCurrentEdge().get(), entityVertex);
+
+        if (entityVertex == null) {
+            AtlasObjectId objId = getObjectId(ctx.getValue());
+
+            entityVertex = context.getDiscoveryContext().getResolvedEntityVertex(objId);
+        }
+
+        if (entityVertex == null) {
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, ctx.getValue().toString());
+        }
+
+        if (ctx.getCurrentEdge() != null) {
+            ret = updateEdge(ctx.getAttributeDef(), ctx.getValue(), ctx.getCurrentEdge(), entityVertex);
         } else if (ctx.getValue() != null) {
-            String edgeLabel = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexPropertyKey());
+            String edgeLabel = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexProperty());
+
             try {
-                result = graphHelper.getOrCreateEdge(ctx.getReferringVertex(), entityVertex, edgeLabel);
+                ret = graphHelper.getOrCreateEdge(ctx.getReferringVertex(), entityVertex, edgeLabel);
             } catch (RepositoryException e) {
                 throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, e);
             }
         }
 
-        return result;
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== mapEntityValue({})", ctx);
+        }
+
+        return ret;
+    }
+
+    private Map<String, Object> mapMapValue(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> mapMapValue({})", ctx);
+        }
+
+        @SuppressWarnings("unchecked")
+        Map<Object, Object> newVal  = (Map<Object, Object>) ctx.getValue();
+        Map<String, Object> newMap  = new HashMap<>();
+        AtlasMapType        mapType = (AtlasMapType) ctx.getAttrType();
+
+        try {
+            AtlasAttribute      attribute   = ctx.getAttribute();
+            List<String>        currentKeys = GraphHelper.getListProperty(ctx.getReferringVertex(), ctx.getVertexProperty());
+            Map<String, Object> currentMap  = new HashMap<>();
+
+            if (CollectionUtils.isNotEmpty(currentKeys)) {
+                for (String key : currentKeys) {
+                    String propertyNameForKey  = GraphHelper.getQualifiedNameForMapKey(ctx.getVertexProperty(), GraphHelper.encodePropertyKey(key));
+                    Object propertyValueForKey = getMapValueProperty(mapType.getValueType(), ctx.getReferringVertex(), propertyNameForKey);
+
+                    currentMap.put(key, propertyValueForKey);
+                }
+            }
+
+            if (MapUtils.isNotEmpty(newVal)) {
+                for (Map.Entry<Object, Object> entry : newVal.entrySet()) {
+                    String    key          = entry.getKey().toString();
+                    String    propertyName = GraphHelper.getQualifiedNameForMapKey(ctx.getVertexProperty(), GraphHelper.encodePropertyKey(key));
+                    AtlasEdge existingEdge = getEdgeIfExists(mapType, currentMap, key);
+
+                    AttributeMutationContext mapCtx =  new AttributeMutationContext(ctx.getOp(), ctx.getReferringVertex(), attribute, entry.getValue(), propertyName, mapType.getValueType(), existingEdge);
+
+                    //Add/Update/Remove property value
+                    Object newEntry = mapCollectionElementsToVertex(mapCtx, context);
+                    setMapValueProperty(mapType.getValueType(), ctx.getReferringVertex(), propertyName, newEntry);
+
+                    newMap.put(key, newEntry);
+                }
+            }
+
+            Map<String, Object> finalMap = removeUnusedMapEntries(attribute, ctx.getReferringVertex(), ctx.getVertexProperty(), currentMap, newMap);
+
+            Set<String> newKeys = new HashSet<>(newMap.keySet());
+            newKeys.addAll(finalMap.keySet());
+
+            // for dereference on way out
+            GraphHelper.setListProperty(ctx.getReferringVertex(), ctx.getVertexProperty(), new ArrayList<>(newKeys));
+        } catch (AtlasException e) {
+            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, e);
+        }
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== mapMapValue({})", ctx);
         }
 
-    @Override
-    public void cleanUp() throws AtlasBaseException {
+        return newMap;
+    }
+
+    public List mapArrayValue(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> mapArrayValue({})", ctx);
+        }
+
+        AtlasAttribute attribute       = ctx.getAttribute();
+        List           newElements     = (List) ctx.getValue();
+        AtlasArrayType arrType         = (AtlasArrayType) attribute.getAttributeType();
+        AtlasType      elementType     = arrType.getElementType();
+        List<Object>   currentElements = getArrayElementsProperty(elementType, ctx.getReferringVertex(), ctx.getVertexProperty());
+
+        List<Object> newElementsCreated = new ArrayList<>();
+
+        if (CollectionUtils.isNotEmpty(newElements)) {
+            for (int index = 0; index < newElements.size(); index++) {
+                AtlasEdge               existingEdge = getEdgeAt(currentElements, index, elementType);
+                AttributeMutationContext arrCtx      =  new AttributeMutationContext(ctx.getOp(), ctx.getReferringVertex(), ctx.getAttribute(), newElements.get(index),
+                                                                                     ctx.getVertexProperty(), elementType, existingEdge);
+
+                Object newEntry = mapCollectionElementsToVertex(arrCtx, context);
+
+                newElementsCreated.add(newEntry);
+            }
+        }
+
+        if (AtlasGraphUtilsV1.isReference(elementType)) {
+            List<AtlasEdge> additionalEdges = removeUnusedArrayEntries(attribute, (List) currentElements, (List) newElementsCreated);
+            newElementsCreated.addAll(additionalEdges);
+        }
+
+        // for dereference on way out
+        setArrayElementsProperty(elementType, ctx.getReferringVertex(), ctx.getVertexProperty(), newElementsCreated);
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== mapArrayValue({})", ctx);
+        }
+
+        return newElementsCreated;
+    }
+
+
+    private AtlasEdge createVertex(AtlasStruct struct, AtlasVertex referringVertex, String edgeLabel, EntityMutationContext context) throws AtlasBaseException {
+        AtlasVertex vertex = createStructVertex(struct);
+
+        mapAttributes(struct, vertex, CREATE, context);
+
+        try {
+            //TODO - Map directly in AtlasGraphUtilsV1
+            return graphHelper.getOrCreateEdge(referringVertex, vertex, edgeLabel);
+        } catch (RepositoryException e) {
+            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, e);
+        }
+    }
+
+    private void updateVertex(AtlasStruct struct, AtlasVertex vertex, EntityMutationContext context) throws AtlasBaseException {
+        mapAttributes(struct, vertex, UPDATE, context);
+    }
+
+    private void updateModificationMetadata(AtlasVertex vertex) {
+        AtlasGraphUtilsV1.setProperty(vertex, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContextV1.get().getRequestTime());
+        GraphHelper.setProperty(vertex, Constants.MODIFIED_BY_KEY, RequestContextV1.get().getUser());
+    }
+
+    private int getEntityVersion(AtlasEntity entity) {
+        Long ret = entity != null ? entity.getVersion() : null;
+
+        return (ret != null) ? ret.intValue() : 0;
+    }
+
+    private AtlasStructType getStructType(String typeName) throws AtlasBaseException {
+        AtlasType objType = typeRegistry.getType(typeName);
+
+        if (!(objType instanceof AtlasStructType)) {
+            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, typeName);
+        }
+
+        return (AtlasStructType)objType;
+    }
+
+    private Object mapCollectionElementsToVertex(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
+        switch(ctx.getAttrType().getTypeCategory()) {
+            case PRIMITIVE:
+            case ENUM:
+                return ctx.getValue();
+
+            case STRUCT:
+                return mapStructValue(ctx, context);
+
+            case ENTITY:
+                AtlasEntityType instanceType = getInstanceType(ctx.getValue());
+                ctx.setElementType(instanceType);
+                return mapEntityValue(ctx, context);
+
+            case MAP:
+            case ARRAY:
+            default:
+                throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, ctx.getAttrType().getTypeCategory().name());
+        }
+    }
+
+    private static AtlasObjectId getObjectId(Object val) throws AtlasBaseException {
+        if (val != null) {
+            if ( val instanceof  AtlasObjectId) {
+                return ((AtlasObjectId) val);
+            } else if (val instanceof AtlasEntity) {
+                return ((AtlasEntity) val).getAtlasObjectId();
+            } else if (val instanceof Map) {
+                AtlasObjectId ret = new AtlasObjectId((Map)val);
+
+                if (ret.isValid()) {
+                    return ret;
+                }
+            }
+
+            throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, val.toString());
+        }
+
+        return null;
+    }
+
+    private static String getGuid(Object val) throws AtlasBaseException {
+        if (val != null) {
+            if ( val instanceof  AtlasObjectId) {
+                return ((AtlasObjectId) val).getGuid();
+            } else if (val instanceof AtlasEntity) {
+                return ((AtlasEntity) val).getGuid();
+            } else if (val instanceof Map) {
+                Object guidVal = ((Map)val).get(AtlasObjectId.KEY_GUID);
+
+                return guidVal != null ? guidVal.toString() : null;
+            }
+        }
+
+        return null;
+    }
+
+    private AtlasEntityType getInstanceType(Object val) throws AtlasBaseException {
+        AtlasEntityType ret = null;
+
+        if (val != null) {
+            String typeName = null;
+
+            if (val instanceof AtlasObjectId) {
+                typeName = ((AtlasObjectId)val).getTypeName();
+            } else if (val instanceof AtlasEntity) {
+                typeName = ((AtlasEntity)val).getTypeName();
+            } else if (val instanceof Map) {
+                Object typeNameVal = ((Map)val).get(AtlasObjectId.KEY_TYPENAME);
+
+                if (typeNameVal != null) {
+                    typeName = typeNameVal.toString();
+                }
+            }
+
+            ret = typeName != null ? typeRegistry.getEntityTypeByName(typeName) : null;
+
+            if (ret == null) {
+                throw new AtlasBaseException(AtlasErrorCode.INVALID_OBJECT_ID, val.toString());
+            }
+        }
+
+        return ret;
+    }
+
+
+    public static Object getMapValueProperty(AtlasType elementType, AtlasVertex vertex, String vertexPropertyName) {
+        if (AtlasGraphUtilsV1.isReference(elementType)) {
+            return vertex.getProperty(vertexPropertyName, AtlasEdge.class);
+        }
+        else {
+            return vertex.getProperty(vertexPropertyName, String.class).toString();
+        }
+    }
+
+    private static void setMapValueProperty(AtlasType elementType, AtlasVertex vertex, String vertexPropertyName, Object value) {
+        if (AtlasGraphUtilsV1.isReference(elementType)) {
+            vertex.setPropertyFromElementId(vertexPropertyName, (AtlasEdge)value);
+        }
+        else {
+            vertex.setProperty(vertexPropertyName, value);
+        }
+    }
+
+    //Remove unused entries from map
+    private Map<String, Object> removeUnusedMapEntries(AtlasAttribute attribute, AtlasVertex vertex, String propertyName,
+                                                       Map<String, Object> currentMap, Map<String, Object> newMap)
+                                                                             throws AtlasException, AtlasBaseException {
+        AtlasMapType        mapType       = (AtlasMapType) attribute.getAttributeType();
+        Map<String, Object> additionalMap = new HashMap<>();
+
+        for (String currentKey : currentMap.keySet()) {
+            boolean shouldDeleteKey = !newMap.containsKey(currentKey);
+
+            if (AtlasGraphUtilsV1.isReference(mapType.getValueType())) {
+                //Delete the edge reference if its not part of new edges created/updated
+                AtlasEdge currentEdge = (AtlasEdge)currentMap.get(currentKey);
+
+                if (!newMap.values().contains(currentEdge)) {
+                    boolean deleted = deleteHandler.deleteEdgeReference(currentEdge, mapType.getValueType().getTypeCategory(), attribute.isOwnedRef(), true);
+
+                    /* TODO: need to review the following 'if' block. Wouldn't this leave deleted keys in the map?
+                     *
+                    if (!deleted) {
+                        additionalMap.put(currentKey, currentEdge);
+                        shouldDeleteKey = false;
+                    }
+                    *
+                    */
+                }
+            }
+
+            if (shouldDeleteKey) {
+                String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(propertyName, GraphHelper.encodePropertyKey(currentKey));
+                GraphHelper.setProperty(vertex, propertyNameForKey, null);
+            }
+        }
+
+        return additionalMap;
+    }
+
+    private static AtlasEdge getEdgeIfExists(AtlasMapType mapType, Map<String, Object> currentMap, String keyStr) {
+        AtlasEdge ret = null;
+
+        if (AtlasGraphUtilsV1.isReference(mapType.getValueType())) {
+            Object val = currentMap.get(keyStr);
+
+            if (val != null) {
+                ret = (AtlasEdge) val;
+            }
+        }
+
+        return ret;
     }
 
     private AtlasEdge updateEdge(AtlasAttributeDef attributeDef, Object value, AtlasEdge currentEdge, final AtlasVertex entityVertex) throws AtlasBaseException {
@@ -132,59 +664,68 @@ public class EntityGraphMapper implements InstanceGraphMapper<AtlasEdge> {
         return newEdge;
     }
 
-    public EntityMutationResponse
-    mapAttributes(EntityMutationContext ctx) throws AtlasBaseException {
-
-        this.context = ctx;
-        structVertexMapper.init(this);
-
-        EntityMutationResponse resp = new EntityMutationResponse();
-        //Map attributes
-        if (ctx.getCreatedEntities() != null) {
-            for (AtlasEntity createdEntity : ctx.getCreatedEntities()) {
-                AtlasVertex vertex = ctx.getVertex(createdEntity);
-                structVertexMapper.mapAttributestoVertex(EntityMutations.EntityOperation.CREATE, ctx.getType(createdEntity), createdEntity, vertex);
-                resp.addEntity(EntityMutations.EntityOperation.CREATE, constructHeader(createdEntity, ctx.getType(createdEntity), vertex));
+    public static List<Object> getArrayElementsProperty(AtlasType elementType, AtlasVertex vertex, String vertexPropertyName) {
+        if (AtlasGraphUtilsV1.isReference(elementType)) {
+            return (List)vertex.getListProperty(vertexPropertyName, AtlasEdge.class);
+        }
+        else {
+            return (List)vertex.getListProperty(vertexPropertyName);
         }
     }
 
-        if (ctx.getUpdatedEntities() != null) {
-            for (AtlasEntity updated : ctx.getUpdatedEntities()) {
-                AtlasVertex vertex = ctx.getVertex(updated);
-                structVertexMapper.mapAttributestoVertex(EntityMutations.EntityOperation.UPDATE, ctx.getType(updated), updated, vertex);
+    private AtlasEdge getEdgeAt(List<Object> currentElements, int index, AtlasType elemType) {
+        AtlasEdge ret = null;
 
-                resp.addEntity(EntityMutations.EntityOperation.UPDATE, constructHeader(updated, ctx.getType(updated), vertex));
+        if (AtlasGraphUtilsV1.isReference(elemType)) {
+            if (currentElements != null && index < currentElements.size()) {
+                ret = (AtlasEdge) currentElements.get(index);
             }
         }
 
-        RequestContextV1 req = RequestContextV1.get();
-        for (AtlasObjectId id : req.getDeletedEntityIds()) {
-            resp.addEntity(EntityMutations.EntityOperation.DELETE, constructHeader(id));
+        return ret;
     }
 
-        return resp;
-    }
+    //Removes unused edges from the old collection, compared to the new collection
+    private List<AtlasEdge> removeUnusedArrayEntries(AtlasAttribute attribute, List<AtlasEdge> currentEntries, List<AtlasEdge> newEntries) throws AtlasBaseException {
+        if (CollectionUtils.isNotEmpty(currentEntries)) {
+            AtlasStructType entityType = attribute.getDefinedInType();
+            AtlasType       entryType  = ((AtlasArrayType)attribute.getAttributeType()).getElementType();
 
+            if (AtlasGraphUtilsV1.isReference(entryType)) {
+                Collection<AtlasEdge> edgesToRemove = CollectionUtils.subtract(currentEntries, newEntries);
 
-    public AtlasObjectId getId(Object value) throws AtlasBaseException {
-        if (value != null) {
-            if ( value instanceof  AtlasObjectId) {
-                return ((AtlasObjectId) value);
-            } else if (value instanceof AtlasEntity) {
-                return ((AtlasEntity) value).getAtlasObjectId();
-            } else if (value instanceof Map) {
-                AtlasObjectId ret = new AtlasObjectId((Map)value);
+                if (CollectionUtils.isNotEmpty(edgesToRemove)) {
+                    List<AtlasEdge> additionalElements = new ArrayList<>();
 
-                if (ret.isValid()) {
-                    return ret;
+                    for (AtlasEdge edge : edgesToRemove) {
+                        boolean deleted = deleteHandler.deleteEdgeReference(edge, entryType.getTypeCategory(), attribute.isOwnedRef(), true);
+
+                        /* TODO: need to review the following 'if' block. Wouldn't this leave deleted elements continue to be in array?
+                         *
+                        if (!deleted) {
+                            additionalElements.add(edge);
+                        }
+                         *
+                         */
+                    }
+
+                    return additionalElements;
+                }
             }
         }
 
-            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, (String) value);
+        return Collections.emptyList();
     }
 
-        return null;
+    private void setArrayElementsProperty(AtlasType elementType, AtlasVertex vertex, String vertexPropertyName, List<Object> values) {
+        if (AtlasGraphUtilsV1.isReference(elementType)) {
+            GraphHelper.setListPropertyFromElementIds(vertex, vertexPropertyName, (List) values);
+        }
+        else {
+            GraphHelper.setProperty(vertex, vertexPropertyName, values);
         }
+    }
+
 
     private AtlasEntityHeader constructHeader(AtlasEntity entity, final AtlasEntityType type, AtlasVertex vertex) {
         //TODO - enhance to return only selective attributes
@@ -211,18 +752,4 @@ public class EntityGraphMapper implements InstanceGraphMapper<AtlasEdge> {
 
         return entity;
     }
-
-    public EntityMutationContext getContext() {
-        return context;
-    }
-
-    public AtlasEntityType getInstanceType(Object val) throws AtlasBaseException {
-        AtlasObjectId guid = getId(val);
-
-        if ( guid != null) {
-            return (AtlasEntityType) getContext().getType(guid);
-        }
-
-        return null;
-    }
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java

@@ -22,6 +22,7 @@ import org.apache.atlas.AtlasException;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasObjectId;
@@ -100,6 +101,22 @@ public final class EntityGraphRetriever {
         return ret;
     }
 
+    public AtlasEntitiesWithExtInfo toAtlasEntitiesWithExtInfo(List<String> guids) throws AtlasBaseException {
+        AtlasEntitiesWithExtInfo ret = new AtlasEntitiesWithExtInfo();
+
+        for (String guid : guids) {
+            AtlasVertex vertex = getEntityVertex(guid);
+
+            AtlasEntity entity = mapVertexToAtlasEntity(vertex, ret);
+
+            ret.addEntity(entity);
+        }
+
+        ret.compact();
+
+        return ret;
+    }
+
     private AtlasVertex getEntityVertex(String guid) throws AtlasBaseException {
         try {
             return graphHelper.getVertexForGUID(guid);
@@ -405,7 +422,7 @@ public final class EntityGraphRetriever {
                         ret = entity.getAtlasObjectId();
                     }
                 } else {
-                    ret = new AtlasObjectId(GraphHelper.getTypeName(referenceVertex), GraphHelper.getGuid(referenceVertex));
+                    ret = new AtlasObjectId(GraphHelper.getGuid(referenceVertex), GraphHelper.getTypeName(referenceVertex));
                 }
             }
         }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityMutationContext.java

@@ -24,6 +24,7 @@ import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasType;
+import org.apache.commons.lang.StringUtils;
 
 import java.util.*;
 
@@ -31,25 +32,29 @@ public class EntityMutationContext {
     private final EntityGraphDiscoveryContext  context;
     private final List<AtlasEntity>            entitiesCreated = new ArrayList<>();
     private final List<AtlasEntity>            entitiesUpdated = new ArrayList<>();
-    private final Map<AtlasObjectId, AtlasEntityType> entityVsType    = new HashMap<>();
-    private final Map<AtlasObjectId, AtlasVertex>     entityVsVertex  = new HashMap<>();
+    private final Map<String, AtlasEntityType> entityVsType    = new HashMap<>();
+    private final Map<String, AtlasVertex>     entityVsVertex  = new HashMap<>();
+    private final Map<String, String>          guidAssignments = new HashMap<>();
 
     public EntityMutationContext(final EntityGraphDiscoveryContext context) {
         this.context = context;
     }
 
-    public void addCreated(AtlasEntity entity, AtlasEntityType type, AtlasVertex atlasVertex) {
-        AtlasObjectId objId = entity.getAtlasObjectId();
+    public void addCreated(String internalGuid, AtlasEntity entity, AtlasEntityType type, AtlasVertex atlasVertex) {
         entitiesCreated.add(entity);
-        entityVsType.put(objId, type);
-        entityVsVertex.put(objId, atlasVertex);
+        entityVsType.put(entity.getGuid(), type);
+        entityVsVertex.put(entity.getGuid(), atlasVertex);
+
+        if (!StringUtils.equals(internalGuid, entity.getGuid())) {
+            guidAssignments.put(internalGuid, entity.getGuid());
+            entityVsVertex.put(internalGuid, atlasVertex);
+        }
     }
 
     public void addUpdated(AtlasEntity entity, AtlasEntityType type, AtlasVertex atlasVertex) {
-        AtlasObjectId objId = entity.getAtlasObjectId();
         entitiesUpdated.add(entity);
-        entityVsType.put(objId, type);
-        entityVsVertex.put(objId, atlasVertex);
+        entityVsType.put(entity.getGuid(), type);
+        entityVsVertex.put(entity.getGuid(), atlasVertex);
     }
 
     public EntityGraphDiscoveryContext getDiscoveryContext() {
@@ -64,21 +69,15 @@ public class EntityMutationContext {
         return entitiesUpdated;
     }
 
-    public AtlasEntityType getType(AtlasEntity entity) {
-        return entityVsType.get(entity.getAtlasObjectId());
+    public Map<String, String> getGuidAssignments() {
+        return guidAssignments;
     }
 
-    public AtlasType getType(AtlasObjectId entityId) {
-        return entityVsType.get(entityId);
+    public AtlasEntityType getType(String guid) {
+        return entityVsType.get(guid);
     }
 
-    public AtlasVertex getVertex(AtlasEntity entity) {
-        return entityVsVertex.get(entity.getAtlasObjectId());
-    }
-
-    public AtlasVertex getVertex(AtlasObjectId entityId) {
-        return entityVsVertex.get(entityId);
-    }
+    public AtlasVertex getVertex(String guid) { return entityVsVertex.get(guid); }
 
 
     @Override

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityStream.java

@@ -28,5 +28,5 @@ public interface EntityStream {
 
     void reset();
 
-    AtlasEntity getById(AtlasObjectId id);
+    AtlasEntity getByGuid(String guid);
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/IDBasedEntityResolver.java

@@ -17,11 +17,9 @@
  */
 package org.apache.atlas.repository.store.graph.v1;
 
-import com.google.common.base.Optional;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
@@ -29,93 +27,53 @@ import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.repository.store.graph.EntityResolver;
 import org.apache.atlas.typesystem.exception.EntityNotFoundException;
 import org.apache.atlas.typesystem.persistence.Id;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
 
 public class IDBasedEntityResolver implements EntityResolver {
+    private static final Logger LOG = LoggerFactory.getLogger(IDBasedEntityResolver.class);
+
 
     private final GraphHelper graphHelper = GraphHelper.getInstance();
-    private final Map<String, AtlasEntity> idToEntityMap = new HashMap<>();
-    private EntityGraphDiscoveryContext    context;
 
-    @Override
-    public void init(EntityGraphDiscoveryContext context) throws AtlasBaseException {
-        this.context = context;
 
-        for (AtlasEntity entity : context.getRootEntities()) {
-            idToEntityMap.put(entity.getGuid(), entity);
-        }
-    }
-
-    public EntityGraphDiscoveryContext resolveEntityReferences() throws AtlasBaseException {
+    public EntityGraphDiscoveryContext resolveEntityReferences(EntityGraphDiscoveryContext context) throws AtlasBaseException {
         if (context == null) {
-            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "Entity resolver not initialized");
+            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "IDBasedEntityResolver.resolveEntityReferences(): context is null");
         }
 
-        List<AtlasObjectId> resolvedReferences = new ArrayList<>();
+        EntityStream entityStream = context.getEntityStream();
 
-        for (AtlasObjectId objId : context.getUnresolvedIds()) {
-            if (objId.isAssignedGuid()) {
-                //validate in graph repo that given guid, typename exists
-                Optional<AtlasVertex> vertex = resolveGuid(objId);
+        for (String guid : context.getReferencedGuids()) {
+            if (AtlasEntity.isAssigned(guid)) { // validate in graph repo that given guid exists
+                AtlasVertex vertex = resolveGuid(guid);
 
-                if (vertex.isPresent()) {
-                    context.addResolvedId(objId, vertex.get());
-                    resolvedReferences.add(objId);
-                }
+                context.addResolvedGuid(guid, vertex);
+            } else  if (entityStream.getByGuid(guid) != null) { //check if entity stream have this reference id
+                context.addLocalGuidReference(guid);
             } else {
-                //check if root references have this temporary id
-               if (!idToEntityMap.containsKey(objId.getGuid()) ) {
-                   throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, objId.toString());
-               }
-                resolvedReferences.add(objId);
-            }
-
-        }
-
-        context.removeUnResolvedIds(resolvedReferences);
-
-        //Resolve root references
-        for (AtlasEntity entity : context.getRootEntities()) {
-            AtlasObjectId objId = entity.getAtlasObjectId();
-
-            if (!context.isResolvedId(objId) && AtlasEntity.isAssigned(entity.getGuid())) {
-                Optional<AtlasVertex> vertex = resolveGuid(objId);
-
-                if (vertex.isPresent()) {
-                    context.addResolvedId(objId, vertex.get());
-                    context.removeUnResolvedId(objId);
-                }
+                throw new AtlasBaseException(AtlasErrorCode.REFERENCED_ENTITY_NOT_FOUND, guid);
             }
         }
 
         return context;
     }
 
-    private Optional<AtlasVertex> resolveGuid(AtlasObjectId objId) throws AtlasBaseException {
+    private AtlasVertex resolveGuid(String guid) throws AtlasBaseException {
         //validate in graph repo that given guid, typename exists
         AtlasVertex vertex = null;
         try {
-            vertex = graphHelper.findVertex(Constants.GUID_PROPERTY_KEY, objId.getGuid(),
-                Constants.TYPE_NAME_PROPERTY_KEY, objId.getTypeName(),
+            vertex = graphHelper.findVertex(Constants.GUID_PROPERTY_KEY, guid,
                                             Constants.STATE_PROPERTY_KEY, Id.EntityState.ACTIVE.name());
         } catch (EntityNotFoundException e) {
             //Ignore
         }
-        if ( vertex != null ) {
-            return Optional.of(vertex);
+
+        if (vertex != null) {
+            return vertex;
         } else {
-            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, objId.getGuid());
-        }
+            throw new AtlasBaseException(AtlasErrorCode.REFERENCED_ENTITY_NOT_FOUND, guid);
         }
-
-    @Override
-    public void cleanUp() throws AtlasBaseException {
-        idToEntityMap.clear();
-        this.context = null;
     }
-
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/InMemoryMapEntityStream.java

@@ -27,14 +27,11 @@ import java.util.Map;
 
 public class InMemoryMapEntityStream implements EntityStream {
 
-    private final Map<AtlasObjectId, AtlasEntity>                 entities = new HashMap<>();
-    private       Iterator<Map.Entry<AtlasObjectId, AtlasEntity>> iterator;
-
-    public InMemoryMapEntityStream(Map<String, AtlasEntity> entityMap) {
-        for (AtlasEntity entity : entityMap.values()) {
-            entities.put(entity.getAtlasObjectId(), entity);
-        }
+    private final Map<String, AtlasEntity>                 entities;
+    private       Iterator<Map.Entry<String, AtlasEntity>> iterator;
 
+    public InMemoryMapEntityStream(Map<String, AtlasEntity> entities) {
+        this.entities = entities;
         this.iterator = entities.entrySet().iterator();
     }
 
@@ -54,7 +51,7 @@ public class InMemoryMapEntityStream implements EntityStream {
     }
 
     @Override
-    public AtlasEntity getById(final AtlasObjectId id) {
-        return entities.get(id);
+    public AtlasEntity getByGuid(final String guid) {
+        return entities.get(guid);
     }
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/InstanceGraphMapper.java

@@ -18,9 +18,6 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.repository.graphdb.AtlasVertex;
-import org.apache.atlas.type.AtlasEntityType;
 
 public interface InstanceGraphMapper<T> {
 
@@ -31,9 +28,5 @@ public interface InstanceGraphMapper<T> {
      * @return the value that was mapped to the vertex
      * @throws AtlasBaseException
      */
-    T toGraph(GraphMutationContext ctx) throws AtlasBaseException;
-
-
-    void cleanUp() throws AtlasBaseException;
-
+    T toGraph(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException;
 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/MapVertexMapper.java

-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.atlas.repository.store.graph.v1;
-
-
-import com.google.common.base.Optional;
-import com.google.inject.Inject;
-import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.AtlasException;
-import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
-import org.apache.atlas.repository.graph.GraphHelper;
-import org.apache.atlas.repository.graphdb.AtlasEdge;
-import org.apache.atlas.repository.graphdb.AtlasVertex;
-import org.apache.atlas.type.AtlasMapType;
-import org.apache.atlas.type.AtlasStructType;
-import org.apache.atlas.type.AtlasType;
-import org.apache.commons.collections.MapUtils;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-
-import static org.apache.atlas.repository.graph.GraphHelper.string;
-
-public class MapVertexMapper implements InstanceGraphMapper<Map> {
-
-    private DeleteHandlerV1 deleteHandler;
-
-    private static final Logger LOG = LoggerFactory.getLogger(MapVertexMapper.class);
-
-    private StructVertexMapper structVertexMapper;
-
-    @Inject
-    public MapVertexMapper(DeleteHandlerV1 deleteHandler) {
-        this.deleteHandler = deleteHandler;
-    }
-
-    void init(StructVertexMapper structVertexMapper) {
-        this.structVertexMapper = structVertexMapper;
-    }
-
-    @Override
-    public Map<String, Object> toGraph(GraphMutationContext ctx) throws AtlasBaseException {
-
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("Mapping instance to vertex {} for map type {}", string(ctx.getReferringVertex()), ctx.getAttrType().getTypeName());
-        }
-
-        @SuppressWarnings("unchecked") Map<Object, Object> newVal =
-            (Map<Object, Object>) ctx.getValue();
-
-        boolean newAttributeEmpty = MapUtils.isEmpty(newVal);
-
-        Map<String, Object> currentMap = new HashMap<>();
-        Map<String, Object> newMap = new HashMap<>();
-
-        AtlasMapType mapType = (AtlasMapType) ctx.getAttrType();
-
-        try {
-            List<String> currentKeys = GraphHelper.getListProperty(ctx.getReferringVertex(), ctx.getVertexPropertyKey());
-            if (currentKeys != null && !currentKeys.isEmpty()) {
-                for (String key : currentKeys) {
-                    String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(ctx.getVertexPropertyKey(), key);
-                    Object propertyValueForKey = getMapValueProperty(mapType.getValueType(), ctx.getReferringVertex(), propertyNameForKey);
-                    currentMap.put(key, propertyValueForKey);
-                }
-            }
-
-            if (!newAttributeEmpty) {
-                for (Map.Entry<Object, Object> entry : newVal.entrySet()) {
-                    String keyStr = entry.getKey().toString();
-                    String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(ctx.getVertexPropertyKey(), keyStr);
-                    Optional<AtlasEdge> existingEdge = getEdgeIfExists(mapType, currentMap, keyStr);
-
-                    GraphMutationContext mapCtx =  new GraphMutationContext.Builder(ctx.getOp(), ctx.getAttribute(), mapType.getValueType(), entry.getValue())
-                        .referringVertex(ctx.getReferringVertex())
-                        .edge(existingEdge)
-                        .vertexProperty(propertyNameForKey).build();
-
-
-                    //Add/Update/Remove property value
-                    Object newEntry = structVertexMapper.mapCollectionElementsToVertex(mapCtx);
-                    MapVertexMapper.setMapValueProperty(mapType.getValueType(), ctx.getReferringVertex(), propertyNameForKey, newEntry);
-
-                    newMap.put(keyStr, newEntry);
-                }
-            }
-
-            Map<String, Object> finalMap =
-                removeUnusedMapEntries(ctx.getParentType(), mapType, ctx.getAttributeDef(), ctx.getReferringVertex(), ctx.getVertexPropertyKey(), currentMap, newMap);
-
-            Set<String> newKeys = new HashSet<>(newMap.keySet());
-            newKeys.addAll(finalMap.keySet());
-
-            // for dereference on way out
-            GraphHelper.setListProperty(ctx.getReferringVertex(), ctx.getVertexPropertyKey(), new ArrayList<>(newKeys));
-        } catch (AtlasException e) {
-            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, e);
-        }
-
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("Map values set in vertex {} {}", mapType.getTypeName(), newMap);
-        }
-
-        return newMap;
-    }
-
-    @Override
-    public void cleanUp() throws AtlasBaseException {
-    }
-
-
-    public static Object getMapValueProperty(AtlasType elementType, AtlasVertex instanceVertex, String propertyName) {
-        String actualPropertyName = GraphHelper.encodePropertyKey(propertyName);
-        if (AtlasGraphUtilsV1.isReference(elementType)) {
-            return instanceVertex.getProperty(actualPropertyName, AtlasEdge.class);
-        }
-        else {
-            return instanceVertex.getProperty(actualPropertyName, String.class).toString();
-        }
-    }
-
-    public static void setMapValueProperty(AtlasType elementType, AtlasVertex instanceVertex, String propertyName, Object value) {
-        String actualPropertyName = GraphHelper.encodePropertyKey(propertyName);
-        if (AtlasGraphUtilsV1.isReference(elementType)) {
-            instanceVertex.setPropertyFromElementId(actualPropertyName, (AtlasEdge)value);
-        }
-        else {
-            instanceVertex.setProperty(actualPropertyName, value);
-        }
-    }
-
-    //Remove unused entries from map
-    private Map<String, Object> removeUnusedMapEntries(
-        AtlasStructType entityType,
-        AtlasMapType mapType, AtlasAttributeDef attributeDef,
-        AtlasVertex instanceVertex, String propertyName,
-        Map<String, Object> currentMap,
-        Map<String, Object> newMap)
-        throws AtlasException, AtlasBaseException {
-
-        Map<String, Object> additionalMap = new HashMap<>();
-        for (String currentKey : currentMap.keySet()) {
-
-            boolean shouldDeleteKey = !newMap.containsKey(currentKey);
-            if (AtlasGraphUtilsV1.isReference(mapType.getValueType())) {
-
-                //Delete the edge reference if its not part of new edges created/updated
-                AtlasEdge currentEdge = (AtlasEdge)currentMap.get(currentKey);
-
-                if (!newMap.values().contains(currentEdge)) {
-                    boolean deleteChildReferences = StructVertexMapper.shouldManageChildReferences(entityType, attributeDef.getName());
-                    boolean deleted =
-                        deleteHandler.deleteEdgeReference(currentEdge, mapType.getValueType().getTypeCategory(), deleteChildReferences, true);
-                    if (!deleted) {
-                        additionalMap.put(currentKey, currentEdge);
-                        shouldDeleteKey = false;
-                    }
-                }
-            }
-
-            if (shouldDeleteKey) {
-                String propertyNameForKey = GraphHelper.getQualifiedNameForMapKey(propertyName, currentKey);
-                GraphHelper.setProperty(instanceVertex, propertyNameForKey, null);
-            }
-        }
-        return additionalMap;
-    }
-    
-    private Optional<AtlasEdge> getEdgeIfExists(AtlasMapType mapType, Map<String, Object> currentMap, String keyStr) {
-        Optional<AtlasEdge> existingEdge = Optional.absent();
-        if ( AtlasGraphUtilsV1.isReference(mapType.getValueType()) ) {
-            if ( currentMap.get(keyStr) != null) {
-                existingEdge = Optional.of((AtlasEdge) currentMap.get(keyStr));
-            }
-        }
-        
-        return existingEdge;
-    }
-}

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/StructVertexMapper.java

-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.atlas.repository.store.graph.v1;
-
-import com.google.common.base.Optional;
-import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.RequestContextV1;
-import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.model.TypeCategory;
-import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasStruct;
-import org.apache.atlas.model.instance.EntityMutations;
-import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
-import org.apache.atlas.repository.Constants;
-import org.apache.atlas.repository.RepositoryException;
-import org.apache.atlas.repository.graph.AtlasGraphProvider;
-import org.apache.atlas.repository.graph.GraphHelper;
-import org.apache.atlas.repository.graphdb.AtlasEdge;
-import org.apache.atlas.repository.graphdb.AtlasGraph;
-import org.apache.atlas.repository.graphdb.AtlasVertex;
-import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasStructType;
-import org.apache.atlas.type.AtlasType;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Map;
-
-public class StructVertexMapper implements InstanceGraphMapper<AtlasEdge> {
-
-    private final AtlasGraph graph;
-
-    private final GraphHelper graphHelper = GraphHelper.getInstance();
-
-    private final MapVertexMapper mapVertexMapper;
-
-    private final ArrayVertexMapper arrVertexMapper;
-
-    private EntityGraphMapper entityVertexMapper;
-
-    private DeleteHandlerV1 deleteHandler;
-
-    private static final Logger LOG = LoggerFactory.getLogger(StructVertexMapper.class);
-
-    public StructVertexMapper(ArrayVertexMapper arrayVertexMapper, MapVertexMapper mapVertexMapper, DeleteHandlerV1 deleteHandler) {
-        this.graph = AtlasGraphProvider.getGraphInstance();;
-        this.mapVertexMapper = mapVertexMapper;
-        this.arrVertexMapper = arrayVertexMapper;
-        this.deleteHandler = deleteHandler;
-    }
-
-    void init(final EntityGraphMapper entityVertexMapper) {
-        this.entityVertexMapper = entityVertexMapper;
-    }
-
-    @Override
-    public AtlasEdge toGraph(GraphMutationContext ctx) throws AtlasBaseException {
-        AtlasEdge ret = null;
-
-        if ( ctx.getCurrentEdge().isPresent() ) {
-            updateVertex(ctx.getParentType(), (AtlasStructType) ctx.getAttrType(), ctx.getAttributeDef(), (AtlasStruct) ctx.getValue(), ctx.getCurrentEdge().get().getInVertex());
-            ret = ctx.getCurrentEdge().get();
-        } else if (ctx.getValue() != null) {
-            String edgeLabel = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexPropertyKey());
-            ret = createVertex(ctx.getParentType(), (AtlasStructType) ctx.getAttrType(), ctx.getAttributeDef(), (AtlasStruct) ctx.getValue(), ctx.getReferringVertex(), edgeLabel);
-        }
-
-        return ret;
-    }
-
-    @Override
-    public void cleanUp() throws AtlasBaseException {
-    }
-
-    public static boolean shouldManageChildReferences(AtlasStructType type, String attributeName) {
-        AtlasStructType.AtlasAttribute attribute = type.getAttribute(attributeName);
-
-        return attribute != null ? attribute.isOwnedRef() : false;
-    }
-
-    /**
-     * Map attributes for entity, struct or trait
-     *
-     * @param op
-     * @param structType
-     * @param struct
-     * @param vertex
-     * @return
-     * @throws AtlasBaseException
-     */
-    public AtlasVertex mapAttributestoVertex(final EntityMutations.EntityOperation op, AtlasStructType structType, AtlasStruct struct, AtlasVertex vertex) throws AtlasBaseException {
-        if (struct.getAttributes() != null) {
-            if (op.equals(EntityMutations.EntityOperation.CREATE)) {
-                final Map<String, AtlasStructType.AtlasAttribute> allAttributes = structType.getAllAttributes();
-                for (String attrName : allAttributes.keySet()) {
-                    Object value = struct.getAttribute(attrName);
-
-                    mapAttribute(op, structType, attrName, value, vertex);
-                }
-            } else if (op.equals(EntityMutations.EntityOperation.UPDATE)) {
-                for (String attrName : struct.getAttributes().keySet()) {
-                    Object value = struct.getAttribute(attrName);
-                    mapAttribute(op, structType, attrName, value, vertex);
-                }
-            }
-            updateModificationMetadata(vertex);
-        }
-        return vertex;
-    }
-
-    private void mapAttribute(final EntityMutations.EntityOperation op, AtlasStructType structType, String attrName, Object value, AtlasVertex vertex) throws AtlasBaseException {
-        AtlasType attributeType = structType.getAttributeType(attrName);
-        if (attributeType != null) {
-            final AtlasStructType.AtlasAttribute attribute = structType.getAttribute(attrName);
-
-            if (value == null) {
-                if ( attribute.getAttributeType().getTypeCategory() == TypeCategory.PRIMITIVE) {
-                    if ( attribute.getAttributeDef().getIsOptional()) {
-                        value = attribute.getAttributeType().createOptionalDefaultValue();
-                    } else {
-                        value = attribute.getAttributeType().createDefaultValue();
-                    }
-                }
-            }
-
-            final String vertexProperty = structType.getQualifiedAttributeName(attrName);
-            GraphMutationContext ctx = new GraphMutationContext.Builder(op, attribute, value)
-                .referringVertex(vertex)
-                .vertexProperty(GraphHelper.encodePropertyKey(vertexProperty)).build();
-            mapToVertexByTypeCategory(ctx);
-        }
-    }
-
-    private void updateModificationMetadata(AtlasVertex vertex) {
-        //Set updated timestamp
-        AtlasGraphUtilsV1.setProperty(vertex, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContextV1.get().getRequestTime());
-        GraphHelper.setProperty(vertex, Constants.MODIFIED_BY_KEY, RequestContextV1.get().getUser());
-    }
-
-
-    protected Object mapToVertexByTypeCategory(GraphMutationContext ctx) throws AtlasBaseException {
-        if (ctx.getOp() == EntityMutations.EntityOperation.CREATE && ctx.getValue() == null) {
-            return null;
-        }
-
-        switch (ctx.getAttrType().getTypeCategory()) {
-        case PRIMITIVE:
-        case ENUM:
-            return primitivesToVertex(ctx);
-        case STRUCT:
-            String edgeLabel = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexPropertyKey());
-            AtlasEdge currentEdge = graphHelper.getEdgeForLabel(ctx.getReferringVertex(), edgeLabel);
-            Optional<AtlasEdge> edge = currentEdge != null ? Optional.of(currentEdge) : Optional.<AtlasEdge>absent();
-            ctx.setExistingEdge(edge);
-            AtlasEdge newEdge = toGraph(ctx);
-
-            if (currentEdge != null && !currentEdge.equals(newEdge)) {
-                deleteHandler.deleteEdgeReference(currentEdge, ctx.getAttrType().getTypeCategory(), false, true);
-            }
-            return newEdge;
-        case ENTITY:
-            edgeLabel = AtlasGraphUtilsV1.getEdgeLabel(ctx.getVertexPropertyKey());
-            currentEdge = graphHelper.getEdgeForLabel(ctx.getReferringVertex(), edgeLabel);
-            AtlasEntityType instanceType = entityVertexMapper.getInstanceType(ctx.getValue());
-            edge = currentEdge != null ? Optional.of(currentEdge) : Optional.<AtlasEdge>absent();
-            ctx.setElementType(instanceType);
-            ctx.setExistingEdge(edge);
-            newEdge = entityVertexMapper.toGraph(ctx);
-
-            if (currentEdge != null && !currentEdge.equals(newEdge)) {
-                deleteHandler.deleteEdgeReference(currentEdge, ctx.getAttrType().getTypeCategory(), shouldManageChildReferences(ctx.getParentType(), ctx.getAttributeDef().getName()), true);
-            }
-            return newEdge;
-        case MAP:
-            return mapVertexMapper.toGraph(ctx);
-        case ARRAY:
-            return arrVertexMapper.toGraph(ctx);
-        default:
-            throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, ctx.getAttrType().getTypeCategory().name());
-        }
-    }
-
-    protected Object primitivesToVertex(GraphMutationContext ctx) {
-        AtlasGraphUtilsV1.setProperty(ctx.getReferringVertex(), ctx.getVertexPropertyKey(), ctx.getValue());
-        return ctx.getValue();
-    }
-
-    private AtlasEdge createVertex(AtlasStructType parentType, AtlasStructType attrType, AtlasAttributeDef attributeDef, AtlasStruct struct, AtlasVertex referringVertex, String edgeLabel) throws AtlasBaseException {
-        AtlasVertex vertex = createVertexTemplate(struct, attrType);
-        mapAttributestoVertex(EntityMutations.EntityOperation.CREATE, attrType, struct, vertex);
-
-        try {
-            //TODO - Map directly in AtlasGraphUtilsV1
-            return graphHelper.getOrCreateEdge(referringVertex, vertex, edgeLabel);
-        } catch (RepositoryException e) {
-            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, e);
-        }
-    }
-
-    private void updateVertex(AtlasStructType parentType, AtlasStructType structAttributeType, AtlasAttributeDef attributeDef, AtlasStruct value, AtlasVertex structVertex) throws AtlasBaseException {
-        mapAttributestoVertex(EntityMutations.EntityOperation.CREATE, structAttributeType, value, structVertex);
-    }
-
-    protected AtlasVertex createVertexTemplate(final AtlasStruct instance, final AtlasStructType structType) {
-        LOG.debug("Creating AtlasVertex for type {}", instance.getTypeName());
-        final AtlasVertex vertexWithoutIdentity = graph.addVertex();
-
-        // add type information
-        AtlasGraphUtilsV1.setProperty(vertexWithoutIdentity, Constants.ENTITY_TYPE_PROPERTY_KEY, instance.getTypeName());
-
-        // add state information
-        AtlasGraphUtilsV1.setProperty(vertexWithoutIdentity, Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
-
-        // add timestamp information
-        AtlasGraphUtilsV1.setProperty(vertexWithoutIdentity, Constants.TIMESTAMP_PROPERTY_KEY, RequestContextV1.get().getRequestTime());
-        AtlasGraphUtilsV1.setProperty(vertexWithoutIdentity, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY,
-            RequestContextV1.get().getRequestTime());
-
-        AtlasGraphUtilsV1.setProperty(vertexWithoutIdentity, Constants.CREATED_BY_KEY, RequestContextV1.get().getUser());
-
-        GraphHelper.setProperty(vertexWithoutIdentity, Constants.MODIFIED_BY_KEY, RequestContextV1.get().getUser());
-
-        return vertexWithoutIdentity;
-    }
-
-    protected Object mapCollectionElementsToVertex(GraphMutationContext ctx) throws AtlasBaseException {
-        switch(ctx.getAttrType().getTypeCategory()) {
-        case PRIMITIVE:
-        case ENUM:
-            return ctx.getValue();
-        case STRUCT:
-            return toGraph(ctx);
-        case ENTITY:
-            AtlasEntityType instanceType = entityVertexMapper.getInstanceType(ctx.getValue());
-            ctx.setElementType(instanceType);
-            return entityVertexMapper.toGraph(ctx);
-        case MAP:
-        case ARRAY:
-        default:
-            throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, ctx.getAttrType().getTypeCategory().name());
-        }
-    }
-}

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/UniqAttrBasedEntityResolver.java

@@ -17,7 +17,6 @@
  */
 package org.apache.atlas.repository.store.graph.v1;
 
-import com.google.common.base.Optional;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
@@ -29,108 +28,53 @@ import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.repository.store.graph.EntityResolver;
 import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasStructType;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.typesystem.exception.EntityNotFoundException;
 import org.apache.commons.collections.MapUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import javax.inject.Inject;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
+
 public class UniqAttrBasedEntityResolver implements EntityResolver {
     private static final Logger LOG = LoggerFactory.getLogger(UniqAttrBasedEntityResolver.class);
 
     private final GraphHelper           graphHelper = GraphHelper.getInstance();
     private final AtlasTypeRegistry     typeRegistry;
-    private EntityGraphDiscoveryContext context;
 
-    @Inject
     public UniqAttrBasedEntityResolver(AtlasTypeRegistry typeRegistry) {
         this.typeRegistry = typeRegistry;
     }
 
     @Override
-    public void init(EntityGraphDiscoveryContext context) throws AtlasBaseException {
-        this.context = context;
-    }
-
-    @Override
-    public EntityGraphDiscoveryContext resolveEntityReferences() throws AtlasBaseException {
+    public EntityGraphDiscoveryContext resolveEntityReferences(EntityGraphDiscoveryContext context) throws AtlasBaseException {
         if (context == null) {
-            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "Unique attribute based entity resolver not initialized");
+            throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "UniqAttrBasedEntityResolver.resolveEntityReferences(): context is null");
         }
 
         //Resolve attribute references
         List<AtlasObjectId> resolvedReferences = new ArrayList<>();
 
-        for (AtlasObjectId entityId : context.getUnresolvedIdsByUniqAttribs()) {
+        for (AtlasObjectId objId : context.getReferencedByUniqAttribs()) {
             //query in graph repo that given unique attribute - check for deleted also?
-            Optional<AtlasVertex> vertex = resolveByUniqueAttribute(entityId);
-
-            if (vertex.isPresent()) {
-                context.addResolvedId(entityId, vertex.get());
-                resolvedReferences.add(entityId);
-            }
-        }
+            AtlasVertex vertex = resolveByUniqueAttribute(objId);
 
-        context.removeUnresolvedIdsByUniqAttribs(resolvedReferences);
-
-        //Resolve root references
-        for (AtlasEntity entity : context.getRootEntities()) {
-            AtlasObjectId entityId = entity.getAtlasObjectId();
-
-            if (!context.isResolvedId(entityId) ) {
-                Optional<AtlasVertex> vertex = resolveByUniqueAttribute(entity);
-
-                if (vertex.isPresent()) {
-                    context.addResolvedId(entityId, vertex.get());
-                    context.removeUnResolvedId(entityId);
-                }
+            if (vertex != null) {
+                context.addResolvedIdByUniqAttribs(objId, vertex);
+                resolvedReferences.add(objId);
+            } else {
+                throw new AtlasBaseException(AtlasErrorCode.REFERENCED_ENTITY_NOT_FOUND, objId.toString());
             }
         }
 
         return context;
     }
 
-    Optional<AtlasVertex> resolveByUniqueAttribute(AtlasEntity entity) throws AtlasBaseException {
-        AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
-
-        if (entityType == null) {
-            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
-        }
-
-        for (AtlasStructType.AtlasAttribute attr : entityType.getAllAttributes().values()) {
-            if (attr.getAttributeDef().getIsUnique()) {
-                Object attrVal = entity.getAttribute(attr.getName());
-
-                if (attrVal == null) {
-                    continue;
-                }
-
-                Optional<AtlasVertex> vertex = findByTypeAndQualifiedName(entityType.getTypeName(), attr.getQualifiedAttributeName(), attrVal);
-
-                if (LOG.isDebugEnabled()) {
-                    LOG.debug("Found vertex by unique attribute : " + attr.getQualifiedAttributeName() + "=" + attrVal);
-                }
-
-                if (!vertex.isPresent()) {
-                    vertex = findBySuperTypeAndQualifiedName(entityType.getTypeName(), attr.getQualifiedAttributeName(), attrVal);
-                }
-
-                if (vertex.isPresent()) {
-                    return vertex;
-                }
-            }
-        }
-
-        return Optional.absent();
-    }
-
-    Optional<AtlasVertex> resolveByUniqueAttribute(AtlasObjectId entityId) throws AtlasBaseException {
+    private AtlasVertex resolveByUniqueAttribute(AtlasObjectId entityId) throws AtlasBaseException {
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entityId.getTypeName());
 
         if (entityType == null) {
@@ -138,77 +82,66 @@ public class UniqAttrBasedEntityResolver implements EntityResolver {
         }
 
         final Map<String, Object> uniqueAttributes = entityId.getUniqueAttributes();
-        if (MapUtils.isNotEmpty(uniqueAttributes)) {
-            for (String attrName : uniqueAttributes.keySet()) {
-                AtlasStructType.AtlasAttribute attr = entityType.getAttribute(attrName);
 
-                if (attr.getAttributeDef().getIsUnique()) {
-                    Object attrVal = uniqueAttributes.get(attr.getName());
+        if (MapUtils.isNotEmpty(uniqueAttributes)) {
+            for (Map.Entry<String, Object> e : uniqueAttributes.entrySet()) {
+                String         attrName  = e.getKey();
+                Object         attrValue = e.getValue();
+                AtlasAttribute attr      = entityType.getAttribute(attrName);
 
-                    if (attrVal == null) {
+                if (attr == null || !attr.getAttributeDef().getIsUnique() || attrValue == null) {
                     continue;
                 }
 
-                    Optional<AtlasVertex> vertex = findByTypeAndQualifiedName(entityId.getTypeName(), attr.getQualifiedAttributeName(), attrVal);
+                AtlasVertex vertex = findByTypeAndQualifiedName(entityId.getTypeName(), attr.getQualifiedName(), attrValue);
 
-                    if (!vertex.isPresent()) {
-                        vertex = findBySuperTypeAndQualifiedName(entityId.getTypeName(), attr.getQualifiedAttributeName(), attrVal);
+                if (vertex == null) {
+                    vertex = findBySuperTypeAndQualifiedName(entityId.getTypeName(), attr.getQualifiedName(), attrValue);
                 }
 
-                    if (vertex.isPresent()) {
+                if (vertex != null) {
                     return vertex;
                 }
             }
         }
-        }
-        return Optional.absent();
+
+        throw new AtlasBaseException(AtlasErrorCode.REFERENCED_ENTITY_NOT_FOUND, entityId.toString());
     }
 
-    Optional<AtlasVertex> findByTypeAndQualifiedName(String typeName, String qualifiedAttrName, Object attrVal) {
+    private AtlasVertex findByTypeAndQualifiedName(String typeName, String qualifiedAttrName, Object attrVal) {
         AtlasVertex vertex = null;
+
         try {
             vertex = graphHelper.findVertex(qualifiedAttrName, attrVal,
                                             Constants.ENTITY_TYPE_PROPERTY_KEY, typeName,
-                Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE
-                    .name());
+                                            Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
 
             if (LOG.isDebugEnabled()) {
-                LOG.debug("Found vertex by unique attribute and type {} {} ", qualifiedAttrName + "=" + attrVal, typeName);
-            }
-            if (vertex != null) {
-                return Optional.of(vertex);
+                LOG.debug("Found vertex by unique attribute and type ({}={}), {} ", qualifiedAttrName, attrVal, typeName);
             }
         } catch (EntityNotFoundException e) {
             //Ignore if not found
         }
-        return Optional.absent();
+
+        return vertex;
     }
 
-    Optional<AtlasVertex> findBySuperTypeAndQualifiedName(String typeName, String qualifiedAttrName, Object attrVal) {
+    private AtlasVertex findBySuperTypeAndQualifiedName(String typeName, String qualifiedAttrName, Object attrVal) {
         AtlasVertex vertex = null;
+
         try {
             vertex = graphHelper.findVertex(qualifiedAttrName, attrVal,
                                             Constants.SUPER_TYPES_PROPERTY_KEY, typeName,
-                Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE
-                    .name());
+                                            Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
 
             if (LOG.isDebugEnabled()) {
-                LOG.debug("Found vertex by unique attribute and supertype {} ", qualifiedAttrName + "=" + attrVal, typeName);
-            }
-            if (vertex != null) {
-                return Optional.of(vertex);
+                LOG.debug("Found vertex by unique attribute and supertype ({}={}), {} ", qualifiedAttrName, attrVal, typeName);
             }
         } catch (EntityNotFoundException e) {
             //Ignore if not found
         }
-        return Optional.absent();
-    }
 
-    @Override
-    public void cleanUp() {
-        //Nothing to cleanup
-        this.context = null;
+        return vertex;
     }
-
 }
 

repository/src/test/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1Test.java

@@ -25,15 +25,20 @@ import org.apache.atlas.TestUtils;
 import org.apache.atlas.TestUtilsV2;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityExtInfo;
+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.model.instance.AtlasStruct;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
+import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
+import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.services.MetadataService;
 import org.apache.atlas.store.AtlasTypeDefStore;
@@ -43,14 +48,8 @@ import org.apache.atlas.type.AtlasStructType;
 import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.type.AtlasTypeUtil;
-import org.apache.atlas.typesystem.IInstance;
-import org.apache.atlas.typesystem.ITypedReferenceableInstance;
-import org.apache.atlas.typesystem.persistence.Id;
-import org.apache.atlas.typesystem.persistence.ReferenceableInstance;
-import org.apache.atlas.typesystem.persistence.StructInstance;
-import org.apache.atlas.typesystem.types.EnumValue;
-import org.apache.atlas.util.AtlasRepositoryConfiguration;
 import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.collections.MapUtils;
 import org.apache.commons.lang.RandomStringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -62,7 +61,6 @@ import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
 import javax.inject.Inject;
-import java.lang.reflect.Constructor;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -70,13 +68,11 @@ import java.util.List;
 import java.util.Map;
 
 import static org.apache.atlas.TestUtils.COLUMNS_ATTR_NAME;
-import static org.apache.atlas.TestUtils.TABLE_TYPE;
 import static org.apache.atlas.TestUtils.randomString;
 import static org.testng.Assert.assertEquals;
 
 @Guice(modules = RepositoryMetadataModule.class)
 public class AtlasEntityStoreV1Test {
-
     private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityStoreV1Test.class);
 
     @Inject
@@ -90,31 +86,34 @@ public class AtlasEntityStoreV1Test {
     @Inject
     MetadataService metadataService;
 
-    private Map<String, AtlasEntity> deptEntityMap;
-    private Map<String, AtlasEntity> dbEntityMap;
-    private Map<String, AtlasEntity> tableEntityMap;
+    @Inject
+    DeleteHandlerV1 deleteHandler;
+
+    private AtlasEntitiesWithExtInfo deptEntity;
+    private AtlasEntityWithExtInfo   dbEntity;
+    private AtlasEntityWithExtInfo   tblEntity;
 
-    private AtlasEntity deptEntity;
-    private AtlasEntity dbEntity;
-    private AtlasEntity tableEntity;
 
     @BeforeClass
     public void setUp() throws Exception {
         metadataService = TestUtils.addSessionCleanupWrapper(metadataService);
         new GraphBackedSearchIndexer(typeRegistry);
-        final AtlasTypesDef deptTypesDef = TestUtilsV2.defineDeptEmployeeTypes();
-        typeDefStore.createTypesDef(deptTypesDef);
 
-        final AtlasTypesDef hiveTypesDef = TestUtilsV2.defineHiveTypes();
-        typeDefStore.createTypesDef(hiveTypesDef);
+        AtlasTypesDef[] testTypesDefs = new AtlasTypesDef[] { TestUtilsV2.defineDeptEmployeeTypes(),
+                                                              TestUtilsV2.defineHiveTypes()
+                                                            };
 
-        deptEntityMap = TestUtilsV2.createDeptEg1();
-        dbEntityMap = TestUtilsV2.createDBEntity();
-        tableEntityMap = TestUtilsV2.createTableEntity(dbEntityMap.keySet().iterator().next());
+        for (AtlasTypesDef typesDef : testTypesDefs) {
+            AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typesDef, typeRegistry);
 
-        deptEntity = deptEntityMap.values().iterator().next();
-        dbEntity = dbEntityMap.values().iterator().next();
-        tableEntity = tableEntityMap.values().iterator().next();
+            if (!typesToCreate.isEmpty()) {
+                typeDefStore.createTypesDef(typesToCreate);
+            }
+        }
+
+        deptEntity = TestUtilsV2.createDeptEg2();
+        dbEntity   = TestUtilsV2.createDBEntityV2();
+        tblEntity  = TestUtilsV2.createTableEntityV2(dbEntity.getEntity());
     }
 
     @AfterClass
@@ -125,209 +124,180 @@ public class AtlasEntityStoreV1Test {
 
     @BeforeTest
     public void init() throws Exception {
-        final Class<? extends DeleteHandlerV1> deleteHandlerImpl = AtlasRepositoryConfiguration.getDeleteHandlerV1Impl();
-        final Constructor<? extends DeleteHandlerV1> deleteHandlerImplConstructor = deleteHandlerImpl.getConstructor(AtlasTypeRegistry.class);
-        DeleteHandlerV1 deleteHandler = deleteHandlerImplConstructor.newInstance(typeRegistry);
-        ArrayVertexMapper arrVertexMapper = new ArrayVertexMapper(deleteHandler);
-        MapVertexMapper mapVertexMapper = new MapVertexMapper(deleteHandler);
-
-
-        entityStore = new AtlasEntityStoreV1(new EntityGraphMapper(arrVertexMapper, mapVertexMapper, deleteHandler));
-        entityStore.init(typeRegistry);
+        entityStore = new AtlasEntityStoreV1(deleteHandler, typeRegistry);
 
         RequestContextV1.clear();
     }
 
     @Test
     public void testCreate() throws Exception {
-        EntityMutationResponse response = entityStore.createOrUpdate(deptEntityMap);
+        init();
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(deptEntity));
+
+        validateMutationResponse(response, EntityOperation.CREATE, 5);
+        AtlasEntityHeader dept1 = response.getFirstCreatedEntityByTypeName(TestUtilsV2.DEPARTMENT_TYPE);
+        validateEntity(deptEntity, getEntityFromStore(dept1), deptEntity.getEntities().get(0));
+
+        final Map<EntityOperation, List<AtlasEntityHeader>> entitiesMutated = response.getEntitiesMutated();
+        List<AtlasEntityHeader> entitiesCreated = entitiesMutated.get(EntityOperation.CREATE);
 
-        validateMutationResponse(response, EntityMutations.EntityOperation.CREATE, 5);
-        AtlasEntityHeader deptEntity = response.getFirstCreatedEntityByTypeName(TestUtilsV2.DEPARTMENT_TYPE);
+        Assert.assertTrue(entitiesCreated.size() >= deptEntity.getEntities().size());
 
-        final Map<EntityMutations.EntityOperation, List<AtlasEntityHeader>> entitiesMutated = response.getEntitiesMutated();
-        List<AtlasEntityHeader> entitiesCreated = entitiesMutated.get(EntityMutations.EntityOperation.CREATE);
+        for (int i = 0; i < deptEntity.getEntities().size(); i++) {
+            AtlasEntity expected = deptEntity.getEntities().get(i);
+            AtlasEntity actual   = getEntityFromStore(entitiesCreated.get(i));
 
-        for (AtlasEntityHeader header : entitiesCreated) {
-            validateAttributes(deptEntityMap, header);
+            validateEntity(deptEntity, actual, expected);
         }
 
         //Create DB
-        EntityMutationResponse dbCreationResponse = entityStore.createOrUpdate(dbEntityMap);
-        validateMutationResponse(dbCreationResponse, EntityMutations.EntityOperation.CREATE, 1);
+        init();
+        EntityMutationResponse dbCreationResponse = entityStore.createOrUpdate(new AtlasEntityStream(dbEntity));
+        validateMutationResponse(dbCreationResponse, EntityOperation.CREATE, 1);
 
-        AtlasEntityHeader dbEntity = dbCreationResponse.getFirstCreatedEntityByTypeName(TestUtilsV2.DATABASE_TYPE);
-        validateAttributes(dbEntityMap, dbEntity);
+        AtlasEntityHeader db1 = dbCreationResponse.getFirstCreatedEntityByTypeName(TestUtilsV2.DATABASE_TYPE);
+        validateEntity(dbEntity, getEntityFromStore(db1));
 
         //Create Table
         //Update DB guid
-        AtlasObjectId dbId = (AtlasObjectId) tableEntity.getAttribute("database");
-        dbId.setGuid(dbEntity.getGuid());
-        tableEntityMap.put(dbId.getGuid(), dbEntityMap.values().iterator().next());
-        tableEntity.setAttribute("database", dbId);
+        AtlasObjectId dbObjectId = (AtlasObjectId) tblEntity.getEntity().getAttribute("database");
+        dbObjectId.setGuid(db1.getGuid());
+        tblEntity.addReferredEntity(dbEntity.getEntity());
 
-        EntityMutationResponse tableCreationResponse = entityStore.createOrUpdate(tableEntityMap);
-        validateMutationResponse(tableCreationResponse, EntityMutations.EntityOperation.CREATE, 1);
+        init();
+        EntityMutationResponse tableCreationResponse = entityStore.createOrUpdate(new AtlasEntityStream(tblEntity));
+        validateMutationResponse(tableCreationResponse, EntityOperation.CREATE, 1);
 
         AtlasEntityHeader tableEntity = tableCreationResponse.getFirstCreatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableEntityMap, tableEntity);
-
+        validateEntity(tblEntity, getEntityFromStore(tableEntity));
     }
 
     @Test(dependsOnMethods = "testCreate")
     public void testArrayOfEntityUpdate() throws Exception {
-        //clear state
-        init();
+        AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
+        List<AtlasObjectId>      columns      = new ArrayList<>();
+        AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);
 
-//        Map<String, AtlasEntity> entityCloneMap = new HashMap<>();
-//        AtlasEntity entityClone = new AtlasEntity(deptEntity);
-//        List<AtlasObjectId> employees = (List<AtlasObjectId>) entityClone.getAttribute("employees");
-//        AtlasEntity entityRemoved = clearSubOrdinates(employees, 1);
-//        entityClone.setAttribute("employees", employees);
-//        EntityMutationResponse response = entityStore.createOrUpdate(entityCloneMap);
-//
-//        validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 5);
-//        AtlasEntityHeader deptEntity = response.getFirstEntityUpdated();
-//        Assert.assertEquals(((List<AtlasEntity>)(((List<AtlasEntity>) deptEntity.getAttribute("employees")).get(1).getAttribute("subordinates"))).size(), 1);
-//
-//        init();
-//        //add  entity back
-//        addSubordinate(employees.get(1), entityRemoved);
-//        response = entityStore.createOrUpdate(entityCloneMap);
-//        validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 5);
-//        deptEntity = response.getFirstEntityUpdated();
-//        validateAttributes(deptEntity);
 
+        AtlasEntity col1 = TestUtilsV2.createColumnEntity(tableEntity);
+        col1.setAttribute(TestUtilsV2.NAME, "col1");
 
-        Map<String, AtlasEntity> tableUpdatedMap = new HashMap<>();
-        tableUpdatedMap.put(dbEntity.getGuid(), dbEntity);
-
-        //test array of class with id
-        final List<AtlasObjectId> columns = new ArrayList<>();
+        AtlasEntity col2 = TestUtilsV2.createColumnEntity(tableEntity);
+        col2.setAttribute(TestUtilsV2.NAME, "col2");
 
-        AtlasEntity col1 = TestUtilsV2.createColumnEntity(tableEntity.getGuid());
-        col1.setAttribute(TestUtilsV2.NAME, "col1");
         columns.add(col1.getAtlasObjectId());
-        tableUpdatedMap.put(col1.getGuid(), col1);
+        columns.add(col2.getAtlasObjectId());
 
-        AtlasEntity tableUpdated = new AtlasEntity(tableEntity);
-        tableUpdated.setAttribute(TestUtilsV2.COLUMNS_ATTR_NAME, columns);
-        tableUpdatedMap.put(tableUpdated.getGuid(), tableUpdated);
+        tableEntity.setAttribute(TestUtilsV2.COLUMNS_ATTR_NAME, columns);
+
+        entitiesInfo.addReferredEntity(dbEntity.getEntity());
+        entitiesInfo.addReferredEntity(col1);
+        entitiesInfo.addReferredEntity(col2);
 
         init();
-        EntityMutationResponse response = entityStore.createOrUpdate(tableUpdatedMap);
-        AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(tableUpdated.getTypeName());
-        validateAttributes(tableUpdatedMap, updatedTable);
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+
+        AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         //Complete update. Add  array elements - col3,col4
-        AtlasEntity col3 = TestUtilsV2.createColumnEntity(tableEntity.getGuid());
-        col1.setAttribute(TestUtilsV2.NAME, "col3");
-        columns.add(col3.getAtlasObjectId());
-        tableUpdatedMap.put(col3.getGuid(), col3);
+        AtlasEntity col3 = TestUtilsV2.createColumnEntity(tableEntity);
+        col3.setAttribute(TestUtilsV2.NAME, "col3");
 
-        AtlasEntity col4 = TestUtilsV2.createColumnEntity(tableEntity.getGuid());
-        col1.setAttribute(TestUtilsV2.NAME, "col4");
+        AtlasEntity col4 = TestUtilsV2.createColumnEntity(tableEntity);
+        col4.setAttribute(TestUtilsV2.NAME, "col4");
+
+        columns.add(col3.getAtlasObjectId());
         columns.add(col4.getAtlasObjectId());
-        tableUpdatedMap.put(col4.getGuid(), col4);
 
-        tableUpdated.setAttribute(COLUMNS_ATTR_NAME, columns);
+        entitiesInfo.addReferredEntity(col3);
+        entitiesInfo.addReferredEntity(col4);
+
         init();
-        response = entityStore.createOrUpdate(tableUpdatedMap);
-        updatedTable = response.getFirstUpdatedEntityByTypeName(tableUpdated.getTypeName());
-        validateAttributes(tableUpdatedMap, updatedTable);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+
+        updatedTable = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         //Swap elements
-        tableUpdatedMap.clear();
         columns.clear();
-        tableUpdated.setAttribute(COLUMNS_ATTR_NAME, columns);
-        tableUpdatedMap.put(tableUpdated.getGuid(), tableUpdated);
-        tableUpdatedMap.put(col3.getGuid(), col3);
-        tableUpdatedMap.put(col4.getGuid(), col4);
         columns.add(col4.getAtlasObjectId());
         columns.add(col3.getAtlasObjectId());
 
         init();
-        response = entityStore.createOrUpdate(tableUpdatedMap);
-        updatedTable = response.getFirstUpdatedEntityByTypeName(tableUpdated.getTypeName());
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+
+        updatedTable = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
         Assert.assertEquals(((List<AtlasObjectId>) updatedTable.getAttribute(COLUMNS_ATTR_NAME)).size(), 2);
 
-        assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 1);  //col1 is deleted
+        assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);  // col1, col2 are deleted
 
         //Update array column to null
-        tableUpdated.setAttribute(COLUMNS_ATTR_NAME, null);
+        tableEntity.setAttribute(COLUMNS_ATTR_NAME, null);
+
         init();
-        response = entityStore.createOrUpdate(tableUpdatedMap);
-        updatedTable = response.getFirstUpdatedEntityByTypeName(tableUpdated.getTypeName());
-        validateAttributes(tableUpdatedMap, updatedTable);
-        assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
 
+        updatedTable = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
+        assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);
     }
 
     @Test(dependsOnMethods = "testCreate")
     public void testUpdateEntityWithMap() throws Exception {
+        AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
+        AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);
+        Map<String, AtlasStruct> partsMap     = new HashMap<>();
+        partsMap.put("part0", new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "test"));
 
-        final Map<String, AtlasEntity> tableCloneMap = new HashMap<>();
-        final AtlasEntity tableClone = new AtlasEntity(tableEntity);
-        final Map<String, AtlasStruct> partsMap = new HashMap<>();
-        partsMap.put("part0", new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE,
-            new HashMap<String, Object>() {{
-                put(TestUtilsV2.NAME, "test");
-            }}));
-
-        
-        tableClone.setAttribute("partitionsMap", partsMap);
-        tableCloneMap.put(tableClone.getGuid(), tableClone);
-
+        tableEntity.setAttribute("partitionsMap", partsMap);
+        entitiesInfo.addReferredEntity(tableEntity);
 
         init();
-        EntityMutationResponse response = entityStore.createOrUpdate(tableCloneMap);
-        final AtlasEntityHeader tableDefinition1 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition1);
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+
+        AtlasEntityHeader tableDefinition1 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(tableDefinition1));
                 
         Assert.assertTrue(partsMap.get("part0").equals(((Map<String, AtlasStruct>) tableDefinition1.getAttribute("partitionsMap")).get("part0")));
 
         //update map - add a map key
-        partsMap.put("part1", new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE,
-            new HashMap<String, Object>() {{
-                put(TestUtilsV2.NAME, "test1");
-            }}));
-        tableClone.setAttribute("partitionsMap", partsMap);
+        partsMap.put("part1", new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "test1"));
+        tableEntity.setAttribute("partitionsMap", partsMap);
 
         init();
-        response = entityStore.createOrUpdate(tableCloneMap);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+
         AtlasEntityHeader tableDefinition2 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition2);
+        validateEntity(entitiesInfo, getEntityFromStore(tableDefinition2));
 
         assertEquals(((Map<String, AtlasStruct>) tableDefinition2.getAttribute("partitionsMap")).size(), 2);
         Assert.assertTrue(partsMap.get("part1").equals(((Map<String, AtlasStruct>) tableDefinition2.getAttribute("partitionsMap")).get("part1")));
 
         //update map - remove a key and add another key
         partsMap.remove("part0");
-        partsMap.put("part2", new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE,
-            new HashMap<String, Object>() {{
-                put(TestUtilsV2.NAME, "test2");
-            }}));
-        tableClone.setAttribute("partitionsMap", partsMap);
+        partsMap.put("part2", new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "test2"));
+        tableEntity.setAttribute("partitionsMap", partsMap);
 
         init();
-        response = entityStore.createOrUpdate(tableCloneMap);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+
         AtlasEntityHeader tableDefinition3 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition3);
+        validateEntity(entitiesInfo, getEntityFromStore(tableDefinition3));
 
         assertEquals(((Map<String, AtlasStruct>) tableDefinition3.getAttribute("partitionsMap")).size(), 2);
         Assert.assertNull(((Map<String, AtlasStruct>) tableDefinition3.getAttribute("partitionsMap")).get("part0"));
         Assert.assertTrue(partsMap.get("part2").equals(((Map<String, AtlasStruct>) tableDefinition3.getAttribute("partitionsMap")).get("part2")));
 
         //update struct value for existing map key
-        init();
         AtlasStruct partition2 = partsMap.get("part2");
         partition2.setAttribute(TestUtilsV2.NAME, "test2Updated");
-        response = entityStore.createOrUpdate(tableCloneMap);
-        final AtlasEntityHeader tableDefinition4 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition4);
 
-        assertEquals(((Map<String, AtlasStruct>) tableDefinition4.getAttribute("partitionsMap")).size(), 2);
-        Assert.assertNull(((Map<String, AtlasStruct>) tableDefinition4.getAttribute("partitionsMap")).get("part0"));
+        init();
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+
+        AtlasEntityHeader tableDefinition4 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(tableDefinition4));
 
         assertEquals(((Map<String, AtlasStruct>) tableDefinition4.getAttribute("partitionsMap")).size(), 2);
         Assert.assertNull(((Map<String, AtlasStruct>) tableDefinition4.getAttribute("partitionsMap")).get("part0"));
@@ -335,209 +305,179 @@ public class AtlasEntityStoreV1Test {
 
         //Test map pointing to a class
 
-        final Map<String, AtlasObjectId> columnsMap = new HashMap<>();
-
-        Map<String, AtlasEntity> col0TypeMap = new HashMap<>();
-        AtlasEntity col0Type = new AtlasEntity(TestUtilsV2.COLUMN_TYPE,
-            new HashMap<String, Object>() {{
-                put(TestUtilsV2.NAME, "test1");
-                put("type", "string");
-                put("table", new AtlasObjectId(TABLE_TYPE, tableDefinition1.getGuid()));
-            }});
-
-        col0TypeMap.put(col0Type.getGuid(), col0Type);
+        AtlasEntity col0 = new AtlasEntity(TestUtilsV2.COLUMN_TYPE, TestUtilsV2.NAME, "test1");
+        col0.setAttribute("type", "string");
+        col0.setAttribute("table", tableEntity.getAtlasObjectId());
 
+        AtlasEntityWithExtInfo col0WithExtendedInfo = new AtlasEntityWithExtInfo(col0);
+        col0WithExtendedInfo.addReferredEntity(tableEntity);
+        col0WithExtendedInfo.addReferredEntity(dbEntity.getEntity());
 
         init();
-        entityStore.createOrUpdate(col0TypeMap);
+        entityStore.createOrUpdate(new AtlasEntityStream(col0WithExtendedInfo));
 
-        Map<String, AtlasEntity> col1TypeMap = new HashMap<>();
-        AtlasEntity col1Type = new AtlasEntity(TestUtils.COLUMN_TYPE,
-            new HashMap<String, Object>() {{
-                put(TestUtilsV2.NAME, "test2");
-                put("type", "string");
-                put("table", new AtlasObjectId(TABLE_TYPE, tableDefinition1.getGuid()));
-            }});
+        AtlasEntity col1 = new AtlasEntity(TestUtils.COLUMN_TYPE, TestUtilsV2.NAME, "test2");
+        col1.setAttribute("type", "string");
+        col1.setAttribute("table", tableEntity.getAtlasObjectId());
 
-        init();
-        col1TypeMap.put(col1Type.getGuid(), col1Type);
-        entityStore.createOrUpdate(col1TypeMap);
+        AtlasEntityWithExtInfo col1WithExtendedInfo = new AtlasEntityWithExtInfo(col1);
+        col1WithExtendedInfo.addReferredEntity(tableEntity);
+        col1WithExtendedInfo.addReferredEntity(dbEntity.getEntity());
 
-        AtlasObjectId col0Id = new AtlasObjectId(col0Type.getTypeName(), col0Type.getGuid());
-        AtlasObjectId col1Id = new AtlasObjectId(col1Type.getTypeName(), col1Type.getGuid());
+        init();
+        entityStore.createOrUpdate(new AtlasEntityStream(col1WithExtendedInfo));
 
-        columnsMap.put("col0", col0Id);
-        columnsMap.put("col1", col1Id);
-        tableCloneMap.put(col0Type.getGuid(), col0Type);
-        tableCloneMap.put(col1Type.getGuid(), col1Type);
+        Map<String, AtlasObjectId> columnsMap = new HashMap<String, AtlasObjectId>();
+        columnsMap.put("col0", col0.getAtlasObjectId());
+        columnsMap.put("col1", col1.getAtlasObjectId());
 
-        tableClone.setAttribute(TestUtils.COLUMNS_MAP, columnsMap);
+        tableEntity.setAttribute(TestUtils.COLUMNS_MAP, columnsMap);
 
+        entitiesInfo.addReferredEntity(col0);
+        entitiesInfo.addReferredEntity(col1);
         init();
-        response = entityStore.createOrUpdate(tableCloneMap);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
         AtlasEntityHeader tableDefinition5 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition5);
+        validateEntity(entitiesInfo, getEntityFromStore(tableDefinition5));
 
         //Swap elements
         columnsMap.clear();
-        columnsMap.put("col0", col1Id);
-        columnsMap.put("col1", col0Id);
+        columnsMap.put("col0", col1.getAtlasObjectId());
+        columnsMap.put("col1", col0.getAtlasObjectId());
 
-        tableClone.setAttribute(TestUtils.COLUMNS_MAP, columnsMap);
+        tableEntity.setAttribute(TestUtils.COLUMNS_MAP, columnsMap);
         init();
-        response = entityStore.createOrUpdate(tableCloneMap);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
         AtlasEntityHeader tableDefinition6 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition6);
+        validateEntity(entitiesInfo, getEntityFromStore(tableDefinition6));
 
         //Drop the first key and change the class type as well to col0
         columnsMap.clear();
-        columnsMap.put("col0", col0Id);
-
+        columnsMap.put("col0", col0.getAtlasObjectId());
         init();
-        response = entityStore.createOrUpdate(tableCloneMap);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
         AtlasEntityHeader tableDefinition7 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition7);
+        validateEntity(entitiesInfo, getEntityFromStore(tableDefinition7));
 
         //Clear state
-        tableClone.setAttribute(TestUtils.COLUMNS_MAP, null);
+        tableEntity.setAttribute(TestUtils.COLUMNS_MAP, null);
         init();
-        response = entityStore.createOrUpdate(tableCloneMap);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
         AtlasEntityHeader tableDefinition8 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition8);
+        validateEntity(entitiesInfo, getEntityFromStore(tableDefinition8));
     }
 
     @Test(dependsOnMethods = "testCreate")
     public void testMapOfPrimitivesUpdate() throws Exception {
-        //clear state
-        init();
+        AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
+        AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);
 
-        Map<String, AtlasEntity> entityCloneMap = new HashMap<>();
-        AtlasEntity entityClone = new AtlasEntity(tableEntity);
-        entityCloneMap.put(entityClone.getGuid(), entityClone);
+        entitiesInfo.addReferredEntity(tableEntity);
 
         //Add a new entry
-        Map<String, String> paramsMap = (Map<String, String>) entityClone.getAttribute("parametersMap");
+        Map<String, String> paramsMap = (Map<String, String>) tableEntity.getAttribute("parametersMap");
         paramsMap.put("newParam", "value");
-        entityClone.setAttribute("parametersMap", paramsMap);
-
-        EntityMutationResponse response = entityStore.createOrUpdate(entityCloneMap);
-        validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
-        AtlasEntityHeader tableEntity = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(entityCloneMap, tableEntity);
-
-        //clear state
         init();
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
+        AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         //Remove an entry
         paramsMap.remove("key1");
-        entityClone.setAttribute("parametersMap", paramsMap);
-
-        response = entityStore.createOrUpdate(entityCloneMap);
+        init();
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
         validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
-        tableEntity = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(entityCloneMap, tableEntity);
+        updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
     }
 
     @Test(dependsOnMethods = "testCreate")
     public void testArrayOfStructs() throws Exception {
         //Modify array of structs
-//        TestUtils.dumpGraph(TestUtils.getGraph());
-        init();
-        final AtlasStruct partition1 = new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE);
-        partition1.setAttribute(TestUtilsV2.NAME, "part1");
-        final AtlasStruct partition2 = new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE);
-        partition2.setAttribute(TestUtilsV2.NAME, "part2");
+        AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
+        AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);
 
-        List<AtlasStruct> partitions = new ArrayList<AtlasStruct>(){{ add(partition1); add(partition2); }};
+        List<AtlasStruct> partitions = new ArrayList<AtlasStruct>(){{
+            add(new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "part1"));
+            add(new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "part2"));
+        }};
         tableEntity.setAttribute("partitions", partitions);
 
-        EntityMutationResponse response = entityStore.createOrUpdate(tableEntityMap);
-        AtlasEntityHeader tableDefinition = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-
-        validateAttributes(tableEntityMap, tableDefinition);
+        init();
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         //add a new element to array of struct
+        partitions.add(new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "part3"));
         init();
-        final AtlasStruct partition3 = new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE);
-        partition3.setAttribute(TestUtilsV2.NAME, "part3");
-        partitions.add(partition3);
-        tableEntity.setAttribute("partitions", partitions);
-        response = entityStore.createOrUpdate(tableEntityMap);
-        tableDefinition = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableEntityMap, tableDefinition);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         //remove one of the struct values
         init();
         partitions.remove(1);
-        tableEntity.setAttribute("partitions", partitions);
-        response = entityStore.createOrUpdate(tableEntityMap);
-        tableDefinition = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableEntityMap, tableDefinition);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         //Update struct value within array of struct
         init();
         partitions.get(0).setAttribute(TestUtilsV2.NAME, "part4");
-        tableEntity.setAttribute("partitions", partitions);
-        response = entityStore.createOrUpdate(tableEntityMap);
-        tableDefinition = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableEntityMap, tableDefinition);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
 
         //add a repeated element to array of struct
+        partitions.add(new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "part4"));
         init();
-        final AtlasStruct partition4 = new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE);
-        partition4.setAttribute(TestUtilsV2.NAME, "part4");
-        partitions.add(partition4);
-        tableEntity.setAttribute("partitions", partitions);
-        response = entityStore.createOrUpdate(tableEntityMap);
-        tableDefinition = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableEntityMap, tableDefinition);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         // Remove all elements. Should set array attribute to null
-        init();
         partitions.clear();
-        tableEntity.setAttribute("partitions", partitions);
-        response = entityStore.createOrUpdate(tableEntityMap);
-        tableDefinition = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableEntityMap, tableDefinition);
+        init();
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
     }
 
 
     @Test(dependsOnMethods = "testCreate")
     public void testStructs() throws Exception {
-        init();
+        AtlasEntity              databaseEntity = dbEntity.getEntity();
+        AtlasEntity              tableEntity    = new AtlasEntity(tblEntity.getEntity());
+        AtlasEntitiesWithExtInfo entitiesInfo   = new AtlasEntitiesWithExtInfo(tableEntity);
 
-        Map<String, AtlasEntity> tableCloneMap = new HashMap<>();
-        AtlasEntity tableClone = new AtlasEntity(tableEntity);
-        AtlasStruct serdeInstance = new AtlasStruct(TestUtils.SERDE_TYPE);
-        serdeInstance.setAttribute(TestUtilsV2.NAME, "serde1Name");
+        AtlasStruct serdeInstance = new AtlasStruct(TestUtils.SERDE_TYPE, TestUtilsV2.NAME, "serde1Name");
         serdeInstance.setAttribute("serde", "test");
         serdeInstance.setAttribute("description", "testDesc");
-        tableClone.setAttribute("serde1", serdeInstance);
-        tableClone.setAttribute("database", new AtlasObjectId(dbEntity.getTypeName(), new HashMap<String, Object>() {{
-            put(TestUtilsV2.NAME, dbEntity.getAttribute(TestUtilsV2.NAME));
-        }}));
+        tableEntity.setAttribute("serde1", serdeInstance);
+        tableEntity.setAttribute("database", new AtlasObjectId(databaseEntity.getTypeName(), TestUtilsV2.NAME, databaseEntity.getAttribute(TestUtilsV2.NAME)));
 
-        tableCloneMap.put(tableClone.getGuid(), tableClone);
-
-        EntityMutationResponse response = entityStore.createOrUpdate(tableCloneMap);
-        AtlasEntityHeader tableDefinition1 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition1);
+        init();
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         //update struct attribute
-        init();
         serdeInstance.setAttribute("serde", "testUpdated");
-        response = entityStore.createOrUpdate(tableCloneMap);
-        AtlasEntityHeader tableDefinition2 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        validateAttributes(tableCloneMap, tableDefinition2);
+        init();
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
 
         //set to null
+        tableEntity.setAttribute("description", null);
         init();
-        tableClone.setAttribute("description", null);
-        response = entityStore.createOrUpdate(tableCloneMap);
-        AtlasEntityHeader tableDefinition3 = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
-        Assert.assertNull(tableDefinition3.getAttribute("description"));
-        validateAttributes(tableCloneMap, tableDefinition3);
+        response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo));
+        updatedTable = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
+        Assert.assertNull(updatedTable.getAttribute("description"));
+        validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
     }
 
 //    private AtlasEntity clearSubOrdinates(List<AtlasObjectId> employees, int index) {
@@ -559,118 +499,25 @@ public class AtlasEntityStoreV1Test {
 //        return subOrdinates.size() - 1;
 //    }
 
-    private void validateMutationResponse(EntityMutationResponse response, EntityMutations.EntityOperation op, int expectedNumCreated) {
-        List<AtlasEntityHeader> entitiesCreated = response.getEntitiesByOperation(op);
-        Assert.assertNotNull(entitiesCreated);
-        Assert.assertEquals(entitiesCreated.size(), expectedNumCreated);
-    }
-
-    private void validateAttributes(Map<String, AtlasEntity> entityMap, AtlasEntityHeader entity) throws AtlasBaseException, AtlasException {
-        //TODO : Use the older API for get until new instance API is ready and validated
-        ITypedReferenceableInstance instance = metadataService.getEntityDefinition(entity.getGuid());
-        assertAttributes(entityMap, entity, instance);
-    }
-
-    private void assertAttributes(Map<String, AtlasEntity> entityMap, AtlasStruct entity, IInstance instance) throws AtlasBaseException, AtlasException {
-        LOG.debug("Asserting type : " + entity.getTypeName());
-        AtlasStructType entityType = (AtlasStructType) typeRegistry.getType(instance.getTypeName());
-        for (String attrName : entity.getAttributes().keySet()) {
-            Object actual = entity.getAttribute(attrName);
-            Object expected = instance.get(attrName);
-
-            AtlasType attrType = entityType.getAttributeType(attrName);
-            assertAttribute(entityMap, actual, expected, attrType, attrName);
-        }
-    }
-
-    private void assertAttribute(Map<String, AtlasEntity> actualEntityMap, Object actual, Object expected, AtlasType attributeType, String attrName) throws AtlasBaseException, AtlasException {
-        LOG.debug("Asserting attribute : " + attrName);
-
-        switch(attributeType.getTypeCategory()) {
-        case ENTITY:
-            if ( expected instanceof Id) {
-                String guid = ((Id) expected)._getId();
-                Assert.assertTrue(AtlasEntity.isAssigned(guid));
-            } else {
-                ReferenceableInstance expectedInstance = (ReferenceableInstance) expected;
-                if (actual instanceof AtlasObjectId) {
-                    AtlasEntity actualEntity = actualEntityMap.get(((AtlasObjectId) actual).getGuid());
-                    if (actualEntity != null) {
-                        assertAttributes(actualEntityMap, actualEntity, expectedInstance);
-                    }
-                } else {
-                    AtlasEntity actualInstance = (AtlasEntity) actual;
-                    if (actualInstance != null) {
-                        assertAttributes(actualEntityMap , actualInstance, expectedInstance);
-                    }
-                }
-            }
-            break;
-        case PRIMITIVE:
-            Assert.assertEquals(actual, expected);
-            break;
-        case ENUM:
-            EnumValue expectedEnumVal = (EnumValue) expected;
-            if ( actual != null) {
-                Assert.assertEquals(actual, expectedEnumVal.value);
-            }
-            break;
-        case MAP:
-            AtlasMapType mapType = (AtlasMapType) attributeType;
-            AtlasType keyType = mapType.getKeyType();
-            AtlasType valueType = mapType.getValueType();
-            Map actualMap = (Map) actual;
-            Map expectedMap = (Map) expected;
-
-            if (expectedMap != null && actualMap != null) {
-                Assert.assertEquals(actualMap.size(), expectedMap.size());
-                for (Object key : actualMap.keySet()) {
-                    assertAttribute(actualEntityMap, actualMap.get(key), expectedMap.get(key), valueType, attrName);
-                }
-            }
-            break;
-        case ARRAY:
-            AtlasArrayType arrType = (AtlasArrayType) attributeType;
-            AtlasType elemType = arrType.getElementType();
-            List actualList = (List) actual;
-            List expectedList = (List) expected;
-
-            if (CollectionUtils.isNotEmpty(actualList)) {
-                //actual list could have deleted entities . Hence size may not match.
-                for (int i = 0; i < actualList.size(); i++) {
-                    assertAttribute(actualEntityMap, actualList.get(i), expectedList.get(i), elemType, attrName);
-                }
-            }
-            break;
-        case STRUCT:
-            StructInstance structInstance = (StructInstance) expected;
-            AtlasStruct newStructVal = (AtlasStruct) actual;
-            assertAttributes(actualEntityMap, newStructVal, structInstance);
-            break;
-        default:
-            Assert.fail("Unknown type category");
-        }
-    }
-
     @Test(dependsOnMethods = "testCreate")
     public void testClassUpdate() throws Exception {
 
         init();
         //Create new db instance
-        final Map<String, AtlasEntity> databaseInstance = TestUtilsV2.createDBEntity();
+        final AtlasEntity databaseInstance = TestUtilsV2.createDBEntity();
 
-        EntityMutationResponse response = entityStore.createOrUpdate(databaseInstance);
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(databaseInstance));
         final AtlasEntityHeader dbCreated = response.getFirstCreatedEntityByTypeName(TestUtilsV2.DATABASE_TYPE);
 
         init();
         Map<String, AtlasEntity> tableCloneMap = new HashMap<>();
-        AtlasEntity tableClone = new AtlasEntity(tableEntity);
-        tableClone.setAttribute("database", new AtlasObjectId(TestUtils.DATABASE_TYPE, dbCreated.getGuid()));
+        AtlasEntity tableClone = new AtlasEntity(tblEntity.getEntity());
+        tableClone.setAttribute("database", new AtlasObjectId(dbCreated.getGuid(), TestUtils.DATABASE_TYPE));
 
-        tableCloneMap.put(dbCreated.getGuid(), databaseInstance.values().iterator().next());
+        tableCloneMap.put(dbCreated.getGuid(), databaseInstance);
         tableCloneMap.put(tableClone.getGuid(), tableClone);
 
-        response = entityStore.createOrUpdate(tableCloneMap);
+        response = entityStore.createOrUpdate(new InMemoryMapEntityStream(tableCloneMap));
         final AtlasEntityHeader tableDefinition = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.TABLE_TYPE);
         Assert.assertNotNull(tableDefinition.getAttribute("database"));
         Assert.assertEquals(((AtlasObjectId) tableDefinition.getAttribute("database")).getGuid(), dbCreated.getGuid());
@@ -679,8 +526,10 @@ public class AtlasEntityStoreV1Test {
     @Test
     public void testCheckOptionalAttrValueRetention() throws Exception {
 
-        Map<String, AtlasEntity> dbEntityMap = TestUtilsV2.createDBEntity();
-        EntityMutationResponse response = entityStore.createOrUpdate(dbEntityMap);
+        AtlasEntity dbEntity = TestUtilsV2.createDBEntity();
+
+        init();
+        EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(dbEntity));
         AtlasEntityHeader firstEntityCreated = response.getFirstCreatedEntityByTypeName(TestUtilsV2.DATABASE_TYPE);
 
         //The optional boolean attribute should have a non-null value
@@ -691,18 +540,17 @@ public class AtlasEntityStoreV1Test {
         Assert.assertNull(firstEntityCreated.getAttribute(paramsAttr));
 
         //Update to true
-        init();
-        AtlasEntity dbEntity = dbEntityMap.values().iterator().next();
         dbEntity.setAttribute(isReplicatedAttr, Boolean.TRUE);
         //Update array
         final HashMap<String, String> params = new HashMap<String, String>() {{ put("param1", "val1"); put("param2", "val2"); }};
         dbEntity.setAttribute(paramsAttr, params);
         //Complete update
 
-        dbEntityMap.put(dbEntity.getGuid(), dbEntity);
-        response = entityStore.createOrUpdate(dbEntityMap);
+        init();
+        response = entityStore.createOrUpdate(new AtlasEntityStream(dbEntity));
         AtlasEntityHeader firstEntityUpdated = response.getFirstUpdatedEntityByTypeName(TestUtilsV2.DATABASE_TYPE);
 
+        Assert.assertNotNull(firstEntityUpdated);
         Assert.assertNotNull(firstEntityUpdated.getAttribute(isReplicatedAttr));
         Assert.assertEquals(firstEntityUpdated.getAttribute(isReplicatedAttr), Boolean.TRUE);
         Assert.assertEquals(firstEntityUpdated.getAttribute(paramsAttr), params);
@@ -742,8 +590,6 @@ public class AtlasEntityStoreV1Test {
         typeDefStore.createTypesDef(atlasTypesDef);
 
         //verify that entity can be created with reserved characters in string value, array value and map key and value
-        Map<String, AtlasEntity> entityCloneMap = new HashMap<>();
-
         AtlasEntity entity = new AtlasEntity();
         entity.setAttribute(strAttrName, randomStrWithReservedChars());
         entity.setAttribute(arrayAttrName, new String[]{randomStrWithReservedChars()});
@@ -751,10 +597,11 @@ public class AtlasEntityStoreV1Test {
             put(randomStrWithReservedChars(), randomStrWithReservedChars());
         }});
 
-        entityCloneMap.put(entity.getGuid(), entity);
-        final EntityMutationResponse response = entityStore.createOrUpdate(entityCloneMap);
+        AtlasEntityWithExtInfo entityWithExtInfo = new AtlasEntityWithExtInfo(entity);
+
+        final EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entityWithExtInfo));
         final AtlasEntityHeader firstEntityCreated = response.getFirstEntityCreated();
-        validateAttributes(entityCloneMap, firstEntityCreated);
+        validateEntity(entityWithExtInfo, getEntityFromStore(firstEntityCreated));
 
 
         //Verify that search with reserved characters works - for string attribute
@@ -765,10 +612,6 @@ public class AtlasEntityStoreV1Test {
 //        assertEquals(response.getJSONArray("rows").length(), 1);
     }
 
-    private String randomStrWithReservedChars() {
-        return randomString() + "\"${}%";
-    }
-
     @Test(expectedExceptions = AtlasBaseException.class)
     public void testCreateRequiredAttrNull() throws Exception {
         //Update required attribute
@@ -777,7 +620,107 @@ public class AtlasEntityStoreV1Test {
         tableEntity.setAttribute(TestUtilsV2.NAME, "table_" + TestUtils.randomString());
         tableCloneMap.put(tableEntity.getGuid(), tableEntity);
 
-        entityStore.createOrUpdate(tableCloneMap);
+        entityStore.createOrUpdate(new InMemoryMapEntityStream(tableCloneMap));
         Assert.fail("Expected exception while creating with required attribute null");
     }
+
+
+    private String randomStrWithReservedChars() {
+        return randomString() + "\"${}%";
+    }
+
+    private void validateMutationResponse(EntityMutationResponse response, EntityMutations.EntityOperation op, int expectedNumCreated) {
+        List<AtlasEntityHeader> entitiesCreated = response.getEntitiesByOperation(op);
+        Assert.assertNotNull(entitiesCreated);
+        Assert.assertEquals(entitiesCreated.size(), expectedNumCreated);
+    }
+
+    private void validateEntity(AtlasEntityExtInfo entityExtInfo, AtlasEntity actual) throws AtlasBaseException, AtlasException {
+        validateEntity(entityExtInfo, actual, entityExtInfo.getEntity(actual.getGuid()));
+    }
+
+    private void validateEntity(AtlasEntityExtInfo entityExtInfo, AtlasStruct actual, AtlasStruct expected) throws AtlasBaseException, AtlasException {
+        if (expected == null) {
+            Assert.assertNull(actual, "expected null instance. Found " + actual);
+
+            return;
+        }
+
+        Assert.assertNotNull(actual, "found null instance");
+
+        AtlasStructType entityType = (AtlasStructType) typeRegistry.getType(actual.getTypeName());
+        for (String attrName : expected.getAttributes().keySet()) {
+            Object expectedVal = expected.getAttribute(attrName);
+            Object actualVal   = actual.getAttribute(attrName);
+
+            AtlasType attrType = entityType.getAttributeType(attrName);
+            validateAttribute(entityExtInfo, actualVal, expectedVal, attrType, attrName);
+        }
+    }
+
+    private void validateAttribute(AtlasEntityExtInfo entityExtInfo, Object actual, Object expected, AtlasType attributeType, String attrName) throws AtlasBaseException, AtlasException {
+        switch(attributeType.getTypeCategory()) {
+            case ENTITY:
+                Assert.assertTrue(actual instanceof AtlasObjectId);
+                String guid = ((AtlasObjectId) actual).getGuid();
+                Assert.assertTrue(AtlasEntity.isAssigned(guid), "expected assigned guid. found " + guid);
+                break;
+
+            case PRIMITIVE:
+            case ENUM:
+                Assert.assertEquals(actual, expected);
+                break;
+
+            case MAP:
+                AtlasMapType mapType     = (AtlasMapType) attributeType;
+                AtlasType    valueType   = mapType.getValueType();
+                Map          actualMap   = (Map) actual;
+                Map          expectedMap = (Map) expected;
+
+                if (MapUtils.isEmpty(expectedMap)) {
+                    Assert.assertEquals(0, (actualMap == null ? 0 : actualMap.size()));
+                } else {
+                    Assert.assertFalse(MapUtils.isEmpty(actualMap));
+                    Assert.assertEquals(actualMap.size(), expectedMap.size());
+
+                    for (Object key : actualMap.keySet()) {
+                        validateAttribute(entityExtInfo, actualMap.get(key), expectedMap.get(key), valueType, attrName);
+                    }
+                }
+                break;
+
+            case ARRAY:
+                AtlasArrayType arrType      = (AtlasArrayType) attributeType;
+                AtlasType      elemType     = arrType.getElementType();
+                List           actualList   = (List) actual;
+                List           expectedList = (List) expected;
+
+                if (CollectionUtils.isEmpty(expectedList)) {
+                    Assert.assertTrue(CollectionUtils.isEmpty(actualList));
+                } else {
+                    Assert.assertFalse(CollectionUtils.isEmpty(actualList));
+                    Assert.assertEquals(actualList.size(), expectedList.size());
+
+                    //actual list could have deleted entities. Hence size may not match.
+                    for (int i = 0; i < actualList.size(); i++) {
+                        validateAttribute(entityExtInfo, actualList.get(i), expectedList.get(i), elemType, attrName);
+                    }
+                }
+                break;
+            case STRUCT:
+                AtlasStruct expectedStruct = (AtlasStruct) expected;
+                AtlasStruct actualStruct   = (AtlasStruct) actual;
+
+                validateEntity(entityExtInfo, actualStruct, expectedStruct);
+                break;
+            default:
+                Assert.fail("Unknown type category");
+        }
+    }
+
+    private AtlasEntity getEntityFromStore(AtlasEntityHeader header) throws AtlasBaseException {
+        AtlasEntityWithExtInfo entity = header != null ? entityStore.getById(header.getGuid()) : null;
+
+        return entity != null ? entity.getEntity() : null;
+    }
 }

repository/src/test/java/org/apache/atlas/service/DefaultMetadataServiceTest.java

@@ -765,7 +765,7 @@ public class DefaultMetadataServiceTest {
         Referenceable dbDef = InstanceSerialization.fromJsonReferenceable(dbDefJson, true);
 
         Assert.assertNotEquals(dbId, (((Id) tableDefinitionActual.get("database"))._getId()));
-        Assert.assertEquals(dbDef.getId()._getId(), (((Id) tableDefinitionActual.get("database"))._getId())); */
+        Assert.assertEquals(dbDef.getObjectId()._getId(), (((Id) tableDefinitionActual.get("database"))._getId())); */
 
     }
 

webapp/src/main/java/org/apache/atlas/examples/QuickStartV2.java

@@ -341,7 +341,7 @@ public class QuickStartV2 {
 
     private AtlasEntity createInstance(AtlasEntity entity, String[] traitNames) throws Exception {
         AtlasEntity ret = null;
-        EntityMutationResponse  response = entitiesClient.createEntity(entity);
+        EntityMutationResponse  response = entitiesClient.createEntity(new AtlasEntityWithExtInfo(entity));
         List<AtlasEntityHeader> entities = response.getEntitiesByOperation(EntityOperation.CREATE);
 
         if (CollectionUtils.isNotEmpty(entities)) {
@@ -422,7 +422,7 @@ public class QuickStartV2 {
         entity.setAttribute("retention", System.currentTimeMillis());
         entity.setAttribute("db", db.getAtlasObjectId());
         entity.setAttribute("sd", sd.getAtlasObjectId());
-        entity.setAttribute("columns", getObjectIds(columns));
+        entity.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));
 
         return createInstance(entity, traitNames);
     }
@@ -563,18 +563,4 @@ public class QuickStartV2 {
         AtlasEntity tableEntity = entitiesClient.getEntityByAttribute(TABLE_TYPE, attributes).getEntity();
         return tableEntity.getGuid();
     }
-
-    private Collection<AtlasObjectId> getObjectIds(Collection<AtlasEntity> entities) {
-        List<AtlasObjectId> ret = new ArrayList<>();
-
-        if (CollectionUtils.isNotEmpty(entities)) {
-            for (AtlasEntity entity : entities) {
-                if (entity != null) {
-                    ret.add(entity.getAtlasObjectId());
-                }
-            }
-        }
-
-        return ret;
-    }
 }

webapp/src/main/java/org/apache/atlas/web/adapters/AtlasEntityFormatConverter.java

@@ -60,11 +60,11 @@ public class AtlasEntityFormatConverter extends AtlasStructFormatConverter {
             if (v1Obj instanceof Id) {
                 Id id = (Id) v1Obj;
 
-                ret = new AtlasObjectId(id.getTypeName(), id._getId());
+                ret = new AtlasObjectId(id._getId(), id.getTypeName());
             } else if (v1Obj instanceof IReferenceableInstance) {
                 IReferenceableInstance entRef = (IReferenceableInstance) v1Obj;
 
-                ret = new AtlasObjectId(entRef.getTypeName(), entRef.getId()._getId());
+                ret = new AtlasObjectId(entRef.getId()._getId(), entRef.getTypeName());
 
                 if (!context.entityExists(ret.getGuid())) {
                     Map<String, Object> v1Attribs = null;

webapp/src/main/java/org/apache/atlas/web/adapters/AtlasObjectIdConverter.java

@@ -51,10 +51,10 @@ AtlasObjectIdConverter extends  AtlasAbstractFormatConverter {
         if (v1Obj != null) {
             if (v1Obj instanceof Id) {
                 Id id = (Id) v1Obj;
-                ret = new AtlasObjectId(id.getTypeName(), id._getId());
+                ret = new AtlasObjectId(id._getId(), id.getTypeName());
             } else if (v1Obj instanceof IReferenceableInstance) {
                 IReferenceableInstance entity = (IReferenceableInstance) v1Obj;
-                ret = new AtlasObjectId(entity.getTypeName(), entity.getId()._getId());
+                ret = new AtlasObjectId(entity.getId()._getId(), entity.getTypeName());
             }
         }
         return ret;

webapp/src/main/java/org/apache/atlas/web/rest/EntitiesREST.java

@@ -20,15 +20,16 @@ package org.apache.atlas.web.rest;
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.AtlasException;
-import org.apache.atlas.CreateUpdateEntitiesResult;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.ClassificationAssociateRequest;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
+import org.apache.atlas.repository.store.graph.v1.AtlasEntityStream;
+import org.apache.atlas.repository.store.graph.v1.EntityStream;
 import org.apache.atlas.services.MetadataService;
-import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedStruct;
 import org.apache.atlas.web.adapters.AtlasInstanceRestAdapters;
 import org.apache.atlas.web.util.Servlets;
@@ -50,13 +51,9 @@ import javax.ws.rs.Produces;
 import javax.ws.rs.QueryParam;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
-import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
-import java.util.Map;
 
 import static org.apache.atlas.web.adapters.AtlasInstanceRestAdapters.toAtlasBaseException;
-import static org.apache.atlas.web.adapters.AtlasInstanceRestAdapters.toEntityMutationResponse;
 
 
 @Path("v2/entities")
@@ -81,101 +78,6 @@ public class EntitiesREST {
         this.restAdapters = restAdapters;
     }
 
-    /*******
-     * Entity Creation/Updation if it already exists in ATLAS
-     * An existing entity is matched by its guid if supplied or by its unique attribute eg: qualifiedName
-     * Any associations like Classifications, Business Terms will have to be handled through the respective APIs
-     *******/
-
-    @POST
-    @Consumes(Servlets.JSON_MEDIA_TYPE)
-    @Produces(Servlets.JSON_MEDIA_TYPE)
-    public EntityMutationResponse createOrUpdate(Map<String, AtlasEntity> entities) throws AtlasBaseException {
-        EntityMutationResponse response = null;
-
-        ITypedReferenceableInstance[] entitiesInOldFormat = restAdapters.getITypedReferenceables(entities.values());
-
-        try {
-            final CreateUpdateEntitiesResult result = metadataService.updateEntities(entitiesInOldFormat);
-            response = toEntityMutationResponse(result);
-        } catch (AtlasException e) {
-            LOG.error("Exception while getting a typed reference for the entity ", e);
-            throw AtlasInstanceRestAdapters.toAtlasBaseException(e);
-        }
-        return response;
-    }
-
-    /*******
-     * Entity Updation - Allows full update of the specified entities.
-     * Any associations like Classifications, Business Terms will have to be handled through the respective APIs
-     * Null updates are supported i.e Set an attribute value to Null if its an optional attribute
-     *******/
-    @PUT
-    @Consumes(Servlets.JSON_MEDIA_TYPE)
-    @Produces(Servlets.JSON_MEDIA_TYPE)
-    public EntityMutationResponse update(Map<String, AtlasEntity> entities) throws AtlasBaseException {
-       return createOrUpdate(entities);
-    }
-
-    @GET
-    @Path("/guids")
-    @Consumes(Servlets.JSON_MEDIA_TYPE)
-    @Produces(Servlets.JSON_MEDIA_TYPE)
-    public AtlasEntity.AtlasEntities getById(@QueryParam("guid") List<String> guids) throws AtlasBaseException {
-
-        if (CollectionUtils.isEmpty(guids)) {
-            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guids);
-        }
-
-        AtlasEntity.AtlasEntities entities = new AtlasEntity.AtlasEntities();
-
-        List<AtlasEntity> entityList = new ArrayList<>();
-
-        for (String guid : guids) {
-            try {
-               ITypedReferenceableInstance ref       = metadataService.getEntityDefinition(guid);
-               Map<String, AtlasEntity>    entityRet = restAdapters.getAtlasEntity(ref);
-
-               addToEntityList(entityList, entityRet.values());
-
-            } catch (AtlasException e) {
-                throw toAtlasBaseException(e);
-            }
-        }
-
-        entities.setList(entityList);
-        return entities;
-    }
-
-    private void addToEntityList(final List<AtlasEntity> entityList, final Collection<AtlasEntity> values) {
-        for (AtlasEntity val : values) {
-            if ( !entityList.contains(val)) {
-                entityList.add(val);
-            }
-        }
-    }
-
-    /*******
-     * Entity Delete
-     *******/
-
-    @DELETE
-    @Path("/guids")
-    @Consumes(Servlets.JSON_MEDIA_TYPE)
-    @Produces(Servlets.JSON_MEDIA_TYPE)
-    public EntityMutationResponse deleteById(@QueryParam("guid") final List<String> guids) throws AtlasBaseException {
-
-        if (CollectionUtils.isEmpty(guids)) {
-            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guids);
-        }
-        try {
-            AtlasClient.EntityResult result = metadataService.deleteEntities(guids);
-            return toEntityMutationResponse(result);
-        } catch (AtlasException e) {
-            throw toAtlasBaseException(e);
-        }
-    }
-
     /**
      * Bulk API to associate a tag to multiple entities
      *

webapp/src/main/java/org/apache/atlas/web/rest/EntityREST.java

@@ -26,21 +26,23 @@ import org.apache.atlas.model.TypeCategory;
 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.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
+import org.apache.atlas.repository.store.graph.v1.AtlasEntityStream;
+import org.apache.atlas.repository.store.graph.v1.EntityStream;
 import org.apache.atlas.services.MetadataService;
 import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.typesystem.IStruct;
-import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedStruct;
 import org.apache.atlas.typesystem.Referenceable;
 import org.apache.atlas.web.adapters.AtlasFormatConverter;
 import org.apache.atlas.web.adapters.AtlasInstanceRestAdapters;
 import org.apache.atlas.web.util.Servlets;
+import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.collections.MapUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
@@ -114,9 +116,23 @@ public class EntityREST {
         AtlasEntityType     entityType = ensureEntityType(typeName);
         Map<String, Object> attributes = getAttributes(servletRequest);
 
-        validateAttributes(entityType, attributes);
+        validateUniqueAttribute(entityType, attributes);
 
-        return entitiesStore.getByUniqueAttribute(entityType, attributes);
+        return entitiesStore.getByUniqueAttributes(entityType, attributes);
+    }
+
+    /**
+     * Create new entity or update existing entity in Atlas.
+     * Existing entity is matched using its unique guid if supplied or by its unique attributes eg: qualifiedName
+     * @param entity
+     * @return EntityMutationResponse
+     * @throws AtlasBaseException
+     */
+    @POST
+    @Consumes(Servlets.JSON_MEDIA_TYPE)
+    @Produces(Servlets.JSON_MEDIA_TYPE)
+    public EntityMutationResponse createOrUpdate(AtlasEntityWithExtInfo entity) throws AtlasBaseException {
+        return entitiesStore.createOrUpdate(new AtlasEntityStream(entity));
     }
 
     /**
@@ -152,18 +168,23 @@ public class EntityREST {
     @PUT
     @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    @Path("/uniqueAttribute/type/{typeName}/attribute/{attrName}")
-    public EntityMutationResponse partialUpdateByUniqueAttribute(@PathParam("typeName") String entityType,
-                                                                 @PathParam("attrName") String attribute,
-                                                                 @QueryParam("value") String value, AtlasEntity entity) throws Exception {
+    @Path("/uniqueAttribute/type/{typeName}")
+    public EntityMutationResponse partialUpdateByUniqueAttribute(@PathParam("typeName") String typeName,
+                                                                 @Context HttpServletRequest servletRequest,
+                                                                 AtlasEntity entity) throws Exception {
+        AtlasEntityType     entityType = ensureEntityType(typeName);
+        Map<String, Object> attributes = getAttributes(servletRequest);
+
+        validateUniqueAttribute(entityType, attributes);
 
-        AtlasEntityType type = (AtlasEntityType) validateType(entityType, TypeCategory.ENTITY);
-        validateUniqueAttribute(type, attribute);
+        // legacy API supports only one unique attribute
+        String attribute = attributes.keySet().toArray(new String[1])[0];
+        String value     = (String)attributes.get(attribute);
 
         AtlasFormatConverter.ConverterContext ctx = new AtlasFormatConverter.ConverterContext();
         ctx.addEntity(entity);
         Referenceable ref = restAdapters.getReferenceable(entity, ctx);
-        CreateUpdateEntitiesResult result = metadataService.updateEntityByUniqueAttribute(entityType, attribute, value, ref);
+        CreateUpdateEntitiesResult result = metadataService.updateEntityByUniqueAttribute(typeName, attribute, value, ref);
         return toEntityMutationResponse(result);
     }
 
@@ -171,46 +192,67 @@ public class EntityREST {
     @DELETE
     @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    @Path("/uniqueAttribute/type/{typeName}/attribute/{attrName}")
+    @Path("/uniqueAttribute/type/{typeName}")
     public EntityMutationResponse deleteByUniqueAttribute(@PathParam("typeName") String typeName,
-        @PathParam("attrName") String attribute,
-        @QueryParam("value") String value) throws Exception {
-
+                                                          @Context HttpServletRequest servletRequest) throws Exception {
         AtlasEntityType     entityType = ensureEntityType(typeName);
+        Map<String, Object> attributes = getAttributes(servletRequest);
 
-        validateUniqueAttribute(entityType, attribute);
+        validateUniqueAttribute(entityType, attributes);
+
+        // legacy API supports only one unique attribute
+        String attribute = attributes.keySet().toArray(new String[1])[0];
+        String value     = (String)attributes.get(attribute);
 
         final AtlasClient.EntityResult result = metadataService.deleteEntityByUniqueAttribute(typeName, attribute, value);
         return toEntityMutationResponse(result);
     }
 
-    /**
-     * Fetch the complete definition of an entity
-     * which is identified by its type and unique attribute  eg: Referenceable.qualifiedName.
-     */
-    @Deprecated
     @GET
-    @Consumes({Servlets.JSON_MEDIA_TYPE, MediaType.APPLICATION_JSON})
+    @Path("/bulk")
+    @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    @Path("/uniqueAttribute/type/{typeName}/attribute/{attrName}")
-    public List<AtlasEntity> getByUniqueAttribute(@PathParam("typeName") String typeName,
-        @PathParam("attrName") String attribute,
-        @QueryParam("value") String value) throws AtlasBaseException {
+    public AtlasEntitiesWithExtInfo getByGuids(@QueryParam("guid") List<String> guids) throws AtlasBaseException {
 
-        List<AtlasEntity> entityList = new ArrayList<>();
-        AtlasEntityType                   entityType = ensureEntityType(typeName);
+        if (CollectionUtils.isEmpty(guids)) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guids);
+        }
 
-        validateUniqueAttribute(entityType, attribute);
+        AtlasEntitiesWithExtInfo entities = entitiesStore.getByIds(guids);
 
-        try {
-            final ITypedReferenceableInstance entityDefinitionReference = metadataService.getEntityDefinitionReference(typeName, attribute, value);
-            Map<String, AtlasEntity> entityRet = restAdapters.getAtlasEntity(entityDefinitionReference);
-            entityList.addAll(entityRet.values());
-        } catch (AtlasException e) {
-            throw toAtlasBaseException(e);
+        return entities;
     }
 
-        return entityList;
+    /**
+     * Create new entity or update existing entity in Atlas.
+     * Existing entity is matched using its unique guid if supplied or by its unique attributes eg: qualifiedName
+     * @param entities
+     * @return EntityMutationResponse
+     * @throws AtlasBaseException
+     */
+    @POST
+    @Path("/bulk")
+    @Consumes(Servlets.JSON_MEDIA_TYPE)
+    @Produces(Servlets.JSON_MEDIA_TYPE)
+    public EntityMutationResponse createOrUpdate(AtlasEntitiesWithExtInfo entities) throws AtlasBaseException {
+
+        EntityStream entityStream = new AtlasEntityStream(entities);
+
+        return entitiesStore.createOrUpdate(entityStream);
+    }
+
+    /*******
+     * Entity Delete
+     *******/
+
+    @DELETE
+    @Path("/bulk")
+    @Consumes(Servlets.JSON_MEDIA_TYPE)
+    @Produces(Servlets.JSON_MEDIA_TYPE)
+    public EntityMutationResponse deleteByGuids(@QueryParam("guid") final List<String> guids) throws AtlasBaseException {
+        EntityMutationResponse ret = entitiesStore.deleteByIds(guids);
+
+        return ret;
     }
 
     /**
@@ -348,18 +390,6 @@ public class EntityREST {
         }
     }
 
-    private AtlasType validateType(String entityType, TypeCategory expectedCategory) throws AtlasBaseException {
-        if ( StringUtils.isEmpty(entityType) ) {
-            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, entityType);
-        }
-
-        AtlasType type = typeRegistry.getType(entityType);
-        if (type.getTypeCategory() != expectedCategory) {
-            throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, type.getTypeCategory().name(), expectedCategory.name());
-        }
-
-        return type;
-    }
 
     private AtlasEntityType ensureEntityType(String typeName) throws AtlasBaseException {
         AtlasEntityType ret = typeRegistry.getEntityTypeByName(typeName);
@@ -381,50 +411,6 @@ public class EntityREST {
         return ret;
     }
 
-    /**
-     * Deprecated method - not used
-     * Validate that attribute is unique attribute
-     * @param entityType     the entity type
-     * @param attributeName  the name of the attribute
-     */
-    private void validateUniqueAttribute(AtlasEntityType entityType, String attributeName) throws AtlasBaseException {
-        AtlasAttributeDef attribute = entityType.getAttributeDef(attributeName);
-
-        if (attribute == null || !attribute.getIsUnique()) {
-            throw new AtlasBaseException(AtlasErrorCode.ATTRIBUTE_UNIQUE_INVALID, entityType.getTypeName(), attributeName);
-        }
-    }
-
-    private List<AtlasEntity> getOrderedEntityList(Map<String, AtlasEntity> entities, String firstItemGuid) {
-        List<AtlasEntity> ret = new ArrayList<>(entities.size());
-
-        for (AtlasEntity entity : entities.values()) {
-            if (StringUtils.equals(entity.getGuid(), firstItemGuid)) {
-                ret.add(0, entity);
-            } else {
-                ret.add(entity);
-            }
-        }
-
-        return ret;
-    }
-
-    private void validateAttributes(AtlasEntityType entityType, Map<String, Object> attributes)
-            throws AtlasBaseException {
-
-        if (MapUtils.isEmpty(attributes)) {
-            throw new AtlasBaseException(AtlasErrorCode.ATTRIBUTE_UNIQUE_INVALID, entityType.getTypeName(), "");
-        }
-
-        for (String attrName : attributes.keySet()) {
-            AtlasAttributeDef attrDef = entityType.getAttributeDef(attrName);
-
-            if (attrDef == null || !attrDef.getIsUnique()) {
-                throw new AtlasBaseException(AtlasErrorCode.ATTRIBUTE_UNIQUE_INVALID, entityType.getTypeName(), attrName);
-            }
-        }
-    }
-
     private Map<String, Object> getAttributes(HttpServletRequest request) {
         Map<String, Object> attributes = new HashMap<>();
 
@@ -443,4 +429,23 @@ public class EntityREST {
 
         return attributes;
     }
+
+    /**
+     * Validate that each attribute given is an unique attribute
+     * @param entityType the entity type
+     * @param attributes attributes
+     */
+    private void validateUniqueAttribute(AtlasEntityType entityType, Map<String, Object> attributes) throws AtlasBaseException {
+        if (MapUtils.isEmpty(attributes)) {
+            throw new AtlasBaseException(AtlasErrorCode.ATTRIBUTE_UNIQUE_INVALID, entityType.getTypeName(), "");
+        }
+
+        for (String attributeName : attributes.keySet()) {
+            AtlasAttributeDef attribute = entityType.getAttributeDef(attributeName);
+
+            if (attribute == null || !attribute.getIsUnique()) {
+                throw new AtlasBaseException(AtlasErrorCode.ATTRIBUTE_UNIQUE_INVALID, entityType.getTypeName(), attributeName);
+            }
+        }
+    }
 }

webapp/src/test/java/org/apache/atlas/web/adapters/TestEntitiesREST.java

@@ -25,6 +25,7 @@ import org.apache.atlas.RequestContext;
 import org.apache.atlas.TestUtilsV2;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasStruct;
@@ -33,7 +34,9 @@ import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
+import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer;
 import org.apache.atlas.store.AtlasTypeDefStore;
+import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.web.rest.EntitiesREST;
 
 import org.apache.atlas.web.rest.EntityREST;
@@ -49,16 +52,21 @@ import org.testng.annotations.Test;
 import javax.inject.Inject;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+
 @Guice(modules = {RepositoryMetadataModule.class})
 public class TestEntitiesREST {
 
     private static final Logger LOG = LoggerFactory.getLogger(TestEntitiesREST.class);
 
     @Inject
+    AtlasTypeRegistry typeRegistry;
+
+    @Inject
     private AtlasTypeDefStore typeStore;
 
     @Inject
@@ -69,10 +77,6 @@ public class TestEntitiesREST {
 
     private List<String> createdGuids = new ArrayList<>();
 
-    private Map<String, AtlasEntity> dbEntityMap;
-
-    private Map<String, AtlasEntity> tableEntityMap;
-
     private AtlasEntity dbEntity;
 
     private AtlasEntity tableEntity;
@@ -81,17 +85,22 @@ public class TestEntitiesREST {
 
     @BeforeClass
     public void setUp() throws Exception {
-        AtlasTypesDef typesDef = TestUtilsV2.defineHiveTypes();
-        typeStore.createTypesDef(typesDef);
-        dbEntityMap = TestUtilsV2.createDBEntity();
-        dbEntity = dbEntityMap.values().iterator().next();
+        AtlasTypesDef[] testTypesDefs = new AtlasTypesDef[] {  TestUtilsV2.defineHiveTypes() };
 
-        tableEntityMap = TestUtilsV2.createTableEntity(dbEntity.getGuid());
-        tableEntity = tableEntityMap.values().iterator().next();
+        for (AtlasTypesDef typesDef : testTypesDefs) {
+            AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typesDef, typeRegistry);
 
-        final AtlasEntity colEntity = TestUtilsV2.createColumnEntity(tableEntity.getGuid());
+            if (!typesToCreate.isEmpty()) {
+                typeStore.createTypesDef(typesToCreate);
+            }
+        }
+
+        dbEntity    = TestUtilsV2.createDBEntity();
+        tableEntity = TestUtilsV2.createTableEntity(dbEntity);
+
+        final AtlasEntity colEntity = TestUtilsV2.createColumnEntity(tableEntity);
         columns = new ArrayList<AtlasEntity>() {{ add(colEntity); }};
-        tableEntity.setAttribute("columns", columns);
+        tableEntity.setAttribute("columns", getObjIdList(columns));
     }
 
     @AfterMethod
@@ -106,11 +115,15 @@ public class TestEntitiesREST {
 
     @Test
     public void testCreateOrUpdateEntities() throws Exception {
-        Map<String, AtlasEntity> entities = new HashMap<>();
-        entities.put(dbEntity.getGuid(), dbEntity);
-        entities.put(tableEntity.getGuid(), tableEntity);
+        AtlasEntitiesWithExtInfo entities = new AtlasEntitiesWithExtInfo();
 
-        EntityMutationResponse response = entitiesREST.createOrUpdate(entities);
+        entities.addEntity(dbEntity);
+        entities.addEntity(tableEntity);
+        for (AtlasEntity column : columns) {
+            entities.addReferredEntity(column);
+        }
+
+        EntityMutationResponse response = entityREST.createOrUpdate(entities);
         List<AtlasEntityHeader> guids = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);
 
         Assert.assertNotNull(guids);
@@ -137,23 +150,24 @@ public class TestEntitiesREST {
     public void testUpdateWithSerializedEntities() throws  Exception {
         //Check with serialization and deserialization of entity attributes for the case
         // where attributes which are de-serialized into a map
-        Map<String, AtlasEntity> dbEntityMap = TestUtilsV2.createDBEntity();
-        AtlasEntity dbEntity = dbEntityMap.values().iterator().next();
-
-        Map<String, AtlasEntity> tableEntityMap = TestUtilsV2.createTableEntity(dbEntity.getGuid());
-        AtlasEntity tableEntity = tableEntityMap.values().iterator().next();
+        AtlasEntity dbEntity    = TestUtilsV2.createDBEntity();
+        AtlasEntity tableEntity = TestUtilsV2.createTableEntity(dbEntity);
 
-        final AtlasEntity colEntity = TestUtilsV2.createColumnEntity(tableEntity.getGuid());
+        final AtlasEntity colEntity = TestUtilsV2.createColumnEntity(tableEntity);
         List<AtlasEntity> columns = new ArrayList<AtlasEntity>() {{ add(colEntity); }};
-        tableEntity.setAttribute("columns", columns);
+        tableEntity.setAttribute("columns", getObjIdList(columns));
 
         AtlasEntity newDBEntity = serDeserEntity(dbEntity);
         AtlasEntity newTableEntity = serDeserEntity(tableEntity);
 
-        Map<String, AtlasEntity> newEntities = new HashMap<>();
-        newEntities.put(newDBEntity.getGuid(), newDBEntity);
-        newEntities.put(newTableEntity.getGuid(), newTableEntity);
-        EntityMutationResponse response2 = entitiesREST.createOrUpdate(newEntities);
+        AtlasEntitiesWithExtInfo newEntities = new AtlasEntitiesWithExtInfo();
+        newEntities.addEntity(newDBEntity);
+        newEntities.addEntity(newTableEntity);
+        for (AtlasEntity column : columns) {
+            newEntities.addReferredEntity(serDeserEntity(column));
+        }
+
+        EntityMutationResponse response2 = entityREST.createOrUpdate(newEntities);
 
         List<AtlasEntityHeader> newGuids = response2.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);
         Assert.assertNotNull(newGuids);
@@ -163,8 +177,8 @@ public class TestEntitiesREST {
     @Test(dependsOnMethods = "testCreateOrUpdateEntities")
     public void testGetEntities() throws Exception {
 
-        final AtlasEntity.AtlasEntities response = entitiesREST.getById(createdGuids);
-        final List<AtlasEntity> entities = response.getList();
+        final AtlasEntitiesWithExtInfo response = entityREST.getByGuids(createdGuids);
+        final List<AtlasEntity> entities = response.getEntities();
 
         Assert.assertNotNull(entities);
         Assert.assertEquals(entities.size(), 3);
@@ -174,7 +188,7 @@ public class TestEntitiesREST {
     @Test(dependsOnMethods = "testGetEntities")
     public void testDeleteEntities() throws Exception {
 
-        final EntityMutationResponse response = entitiesREST.deleteById(createdGuids);
+        final EntityMutationResponse response = entityREST.deleteByGuids(createdGuids);
         final List<AtlasEntityHeader> entities = response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE);
 
         Assert.assertNotNull(entities);
@@ -247,4 +261,14 @@ public class TestEntitiesREST {
         AtlasEntity newEntity = mapper.readValue(entityJson, AtlasEntity.class);
         return newEntity;
     }
+
+    private List<AtlasObjectId> getObjIdList(Collection<AtlasEntity> entities) {
+        List<AtlasObjectId> ret = new ArrayList<>();
+
+        for (AtlasEntity entity : entities) {
+            ret.add(entity.getAtlasObjectId());
+        }
+
+        return ret;
+    }
 }

webapp/src/test/java/org/apache/atlas/web/adapters/TestEntityREST.java

@@ -22,6 +22,7 @@ import org.apache.atlas.RequestContext;
 import org.apache.atlas.TestUtilsV2;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.EntityMutationResponse;
@@ -29,8 +30,8 @@ import org.apache.atlas.model.instance.EntityMutations;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.store.AtlasTypeDefStore;
-import org.apache.atlas.web.rest.EntitiesREST;
 import org.apache.atlas.web.rest.EntityREST;
+import org.mockito.Mockito;
 import org.testng.Assert;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeClass;
@@ -39,6 +40,7 @@ import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
 import javax.inject.Inject;
+import javax.servlet.http.HttpServletRequest;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -53,9 +55,6 @@ public class TestEntityREST {
     @Inject
     private EntityREST entityREST;
 
-    @Inject
-    private EntitiesREST entitiesREST;
-
     private AtlasEntity dbEntity;
 
     private AtlasClassification testClassification;
@@ -77,12 +76,9 @@ public class TestEntityREST {
     }
 
     private void createTestEntity() throws Exception {
-        Map<String, AtlasEntity> dbEntityMap = TestUtilsV2.createDBEntity();
-        AtlasEntity              dbEntity    = dbEntityMap.values().iterator().next();
+        AtlasEntity dbEntity = TestUtilsV2.createDBEntity();
 
-        dbEntityMap.put(dbEntity.getGuid(), dbEntity);
-
-        final EntityMutationResponse response = entitiesREST.createOrUpdate(dbEntityMap);
+        final EntityMutationResponse response = entityREST.createOrUpdate(new AtlasEntitiesWithExtInfo(dbEntity));
 
         Assert.assertNotNull(response);
         List<AtlasEntityHeader> entitiesMutated = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);
@@ -158,9 +154,8 @@ public class TestEntityREST {
 
     @Test
     public void  testUpdateGetDeleteEntityByUniqueAttribute() throws Exception {
-        Map<String, AtlasEntity> dbEntityMap = TestUtilsV2.createDBEntity();
-        AtlasEntity              dbEntity    = dbEntityMap.values().iterator().next();
-        EntityMutationResponse   response    = entitiesREST.createOrUpdate(dbEntityMap);
+        AtlasEntity            dbEntity = TestUtilsV2.createDBEntity();
+        EntityMutationResponse response = entityREST.createOrUpdate(new AtlasEntitiesWithExtInfo(dbEntity));
         String                 dbGuid   = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).get(0).getGuid();
 
         Assert.assertTrue(AtlasEntity.isAssigned(dbGuid));
@@ -170,21 +165,35 @@ public class TestEntityREST {
 
         dbEntity.setAttribute(TestUtilsV2.NAME, updatedDBName);
 
-        response = entityREST.partialUpdateByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, TestUtilsV2.NAME, prevDBName, dbEntity);
+        response = entityREST.partialUpdateByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, toHttpServletRequest(TestUtilsV2.NAME, prevDBName), dbEntity);
         Assert.assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).get(0).getGuid(), dbGuid);
 
         //Get By unique attribute
-        List<AtlasEntity> entities = entityREST.getByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, TestUtilsV2.NAME, updatedDBName);
-        Assert.assertNotNull(entities);
-        Assert.assertNotNull(entities.get(0).getGuid());
-        Assert.assertEquals(entities.get(0).getGuid(), dbGuid);
-        TestEntitiesREST.verifyAttributes(entities.get(0).getAttributes(), dbEntity.getAttributes());
+        AtlasEntityWithExtInfo entity = entityREST.getByUniqueAttributes(TestUtilsV2.DATABASE_TYPE, toHttpServletRequest(TestUtilsV2.NAME, updatedDBName));
+        Assert.assertNotNull(entity);
+        Assert.assertNotNull(entity.getEntity().getGuid());
+        Assert.assertEquals(entity.getEntity().getGuid(), dbGuid);
+        TestEntitiesREST.verifyAttributes(entity.getEntity().getAttributes(), dbEntity.getAttributes());
 
-        final EntityMutationResponse deleteResponse = entityREST.deleteByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, TestUtilsV2.NAME, (String) dbEntity.getAttribute(TestUtilsV2.NAME));
+        final EntityMutationResponse deleteResponse = entityREST.deleteByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, toHttpServletRequest(TestUtilsV2.NAME, (String) dbEntity.getAttribute(TestUtilsV2.NAME)));
 
         Assert.assertNotNull(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE));
         Assert.assertEquals(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 1);
         Assert.assertEquals(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).get(0).getGuid(), dbGuid);
     }
 
+    private HttpServletRequest toHttpServletRequest(String attrName, String attrValue) {
+        HttpServletRequest    request   = Mockito.mock(HttpServletRequest.class);
+        Map<String, String[]> paramsMap = toParametersMap(EntityREST.PREFIX_ATTR + attrName, attrValue);
+
+        Mockito.when(request.getParameterMap()).thenReturn(paramsMap);
+
+        return request;
+    }
+
+    private Map<String, String[]> toParametersMap(final String name, final String value) {
+        return new HashMap<String, String[]>() {{
+            put(name, new String[] { value });
+        }};
+    }
 }

webapp/src/test/java/org/apache/atlas/web/resources/BaseResourceIT.java

@@ -43,6 +43,7 @@ import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasStruct;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.typedef.AtlasClassificationDef;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
 import org.apache.atlas.model.typedef.AtlasEnumDef;
@@ -74,6 +75,7 @@ import org.apache.atlas.typesystem.types.TypeUtils;
 import org.apache.atlas.typesystem.types.utils.TypesUtil;
 import org.apache.atlas.utils.AuthenticationUtil;
 import org.apache.atlas.utils.ParamChecker;
+import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.configuration.Configuration;
 import org.apache.commons.lang.RandomStringUtils;
 import org.codehaus.jettison.json.JSONArray;
@@ -286,13 +288,13 @@ public abstract class BaseResourceIT {
         EntityMutationResponse entity = null;
         try {
             if (!update) {
-                entity = entitiesClientV2.createEntity(atlasEntity);
+                entity = entitiesClientV2.createEntity(new AtlasEntityWithExtInfo(atlasEntity));
                 assertNotNull(entity);
                 assertNotNull(entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
                 assertTrue(entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).size() > 0);
                 return entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).get(0);
             } else {
-                entity = entitiesClientV2.updateEntity(atlasEntity);
+                entity = entitiesClientV2.updateEntity(new AtlasEntityWithExtInfo(atlasEntity));
                 assertNotNull(entity);
                 assertNotNull(entity.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
                 assertTrue(entity.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).size() > 0);

webapp/src/test/java/org/apache/atlas/web/resources/EntityV2JerseyResourceIT.java

@@ -34,9 +34,11 @@ import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasServiceException;
 import org.apache.atlas.EntityAuditEvent;
 import org.apache.atlas.model.instance.AtlasClassification;
-import org.apache.atlas.model.instance.AtlasClassification.AtlasClassifications;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
+import org.apache.atlas.model.instance.AtlasObjectId;
+import org.apache.atlas.model.instance.AtlasClassification.AtlasClassifications;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
 import org.apache.atlas.model.typedef.AtlasClassificationDef;
@@ -106,54 +108,57 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
 
     @Test
     public void testCreateNestedEntities() throws Exception {
+        AtlasEntity.AtlasEntitiesWithExtInfo entities = new AtlasEntity.AtlasEntitiesWithExtInfo();
 
-        AtlasEntity databaseInstance = new AtlasEntity(DATABASE_TYPE_V2);
+        AtlasEntity databaseInstance = new AtlasEntity(DATABASE_TYPE_V2, "name", "db1");
         databaseInstance.setAttribute("name", "db1");
         databaseInstance.setAttribute("description", "foo database");
         databaseInstance.setAttribute("owner", "user1");
         databaseInstance.setAttribute("locationUri", "/tmp");
         databaseInstance.setAttribute("createTime",1000);
+        entities.addEntity(databaseInstance);
 
         int nTables = 5;
         int colsPerTable=3;
-        List<AtlasEntity> tables = new ArrayList<>();
-        List<AtlasEntity> allColumns = new ArrayList<>();
 
         for(int i = 0; i < nTables; i++) {
             String tableName = "db1-table-" + i;
 
-            AtlasEntity tableInstance =
-                    new AtlasEntity(HIVE_TABLE_TYPE_V2);
-            tableInstance.setAttribute("name", tableName);
+            AtlasEntity tableInstance = new AtlasEntity(HIVE_TABLE_TYPE_V2, "name", tableName);
             tableInstance.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName);
-            tableInstance.setAttribute("db", databaseInstance);
+            tableInstance.setAttribute("db", databaseInstance.getAtlasObjectId());
             tableInstance.setAttribute("description", tableName + " table");
-            tables.add(tableInstance);
+            entities.addEntity(tableInstance);
 
-            List<AtlasEntity> columns = new ArrayList<>();
+            List<AtlasObjectId> columns = new ArrayList<>();
             for(int j = 0; j < colsPerTable; j++) {
                 AtlasEntity columnInstance = new AtlasEntity(COLUMN_TYPE_V2);
                 columnInstance.setAttribute("name", tableName + "-col-" + j);
                 columnInstance.setAttribute("dataType", "String");
                 columnInstance.setAttribute("comment", "column " + j + " for table " + i);
-                allColumns.add(columnInstance);
-                columns.add(columnInstance);
+
+                columns.add(columnInstance.getAtlasObjectId());
+
+                entities.addReferredEntity(columnInstance);
             }
-            tableInstance.setAttribute("columns", ImmutableList.builder().addAll(columns).build());
+            tableInstance.setAttribute("columns", columns);
         }
 
         //Create the tables.  The database and columns should be created automatically, since
         //the tables reference them.
-        EntityMutationResponse response = entitiesClientV2.createEntities(tables);
+
+        EntityMutationResponse response = entitiesClientV2.createEntities(entities);
         Assert.assertNotNull(response);
 
         Map<String,String> guidsCreated = response.getGuidAssignments();
         assertEquals(guidsCreated.size(), nTables * colsPerTable + nTables + 1);
         assertNotNull(guidsCreated.get(databaseInstance.getGuid()));
-        for(AtlasEntity r : allColumns) {
+
+        for(AtlasEntity r : entities.getEntities()) {
             assertNotNull(guidsCreated.get(r.getGuid()));
         }
-        for(AtlasEntity r : tables) {
+
+        for(AtlasEntity r : entities.getReferredEntities().values()) {
             assertNotNull(guidsCreated.get(r.getGuid()));
         }
     }
@@ -189,7 +194,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         AtlasEntity hiveTableInstanceV2 = createHiveTableInstanceV2(hiveDBInstanceV2, tableName);
         hiveTableInstanceV2.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName);
 
-        EntityMutationResponse entity = entitiesClientV2.createEntity(hiveTableInstanceV2);
+        EntityMutationResponse entity = entitiesClientV2.createEntity(new AtlasEntityWithExtInfo(hiveTableInstanceV2));
         assertNotNull(entity);
         assertNotNull(entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
         results = searchByDSL(String.format("%s where name='%s'", DATABASE_TYPE_V2, DATABASE_NAME));
@@ -226,7 +231,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         //create entity for the type
         AtlasEntity instance = new AtlasEntity(entityDef.getName());
         instance.setAttribute("name", randomString());
-        EntityMutationResponse mutationResponse = entitiesClientV2.createEntity(instance);
+        EntityMutationResponse mutationResponse = entitiesClientV2.createEntity(new AtlasEntityWithExtInfo(instance));
         assertNotNull(mutationResponse);
         assertNotNull(mutationResponse.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
         assertEquals(mutationResponse.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).size(),1 );
@@ -272,9 +277,8 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         AtlasEntity hiveDB = createHiveDB();
         String qualifiedName = (String) hiveDB.getAttribute(NAME);
         //get entity by attribute
-        Map<String, String> attributes = new HashMap<>();
-        attributes.put(NAME, qualifiedName);
-        AtlasEntity byAttribute = entitiesClientV2.getEntityByAttribute(DATABASE_TYPE_V2, attributes).getEntity();
+
+        AtlasEntity byAttribute = entitiesClientV2.getEntityByAttribute(DATABASE_TYPE_V2, toMap(NAME, qualifiedName)).getEntity();
         assertEquals(byAttribute.getTypeName(), DATABASE_TYPE_V2);
         assertEquals(byAttribute.getAttribute(NAME), qualifiedName);
     }
@@ -407,7 +411,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
 
         AtlasEntity entityByGuid = getEntityByGuid(guid);
         entityByGuid.setAttribute(property, value);
-        EntityMutationResponse response = entitiesClientV2.updateEntity(entityByGuid);
+        EntityMutationResponse response = entitiesClientV2.updateEntity(new AtlasEntityWithExtInfo(entityByGuid));
         assertNotNull(response);
         assertNotNull(response.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
     }
@@ -624,10 +628,10 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         AtlasEntity hiveTable = createHiveTable();
         AtlasEntity tableUpdated = hiveTable;
 
-        hiveTable.setAttribute("columns", columns);
+        hiveTable.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));
 
         LOG.debug("Updating entity= " + tableUpdated);
-        EntityMutationResponse updateResult = entitiesClientV2.updateEntity(tableUpdated);
+        EntityMutationResponse updateResult = entitiesClientV2.updateEntity(new AtlasEntityWithExtInfo(tableUpdated));
         assertNotNull(updateResult);
         assertNotNull(updateResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
         assertTrue(updateResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).size() > 0);
@@ -642,11 +646,12 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         ref = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values);
         columns.set(0, ref);
         tableUpdated = hiveTable;
-        tableUpdated.setAttribute("columns", columns);
+        tableUpdated.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));
 
         LOG.debug("Updating entity= " + tableUpdated);
-        EntityMutationResponse updateResponse = entitiesClientV2.updateEntityByAttribute(BaseResourceIT.HIVE_TABLE_TYPE_V2, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
-                (String) hiveTable.getAttribute("name"), tableUpdated);
+        Map<String, String> uniqAttributes = new HashMap<>();
+        uniqAttributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, (String)hiveTable.getAttribute("name"));
+        EntityMutationResponse updateResponse = entitiesClientV2.updateEntityByAttribute(BaseResourceIT.HIVE_TABLE_TYPE_V2, uniqAttributes, tableUpdated);
         assertNotNull(updateResponse);
         assertNotNull(updateResponse.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
         assertTrue(updateResponse.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).size() > 0);
@@ -680,8 +685,8 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         columns.add(ref1);
         columns.add(ref2);
         AtlasEntity hiveTable = createHiveTable();
-        hiveTable.setAttribute("columns", columns);
-        EntityMutationResponse updateEntityResult = entitiesClientV2.updateEntity(hiveTable);
+        hiveTable.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));
+        EntityMutationResponse updateEntityResult = entitiesClientV2.updateEntity(new AtlasEntityWithExtInfo(hiveTable));
         assertNotNull(updateEntityResult);
         assertNotNull(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
         assertNotNull(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
@@ -690,7 +695,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         assertEquals(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).size(), 2);
 
         AtlasEntity entityByGuid = getEntityByGuid(hiveTable.getGuid());
-        List<AtlasEntity> refs = (List<AtlasEntity>) entityByGuid.getAttribute("columns");
+        List<AtlasObjectId> refs = (List<AtlasObjectId>) entityByGuid.getAttribute("columns");
         assertEquals(refs.size(), 2);
     }
 
@@ -713,7 +718,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         AtlasEntityHeader entity2Header = createEntity(db2);
 
         // Delete the database entities
-        EntityMutationResponse deleteResponse = entitiesClientV2.deleteEntityByGuid(ImmutableList.of(entity1Header.getGuid(), entity2Header.getGuid()));
+        EntityMutationResponse deleteResponse = entitiesClientV2.deleteEntitiesByGuids(ImmutableList.of(entity1Header.getGuid(), entity2Header.getGuid()));
 
         // Verify that deleteEntities() response has database entity guids
         assertNotNull(deleteResponse);
@@ -727,10 +732,9 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
     public void testDeleteEntityByUniqAttribute() throws Exception {
         // Create database entity
         AtlasEntity hiveDB = createHiveDB(DATABASE_NAME + random());
-        String qualifiedName = (String) hiveDB.getAttribute(NAME);
 
         // Delete the database entity
-        EntityMutationResponse deleteResponse = entitiesClientV2.deleteEntityByAttribute(DATABASE_TYPE_V2, NAME, qualifiedName);
+        EntityMutationResponse deleteResponse = entitiesClientV2.deleteEntityByAttribute(DATABASE_TYPE_V2, toMap(NAME, (String) hiveDB.getAttribute(NAME)));
 
         // Verify that deleteEntities() response has database entity guids
         assertNotNull(deleteResponse);
@@ -740,4 +744,9 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         // Verify entities were deleted from the repository.
     }
 
+    private Map<String, String> toMap(final String name, final String value) {
+        return new HashMap<String, String>() {{
+            put(name, value);
+        }};
+    }
 }