ATLAS-1522: V2 entity API changes to accept only AtlasObjectId for child references

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

@@ -32,7 +32,9 @@ import org.apache.hadoop.security.UserGroupInformation;
 import javax.ws.rs.HttpMethod;
 import javax.ws.rs.core.MultivaluedMap;
 import javax.ws.rs.core.Response;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import static org.apache.atlas.model.instance.AtlasEntity.AtlasEntities;
 
@@ -43,12 +45,11 @@ public class AtlasEntitiesClientV2 extends AtlasBaseClient {
 
     private static final APIInfo GET_ENTITY_BY_GUID = new APIInfo(ENTITY_API + "guid/", HttpMethod.GET, Response.Status.OK);
     private static final APIInfo GET_ENTITY_WITH_ASSOCIATION_BY_GUID = new APIInfo(ENTITY_API + "guid/%s/associations", HttpMethod.GET, Response.Status.OK);
-    private static final APIInfo CREATE_ENTITY = new APIInfo(ENTITY_API, HttpMethod.POST, 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_ATTRIBUTE = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s/attribute/%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 UPDATE_ENTITY_BY_GUID = new APIInfo(ENTITY_API + "guid/", HttpMethod.PUT, 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 GET_CLASSIFICATIONS = new APIInfo(ENTITY_API + "guid/%s/classifications", HttpMethod.GET, Response.Status.OK);
@@ -113,16 +114,12 @@ public class AtlasEntitiesClientV2 extends AtlasBaseClient {
         return callAPI(formatPathForPathParams(DELETE_ENTITY_BY_ATTRIBUTE, type, attribute), null, EntityMutationResponse.class, queryParams);
     }
 
-    public EntityMutationResponse createEntity(AtlasEntity atlasEntity) throws AtlasServiceException {
-        return callAPI(CREATE_ENTITY, atlasEntity, EntityMutationResponse.class);
+    public EntityMutationResponse createEntity(final AtlasEntity atlasEntity) throws AtlasServiceException {
+        return callAPI(CREATE_ENTITY, new HashMap<String, AtlasEntity>() {{ put(atlasEntity.getGuid(), atlasEntity); }}, EntityMutationResponse.class);
     }
 
-    public EntityMutationResponse updateEntity(AtlasEntity atlasEntity) throws AtlasServiceException {
-        return callAPI(UPDATE_ENTITY, atlasEntity, EntityMutationResponse.class);
-    }
-
-    public EntityMutationResponse updateEntity(String guid, AtlasEntity atlasEntity) throws AtlasServiceException {
-        return callAPI(UPDATE_ENTITY_BY_GUID, atlasEntity, EntityMutationResponse.class, guid);
+    public EntityMutationResponse updateEntity(final AtlasEntity atlasEntity) throws AtlasServiceException {
+        return callAPI(UPDATE_ENTITY, new HashMap<String, AtlasEntity>() {{ put(atlasEntity.getGuid(), atlasEntity); }}, EntityMutationResponse.class);
     }
 
     public AtlasEntity deleteEntityByGuid(String guid) throws AtlasServiceException {
@@ -159,15 +156,11 @@ public class AtlasEntitiesClientV2 extends AtlasBaseClient {
         return null;
     }
 
-    public List<AtlasEntity> createEntities(List<AtlasEntity> atlasEntities) throws AtlasServiceException {
+    public List<AtlasEntity> createEntities(Map<String, AtlasEntity> atlasEntities) throws AtlasServiceException {
         return (List<AtlasEntity>)callAPI(CREATE_ENTITIES, atlasEntities, List.class);
     }
 
-    public List<AtlasEntity> updateEntities(List<AtlasEntity> atlasEntities) throws AtlasServiceException {
+    public List<AtlasEntity> updateEntities(Map<String, AtlasEntity> atlasEntities) throws AtlasServiceException {
         return (List<AtlasEntity>)callAPI(UPDATE_ENTITIES, atlasEntities, List.class);
     }
-
-    public AtlasEntity.AtlasEntities searchEntities(SearchFilter searchFilter) throws AtlasServiceException {
-        return callAPI(GET_ENTITIES, AtlasEntity.AtlasEntities.class, searchFilter.getParams());
-    }
 }

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

@@ -80,6 +80,8 @@ public enum AtlasErrorCode {
 
     INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND(400, "ATLAS40018E", "Instance {0} with unique attribute {1} does not exist"),
 
+    UNRESOLVED_REFERENCES_FOUND(400, "ATLAS40010E", "Unresolved references: byId={0}; byUniqueAttributes={1}"),
+
     UNKNOWN_ATTRIBUTE(400, "ATLAS40019E", "Attribute {0} not found for type {1}");
 
     private String errorCode;

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

@@ -161,6 +161,41 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
         this.version = version;
     }
 
+    @JsonIgnore
+    public boolean isUnassigned() {
+        return isUnAssigned(guid);
+    }
+
+    @JsonIgnore
+    public boolean isAssigned() {
+        return isAssigned(guid);
+    }
+
+    @JsonIgnore
+    public static boolean isAssigned(String guid) {
+        try {
+            UUID.fromString(guid);
+        } catch (IllegalArgumentException e) {
+            return false;
+        }
+
+        return true;
+    }
+
+    @JsonIgnore
+    public static boolean isUnAssigned(String guid) {
+        return guid != null && guid.length() > 0 && guid.charAt(0) == '-';
+    }
+
+    private static String nextInternalId() {
+        return "-" + Long.toString(s_nextId.getAndIncrement());
+    }
+
+    @JsonIgnore
+    public AtlasObjectId getAtlasObjectId() {
+        return new AtlasObjectId(getTypeName(), getGuid());
+    }
+
     @Override
     public StringBuilder toString(StringBuilder sb) {
         if (sb == null) {
@@ -233,44 +268,4 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
             super(list, startIndex, pageSize, totalCount, sortType, sortBy);
         }
     }
-
-    @JsonIgnore
-    public boolean validate(String id) {
-        try {
-            long l = Long.parseLong(id);
-            return l < 0;
-        } catch (NumberFormatException ne) {
-            return false;
-        }
-    }
-
-    @JsonIgnore
-    public boolean isUnassigned() {
-        return isUnAssigned(guid);
-    }
-
-    @JsonIgnore
-    public boolean isAssigned() {
-        return isAssigned(guid);
-    }
-
-    @JsonIgnore
-    public static boolean isAssigned(String guid) {
-        try {
-            UUID.fromString(guid);
-        } catch (IllegalArgumentException e) {
-            return false;
-        }
-
-        return true;
-    }
-
-    @JsonIgnore
-    public static boolean isUnAssigned(String guid) {
-        return guid != null && guid.length() > 0 && guid.charAt(0) == '-';
-    }
-
-    private String nextInternalId() {
-        return "-" + Long.toString(s_nextId.getAndIncrement());
-    }
 }

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

@@ -73,6 +73,10 @@ public class AtlasEntityWithAssociations extends AtlasEntity implements Serializ
         setClassifications(other != null ? other.getClassifications() : null);
     }
 
+    public AtlasEntityWithAssociations(AtlasEntity other) {
+        super(other);
+    }
+
     @Override
     public StringBuilder toString(StringBuilder sb) {
         if (sb == null) {

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

@@ -30,10 +30,13 @@ import javax.xml.bind.annotation.XmlSeeAlso;
 import org.apache.atlas.model.PList;
 import org.apache.atlas.model.SearchFilter.SortType;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
+import org.apache.commons.collections.MapUtils;
+import org.apache.commons.lang.StringUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
 import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONLY;
 import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.NONE;
 
+import org.codehaus.jackson.annotate.JsonIgnore;
 import org.codehaus.jackson.annotate.JsonIgnoreProperties;
 import org.codehaus.jackson.map.annotate.JsonSerialize;
 
@@ -56,6 +59,12 @@ public class AtlasObjectId  implements Serializable {
     private String              guid;
     private Map<String, Object> uniqueAttributes;
 
+    @JsonIgnore
+    private boolean isAssignedGuid = false;
+
+    @JsonIgnore
+    private boolean isUnAssignedGuid = false;
+
     public AtlasObjectId() {
         this(null, null, null);
     }
@@ -120,6 +129,10 @@ public class AtlasObjectId  implements Serializable {
 
     public void setGuid(String guid) {
         this.guid = guid;
+        if ( guid != null) {
+            this.isAssignedGuid = AtlasEntity.isAssigned(guid);
+            this.isUnAssignedGuid = AtlasEntity.isUnAssigned(guid);
+        }
     }
 
     public Map<String, Object> getUniqueAttributes() {
@@ -130,6 +143,31 @@ public class AtlasObjectId  implements Serializable {
         this.uniqueAttributes = uniqueAttributes;
     }
 
+    @JsonIgnore
+    public boolean isAssignedGuid() {
+        return isAssignedGuid;
+    }
+
+    @JsonIgnore
+    public boolean isUnAssignedGuid() {
+        return isUnAssignedGuid;
+    }
+
+    @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;
+        }
+
+        return true;
+    }
+
     public StringBuilder toString(StringBuilder sb) {
         if (sb == null) {
             sb = new StringBuilder();

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

@@ -89,6 +89,34 @@ public class EntityMutationResponse {
         return null;
     }
 
+    @JsonIgnore
+    public AtlasEntityHeader getFirstCreatedEntityByTypeName(String typeName) {
+        final List<AtlasEntityHeader> entitiesByOperation = getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);
+        if ( entitiesByOperation != null && entitiesByOperation.size() > 0) {
+            for (AtlasEntityHeader header : entitiesByOperation) {
+                if ( header.getTypeName().equals(typeName)) {
+                    return header;
+                }
+            }
+        }
+
+        return null;
+    }
+
+    @JsonIgnore
+    public AtlasEntityHeader getFirstUpdatedEntityByTypeName(String typeName) {
+        final List<AtlasEntityHeader> entitiesByOperation = getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE);
+        if ( entitiesByOperation != null && entitiesByOperation.size() > 0) {
+            for (AtlasEntityHeader header : entitiesByOperation) {
+                if ( header.getTypeName().equals(typeName)) {
+                    return header;
+                }
+            }
+        }
+
+        return null;
+    }
+
     public void addEntity(EntityMutations.EntityOperation op, AtlasEntityHeader header) {
         if (entitiesMutated == null) {
             entitiesMutated = new HashMap<>();

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

@@ -193,7 +193,7 @@ public class AtlasEntityType extends AtlasStructType {
         if (obj != null) {
             if (obj instanceof AtlasObjectId) {
                 AtlasObjectId objId = (AtlasObjectId ) obj;
-                return validateAtlasObjectId(objId);
+                return isAssignableFrom(objId);
             } else {
                 for (AtlasEntityType superType : superTypes) {
                     if (!superType.isValidValue(obj)) {
@@ -240,7 +240,7 @@ public class AtlasEntityType extends AtlasStructType {
         if (obj != null) {
             if (obj instanceof AtlasObjectId) {
                 AtlasObjectId objId = (AtlasObjectId ) obj;
-                return validateAtlasObjectId(objId);
+                return isAssignableFrom(objId);
             }
 
             for (AtlasEntityType superType : superTypes) {
@@ -406,16 +406,10 @@ public class AtlasEntityType extends AtlasStructType {
         return ret == null ? Collections.<String, AtlasAttribute>emptyMap() : ret;
     }
 
-    private boolean validateAtlasObjectId(AtlasObjectId objId) {
-        if (StringUtils.isEmpty(objId.getTypeName()) || StringUtils.isEmpty(objId.getGuid())) {
-            return false;
-        } else {
-            String typeName = objId.getTypeName();
-            if (!typeName.equals(getTypeName()) && !isSuperTypeOf(typeName)) {
-                return false;
-            }
-        }
-        return AtlasEntity.isAssigned(objId.getGuid()) || AtlasEntity.isUnAssigned((objId.getGuid()));
+    boolean isAssignableFrom(AtlasObjectId objId) {
+        boolean ret = objId.isValid() && (StringUtils.equals(objId.getTypeName(), getTypeName()) || isSuperTypeOf(objId.getTypeName()));
+
+        return ret;
     }
 
     public static class ForeignKeyReference {

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

@@ -34,7 +34,6 @@ import org.slf4j.LoggerFactory;
 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;
@@ -481,6 +480,22 @@ public class AtlasStructType extends AtlasType {
             return qualifiedName;
         }
 
+        /*
+         * "isContainedAttribute" can not be computed and cached in the constructor - as structType is not fully
+         * populated at the time AtlasAttribute object is constructed.
+         */
+        public boolean isContainedAttribute() {
+            if ( structType.isForeignKeyOnDeleteActionUpdate(attributeDef.getName()) ) {
+                return true;
+            }
+
+            if ( structType instanceof AtlasEntityType) {
+                return ((AtlasEntityType) structType).isMappedFromRefAttribute(attributeDef.getName());
+            }
+
+            return false;
+        }
+
         public static String getQualifiedAttributeName(AtlasStructDef structDef, String attrName) {
             final String typeName = structDef.getName();
             return attrName.contains(".") ? attrName : String.format("%s.%s", typeName, attrName);

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

@@ -284,7 +284,9 @@ public final class TestUtilsV2 {
     public static final String PERSON_TYPE = "Person";
     public static final String EMPLOYEE_TYPE = "Employee";
 
-    public static AtlasEntity createDeptEg1() {
+    public static Map<String, AtlasEntity> createDeptEg1() {
+        Map<String, AtlasEntity> deptEmpEntities = new HashMap<>();
+
         AtlasEntity hrDept = new AtlasEntity(DEPARTMENT_TYPE);
         AtlasEntity john = new AtlasEntity(EMPLOYEE_TYPE);
 
@@ -329,7 +331,8 @@ public final class TestUtilsV2 {
         julius.setAttribute("address", juliusAddr);
         julius.setAttribute("subordinates", ImmutableList.of());
 
-        AtlasObjectId janeId = new AtlasObjectId(jane.getTypeName(), jane.getGuid());
+        AtlasObjectId janeId = jane.getAtlasObjectId();
+        AtlasObjectId johnId = john.getAtlasObjectId();
 
         //TODO - Change to MANAGER_TYPE for JULIUS
         AtlasObjectId maxId = new AtlasObjectId(EMPLOYEE_TYPE, max.getGuid());
@@ -355,15 +358,20 @@ public final class TestUtilsV2 {
 
         john.setAttribute("manager", janeId);
         john.setAttribute("mentor", maxId);
-        hrDept.setAttribute("employees", ImmutableList.of(john, jane, julius, max));
+        hrDept.setAttribute("employees", ImmutableList.of(johnId, janeId, juliusId, maxId));
 
-        jane.setAttribute("subordinates", ImmutableList.of(john, max));
+        jane.setAttribute("subordinates", ImmutableList.of(johnId, maxId));
 
 //        Map<String, Integer> secClearanceLevelMap = new HashMap<>();
 //        secClearanceLevelMap.put("level", 1);
 //        jane.setAttribute("SecurityClearance", secClearanceLevelMap);
 
-        return hrDept;
+        deptEmpEntities.put(jane.getGuid(), jane);
+        deptEmpEntities.put(john.getGuid(), john);
+        deptEmpEntities.put(julius.getGuid(), julius);
+        deptEmpEntities.put(max.getGuid(), max);
+        deptEmpEntities.put(deptId.getGuid(), hrDept);
+        return deptEmpEntities;
     }
 
     public static final String DATABASE_TYPE = "hive_database";
@@ -661,15 +669,20 @@ public final class TestUtilsV2 {
         return RandomStringUtils.randomAlphanumeric(10);
     }
 
-    public static AtlasEntity createDBEntity() {
+    public static Map<String, AtlasEntity> createDBEntity() {
+        Map<String, AtlasEntity> ret = new HashMap<>();
         AtlasEntity entity = new AtlasEntity(DATABASE_TYPE);
         String dbName = RandomStringUtils.randomAlphanumeric(10);
         entity.setAttribute(NAME, dbName);
         entity.setAttribute("description", "us db");
-        return entity;
+
+        ret.put(entity.getGuid(), entity);
+        return ret;
     }
 
-    public static AtlasEntity createTableEntity(String dbId) {
+    public static Map<String, AtlasEntity> createTableEntity(String dbId) {
+        Map<String, AtlasEntity> ret = new HashMap<>();
+
         AtlasEntity entity = new AtlasEntity(TABLE_TYPE);
         String tableName = RandomStringUtils.randomAlphanumeric(10);
         entity.setAttribute(NAME, tableName);
@@ -688,10 +701,13 @@ public final class TestUtilsV2 {
         entity.setAttribute("parametersMap", new java.util.HashMap<String, String>() {{
             put("key1", "value1");
         }});
-        return entity;
+
+        ret.put(entity.getGuid(), entity);
+        return ret;
     }
 
     public static AtlasEntity createColumnEntity(String tableId) {
+
         AtlasEntity entity = new AtlasEntity(COLUMN_TYPE);
         entity.setAttribute(NAME, RandomStringUtils.randomAlphanumeric(10));
         entity.setAttribute("type", "VARCHAR(32)");

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

@@ -28,6 +28,7 @@ import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.type.AtlasTypeRegistry;
 
 import java.util.List;
+import java.util.Map;
 
 /**
  * Persistence/Retrieval API for AtlasEntity
@@ -37,23 +38,8 @@ public interface AtlasEntityStore {
     /**
      * Initialization
      */
-    void init(AtlasTypeRegistry typeRegistry, EntityGraphDiscovery graphDiscovery) throws AtlasBaseException;
+    void init(AtlasTypeRegistry typeRegistry) throws AtlasBaseException;
 
-    /**
-     * Create or update an entity if it already exists.
-     * @param entity
-     * @return
-     */
-    EntityMutationResponse createOrUpdate(AtlasEntity entity) throws AtlasBaseException;
-
-
-    /**
-     * Update entity identified by its guid
-     * @param guid
-     * @param entity
-     * @return
-     */
-    EntityMutationResponse updateById(String guid, AtlasEntity entity);
 
     /**
      *
@@ -70,25 +56,14 @@ public interface AtlasEntityStore {
      */
     EntityMutationResponse deleteById(String guid);
 
-
     /**
-     * Create or update a list of entities
-     * @param entities List of AtlasEntity objects that need to be created
+     * 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(List<AtlasEntity> entities) throws AtlasBaseException;
-
-    /**
-     *
-     * Provides list of updated entity guids including any child entities
-     * @param guid
-     * @param entity
-     * @return
-     * @throws AtlasBaseException
-     */
-    EntityMutationResponse updateByIds(String guid, AtlasEntity entity) throws AtlasBaseException;
+    EntityMutationResponse createOrUpdate(Map<String, AtlasEntity> entities) throws AtlasBaseException;
 
     /**
      * Batch GET to retrieve entities by their ID
@@ -145,14 +120,6 @@ public interface AtlasEntityStore {
     EntityMutationResponse deleteByUniqueAttribute(String typeName, String attributeName, String attributeValue) throws AtlasBaseException;
 
     /**
-     * Compose any type of mutation op - EntityMutation.EntityOperation - CREATE_OR_UPDATE, PARTIAL_UPDATE, DELETE etc in a single transaction
-     * @param mutations
-     * @return
-     * @throws AtlasBaseException
-     */
-    EntityMutationResponse batchMutate(EntityMutations mutations) throws AtlasBaseException;
-
-    /**
      * Add classification(s)
      */
     void addClassifications(String guid, List<AtlasClassification> classification) throws AtlasBaseException;
@@ -168,13 +135,4 @@ public interface AtlasEntityStore {
      */
     void deleteClassifications(String guid, List<String> classificationNames) throws AtlasBaseException;
 
-    /**
-     *
-     * Search by AND filters like typename, pre-defined attribute(s) eg: name, qualifiedName
-     * @param searchFilter
-     * @return
-     * @throws AtlasBaseException
-     */
-    AtlasEntity.AtlasEntities searchEntities(SearchFilter searchFilter) throws AtlasBaseException;
-
 }

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

@@ -17,106 +17,145 @@
  */
 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.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.HashSet;
 import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
 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);
 
-    /**
-     *  Keeps track of all the entities that need to be created/updated including its child entities *
-     */
-    private Set<AtlasEntity> rootEntities = new LinkedHashSet<>();
+    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<>();
 
-    //Key is a transient id/guid
-    /**
-     * These references have been resolved using a unique identifier like guid or a qualified name etc in Atlas repository
-     */
-    private Map<String, AtlasVertex> repositoryResolvedReferences = new LinkedHashMap<>();
-
-    /**
-     * Unresolved entity references
-     */
-    private List<AtlasEntity> unresolvedEntityReferences = new ArrayList<>();
+    public EntityGraphDiscoveryContext(AtlasTypeRegistry typeRegistry) {
+        this.typeRegistry = typeRegistry;
+    }
 
-    /**
-     * Unresolved entity id references
-     */
-    private Set<AtlasObjectId> unresolvedIdReferences = new HashSet<>();
 
-    public void addRepositoryResolvedReference(AtlasObjectId id, AtlasVertex vertex) {
-        repositoryResolvedReferences.put(id.getGuid(), vertex);
+    public Collection<AtlasEntity> getRootEntities() {
+        return rootEntities;
     }
 
-    public void addUnResolvedEntityReference(AtlasEntity entity) {
-        this.unresolvedEntityReferences.add(entity);
+    public Map<AtlasObjectId, AtlasVertex> getResolvedIds() {
+        return resolvedIds;
     }
 
-    public void addUnResolvedIdReference(AtlasEntityType entityType, String id) {
-        this.unresolvedIdReferences.add(new AtlasObjectId(entityType.getTypeName(), id));
+    public Set<AtlasObjectId> getUnresolvedIds() {
+        return unresolvedIds;
     }
 
-    public Set<AtlasObjectId> getUnresolvedIdReferences() {
-        return unresolvedIdReferences;
+    public List<AtlasObjectId> getUnresolvedIdsByUniqAttribs() {
+        return unresolvedIdsByUniqAttribs;
     }
 
-    public boolean isResolved(String guid) {
-        return repositoryResolvedReferences.containsKey(guid);
-    }
 
-    public AtlasVertex getResolvedReference(AtlasObjectId ref) {
-        return repositoryResolvedReferences.get(ref.getGuid());
+    public void addRootEntity(AtlasEntity rootEntity) {
+        this.rootEntities.add(rootEntity);
     }
 
-    public Map<String, AtlasVertex> getRepositoryResolvedReferences() {
-        return repositoryResolvedReferences;
-    }
 
-    public AtlasVertex getResolvedReference(String id) {
-        return repositoryResolvedReferences.get(id);
+    public void addResolvedId(AtlasObjectId objId, AtlasVertex vertex) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("addResolvedId({})", objId);
+        }
+
+        resolvedIds.put(objId, vertex);
     }
 
-    public List<AtlasEntity> getUnResolvedEntityReferences() {
-        return unresolvedEntityReferences;
+    public boolean removeUnResolvedId(AtlasObjectId objId) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("removeUnResolvedId({})", objId);
+        }
+
+        return unresolvedIds.remove(objId);
     }
 
-    public void addRootEntity(AtlasEntity rootEntity) {
-        this.rootEntities.add(rootEntity);
+
+    public void addUnResolvedId(AtlasObjectId objId) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("addUnResolvedId({})", objId);
+        }
+
+        this.unresolvedIds.add(objId);
     }
 
-    public Collection<AtlasEntity> getRootEntities() {
-        return rootEntities;
+    public boolean removeUnResolvedIds(List<AtlasObjectId> objIds) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("removeUnResolvedIds({})", objIds);
+        }
+
+        return unresolvedIds.removeAll(objIds);
     }
 
-    public boolean removeUnResolvedEntityReference(final AtlasEntity entity) {
-        return unresolvedEntityReferences.remove(entity);
+
+    public void addUnresolvedIdByUniqAttribs(AtlasObjectId objId) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("addUnresolvedIdByUniqAttribs({})", objId);
+        }
+
+        this.unresolvedIdsByUniqAttribs.add(objId);
     }
 
-    public boolean removeUnResolvedEntityReferences(final List<AtlasEntity> entities) {
-        return unresolvedEntityReferences.removeAll(entities);
+    public boolean removeUnresolvedIdsByUniqAttribs(List<AtlasObjectId> objIds) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("removeUnresolvedIdsByUniqAttribs({})", objIds);
+        }
+
+        return unresolvedIdsByUniqAttribs.removeAll(objIds);
     }
 
-    public boolean removeUnResolvedIdReferences(final List<AtlasObjectId> entities) {
-        return unresolvedIdReferences.removeAll(entities);
+    public boolean hasUnresolvedReferences() {
+        return unresolvedIdsByUniqAttribs.size() > 0 || unresolvedIds.size() > 0;
     }
 
-    public boolean removeUnResolvedIdReference(final AtlasObjectId entity) {
-        return unresolvedIdReferences.remove(entity);
+    public boolean isResolvedId(AtlasObjectId id) {
+        return resolvedIds.containsKey(id);
     }
 
-    public boolean hasUnresolvedReferences() {
-        return unresolvedEntityReferences.size() > 0 || unresolvedIdReferences.size() > 0;
+    public AtlasVertex getResolvedEntityVertex(AtlasObjectId ref) throws AtlasBaseException {
+        AtlasVertex vertex = resolvedIds.get(ref);
+
+        // 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 Set<String>     allSubTypes = entityType.getAllSubTypes();
+
+            for (String subType : allSubTypes) {
+                AtlasObjectId subTypeObjId = new AtlasObjectId(subType, ref.getGuid(), ref.getUniqueAttributes());
+
+                vertex = resolvedIds.get(subTypeObjId);
+
+                if (vertex != null) {
+                    resolvedIds.put(ref, vertex);
+                    break;
+                }
+            }
+        }
+
+        if (vertex == null) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS,
+                                         " : Could not find an entity with " + ref.toString());
+        }
+
+        return vertex;
     }
 
     @Override
@@ -130,15 +169,15 @@ public final class EntityGraphDiscoveryContext {
         } else {
             EntityGraphDiscoveryContext ctx = (EntityGraphDiscoveryContext) obj;
             return Objects.equals(rootEntities, ctx.getRootEntities()) &&
-                Objects.equals(repositoryResolvedReferences, ctx.getRepositoryResolvedReferences()) &&
-                Objects.equals(unresolvedEntityReferences, ctx.getUnResolvedEntityReferences()) &&
-                Objects.equals(unresolvedIdReferences, ctx.getUnresolvedIdReferences());
+                Objects.equals(resolvedIds, ctx.getResolvedIds()) &&
+                Objects.equals(unresolvedIdsByUniqAttribs, ctx.getUnresolvedIdsByUniqAttribs()) &&
+                Objects.equals(unresolvedIds, ctx.getUnresolvedIds());
         }
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(rootEntities, repositoryResolvedReferences, unresolvedEntityReferences, unresolvedIdReferences);
+        return Objects.hash(rootEntities, resolvedIds, unresolvedIdsByUniqAttribs, unresolvedIds);
     }
 
     public StringBuilder toString(StringBuilder sb) {
@@ -148,9 +187,9 @@ public final class EntityGraphDiscoveryContext {
 
         sb.append("EntityGraphDiscoveryCtx{");
         sb.append("rootEntities='").append(rootEntities).append('\'');
-        sb.append(", repositoryResolvedReferences=").append(repositoryResolvedReferences);
-        sb.append(", unresolvedEntityReferences='").append(unresolvedEntityReferences).append('\'');
-        sb.append(", unresolvedIdReferences='").append(unresolvedIdReferences).append('\'');
+        sb.append(", resolvedIds=").append(resolvedIds);
+        sb.append(", unresolvedIdsByUniqAttribs='").append(unresolvedIdsByUniqAttribs).append('\'');
+        sb.append(", unresolvedIds='").append(unresolvedIds).append('\'');
         sb.append('}');
 
         return sb;
@@ -163,8 +202,8 @@ public final class EntityGraphDiscoveryContext {
 
     public void cleanUp() {
         rootEntities.clear();
-        unresolvedEntityReferences.clear();
-        repositoryResolvedReferences.clear();
-        unresolvedIdReferences.clear();
+        unresolvedIdsByUniqAttribs.clear();
+        resolvedIds.clear();
+        unresolvedIds.clear();
     }
 }

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

@@ -41,6 +41,7 @@ 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;
@@ -49,17 +50,15 @@ import com.google.inject.Provider;
 
 public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
 
-    private AtlasTypeRegistry typeRegistry;
-
-    private Set<String> processedIds = new HashSet<>();
-
-    private EntityGraphDiscoveryContext discoveredEntities = new EntityGraphDiscoveryContext();
-
-    private final Collection<EntityResolver> entityResolvers = new LinkedHashSet<>();
+    private final AtlasTypeRegistry           typeRegistry;
+    private final EntityGraphDiscoveryContext discoveredEntities;
+    private final Set<String>                 processedIds    = new HashSet<>();
+    private final Collection<EntityResolver>  entityResolvers = new LinkedHashSet<>();
 
     @Inject
-    public AtlasEntityGraphDiscoveryV1(AtlasTypeRegistry typeRegistry, final Collection<Provider<EntityResolver>> entityResolverProviders) {
-        this.typeRegistry = typeRegistry;
+    public AtlasEntityGraphDiscoveryV1(AtlasTypeRegistry typeRegistry, Collection<Provider<EntityResolver>> entityResolverProviders) {
+        this.typeRegistry       = typeRegistry;
+        this.discoveredEntities = new EntityGraphDiscoveryContext(typeRegistry);
 
         for (Provider<EntityResolver> entityResolverProvider : entityResolverProviders) {
              entityResolvers.add(entityResolverProvider.get());
@@ -67,8 +66,9 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
     }
 
     @VisibleForTesting
-    public AtlasEntityGraphDiscoveryV1(AtlasTypeRegistry typeRegistry, final List<EntityResolver> entityResolvers) {
-        this.typeRegistry = typeRegistry;
+    public AtlasEntityGraphDiscoveryV1(AtlasTypeRegistry typeRegistry, List<EntityResolver> entityResolvers) {
+        this.typeRegistry       = typeRegistry;
+        this.discoveredEntities = new EntityGraphDiscoveryContext(typeRegistry);
 
         for (EntityResolver entityResolver : entityResolvers) {
             this.entityResolvers.add(entityResolver);
@@ -96,22 +96,14 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
     public void cleanUp() throws AtlasBaseException {
         processedIds.clear();
         discoveredEntities.cleanUp();
-        final Collection<EntityResolver> entityResolvers = this.entityResolvers;
+
         for (EntityResolver resolver : entityResolvers) {
             resolver.cleanUp();
         }
     }
 
 
-    protected void resolveReferences() throws AtlasBaseException {
-        for (EntityResolver resolver : entityResolvers ) {
-            resolver.init(discoveredEntities);
-            resolver.resolveEntityReferences();
-        }
-    }
-
-
-    protected void discover(final List<AtlasEntity> entities) throws AtlasBaseException {
+    protected void discover(List<AtlasEntity> entities) throws AtlasBaseException {
         for (AtlasEntity entity : entities) {
             AtlasEntityType type = typeRegistry.getEntityTypeByName(entity.getTypeName());
 
@@ -120,27 +112,41 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
             }
 
             discoveredEntities.addRootEntity(entity);
+
             walkEntityGraph(type, entity);
         }
     }
 
-    private void visitReference(AtlasEntityType type, Object entity, boolean isManagedEntity) throws AtlasBaseException {
+    protected void resolveReferences() throws AtlasBaseException {
+        for (EntityResolver resolver : entityResolvers) {
+            resolver.init(discoveredEntities);
+
+            resolver.resolveEntityReferences();
+        }
+
+        if (discoveredEntities.hasUnresolvedReferences()) {
+            throw new AtlasBaseException(AtlasErrorCode.UNRESOLVED_REFERENCES_FOUND,
+                                                         discoveredEntities.getUnresolvedIds().toString(),
+                                                         discoveredEntities.getUnresolvedIdsByUniqAttribs().toString());
+        }
+    }
+
+    private void visitReference(AtlasEntityType type, Object entity) throws AtlasBaseException {
         if (entity != null) {
             if (entity instanceof AtlasObjectId) {
-                final String guid = ((AtlasObjectId) entity).getGuid();
-                discoveredEntities.addUnResolvedIdReference(type, guid);
-            } else if (entity instanceof AtlasEntity) {
-                AtlasEntity entityObj = (AtlasEntity) entity;
-                if (isManagedEntity) {
-                    if (!processedIds.contains(entityObj.getGuid())) {
-                        processedIds.add(entityObj.getGuid());
-
-                        discoveredEntities.addRootEntity(entityObj);
-                        visitStruct(type, entityObj);
-                    }
+                AtlasObjectId objId = (AtlasObjectId)entity;
+
+                if (!objId.isValid()) {
+                    throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, "Invalid object id " + objId);
+                }
+
+                if (!StringUtils.isEmpty(objId.getGuid()) && (objId.isAssignedGuid() || objId.isUnAssignedGuid())) {
+                    discoveredEntities.addUnResolvedId(objId);
                 } else {
-                    discoveredEntities.addUnResolvedEntityReference(entityObj);
+                    discoveredEntities.addUnresolvedIdByUniqAttribs(objId);
                 }
+            } 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());
             } else {
                 throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, "Invalid object type " + entity.getClass());
             }
@@ -154,26 +160,18 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
             }
             if (attrType.getTypeCategory() == TypeCategory.ARRAY) {
                 AtlasArrayType arrayType = (AtlasArrayType) attrType;
-                AtlasType elemType = arrayType.getElementType();
+                AtlasType      elemType  = arrayType.getElementType();
+
                 visitCollectionReferences(parentType, attrType, attrDef, elemType, val);
             } else if (attrType.getTypeCategory() == TypeCategory.MAP) {
-                AtlasType keyType = ((AtlasMapType) attrType).getKeyType();
+                AtlasType keyType   = ((AtlasMapType) attrType).getKeyType();
                 AtlasType valueType = ((AtlasMapType) attrType).getValueType();
+
                 visitMapReferences(parentType, attrType, attrDef, keyType, valueType, val);
             } else if (attrType.getTypeCategory() == TypeCategory.STRUCT) {
-                visitStruct(attrType, val);
+                visitStruct((AtlasStructType)attrType, val);
             } else if (attrType.getTypeCategory() == TypeCategory.ENTITY) {
-                if ( val instanceof AtlasObjectId) {
-                    visitReference((AtlasEntityType) attrType,  val, false);
-                } else if ( val instanceof AtlasEntity ) {
-                    //TODO - Change this to foreign key checks after changes in the model
-                   if ((parentType instanceof AtlasEntityType) &&
-                       ((AtlasEntityType)parentType).isMappedFromRefAttribute(attrDef.getName())) {
-                       visitReference((AtlasEntityType) attrType,  val, true);
-                   } else {
-                       visitReference((AtlasEntityType) attrType,  val, false);
-                   }
-                }
+                visitReference((AtlasEntityType) attrType,  val);
             }
         }
     }
@@ -184,7 +182,6 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
         }
 
         if (val != null) {
-
             if (Map.class.isAssignableFrom(val.getClass())) {
                 Iterator<Map.Entry> it = ((Map) val).entrySet().iterator();
                 while (it.hasNext()) {
@@ -197,7 +194,6 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
     }
 
     void visitCollectionReferences(final AtlasStructType parentType, final AtlasType attrType, final AtlasAttributeDef attrDef, AtlasType elemType, Object val) throws AtlasBaseException {
-
         if (isPrimitive(elemType.getTypeCategory())) {
             return;
         }
@@ -220,24 +216,22 @@ public class AtlasEntityGraphDiscoveryV1 implements EntityGraphDiscovery {
         }
     }
 
-    void visitStruct(AtlasType type, Object val) throws AtlasBaseException {
-
-        if (val == null || !(val instanceof AtlasStruct)) {
+    void visitStruct(AtlasStructType structType, Object val) throws AtlasBaseException {
+        if (structType == null) {
             return;
         }
 
-        AtlasStructType structType = (AtlasStructType) type;
-
         for (AtlasStructType.AtlasAttribute attribute : structType.getAllAttributes().values()) {
             AtlasType attrType = attribute.getAttributeType();
-            Object attrVal = ((AtlasStruct) val).getAttribute(attribute.getName());
+            Object    attrVal  = ((AtlasStruct) val).getAttribute(attribute.getName());
+
             visitAttribute(structType, attrType, attribute.getAttributeDef(), attrVal);
         }
     }
 
 
-    void walkEntityGraph(AtlasType type, AtlasEntity entity) throws AtlasBaseException {
-        visitStruct(type, entity);
+    void walkEntityGraph(AtlasEntityType entityType, AtlasEntity entity) throws AtlasBaseException {
+        visitStruct(entityType, entity);
     }
 
 

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

@@ -25,31 +25,29 @@ 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.SearchFilter;
 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.AtlasEntityWithAssociations;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.EntityMutationResponse;
-import org.apache.atlas.model.instance.EntityMutations;
 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 com.google.common.annotations.VisibleForTesting;
-import com.google.inject.Inject;
+import java.util.Map;
 
+import com.google.inject.Inject;
 
 
 public class AtlasEntityStoreV1 implements AtlasEntityStore {
 
-    protected EntityGraphDiscovery graphDiscoverer;
     protected AtlasTypeRegistry typeRegistry;
 
     private EntityGraphMapper graphMapper;
@@ -62,22 +60,11 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Inject
-    public void init(AtlasTypeRegistry typeRegistry, EntityGraphDiscovery graphDiscoverer) throws AtlasBaseException {
-        this.graphDiscoverer = graphDiscoverer;
+    public void init(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
         this.typeRegistry = typeRegistry;
     }
 
     @Override
-    public EntityMutationResponse createOrUpdate(final AtlasEntity entity) throws AtlasBaseException {
-        return createOrUpdate(new ArrayList<AtlasEntity>() {{ add(entity); }});
-    }
-
-    @Override
-    public EntityMutationResponse updateById(final String guid, final AtlasEntity entity) {
-        return null;
-    }
-
-    @Override
     public AtlasEntity getById(final String guid) {
         return null;
     }
@@ -89,7 +76,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
 
     @Override
     @GraphTransaction
-    public EntityMutationResponse createOrUpdate(final List<AtlasEntity> entities) throws AtlasBaseException {
+    public EntityMutationResponse createOrUpdate(final Map<String, AtlasEntity> entities) throws AtlasBaseException {
 
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> AtlasEntityStoreV1.createOrUpdate({}, {})", entities);
@@ -109,11 +96,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Override
-    public EntityMutationResponse updateByIds(final String guid, final AtlasEntity entity) throws AtlasBaseException {
-        return null;
-    }
-
-    @Override
     public AtlasEntity.AtlasEntities getByIds(final List<String> guid) throws AtlasBaseException {
         return null;
     }
@@ -144,12 +126,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Override
-    public EntityMutationResponse batchMutate(final EntityMutations mutations) throws AtlasBaseException {
-        return null;
-    }
-
-
-    @Override
     public void addClassifications(final String guid, final List<AtlasClassification> classification) throws AtlasBaseException {
 
     }
@@ -164,42 +140,49 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
 
     }
 
-    @Override
-    public AtlasEntity.AtlasEntities searchEntities(final SearchFilter searchFilter) throws AtlasBaseException {
-        // TODO: Add checks here to ensure that typename and supertype are mandatory in the request
-        return null;
-    }
-
     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);
-        for (AtlasEntity entity : discoveredEntities.getRootEntities()) {
+        EntityMutationContext       context            = new EntityMutationContext(discoveredEntities);
 
-            AtlasVertex vertex = null;
+        for (AtlasEntity entity : discoveredEntities.getRootEntities()) {
             if (LOG.isDebugEnabled()) {
-                LOG.debug("<== AtlasEntityStoreV1.preCreateOrUpdate({}): {}", entity);
+                LOG.debug("==> AtlasEntityStoreV1.preCreateOrUpdate({}): {}", entity);
             }
 
             AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
 
-            if ( entityType == null) {
+            if (entityType == null) {
                 throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
             }
 
-            if ( discoveredEntities.isResolved(entity.getGuid()) ) {
-                vertex = discoveredEntities.getResolvedReference(entity.getGuid());
+            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));
             } else {
                 //Create vertices which do not exist in the repository
                 vertex = graphMapper.createVertexTemplate(entity, entityType);
+
                 context.addCreated(entity, entityType, vertex);
-                discoveredEntities.addRepositoryResolvedReference(new AtlasObjectId(entityType.getTypeName(), entity.getGuid()), vertex);
+
+                discoveredEntities.addResolvedId(objId, vertex);
+                discoveredEntities.removeUnResolvedId(objId);
 
                 String guid = AtlasGraphUtilsV1.getIdFromVertex(vertex);
+
                 RequestContextV1.get().recordEntityCreate(new AtlasObjectId(entityType.getTypeName(), guid));
             }
 
@@ -211,12 +194,21 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
         return context;
     }
 
-    private List<AtlasEntity> validateAndNormalize(final List<AtlasEntity> entities) throws AtlasBaseException {
-
+    private List<AtlasEntity> validateAndNormalize(final Map<String, AtlasEntity> entities) throws AtlasBaseException {
         List<AtlasEntity> normalizedEntities = new ArrayList<>();
-        List<String> messages = 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);
+            }
 
-        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());
@@ -227,11 +219,9 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
             if ( !messages.isEmpty()) {
                 throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, messages);
             }
+
             AtlasEntity normalizedEntity = (AtlasEntity) type.getNormalizedValue(entity);
-            if ( normalizedEntity == null) {
-                //TODO - Fix this. Should not come here. Should ideally fail above
-                throw new AtlasBaseException(AtlasErrorCode.INSTANCE_CRUD_INVALID_PARAMS, "Failed to validate entity");
-            }
+
             normalizedEntities.add(normalizedEntity);
         }
 
@@ -239,6 +229,5 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     public void cleanUp() throws AtlasBaseException {
-        this.graphDiscoverer.cleanUp();
     }
 }

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

@@ -25,6 +25,7 @@ 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.AtlasObjectId;
+import org.apache.atlas.model.typedef.AtlasStructDef;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graph.AtlasEdgeLabel;

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

@@ -88,8 +88,8 @@ public class EntityGraphMapper implements InstanceGraphMapper<AtlasEdge> {
     public AtlasEdge toGraph(GraphMutationContext ctx) throws AtlasBaseException {
         AtlasEdge result = null;
 
-        String guid = getId(ctx.getValue());
-        AtlasVertex entityVertex = context.getDiscoveryContext().getResolvedReference(guid);
+        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);
         } else if (ctx.getValue() != null) {
@@ -166,12 +166,18 @@ public class EntityGraphMapper implements InstanceGraphMapper<AtlasEdge> {
     }
 
 
-    public String getId(Object value) throws AtlasBaseException {
-        if ( value != null) {
+    public AtlasObjectId getId(Object value) throws AtlasBaseException {
+        if (value != null) {
             if ( value instanceof  AtlasObjectId) {
-                return ((AtlasObjectId) value).getGuid();
+                return ((AtlasObjectId) value);
             } else if (value instanceof AtlasEntity) {
-                return ((AtlasEntity) value).getGuid();
+                return ((AtlasEntity) value).getAtlasObjectId();
+            } else if (value instanceof Map) {
+                AtlasObjectId ret = new AtlasObjectId((Map)value);
+
+                if (ret.isValid()) {
+                    return ret;
+                }
             }
 
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, (String) value);
@@ -211,7 +217,7 @@ public class EntityGraphMapper implements InstanceGraphMapper<AtlasEdge> {
     }
 
     public AtlasEntityType getInstanceType(Object val) throws AtlasBaseException {
-        String guid = getId(val);
+        AtlasObjectId guid = getId(val);
 
         if ( guid != null) {
             return (AtlasEntityType) getContext().getType(guid);

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

@@ -19,40 +19,41 @@ package org.apache.atlas.repository.store.graph.v1;
 
 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.EntityGraphDiscoveryContext;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasType;
 
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 public class EntityMutationContext {
-
-    private List<AtlasEntity> entitiesCreated  = new ArrayList<>();
-    private List<AtlasEntity> entitiesUpdated  = new ArrayList<>();
-
-    private EntityGraphDiscoveryContext context;
-    private Map<String, AtlasEntityType> entityVsType = new HashMap<>();
-    private Map<String, AtlasVertex> entityVsVertex = new HashMap<>();
+    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<>();
 
     public EntityMutationContext(final EntityGraphDiscoveryContext context) {
         this.context = context;
     }
 
     public void addCreated(AtlasEntity entity, AtlasEntityType type, AtlasVertex atlasVertex) {
+        AtlasObjectId objId = entity.getAtlasObjectId();
         entitiesCreated.add(entity);
-        entityVsVertex.put(entity.getGuid(), atlasVertex);
-        entityVsType.put(entity.getGuid(), type);
+        entityVsType.put(objId, type);
+        entityVsVertex.put(objId, atlasVertex);
     }
 
     public void addUpdated(AtlasEntity entity, AtlasEntityType type, AtlasVertex atlasVertex) {
+        AtlasObjectId objId = entity.getAtlasObjectId();
         entitiesUpdated.add(entity);
-        entityVsVertex.put(entity.getGuid(), atlasVertex);
-        entityVsType.put(entity.getGuid(), type);
+        entityVsType.put(objId, type);
+        entityVsVertex.put(objId, atlasVertex);
+    }
+
+    public EntityGraphDiscoveryContext getDiscoveryContext() {
+        return this.context;
     }
 
     public Collection<AtlasEntity> getCreatedEntities() {
@@ -64,26 +65,21 @@ public class EntityMutationContext {
     }
 
     public AtlasEntityType getType(AtlasEntity entity) {
-        return entityVsType.get(entity.getGuid());
+        return entityVsType.get(entity.getAtlasObjectId());
     }
 
-    public AtlasType getType(String entityId) {
+    public AtlasType getType(AtlasObjectId entityId) {
         return entityVsType.get(entityId);
     }
 
     public AtlasVertex getVertex(AtlasEntity entity) {
-        return entityVsVertex.get(entity.getGuid());
+        return entityVsVertex.get(entity.getAtlasObjectId());
     }
 
-    public AtlasVertex getVertex(String entityId) {
+    public AtlasVertex getVertex(AtlasObjectId entityId) {
         return entityVsVertex.get(entityId);
     }
 
-    public EntityGraphDiscoveryContext getDiscoveryContext() {
-        return this.context;
-    }
-
-    //TODO - equals/hashCode/toString
 
     @Override
     public boolean equals(final Object o) {
@@ -92,32 +88,29 @@ public class EntityMutationContext {
 
         final EntityMutationContext that = (EntityMutationContext) o;
 
-        if (entitiesCreated != null ? !entitiesCreated.equals(that.entitiesCreated) : that.entitiesCreated != null)
-            return false;
-        if (entitiesUpdated != null ? !entitiesUpdated.equals(that.entitiesUpdated) : that.entitiesUpdated != null)
-            return false;
-        if (context != null ? !context.equals(that.context) : that.context != null) return false;
-        if (entityVsType != null ? !entityVsType.equals(that.entityVsType) : that.entityVsType != null) return false;
-        return !(entityVsVertex != null ? !entityVsVertex.equals(that.entityVsVertex) : that.entityVsVertex != null);
-
+        return Objects.equals(context, that.context) &&
+               Objects.equals(entitiesCreated, that.entitiesCreated) &&
+               Objects.equals(entitiesUpdated, that.entitiesUpdated) &&
+               Objects.equals(entityVsType, that.entityVsType) &&
+               Objects.equals(entityVsVertex, that.entityVsVertex);
     }
 
     @Override
     public int hashCode() {
-        int result = entitiesCreated != null ? entitiesCreated.hashCode() : 0;
-        result = 31 * result + (entitiesUpdated != null ? entitiesUpdated.hashCode() : 0);
-        result = 31 * result + (context != null ? context.hashCode() : 0);
-        result = 31 * result + (entityVsType != null ? entityVsType.hashCode() : 0);
-        result = 31 * result + (entityVsVertex != null ? entityVsVertex.hashCode() : 0);
+        int result = (context != null ? context.hashCode() : 0);
+        result = 31 * result + entitiesCreated.hashCode();
+        result = 31 * result + entitiesUpdated.hashCode();
+        result = 31 * result + entityVsType.hashCode();
+        result = 31 * result + entityVsVertex.hashCode();
         return result;
     }
 
     @Override
     public String toString() {
         return "EntityMutationContext{" +
-            "entitiesCreated=" + entitiesCreated +
+            "context=" + context +
+            ", entitiesCreated=" + entitiesCreated +
             ", entitiesUpdated=" + entitiesUpdated +
-            ", context=" + context +
             ", entityVsType=" + entityVsType +
             ", entityVsVertex=" + entityVsVertex +
             '}';

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityStream.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 org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasObjectId;
+
+public interface EntityStream {
+
+    boolean hasNext();
+
+    AtlasEntity next();
+
+    void reset();
+
+    AtlasEntity getById(AtlasObjectId id);
+}

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

@@ -37,67 +37,70 @@ import java.util.Map;
 
 public class IDBasedEntityResolver implements EntityResolver {
 
-    private Map<String, AtlasEntity> idToEntityMap = new HashMap<>();
-
-    private final GraphHelper graphHelper = GraphHelper.getInstance();
-
-    private EntityGraphDiscoveryContext context;
+    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 {
-
-        if ( context == null) {
+        if (context == null) {
             throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "Entity resolver not initialized");
         }
 
         List<AtlasObjectId> resolvedReferences = new ArrayList<>();
 
-        for (AtlasObjectId typeIdPair : context.getUnresolvedIdReferences()) {
-            if ( AtlasEntity.isAssigned(typeIdPair.getGuid())) {
+        for (AtlasObjectId objId : context.getUnresolvedIds()) {
+            if (objId.isAssignedGuid()) {
                 //validate in graph repo that given guid, typename exists
-                Optional<AtlasVertex> vertex = resolveGuid(typeIdPair);
+                Optional<AtlasVertex> vertex = resolveGuid(objId);
 
-                if ( vertex.isPresent() ) {
-                    context.addRepositoryResolvedReference(typeIdPair, vertex.get());
-                    resolvedReferences.add(typeIdPair);
+                if (vertex.isPresent()) {
+                    context.addResolvedId(objId, vertex.get());
+                    resolvedReferences.add(objId);
                 }
             } else {
                 //check if root references have this temporary id
-               if (!idToEntityMap.containsKey(typeIdPair.getGuid()) ) {
-                   throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, "Could not find an entity with the specified id " + typeIdPair + " in the request");
+               if (!idToEntityMap.containsKey(objId.getGuid()) ) {
+                   throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, objId.toString());
                }
+                resolvedReferences.add(objId);
             }
+
         }
 
-        context.removeUnResolvedIdReferences(resolvedReferences);
+        context.removeUnResolvedIds(resolvedReferences);
 
         //Resolve root references
         for (AtlasEntity entity : context.getRootEntities()) {
-            if ( !context.isResolved(entity.getGuid()) && AtlasEntity.isAssigned(entity.getGuid())) {
-                AtlasObjectId typeIdPair = new AtlasObjectId(entity.getTypeName(), entity.getGuid());
-                Optional<AtlasVertex> vertex = resolveGuid(typeIdPair);
+            AtlasObjectId objId = entity.getAtlasObjectId();
+
+            if (!context.isResolvedId(objId) && AtlasEntity.isAssigned(entity.getGuid())) {
+                Optional<AtlasVertex> vertex = resolveGuid(objId);
+
                 if (vertex.isPresent()) {
-                    context.addRepositoryResolvedReference(typeIdPair, vertex.get());
-                    context.removeUnResolvedIdReference(typeIdPair);
+                    context.addResolvedId(objId, vertex.get());
+                    context.removeUnResolvedId(objId);
                 }
             }
         }
+
         return context;
     }
 
-    private Optional<AtlasVertex> resolveGuid(AtlasObjectId typeIdPair) throws AtlasBaseException {
+    private Optional<AtlasVertex> resolveGuid(AtlasObjectId objId) throws AtlasBaseException {
         //validate in graph repo that given guid, typename exists
         AtlasVertex vertex = null;
         try {
-            vertex = graphHelper.findVertex(Constants.GUID_PROPERTY_KEY, typeIdPair.getGuid(),
-                Constants.TYPE_NAME_PROPERTY_KEY, typeIdPair.getTypeName(),
+            vertex = graphHelper.findVertex(Constants.GUID_PROPERTY_KEY, objId.getGuid(),
+                Constants.TYPE_NAME_PROPERTY_KEY, objId.getTypeName(),
                 Constants.STATE_PROPERTY_KEY, Id.EntityState.ACTIVE.name());
         } catch (EntityNotFoundException e) {
             //Ignore
@@ -105,7 +108,7 @@ public class IDBasedEntityResolver implements EntityResolver {
         if ( vertex != null ) {
             return Optional.of(vertex);
         } else {
-            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, typeIdPair.getGuid());
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, objId.getGuid());
         }
     }
 

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/InMemoryMapEntityStream.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 org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasObjectId;
+
+import java.util.HashMap;
+import java.util.Iterator;
+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);
+        }
+
+        this.iterator = entities.entrySet().iterator();
+    }
+
+    @Override
+    public boolean hasNext() {
+        return iterator.hasNext();
+    }
+
+    @Override
+    public AtlasEntity next() {
+        return iterator.hasNext() ? iterator.next().getValue() : null;
+    }
+
+    @Override
+    public void reset() {
+        iterator = entities.entrySet().iterator();
+    }
+
+    @Override
+    public AtlasEntity getById(final AtlasObjectId id) {
+        return entities.get(id);
+    }
+}

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

@@ -23,7 +23,6 @@ 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.AtlasObjectId;
-import org.apache.atlas.model.typedef.AtlasStructDef;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
@@ -33,21 +32,20 @@ import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasStructType;
 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 AtlasTypeRegistry typeRegistry;
-
-    private final GraphHelper graphHelper = GraphHelper.getInstance();
-
+    private final GraphHelper           graphHelper = GraphHelper.getInstance();
+    private final AtlasTypeRegistry     typeRegistry;
     private EntityGraphDiscoveryContext context;
 
     @Inject
@@ -56,41 +54,41 @@ public class UniqAttrBasedEntityResolver implements EntityResolver {
     }
 
     @Override
-    public void init(EntityGraphDiscoveryContext entities) throws AtlasBaseException {
-        this.context = entities;
+    public void init(EntityGraphDiscoveryContext context) throws AtlasBaseException {
+        this.context = context;
     }
 
     @Override
     public EntityGraphDiscoveryContext resolveEntityReferences() throws AtlasBaseException {
-
-        if ( context == null) {
+        if (context == null) {
             throw new AtlasBaseException(AtlasErrorCode.INTERNAL_ERROR, "Unique attribute based entity resolver not initialized");
         }
 
         //Resolve attribute references
-        List<AtlasEntity> resolvedReferences = new ArrayList<>();
+        List<AtlasObjectId> resolvedReferences = new ArrayList<>();
 
-        for (AtlasEntity entity : context.getUnResolvedEntityReferences()) {
+        for (AtlasObjectId entityId : context.getUnresolvedIdsByUniqAttribs()) {
             //query in graph repo that given unique attribute - check for deleted also?
-            Optional<AtlasVertex> vertex = resolveByUniqueAttribute(entity);
+            Optional<AtlasVertex> vertex = resolveByUniqueAttribute(entityId);
+
             if (vertex.isPresent()) {
-                context.addRepositoryResolvedReference(new AtlasObjectId(entity.getTypeName(), entity.getGuid()), vertex.get());
-                resolvedReferences.add(entity);
+                context.addResolvedId(entityId, vertex.get());
+                resolvedReferences.add(entityId);
             }
         }
 
-        context.removeUnResolvedEntityReferences(resolvedReferences);
-
-        if (context.getUnResolvedEntityReferences().size() > 0) {
-            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, context.getUnResolvedEntityReferences().toString());
-        }
+        context.removeUnresolvedIdsByUniqAttribs(resolvedReferences);
 
         //Resolve root references
         for (AtlasEntity entity : context.getRootEntities()) {
-            if ( !context.isResolved(entity.getGuid()) ) {
+            AtlasObjectId entityId = entity.getAtlasObjectId();
+
+            if (!context.isResolvedId(entityId) ) {
                 Optional<AtlasVertex> vertex = resolveByUniqueAttribute(entity);
+
                 if (vertex.isPresent()) {
-                    context.addRepositoryResolvedReference(new AtlasObjectId(entity.getTypeName(), entity.getGuid()), vertex.get());
+                    context.addResolvedId(entityId, vertex.get());
+                    context.removeUnResolvedId(entityId);
                 }
             }
         }
@@ -108,23 +106,57 @@ public class UniqAttrBasedEntityResolver implements EntityResolver {
         for (AtlasStructType.AtlasAttribute attr : entityType.getAllAttributes().values()) {
             if (attr.getAttributeDef().getIsUnique()) {
                 Object attrVal = entity.getAttribute(attr.getName());
-                if (attrVal != null) {
-                    String qualifiedAttrName = attr.getQualifiedAttributeName();
-                    AtlasVertex vertex = null;
-                    try {
-                        vertex = graphHelper.findVertex(qualifiedAttrName, attrVal,
-                            Constants.ENTITY_TYPE_PROPERTY_KEY, entityType.getTypeName(),
-                            Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE
-                                .name());
-
-                        if (LOG.isDebugEnabled()) {
-                            LOG.debug("Found vertex by unique attribute : " + qualifiedAttrName + "=" + attrVal);
-                        }
-                        if (vertex != null) {
-                            return Optional.of(vertex);
-                        }
-                    } catch (EntityNotFoundException e) {
-                        //Ignore if not found
+
+                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 {
+        AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entityId.getTypeName());
+
+        if (entityType == null) {
+            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entityId.getTypeName());
+        }
+
+        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 (attrVal == null) {
+                        continue;
+                    }
+
+                    Optional<AtlasVertex> vertex = findByTypeAndQualifiedName(entityId.getTypeName(), attr.getQualifiedAttributeName(), attrVal);
+
+                    if (!vertex.isPresent()) {
+                        vertex = findBySuperTypeAndQualifiedName(entityId.getTypeName(), attr.getQualifiedAttributeName(), attrVal);
+                    }
+
+                    if (vertex.isPresent()) {
+                        return vertex;
                     }
                 }
             }
@@ -132,6 +164,46 @@ public class UniqAttrBasedEntityResolver implements EntityResolver {
         return Optional.absent();
     }
 
+    Optional<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());
+
+            if (LOG.isDebugEnabled()) {
+                LOG.debug("Found vertex by unique attribute and type {} {} ", qualifiedAttrName + "=" + attrVal, typeName);
+            }
+            if (vertex != null) {
+                return Optional.of(vertex);
+            }
+        } catch (EntityNotFoundException e) {
+            //Ignore if not found
+        }
+        return Optional.absent();
+    }
+
+    Optional<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());
+
+            if (LOG.isDebugEnabled()) {
+                LOG.debug("Found vertex by unique attribute and supertype {} ", qualifiedAttrName + "=" + attrVal, typeName);
+            }
+            if (vertex != null) {
+                return Optional.of(vertex);
+            }
+        } catch (EntityNotFoundException e) {
+            //Ignore if not found
+        }
+        return Optional.absent();
+    }
+
     @Override
     public void cleanUp() {
         //Nothing to cleanup

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

@@ -38,7 +38,7 @@ public class AtlasArrayFormatConverter extends AtlasAbstractFormatConverter {
     }
 
     @Override
-    public Collection fromV1ToV2(Object v1Obj, AtlasType type) throws AtlasBaseException {
+    public Collection fromV1ToV2(Object v1Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         Collection ret = null;
 
         if (v1Obj != null) {
@@ -57,7 +57,7 @@ public class AtlasArrayFormatConverter extends AtlasAbstractFormatConverter {
             Collection           v1List        = (Collection) v1Obj;
 
             for (Object v1Elem : v1List) {
-                Object convertedVal = elemConverter.fromV1ToV2(v1Elem, elemType);
+                Object convertedVal = elemConverter.fromV1ToV2(v1Elem, elemType, ctx);
 
                 ret.add(convertedVal);
             }
@@ -67,7 +67,7 @@ public class AtlasArrayFormatConverter extends AtlasAbstractFormatConverter {
     }
 
     @Override
-    public Collection fromV2ToV1(Object v2Obj, AtlasType type) throws AtlasBaseException {
+    public Collection fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         Collection ret = null;
 
         if (v2Obj != null) {
@@ -86,7 +86,7 @@ public class AtlasArrayFormatConverter extends AtlasAbstractFormatConverter {
             Collection           v2List        = (Collection) v2Obj;
 
             for (Object v2Elem : v2List) {
-                Object convertedVal = elemConverter.fromV2ToV1(v2Elem, elemType);
+                Object convertedVal = elemConverter.fromV2ToV1(v2Elem, elemType, ctx);
 
                 ret.add(convertedVal);
             }

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

@@ -40,7 +40,7 @@ public class AtlasClassificationFormatConverter extends AtlasStructFormatConvert
     }
 
     @Override
-    public AtlasClassification fromV1ToV2(Object v1Obj, AtlasType type) throws AtlasBaseException {
+    public AtlasClassification fromV1ToV2(Object v1Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         AtlasClassification ret = null;
 
         if (v1Obj != null) {
@@ -51,7 +51,7 @@ public class AtlasClassificationFormatConverter extends AtlasStructFormatConvert
                 final Map v1Attribs = (Map) v1Map.get(ATTRIBUTES_PROPERTY_KEY);
 
                 if (MapUtils.isNotEmpty(v1Attribs)) {
-                    ret = new AtlasClassification(type.getTypeName(), fromV1ToV2(classificationType, v1Attribs));
+                    ret = new AtlasClassification(type.getTypeName(), fromV1ToV2(classificationType, v1Attribs, ctx));
                 } else {
                     ret = new AtlasClassification(type.getTypeName());
                 }
@@ -65,7 +65,7 @@ public class AtlasClassificationFormatConverter extends AtlasStructFormatConvert
                     LOG.error("IStruct.getValuesMap() failed", excp);
                 }
 
-                ret = new AtlasClassification(type.getTypeName(), fromV1ToV2(classificationType, v1Attribs));
+                ret = new AtlasClassification(type.getTypeName(), fromV1ToV2(classificationType, v1Attribs, ctx));
             } else {
                 throw new AtlasBaseException(AtlasErrorCode.UNEXPECTED_TYPE, "Map or IStruct",
                                              v1Obj.getClass().getCanonicalName());

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

@@ -53,7 +53,7 @@ public class AtlasEntityFormatConverter extends AtlasStructFormatConverter {
     }
 
     @Override
-    public Object fromV1ToV2(Object v1Obj, AtlasType type) throws AtlasBaseException {
+    public Object fromV1ToV2(Object v1Obj, AtlasType type, ConverterContext context) throws AtlasBaseException {
         Object ret = null;
 
         if (v1Obj != null) {
@@ -67,13 +67,15 @@ public class AtlasEntityFormatConverter extends AtlasStructFormatConverter {
                 IReferenceableInstance entity    = (IReferenceableInstance) v1Obj;
                 Map<String, Object>    v1Attribs = null;
 
+                ret = new AtlasObjectId(entity.getTypeName(), entity.getId()._getId());
+
                 try {
                     v1Attribs = entity.getValuesMap();
                 } catch (AtlasException excp) {
                     LOG.error("IReferenceableInstance.getValuesMap() failed", excp);
                 }
 
-                AtlasEntityWithAssociations ret1 =  new AtlasEntityWithAssociations(entity.getTypeName(), super.fromV1ToV2(entityType, v1Attribs));
+                AtlasEntityWithAssociations ret1 =  new AtlasEntityWithAssociations(entity.getTypeName(), super.fromV1ToV2(entityType, v1Attribs, context));
                 ret1.setGuid(entity.getId()._getId());
                 ret1.setStatus(convertState(entity.getId().getState()));
                 AtlasSystemAttributes systemAttributes = entity.getSystemAttributes();
@@ -90,16 +92,17 @@ public class AtlasEntityFormatConverter extends AtlasStructFormatConverter {
                     for (String traitName : entity.getTraits()) {
                         IStruct             trait          = entity.getTrait(traitName);
                         AtlasType           classifiType   = typeRegistry.getType(traitName);
-                        AtlasClassification classification = (AtlasClassification) traitConverter.fromV1ToV2(trait, classifiType);
+                        AtlasClassification classification = (AtlasClassification) traitConverter.fromV1ToV2(trait, classifiType, context);
 
                         classifications.add(classification);
                     }
 
                     ret1.setClassifications(classifications);
                 }
-                ret = ret1;
+
+                context.addEntity(ret1);
             } else {
-                throw new AtlasBaseException(AtlasErrorCode.UNEXPECTED_TYPE, "Id or IReferenceableInstance",
+                throw new AtlasBaseException(AtlasErrorCode.UNEXPECTED_TYPE, "IReferenceableInstance",
                                              v1Obj.getClass().getCanonicalName());
             }
         }
@@ -116,7 +119,7 @@ public class AtlasEntityFormatConverter extends AtlasStructFormatConverter {
     }
 
     @Override
-    public Object fromV2ToV1(Object v2Obj, AtlasType type) throws AtlasBaseException {
+    public Object fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext context) throws AtlasBaseException {
         Object ret = null;
 
         if (v2Obj != null) {
@@ -136,15 +139,21 @@ public class AtlasEntityFormatConverter extends AtlasStructFormatConverter {
                 if (MapUtils.isEmpty(v2Attribs)) {
                     ret = new Id(idStr, 0, typeName);
                 } else {
-                    ret = new Referenceable(idStr, typeName, super.fromV2ToV1(entityType, v2Attribs));
+                    ret = new Referenceable(idStr, typeName, super.fromV2ToV1(entityType, v2Attribs, context));
                 }
             } else if (v2Obj instanceof AtlasEntity) {
                 AtlasEntity entity = (AtlasEntity) v2Obj;
 
                 ret = new Referenceable(entity.getGuid(), entity.getTypeName(),
-                                        fromV2ToV1(entityType, entity.getAttributes()));
+                                        fromV2ToV1(entityType, entity.getAttributes(), context));
+
             } else if (v2Obj instanceof AtlasObjectId) { // transient-id
-                ret = new Referenceable(((AtlasObjectId) v2Obj).getGuid(), type.getTypeName(), null);
+                AtlasEntity entity = context.getById(((AtlasObjectId) v2Obj).getGuid());
+                if ( entity == null) {
+                    throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "Could not find entity ",
+                        v2Obj.toString());
+                }
+                ret = this.fromV2ToV1(entity, typeRegistry.getType(((AtlasObjectId) v2Obj).getTypeName()), context);
             } else {
                 throw new AtlasBaseException(AtlasErrorCode.UNEXPECTED_TYPE, "Map or AtlasEntity or String",
                                              v2Obj.getClass().getCanonicalName());

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

@@ -30,12 +30,12 @@ public class AtlasEnumFormatConverter extends AtlasAbstractFormatConverter {
     }
 
     @Override
-    public Object fromV1ToV2(Object v1Obj, AtlasType type) throws AtlasBaseException {
+    public Object fromV1ToV2(Object v1Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         return type.getNormalizedValue(v1Obj);
     }
 
     @Override
-    public Object fromV2ToV1(Object v2Obj, AtlasType type) throws AtlasBaseException {
+    public Object fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         return type.getNormalizedValue(v2Obj);
     }
 }

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

@@ -20,12 +20,62 @@ package org.apache.atlas.web.adapters;
 
 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.AtlasEntityWithAssociations;
 import org.apache.atlas.type.AtlasType;
 
+import java.util.HashMap;
+import java.util.Map;
+
 public interface AtlasFormatConverter {
-    Object fromV1ToV2(Object v1Obj, AtlasType type) throws AtlasBaseException;
+    Object fromV1ToV2(Object v1Obj, AtlasType type, ConverterContext context) throws AtlasBaseException;
 
-    Object fromV2ToV1(Object v2Obj, AtlasType type) throws AtlasBaseException;
+    Object fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext context) throws AtlasBaseException;
 
     TypeCategory getTypeCategory();
+
+    public static class ConverterContext {
+
+        private Map<String, AtlasEntityWithAssociations> entities = null;
+
+        public void addEntity(AtlasEntityWithAssociations entity) {
+            if (entities == null) {
+                entities = new HashMap<>();
+            }
+            entities.put(entity.getGuid(), entity);
+        }
+
+        public void addEntity(AtlasEntity entity) {
+            if (entities == null) {
+                entities = new HashMap<>();
+            }
+            entities.put(entity.getGuid(), new AtlasEntityWithAssociations(entity));
+        }
+
+        public boolean exists(AtlasEntityWithAssociations entity) {
+            return entities != null ? entities.containsKey(entity.getGuid()) : false;
+        }
+
+        public AtlasEntity getById(String guid) {
+
+            if( entities != null) {
+                return entities.get(guid);
+            }
+
+            return null;
+        }
+
+        public Map<String, AtlasEntityWithAssociations> getEntities() {
+            return entities;
+        }
+
+        public void addEntities(Map<String, AtlasEntity> entities) {
+            if (this.entities == null) {
+                this.entities = new HashMap<>(entities.size());
+            }
+            for (String entityId : entities.keySet()) {
+                this.entities.put(entityId, new AtlasEntityWithAssociations(entities.get(entityId)));
+            }
+        }
+    }
 }

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

@@ -47,7 +47,11 @@ import org.apache.atlas.typesystem.exception.TypeNotFoundException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 
 @Singleton
 public class AtlasInstanceRestAdapters {
@@ -63,19 +67,24 @@ public class AtlasInstanceRestAdapters {
     @Inject
     private MetadataService metadataService;
 
-    public ITypedReferenceableInstance[] getITypedReferenceables(List<AtlasEntity> entities) throws AtlasBaseException {
+    public ITypedReferenceableInstance[] getITypedReferenceables(Collection<AtlasEntity> entities) throws AtlasBaseException {
         ITypedReferenceableInstance[] entitiesInOldFormat = new ITypedReferenceableInstance[entities.size()];
 
+        AtlasFormatConverter.ConverterContext ctx = new AtlasFormatConverter.ConverterContext();
+        for(Iterator<AtlasEntity> i = entities.iterator(); i.hasNext(); ) {
+            ctx.addEntity(i.next());
+        }
+
+        Iterator<AtlasEntity> entityIterator = entities.iterator();
         for (int i = 0; i < entities.size(); i++) {
-            ITypedReferenceableInstance typedInstance = getITypedReferenceable(entities.get(i));
+            ITypedReferenceableInstance typedInstance = getITypedReferenceable(entityIterator.next(), ctx);
             entitiesInOldFormat[i] = typedInstance;
         }
-
         return entitiesInOldFormat;
     }
 
-    public ITypedReferenceableInstance getITypedReferenceable(AtlasEntity entity) throws AtlasBaseException {
-        Referenceable ref = getReferenceable(entity);
+    public ITypedReferenceableInstance getITypedReferenceable(AtlasEntity entity, AtlasFormatConverter.ConverterContext ctx) throws AtlasBaseException {
+        Referenceable ref = getReferenceable(entity, ctx);
 
         try {
             return metadataService.getTypedReferenceableInstance(ref);
@@ -85,10 +94,10 @@ public class AtlasInstanceRestAdapters {
         }
     }
 
-    public Referenceable getReferenceable(AtlasEntity entity) throws AtlasBaseException {
+    public Referenceable getReferenceable(AtlasEntity entity, final AtlasFormatConverter.ConverterContext ctx) throws AtlasBaseException {
         AtlasFormatConverter converter  = instanceFormatters.getConverter(TypeCategory.ENTITY);
         AtlasType            entityType = typeRegistry.getType(entity.getTypeName());
-        Referenceable        ref        = (Referenceable)converter.fromV2ToV1(entity, entityType);
+        Referenceable        ref        = (Referenceable)converter.fromV2ToV1(entity, entityType, ctx);
 
         return ref;
     }
@@ -96,7 +105,7 @@ public class AtlasInstanceRestAdapters {
     public ITypedStruct getTrait(AtlasClassification classification) throws AtlasBaseException {
         AtlasFormatConverter converter          = instanceFormatters.getConverter(TypeCategory.CLASSIFICATION);
         AtlasType            classificationType = typeRegistry.getType(classification.getTypeName());
-        Struct               trait               = (Struct)converter.fromV2ToV1(classification, classificationType);
+        Struct               trait               = (Struct)converter.fromV2ToV1(classification, classificationType, new AtlasFormatConverter.ConverterContext());
 
         try {
             return metadataService.createTraitInstance(trait);
@@ -112,21 +121,22 @@ public class AtlasInstanceRestAdapters {
         if (classificationType == null) {
             throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.CLASSIFICATION.name(), classification.getTypeName());
         }
-        AtlasClassification  ret                = (AtlasClassification)converter.fromV1ToV2(classification, classificationType);
+        AtlasClassification  ret                = (AtlasClassification)converter.fromV1ToV2(classification, classificationType, new AtlasFormatConverter.ConverterContext());
 
         return ret;
     }
 
-    public AtlasEntityWithAssociations getAtlasEntity(IReferenceableInstance referenceable) throws AtlasBaseException {
+    public Map<String, AtlasEntityWithAssociations> getAtlasEntity(IReferenceableInstance referenceable) throws AtlasBaseException {
+
         AtlasFormatConverter converter  = instanceFormatters.getConverter(TypeCategory.ENTITY);
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName(referenceable.getTypeName());
         if (entityType == null) {
             throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), referenceable.getTypeName());
         }
 
-        AtlasEntityWithAssociations ret = (AtlasEntityWithAssociations)converter.fromV1ToV2(referenceable, entityType);
-
-        return ret;
+        AtlasFormatConverter.ConverterContext ctx = new AtlasFormatConverter.ConverterContext();
+        converter.fromV1ToV2(referenceable, entityType, ctx);
+        return ctx.getEntities();
     }
 
 

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

@@ -35,7 +35,7 @@ public class AtlasMapFormatConverter extends AtlasAbstractFormatConverter {
     }
 
     @Override
-    public Map fromV1ToV2(Object v1Obj, AtlasType type) throws AtlasBaseException {
+    public Map fromV1ToV2(Object v1Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         Map ret = null;
 
         if (v1Obj != null) {
@@ -52,8 +52,8 @@ public class AtlasMapFormatConverter extends AtlasAbstractFormatConverter {
                 for (Object key : v1Map.keySet()) {
                     Object value = v1Map.get(key);
 
-                    Object v2Key   = keyConverter.fromV1ToV2(key, keyType);
-                    Object v2Value = valueConverter.fromV1ToV2(value, valueType);
+                    Object v2Key   = keyConverter.fromV1ToV2(key, keyType, ctx);
+                    Object v2Value = valueConverter.fromV1ToV2(value, valueType, ctx);
 
                     ret.put(v2Key, v2Value);
                 }
@@ -67,7 +67,7 @@ public class AtlasMapFormatConverter extends AtlasAbstractFormatConverter {
     }
 
     @Override
-    public Map fromV2ToV1(Object v2Obj, AtlasType type) throws AtlasBaseException {
+    public Map fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         Map ret = null;
 
         if (v2Obj != null) {
@@ -84,8 +84,8 @@ public class AtlasMapFormatConverter extends AtlasAbstractFormatConverter {
                 for (Object key : v1Map.keySet()) {
                     Object value = v1Map.get(key);
 
-                    Object v2Key   = keyConverter.fromV2ToV1(key, keyType);
-                    Object v2Value = valueConverter.fromV2ToV1(value, valueType);
+                    Object v2Key   = keyConverter.fromV2ToV1(key, keyType, ctx);
+                    Object v2Value = valueConverter.fromV2ToV1(value, valueType, ctx);
 
                     ret.put(v2Key, v2Value);
                 }

webapp/src/main/java/org/apache/atlas/web/adapters/AtlasObjectIdConverter.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.web.adapters;
+
+
+import org.apache.atlas.AtlasErrorCode;
+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.AtlasObjectId;
+import org.apache.atlas.type.AtlasType;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.atlas.typesystem.IReferenceableInstance;
+import org.apache.atlas.typesystem.Referenceable;
+import org.apache.atlas.typesystem.persistence.Id;
+import org.apache.commons.collections.MapUtils;
+import org.apache.commons.lang3.StringUtils;
+
+import java.util.Map;
+
+public class
+AtlasObjectIdConverter extends  AtlasAbstractFormatConverter {
+
+    public AtlasObjectIdConverter(AtlasFormatConverters registry, AtlasTypeRegistry typeRegistry) {
+        this(registry, typeRegistry, TypeCategory.OBJECT_ID_TYPE);
+    }
+
+    protected AtlasObjectIdConverter(AtlasFormatConverters registry, AtlasTypeRegistry typeRegistry, TypeCategory typeCategory) {
+        super(registry, typeRegistry, typeCategory);
+    }
+
+    @Override
+    public Object fromV1ToV2(Object v1Obj, AtlasType type, AtlasFormatConverter.ConverterContext converterContext) throws AtlasBaseException {
+        Object ret = null;
+
+        if (v1Obj != null) {
+            if (v1Obj instanceof Id) {
+                Id id = (Id) v1Obj;
+                ret = new AtlasObjectId(id.getTypeName(), id._getId());
+            } else if (v1Obj instanceof IReferenceableInstance) {
+                IReferenceableInstance entity = (IReferenceableInstance) v1Obj;
+                ret = new AtlasObjectId(entity.getTypeName(), entity.getId()._getId());
+            }
+        }
+        return ret;
+    }
+
+    @Override
+    public Object fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext converterContext) throws AtlasBaseException {
+        Id ret = null;
+
+        if (v2Obj != null) {
+
+            if (v2Obj instanceof Map) {
+                Map    v2Map    = (Map) v2Obj;
+                String idStr    = (String)v2Map.get(AtlasObjectId.KEY_GUID);
+                String typeName = type.getTypeName();
+
+                if (StringUtils.isEmpty(idStr)) {
+                    throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND);
+                }
+
+                ret = new Id(idStr, 0, typeName);
+            } else if (v2Obj instanceof AtlasObjectId) { // transient-id
+                ret = new Id(((AtlasObjectId) v2Obj).getGuid(), 0, type.getTypeName());
+            } else if (v2Obj instanceof AtlasEntity) {
+                AtlasEntity entity = (AtlasEntity) v2Obj;
+                ret = new Id(((AtlasObjectId) v2Obj).getGuid(), 0, type.getTypeName());
+            } else {
+                throw new AtlasBaseException(AtlasErrorCode.TYPE_CATEGORY_INVALID, type.getTypeCategory().name());
+            }
+        }
+        return ret;
+    }
+}

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

@@ -30,12 +30,12 @@ public class AtlasPrimitiveFormatConverter extends AtlasAbstractFormatConverter
     }
 
     @Override
-    public Object fromV1ToV2(Object v1Obj, AtlasType type) throws AtlasBaseException {
+    public Object fromV1ToV2(Object v1Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         return type.getNormalizedValue(v1Obj);
     }
 
     @Override
-    public Object fromV2ToV1(Object v2Obj, AtlasType type) throws AtlasBaseException {
+    public Object fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext ctx) throws AtlasBaseException {
         return type.getNormalizedValue(v2Obj);
     }
 }

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

@@ -50,7 +50,7 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
     }
 
     @Override
-    public Object fromV1ToV2(Object v1Obj, AtlasType type) throws AtlasBaseException {
+    public Object fromV1ToV2(Object v1Obj, AtlasType type, ConverterContext converterContext) throws AtlasBaseException {
         AtlasStruct ret = null;
 
         if (v1Obj != null) {
@@ -61,7 +61,7 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
                 final Map v1Attribs = (Map) v1Map.get(ATTRIBUTES_PROPERTY_KEY);
 
                 if (MapUtils.isNotEmpty(v1Attribs)) {
-                    ret = new AtlasStruct(type.getTypeName(), fromV1ToV2(structType, v1Attribs));
+                    ret = new AtlasStruct(type.getTypeName(), fromV1ToV2(structType, v1Attribs, converterContext));
                 } else {
                     ret = new AtlasStruct(type.getTypeName());
                 }
@@ -75,7 +75,7 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
                     LOG.error("IStruct.getValuesMap() failed", excp);
                 }
 
-                ret = new AtlasStruct(type.getTypeName(), fromV1ToV2(structType, v1Attribs));
+                ret = new AtlasStruct(type.getTypeName(), fromV1ToV2(structType, v1Attribs, converterContext));
             } else {
                 throw new AtlasBaseException(AtlasErrorCode.UNEXPECTED_TYPE, "Map or IStruct", v1Obj.getClass().getCanonicalName());
             }
@@ -85,7 +85,7 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
     }
 
     @Override
-    public Object fromV2ToV1(Object v2Obj, AtlasType type) throws AtlasBaseException {
+    public Object fromV2ToV1(Object v2Obj, AtlasType type, ConverterContext converterContext) throws AtlasBaseException {
         Struct ret = null;
 
         if (v2Obj != null) {
@@ -102,14 +102,14 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
                 }
 
                 if (MapUtils.isNotEmpty(v2Attribs)) {
-                    ret = new Struct(type.getTypeName(), fromV2ToV1(structType, v2Attribs));
+                    ret = new Struct(type.getTypeName(), fromV2ToV1(structType, v2Attribs, converterContext));
                 } else {
                     ret = new Struct(type.getTypeName());
                 }
             } else if (v2Obj instanceof AtlasStruct) {
                 AtlasStruct struct = (AtlasStruct) v2Obj;
 
-                ret = new Struct(type.getTypeName(), fromV2ToV1(structType, struct.getAttributes()));
+                ret = new Struct(type.getTypeName(), fromV2ToV1(structType, struct.getAttributes(), converterContext));
             } else {
                 throw new AtlasBaseException(AtlasErrorCode.UNEXPECTED_TYPE, "Map or AtlasStruct", v2Obj.getClass().getCanonicalName());
             }
@@ -118,7 +118,7 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
         return ret;
     }
 
-    protected Map<String, Object> fromV2ToV1(AtlasStructType structType, Map attributes) throws AtlasBaseException {
+    protected Map<String, Object> fromV2ToV1(AtlasStructType structType, Map attributes, ConverterContext context) throws AtlasBaseException {
         Map<String, Object> ret = null;
 
         if (MapUtils.isNotEmpty(attributes)) {
@@ -132,11 +132,17 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
                     continue;
                 }
 
-                AtlasFormatConverter attrConverter = converterRegistry.getConverter(attrType.getTypeCategory());
-
                 Object v2Value = attributes.get(attr.getName());
-                Object v1Value = attrConverter.fromV2ToV1(v2Value, attrType);
+                Object v1Value = null;
 
+                AtlasFormatConverter attrConverter = null;
+                if (attrType.getTypeCategory() == TypeCategory.ENTITY && !attr.isContainedAttribute()) {
+                    attrConverter = new AtlasObjectIdConverter(converterRegistry, typeRegistry);
+                    v1Value = attrConverter.fromV2ToV1(v2Value, attrType, context);
+                } else {
+                    attrConverter = converterRegistry.getConverter(attrType.getTypeCategory());
+                    v1Value = attrConverter.fromV2ToV1(v2Value, attrType, context);
+                }
                 ret.put(attr.getName(), v1Value);
             }
         }
@@ -144,7 +150,7 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
         return ret;
     }
 
-    protected Map<String, Object> fromV1ToV2(AtlasStructType structType, Map attributes) throws AtlasBaseException {
+    protected Map<String, Object> fromV1ToV2(AtlasStructType structType, Map attributes, ConverterContext context) throws AtlasBaseException {
         Map<String, Object> ret = null;
 
         if (MapUtils.isNotEmpty(attributes)) {
@@ -152,11 +158,23 @@ public class AtlasStructFormatConverter extends AtlasAbstractFormatConverter {
 
             for (AtlasStructType.AtlasAttribute attr : structType.getAllAttributes().values()) {
                 AtlasType            attrType      = attr.getAttributeType();
-                AtlasFormatConverter attrConverter = converterRegistry.getConverter(attrType.getTypeCategory());
+
+                if (attrType == null) {
+                    LOG.warn("ignored attribute {}.{}: failed to find AtlasType", structType.getTypeName(), attr.getName());
+                    continue;
+                }
 
                 Object v1Value = attributes.get(attr.getName());
-                Object v2Value = attrConverter.fromV1ToV2(v1Value, attrType);
+                Object v2Value = null;
 
+                AtlasFormatConverter attrConverter = null;
+                if (attrType.getTypeCategory() == TypeCategory.ENTITY && !attr.isContainedAttribute()) {
+                    attrConverter = new AtlasObjectIdConverter(converterRegistry, typeRegistry);
+                    v2Value = attrConverter.fromV1ToV2(v1Value, attrType, context);
+                } else {
+                    attrConverter = converterRegistry.getConverter(attrType.getTypeCategory());
+                    v2Value = attrConverter.fromV1ToV2(v1Value, attrType, context);
+                }
                 ret.put(attr.getAttributeDef().getName(), v2Value);
             }
         }

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

@@ -25,6 +25,7 @@ import org.apache.atlas.model.SearchFilter;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
+import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
 import org.apache.atlas.model.instance.ClassificationAssociateRequest;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
@@ -53,6 +54,7 @@ import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
@@ -92,9 +94,10 @@ public class EntitiesREST {
     @POST
     @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    public EntityMutationResponse createOrUpdate(List<AtlasEntity> entities) throws AtlasBaseException {
+    public EntityMutationResponse createOrUpdate(Map<String, AtlasEntity> entities) throws AtlasBaseException {
         EntityMutationResponse response = null;
-        ITypedReferenceableInstance[] entitiesInOldFormat = restAdapters.getITypedReferenceables(entities);
+
+        ITypedReferenceableInstance[] entitiesInOldFormat = restAdapters.getITypedReferenceables(entities.values());
 
         try {
             final AtlasClient.EntityResult result = metadataService.updateEntities(entitiesInOldFormat);
@@ -114,7 +117,7 @@ public class EntitiesREST {
     @PUT
     @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    public EntityMutationResponse update(List<AtlasEntity> entities) throws AtlasBaseException {
+    public EntityMutationResponse update(Map<String, AtlasEntity> entities) throws AtlasBaseException {
        return createOrUpdate(entities);
     }
 
@@ -135,8 +138,10 @@ public class EntitiesREST {
         for (String guid : guids) {
             try {
                ITypedReferenceableInstance ref = metadataService.getEntityDefinition(guid);
-               AtlasEntity entity = restAdapters.getAtlasEntity(ref);
-               entityList.add(entity);
+               Map<String, AtlasEntityWithAssociations> entityRet = restAdapters.getAtlasEntity(ref);
+
+               addToEntityList(entityList, entityRet.values());
+
             } catch (AtlasException e) {
                 throw toAtlasBaseException(e);
             }
@@ -146,6 +151,14 @@ public class EntitiesREST {
         return entities;
     }
 
+    private void addToEntityList(final List<AtlasEntity> entityList, final Collection<AtlasEntityWithAssociations> values) {
+        for (AtlasEntityWithAssociations val : values) {
+            if ( !entityList.contains(val)) {
+                entityList.add(val);
+            }
+        }
+    }
+
     /*******
      * Entity Delete
      *******/
@@ -168,28 +181,6 @@ public class EntitiesREST {
     }
 
     /**
-     * Bulk retrieval API for searching on entities by certain predefined attributes ( typeName, superType, name, qualifiedName etc) + optional user defined attributes
-     *
-     * @throws AtlasBaseException
-     */
-    @GET
-    @Produces(Servlets.JSON_MEDIA_TYPE)
-    public AtlasEntityHeader.AtlasEntityHeaders searchEntities() throws AtlasBaseException {
-        SearchFilter searchFilter = getSearchFilter();
-        AtlasEntity.AtlasEntities atlasEntities = entitiesStore.searchEntities(searchFilter);
-        AtlasEntityHeader.AtlasEntityHeaders entityHeaders = new AtlasEntityHeader.AtlasEntityHeaders();
-        entityHeaders.setList(new LinkedList<AtlasEntityHeader>());
-
-        if (atlasEntities != null) {
-            for (AtlasEntity atlasEntity : atlasEntities.getList()) {
-                entityHeaders.getList().add(new AtlasEntityHeader(atlasEntity.getTypeName(), atlasEntity.getAttributes()));
-            }
-        }
-
-        return entityHeaders;
-    }
-
-    /**
      * Bulk API to associate a tag to multiple entities
      *
      */
@@ -219,15 +210,4 @@ public class EntitiesREST {
             throw toAtlasBaseException(e);
         }
     }
-
-    private SearchFilter getSearchFilter() {
-        SearchFilter searchFilter = new SearchFilter();
-        if (null != httpServletRequest && null != httpServletRequest.getParameterMap()) {
-            for (Map.Entry<String, String[]> entry : httpServletRequest.getParameterMap().entrySet()) {
-                searchFilter.setParam(entry.getKey(), Arrays.asList(entry.getValue()));
-            }
-        }
-        return searchFilter;
-    }
-
 }

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

@@ -35,6 +35,7 @@ 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.lang3.StringUtils;
@@ -47,6 +48,7 @@ import javax.ws.rs.*;
 import javax.ws.rs.core.MediaType;
 import java.util.ArrayList;
 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;
@@ -74,45 +76,6 @@ public class EntityREST {
     }
 
     /**
-     * Create or Update an entity if it  already exists
-     *
-     * @param entity The updated entity
-     * @return
-     */
-    @POST
-    @Consumes({Servlets.JSON_MEDIA_TYPE, MediaType.APPLICATION_JSON})
-    @Produces(Servlets.JSON_MEDIA_TYPE)
-    public EntityMutationResponse createOrUpdate(final AtlasEntity entity) throws AtlasBaseException {
-        EntityMutationResponse response = null;
-        ITypedReferenceableInstance[] entitiesInOldFormat = restAdapters.getITypedReferenceables(new ArrayList<AtlasEntity>() {{ add(entity); }});
-
-        try {
-            final AtlasClient.EntityResult 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;
-    }
-
-    /**
-     * Complete Update of an entity identified by its GUID
-     *
-     * @param guid
-     * @param entity The updated entity
-     * @return
-     */
-    @PUT
-    @Path("guid/{guid}")
-    @Consumes({Servlets.JSON_MEDIA_TYPE, MediaType.APPLICATION_JSON})
-    @Produces(Servlets.JSON_MEDIA_TYPE)
-    public EntityMutationResponse updateByGuid(@PathParam("guid") String guid, AtlasEntity entity, @DefaultValue("false") @QueryParam("partialUpdate") boolean partialUpdate) throws AtlasBaseException {
-        return createOrUpdate(entity);
-    }
-
-
-    /**
      * Fetch the complete definition of an entity given its GUID.
      *
      * @param guid GUID for the entity
@@ -120,10 +83,14 @@ public class EntityREST {
     @GET
     @Path("/guid/{guid}")
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    public AtlasEntity getById(@PathParam("guid") String guid) throws AtlasBaseException {
+    public List<AtlasEntityWithAssociations> getById(@PathParam("guid") String guid) throws AtlasBaseException {
+        List<AtlasEntityWithAssociations> entityList = new ArrayList<>();
+
         try {
             ITypedReferenceableInstance ref = metadataService.getEntityDefinition(guid);
-            return restAdapters.getAtlasEntity(ref);
+            Map<String, AtlasEntityWithAssociations> entityRet = restAdapters.getAtlasEntity(ref);
+            entityList.addAll(entityRet.values());
+            return entityList;
         } catch (AtlasException e) {
             throw toAtlasBaseException(e);
         }
@@ -138,10 +105,14 @@ public class EntityREST {
     @GET
     @Path("/guid/{guid}/associations")
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    public AtlasEntityWithAssociations getWithAssociationsByGuid(@PathParam("guid") String guid) throws AtlasBaseException {
+    public List<AtlasEntityWithAssociations> getWithAssociationsByGuid(@PathParam("guid") String guid) throws AtlasBaseException {
+
+        List<AtlasEntityWithAssociations> entityList = new ArrayList<>();
         try {
             ITypedReferenceableInstance ref = metadataService.getEntityDefinition(guid);
-            return restAdapters.getAtlasEntity(ref);
+            Map<String, AtlasEntityWithAssociations> entityRet = restAdapters.getAtlasEntity(ref);
+            entityList.addAll(entityRet.values());
+            return entityList;
         } catch (AtlasException e) {
             throw toAtlasBaseException(e);
         }
@@ -190,7 +161,9 @@ public class EntityREST {
         AtlasEntityType type = (AtlasEntityType) validateType(entityType, TypeCategory.ENTITY);
         validateUniqueAttribute(type, attribute);
 
-        Referenceable ref = restAdapters.getReferenceable(entity);
+        AtlasFormatConverter.ConverterContext ctx = new AtlasFormatConverter.ConverterContext();
+        ctx.addEntity(entity);
+        Referenceable ref = restAdapters.getReferenceable(entity, ctx);
         AtlasClient.EntityResult result = metadataService.updateEntityByUniqueAttribute(entityType, attribute, value, ref);
         return toEntityMutationResponse(result);
     }
@@ -220,19 +193,23 @@ public class EntityREST {
     @Consumes({Servlets.JSON_MEDIA_TYPE, MediaType.APPLICATION_JSON})
     @Produces(Servlets.JSON_MEDIA_TYPE)
     @Path("/uniqueAttribute/type/{typeName}/attribute/{attrName}")
-    public AtlasEntity getByUniqueAttribute(@PathParam("typeName") String entityType,
+    public List<AtlasEntityWithAssociations> getByUniqueAttribute(@PathParam("typeName") String entityType,
         @PathParam("attrName") String attribute,
         @QueryParam("value") String value) throws AtlasBaseException {
 
+        List<AtlasEntityWithAssociations> entityList = new ArrayList<>();
         AtlasEntityType type = (AtlasEntityType) validateType(entityType, TypeCategory.ENTITY);
         validateUniqueAttribute(type, attribute);
 
         try {
             final ITypedReferenceableInstance entityDefinitionReference = metadataService.getEntityDefinitionReference(entityType, attribute, value);
-            return restAdapters.getAtlasEntity(entityDefinitionReference);
+            Map<String, AtlasEntityWithAssociations> entityRet = restAdapters.getAtlasEntity(entityDefinitionReference);
+            entityList.addAll(entityRet.values());
         } catch (AtlasException e) {
             throw toAtlasBaseException(e);
         }
+
+        return entityList;
     }
 
 

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

@@ -26,6 +26,7 @@ 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.AtlasEntityHeader;
+import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasStruct;
 import org.apache.atlas.model.instance.ClassificationAssociateRequest;
 import org.apache.atlas.model.instance.EntityMutationResponse;
@@ -68,6 +69,10 @@ 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;
@@ -78,9 +83,12 @@ public class TestEntitiesREST {
     public void setUp() throws Exception {
         AtlasTypesDef typesDef = TestUtilsV2.defineHiveTypes();
         typeStore.createTypesDef(typesDef);
-        dbEntity = TestUtilsV2.createDBEntity();
+        dbEntityMap = TestUtilsV2.createDBEntity();
+        dbEntity = dbEntityMap.values().iterator().next();
+
+        tableEntityMap = TestUtilsV2.createTableEntity(dbEntity.getGuid());
+        tableEntity = tableEntityMap.values().iterator().next();
 
-        tableEntity = TestUtilsV2.createTableEntity(dbEntity.getGuid());
         final AtlasEntity colEntity = TestUtilsV2.createColumnEntity(tableEntity.getGuid());
         columns = new ArrayList<AtlasEntity>() {{ add(colEntity); }};
         tableEntity.setAttribute("columns", columns);
@@ -98,9 +106,9 @@ public class TestEntitiesREST {
 
     @Test
     public void testCreateOrUpdateEntities() throws Exception {
-        List<AtlasEntity> entities = new ArrayList<AtlasEntity>();
-        entities.add(dbEntity);
-        entities.add(tableEntity);
+        Map<String, AtlasEntity> entities = new HashMap<>();
+        entities.put(dbEntity.getGuid(), dbEntity);
+        entities.put(tableEntity.getGuid(), tableEntity);
 
         EntityMutationResponse response = entitiesREST.createOrUpdate(entities);
         List<AtlasEntityHeader> guids = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);
@@ -129,9 +137,12 @@ 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
-        AtlasEntity dbEntity = TestUtilsV2.createDBEntity();
+        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 tableEntity = TestUtilsV2.createTableEntity(dbEntity.getGuid());
         final AtlasEntity colEntity = TestUtilsV2.createColumnEntity(tableEntity.getGuid());
         List<AtlasEntity> columns = new ArrayList<AtlasEntity>() {{ add(colEntity); }};
         tableEntity.setAttribute("columns", columns);
@@ -139,9 +150,9 @@ public class TestEntitiesREST {
         AtlasEntity newDBEntity = serDeserEntity(dbEntity);
         AtlasEntity newTableEntity = serDeserEntity(tableEntity);
 
-        List<AtlasEntity> newEntities = new ArrayList<AtlasEntity>();
-        newEntities.add(newDBEntity);
-        newEntities.add(newTableEntity);
+        Map<String, AtlasEntity> newEntities = new HashMap<>();
+        newEntities.put(newDBEntity.getGuid(), newDBEntity);
+        newEntities.put(newTableEntity.getGuid(), newTableEntity);
         EntityMutationResponse response2 = entitiesREST.createOrUpdate(newEntities);
 
         List<AtlasEntityHeader> newGuids = response2.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);
@@ -211,7 +222,7 @@ public class TestEntitiesREST {
             //date
             Assert.assertEquals(tableEntity.getAttribute("created"), retrievedTableEntity.getAttribute("created"));
             //array of Ids
-            Assert.assertEquals(((List<AtlasEntity>) retrievedTableEntity.getAttribute("columns")).get(0).getGuid(), retrievedColumnEntity.getGuid());
+            Assert.assertEquals(((List<AtlasObjectId>) retrievedTableEntity.getAttribute("columns")).get(0).getGuid(), retrievedColumnEntity.getGuid());
             //array of structs
             Assert.assertEquals(((List<AtlasStruct>) retrievedTableEntity.getAttribute("partitions")), tableEntity.getAttribute("partitions"));
         }

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

@@ -29,6 +29,7 @@ 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.testng.Assert;
 import org.testng.annotations.AfterMethod;
@@ -41,6 +42,7 @@ import javax.inject.Inject;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 @Guice(modules = {RepositoryMetadataModule.class})
 public class TestEntityREST {
@@ -51,6 +53,9 @@ public class TestEntityREST {
     @Inject
     private EntityREST entityREST;
 
+    @Inject
+    private EntitiesREST entitiesREST;
+
     private AtlasEntity dbEntity;
 
     private String dbGuid;
@@ -61,7 +66,8 @@ public class TestEntityREST {
     public void setUp() throws Exception {
         AtlasTypesDef typesDef = TestUtilsV2.defineHiveTypes();
         typeStore.createTypesDef(typesDef);
-        dbEntity = TestUtilsV2.createDBEntity();
+        Map<String, AtlasEntity> dbEntityMap = TestUtilsV2.createDBEntity();
+        dbEntity = dbEntityMap.values().iterator().next();
     }
 
     @AfterClass
@@ -74,9 +80,10 @@ public class TestEntityREST {
         RequestContext.clear();
     }
 
-    @Test
-    public void testCreateOrUpdateEntity() throws Exception {
-        final EntityMutationResponse response = entityREST.createOrUpdate(dbEntity);
+    public void createOrUpdateEntity() throws Exception {
+        Map<String, AtlasEntity> dbEntityMap = new HashMap<>();
+        dbEntityMap.put(dbEntity.getGuid(), dbEntity);
+        final EntityMutationResponse response = entitiesREST.createOrUpdate(dbEntityMap);
 
         Assert.assertNotNull(response);
         List<AtlasEntityHeader> entitiesMutated = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);
@@ -88,18 +95,19 @@ public class TestEntityREST {
         Assert.assertEquals(entitiesMutated.size(), 1);
     }
 
-    @Test(dependsOnMethods = "testCreateOrUpdateEntity")
+    @Test
     public void testGetEntityById() throws Exception {
-
-        final AtlasEntity response = entityREST.getById(dbGuid);
+        createOrUpdateEntity();
+        final List<AtlasEntityWithAssociations> response = entityREST.getById(dbGuid);
 
         Assert.assertNotNull(response);
-        TestEntitiesREST.verifyAttributes(response.getAttributes(), dbEntity.getAttributes());
+        TestEntitiesREST.verifyAttributes(response.get(0).getAttributes(), dbEntity.getAttributes());
     }
 
-    @Test(dependsOnMethods = "testCreateOrUpdateEntity")
+    @Test
     public void  testAddAndGetClassification() throws Exception {
 
+        createOrUpdateEntity();
         List<AtlasClassification> classifications = new ArrayList<>();
         testClassification = new AtlasClassification(TestUtilsV2.CLASSIFICATION, new HashMap<String, Object>() {{ put("tag", "tagName"); }});
         classifications.add(testClassification);
@@ -122,8 +130,8 @@ public class TestEntityREST {
     @Test(dependsOnMethods = "testAddAndGetClassification")
     public void  testGetEntityWithAssociations() throws Exception {
 
-        AtlasEntityWithAssociations entity = entityREST.getWithAssociationsByGuid(dbGuid);
-        final List<AtlasClassification> retrievedClassifications = entity.getClassifications();
+        List<AtlasEntityWithAssociations> entity = entityREST.getWithAssociationsByGuid(dbGuid);
+        final List<AtlasClassification> retrievedClassifications = entity.get(0).getClassifications();
 
         Assert.assertNotNull(retrievedClassifications);
         Assert.assertEquals(new ArrayList<AtlasClassification>() {{ add(testClassification); }}, retrievedClassifications);
@@ -151,8 +159,8 @@ public class TestEntityREST {
     @Test
     public void  testUpdateGetDeleteEntityByUniqueAttribute() throws Exception {
 
-        AtlasEntity dbEntity = TestUtilsV2.createDBEntity();
-        entityREST.createOrUpdate(dbEntity);
+        Map<String, AtlasEntity> dbEntityMap = TestUtilsV2.createDBEntity();
+        entitiesREST.createOrUpdate(dbEntityMap);
 
         final String prevDBName = (String) dbEntity.getAttribute(TestUtilsV2.NAME);
         final String updatedDBName = "updatedDBName";
@@ -164,11 +172,11 @@ public class TestEntityREST {
         Assert.assertTrue(AtlasEntity.isAssigned(dbGuid));
 
         //Get By unique attribute
-        AtlasEntity entity = entityREST.getByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, TestUtilsV2.NAME, updatedDBName);
-        Assert.assertNotNull(entity);
-        Assert.assertNotNull(entity.getGuid());
-        Assert.assertEquals(entity.getGuid(), dbGuid);
-        TestEntitiesREST.verifyAttributes(entity.getAttributes(), dbEntity.getAttributes());
+        List<AtlasEntityWithAssociations> 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());
 
         final EntityMutationResponse deleteResponse = entityREST.deleteByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, TestUtilsV2.NAME, (String) dbEntity.getAttribute(TestUtilsV2.NAME));
 

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

@@ -564,7 +564,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         hiveTable.setAttribute("columns", columns);
 
         LOG.debug("Updating entity= " + tableUpdated);
-        EntityMutationResponse updateResult = entitiesClientV2.updateEntity(hiveTable.getGuid(), tableUpdated);
+        EntityMutationResponse updateResult = entitiesClientV2.updateEntity(tableUpdated);
         assertNotNull(updateResult);
         assertNotNull(updateResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
         assertTrue(updateResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).size() > 0);