ATLAS-1523: v2 implementation of get-entity by-guid/unique-attributes

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

@@ -30,8 +30,9 @@ import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasClassification.AtlasClassifications;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntities;
-import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.EntityMutationResponse;
+import org.apache.commons.collections.MapUtils;
 import org.apache.commons.configuration.Configuration;
 import org.apache.hadoop.security.UserGroupInformation;
 
@@ -42,15 +43,14 @@ import com.sun.jersey.core.util.MultivaluedMapImpl;
 
 public class AtlasEntitiesClientV2 extends AtlasBaseClient {
 
-    private static final GenericType<List<AtlasEntityWithAssociations>> ENTITY_WITH_ASSOCIATIONS_LIST_TYPE = new GenericType<List<AtlasEntityWithAssociations>>(){};
+    private static final GenericType<List<AtlasEntity>> ENTITY_LIST_TYPE = new GenericType<List<AtlasEntity>>(){};
     public static final String ENTITY_API = BASE_URI + "v2/entity/";
     public static final String ENTITIES_API = BASE_URI + "v2/entities/";
 
     private static final APIInfo GET_ENTITY_BY_GUID = new APIInfo(ENTITY_API + "guid/", HttpMethod.GET, Response.Status.OK);
-    private static final APIInfo 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(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 GET_ENTITY_BY_ATTRIBUTE = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s", HttpMethod.GET, Response.Status.OK);
     private static final APIInfo UPDATE_ENTITY_BY_ATTRIBUTE = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s/attribute/%s", HttpMethod.PUT, Response.Status.OK);
     private static final APIInfo DELETE_ENTITY_BY_ATTRIBUTE = new APIInfo(ENTITY_API + "uniqueAttribute/type/%s/attribute/%s", HttpMethod.DELETE, Response.Status.OK);
     private static final APIInfo DELETE_ENTITY_BY_GUID = new APIInfo(ENTITY_API + "guid/", HttpMethod.DELETE, Response.Status.OK);
@@ -87,25 +87,21 @@ public class AtlasEntitiesClientV2 extends AtlasBaseClient {
         super(service, configuration);
     }
 
-    public List<AtlasEntityWithAssociations> getEntityByGuid(String guid) throws AtlasServiceException {
+    public AtlasEntityWithExtInfo getEntityByGuid(String guid) throws AtlasServiceException {
 
-        return callAPI(GET_ENTITY_BY_GUID, null, ENTITY_WITH_ASSOCIATIONS_LIST_TYPE, guid);
+        return callAPI(GET_ENTITY_BY_GUID, null, AtlasEntityWithExtInfo.class, guid);
     }
 
     public AtlasEntities getEntityByGuids(List<String> guids) throws AtlasServiceException {
         return callAPI(GET_ENTITY_BY_GUID, AtlasEntities.class, "guid", guids);
     }
 
-    public List<AtlasEntityWithAssociations> getEntityWithAssociationByGuid(String guid) throws AtlasServiceException {
+    public static final String PREFIX_ATTR = "attr:";
 
-        return callAPI(formatPathForPathParams(GET_ENTITY_WITH_ASSOCIATION_BY_GUID, guid), null, ENTITY_WITH_ASSOCIATIONS_LIST_TYPE);
-    }
-
-    public List<AtlasEntityWithAssociations> getEntityByAttribute(String type, String attribute, String value) throws AtlasServiceException {
+    public AtlasEntityWithExtInfo getEntityByAttribute(String type, Map<String, String> attributes) throws AtlasServiceException {
+        MultivaluedMap<String, String> queryParams = attributesToQueryParams(attributes);
 
-        MultivaluedMap<String, String> queryParams = new MultivaluedMapImpl();
-        queryParams.add("value", value);
-        return callAPI(formatPathForPathParams(GET_ENTITY_BY_ATTRIBUTE, type, attribute), ENTITY_WITH_ASSOCIATIONS_LIST_TYPE, queryParams);
+        return callAPI(formatPathForPathParams(GET_ENTITY_BY_ATTRIBUTE, type), AtlasEntityWithExtInfo.class, queryParams);
     }
 
     public EntityMutationResponse updateEntityByAttribute(String type, String attribute, String value, AtlasEntity entity) throws AtlasServiceException {
@@ -182,4 +178,23 @@ public class AtlasEntitiesClientV2 extends AtlasBaseClient {
     public AtlasEntity.AtlasEntities searchEntities(SearchFilter searchFilter) throws AtlasServiceException {
         return callAPI(GET_ENTITIES, AtlasEntity.AtlasEntities.class, searchFilter.getParams());
     }
+
+    private MultivaluedMap<String, String> attributesToQueryParams(Map<String, String> attributes) {
+        return attributesToQueryParams(attributes, null);
+    }
+
+    private MultivaluedMap<String, String> attributesToQueryParams(Map<String, String>            attributes,
+                                                                   MultivaluedMap<String, String> queryParams) {
+        if (queryParams == null) {
+            queryParams = new MultivaluedMapImpl();
+        }
+
+        if (MapUtils.isNotEmpty(attributes)) {
+            for (Map.Entry<String, String> e : attributes.entrySet()) {
+                queryParams.putSingle(PREFIX_ATTR + e.getKey(), e.getValue());
+            }
+        }
+
+        return queryParams;
+    }
 }

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

@@ -59,6 +59,7 @@ public enum AtlasErrorCode {
     CONSTRAINT_INVERSE_REF_INVERSE_ATTRIBUTE_NON_EXISTING(400, "ATLAS40028E", "{0}.{1}: invalid {2} constraint. Inverse attribute {3}.{4} does not exist"),
     CONSTRAINT_INVERSE_REF_INVERSE_ATTRIBUTE_INVALID_TYPE(400, "ATLAS40029E", "{0}.{1}: invalid {2} constraint. Inverse attribute {3}.{4} is not an entity type"),
     CONSTRAINT_OWNED_REF_ATTRIBUTE_INVALID_TYPE(400, "ATLAS40030E", "{0}.{1}: invalid {2} constraint. Attribute {3} is not an entity type"),
+    CANNOT_MAP_ATTRIBUTE(400, "ATLAS40027E", "cannot map attribute: {0} of type: {1} from vertex"),
 
     // All Not found enums go here
     TYPE_NAME_NOT_FOUND(404, "ATLAS4041E", "Given typename {0} was invalid"),
@@ -82,7 +83,7 @@ public enum AtlasErrorCode {
     INDEX_ROLLBACK_FAILED(500, "ATLAS5003E", "Index rollback failed for {0}"),
     FAILED_TO_OBTAIN_TYPE_UPDATE_LOCK(500, "ATLAS5004E", "Failed to get the lock; another type update might be in progress. Please try again"),
 
-    INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND(400, "ATLAS40018E", "Instance {0} with unique attribute {1} does not exist"),
+    INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND(400, "ATLAS40018E", "Instance {0} with unique attribute {1}={2} does not exist"),
 
     UNRESOLVED_REFERENCES_FOUND(400, "ATLAS40010E", "Unresolved references: byId={0}; byUniqueAttributes={1}"),
 

intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java

@@ -17,7 +17,7 @@
  */
 package org.apache.atlas.model.discovery;
 
-import org.apache.atlas.model.instance.AtlasEntityHeaderWithAssociations;
+import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang.StringUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
@@ -44,7 +44,7 @@ import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONL
 public class AtlasSearchResult implements Serializable {
     private String                    queryText;
     private AtlasQueryType            queryType;
-    private List<AtlasEntityHeaderWithAssociations> entities;
+    private List<AtlasEntityHeader>   entities;
     private AttributeSearchResult     attributes;
     private List<AtlasFullTextResult> fullTextResult;
 
@@ -66,9 +66,9 @@ public class AtlasSearchResult implements Serializable {
 
     public void setQueryType(AtlasQueryType queryType) { this.queryType = queryType; }
 
-    public List<AtlasEntityHeaderWithAssociations> getEntities() { return entities; }
+    public List<AtlasEntityHeader> getEntities() { return entities; }
 
-    public void setEntities(List<AtlasEntityHeaderWithAssociations> entities) { this.entities = entities; }
+    public void setEntities(List<AtlasEntityHeader> entities) { this.entities = entities; }
 
     public AttributeSearchResult getAttributes() { return attributes; }
 
@@ -104,7 +104,7 @@ public class AtlasSearchResult implements Serializable {
                 '}';
     }
 
-    public void addEntity(AtlasEntityHeaderWithAssociations newEntity) {
+    public void addEntity(AtlasEntityHeader newEntity) {
         if (entities == null) {
             entities = new ArrayList<>();
         }
@@ -117,13 +117,13 @@ public class AtlasSearchResult implements Serializable {
         }
     }
 
-    public void removeEntity(AtlasEntityHeaderWithAssociations entity) {
-        List<AtlasEntityHeaderWithAssociations> entities = this.entities;
+    public void removeEntity(AtlasEntityHeader entity) {
+        List<AtlasEntityHeader> entities = this.entities;
 
         if (CollectionUtils.isNotEmpty(entities)) {
-            Iterator<AtlasEntityHeaderWithAssociations> iter = entities.iterator();
+            Iterator<AtlasEntityHeader> iter = entities.iterator();
             while (iter.hasNext()) {
-                AtlasEntityHeaderWithAssociations currEntity = iter.next();
+                AtlasEntityHeader currEntity = iter.next();
                 if (StringUtils.equals(currEntity.getGuid(), entity.getGuid())) {
                     iter.remove();
                 }
@@ -184,19 +184,19 @@ public class AtlasSearchResult implements Serializable {
     @XmlRootElement
     @XmlAccessorType(XmlAccessType.PROPERTY)
     public static class AtlasFullTextResult {
-        AtlasEntityHeaderWithAssociations entity;
+        AtlasEntityHeader entity;
         Double            score;
 
         public AtlasFullTextResult() {}
 
-        public AtlasFullTextResult(AtlasEntityHeaderWithAssociations entity, Double score) {
+        public AtlasFullTextResult(AtlasEntityHeader entity, Double score) {
             this.entity = entity;
             this.score  = score;
         }
 
-        public AtlasEntityHeaderWithAssociations getEntity() { return entity; }
+        public AtlasEntityHeader getEntity() { return entity; }
 
-        public void setEntity(AtlasEntityHeaderWithAssociations entity) { this.entity = entity; }
+        public void setEntity(AtlasEntityHeader entity) { this.entity = entity; }
 
         public Double getScore() { return score; }
 

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

@@ -19,7 +19,10 @@ package org.apache.atlas.model.instance;
 
 import org.apache.atlas.model.PList;
 import org.apache.atlas.model.SearchFilter.SortType;
+import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang.StringUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
 import org.codehaus.jackson.annotate.JsonIgnore;
 import org.codehaus.jackson.annotate.JsonIgnoreProperties;
@@ -30,7 +33,9 @@ import javax.xml.bind.annotation.XmlAccessorType;
 import javax.xml.bind.annotation.XmlRootElement;
 import javax.xml.bind.annotation.XmlSeeAlso;
 import java.io.Serializable;
+import java.util.ArrayList;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -65,6 +70,8 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
     private Date   updateTime = null;
     private Long   version    = new Long(0);
 
+    private List<AtlasClassification> classifications;
+
     @JsonIgnore
     private static AtomicLong s_nextId = new AtomicLong(System.nanoTime());
 
@@ -89,6 +96,7 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
         setUpdatedBy(null);
         setCreateTime(null);
         setUpdateTime(null);
+        setClassifications(null);
     }
 
     public AtlasEntity(AtlasEntity other) {
@@ -102,6 +110,7 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
             setCreateTime(other.getCreateTime());
             setUpdateTime(other.getUpdateTime());
             setVersion(other.getVersion());
+            setClassifications(other.getClassifications());
         }
     }
 
@@ -161,6 +170,10 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
         this.version = version;
     }
 
+    public List<AtlasClassification> getClassifications() { return classifications; }
+
+    public void setClassifications(List<AtlasClassification> classifications) { this.classifications = classifications; }
+
     @JsonIgnore
     public boolean isUnassigned() {
         return isUnAssigned(guid);
@@ -210,6 +223,9 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
         dumpDateField(", createTime=", createTime, sb);
         dumpDateField(", updateTime=", updateTime, sb);
         sb.append(", version=").append(version);
+        sb.append(", classifications=[");
+        AtlasBaseTypeDef.dumpObjects(classifications, sb);
+        sb.append(']');
         sb.append(", ");
         super.toString(sb);
         sb.append('}');
@@ -230,12 +246,14 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
                 Objects.equals(updatedBy, that.updatedBy) &&
                 Objects.equals(createTime, that.createTime) &&
                 Objects.equals(updateTime, that.updateTime) &&
-                Objects.equals(version, that.version);
+                Objects.equals(version, that.version) &&
+                Objects.equals(classifications, that.classifications);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(super.hashCode(), guid, status, createdBy, updatedBy, createTime, updateTime, version);
+        return Objects.hash(super.hashCode(), guid, status, createdBy, updatedBy, createTime, updateTime, version,
+                            classifications);
     }
 
     @Override
@@ -244,6 +262,339 @@ public class AtlasEntity extends AtlasStruct implements Serializable {
     }
 
     /**
+     * An instance of an entity along with extended info - like hive_table, hive_database.
+     */
+    @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
+    @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
+    @JsonIgnoreProperties(ignoreUnknown=true)
+    @XmlRootElement
+    @XmlAccessorType(XmlAccessType.PROPERTY)
+    public static class AtlasEntityExtInfo implements Serializable {
+        private static final long serialVersionUID = 1L;
+
+        private Map<String, AtlasEntity> referredEntities;
+
+
+        public AtlasEntityExtInfo() {
+            setReferredEntities(null);
+        }
+
+        public AtlasEntityExtInfo(Map<String, AtlasEntity> referredEntities) {
+            setReferredEntities(referredEntities);
+        }
+
+        public AtlasEntityExtInfo(AtlasEntityExtInfo other) {
+            if (other != null) {
+                setReferredEntities(other.getReferredEntities());
+            }
+        }
+
+        public Map<String, AtlasEntity> getReferredEntities() { return referredEntities; }
+
+        public void setReferredEntities(Map<String, AtlasEntity> referredEntities) { this.referredEntities = referredEntities; }
+
+        @JsonIgnore
+        public final void addReferredEntity(String guid, AtlasEntity entity) {
+            Map<String, AtlasEntity> r = this.referredEntities;
+
+            if (r == null) {
+                r = new HashMap<>();
+
+                this.referredEntities = r;
+            }
+
+            if (guid != null) {
+                r.put(guid, entity);
+            }
+        }
+
+        @JsonIgnore
+        public final AtlasEntity removeReferredEntity(String guid) {
+            Map<String, AtlasEntity> r = this.referredEntities;
+
+            return r != null && guid != null ? r.remove(guid) : null;
+        }
+
+        @JsonIgnore
+        public final AtlasEntity getReferredEntity(String guid) {
+            Map<String, AtlasEntity> r = this.referredEntities;
+
+            return r != null && guid != null ? r.get(guid) : null;
+        }
+
+        @JsonIgnore
+        public AtlasEntity getEntity(String guid) {
+            return getReferredEntity(guid);
+        }
+
+        @JsonIgnore
+        public AtlasEntity removeEntity(String guid) {
+            Map<String, AtlasEntity> r = this.referredEntities;
+
+            return r != null && guid != null ? r.remove(guid) : null;
+        }
+
+        public void updateEntityGuid(String oldGuid, String newGuid) {
+            AtlasEntity entity = getEntity(oldGuid);
+
+            if (entity != null) {
+                entity.setGuid(newGuid);
+
+                if(removeEntity(oldGuid) != null) {
+                    addReferredEntity(newGuid, entity);
+                }
+            }
+        }
+
+        public boolean hasEntity(String guid) {
+            return getEntity(guid) != null;
+        }
+
+        public void compact() {
+            // for derived classes to implement their own logic
+        }
+
+        public StringBuilder toString(StringBuilder sb) {
+            if (sb == null) {
+                sb = new StringBuilder();
+            }
+
+            sb.append("AtlasEntityExtInfo{");
+            sb.append("referredEntities={");
+            AtlasBaseTypeDef.dumpObjects(referredEntities, sb);
+            sb.append("}");
+            sb.append("}");
+
+            return sb;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+
+            AtlasEntityExtInfo that = (AtlasEntityExtInfo) o;
+            return Objects.equals(referredEntities, that.referredEntities);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(referredEntities);
+        }
+
+        @Override
+        public String toString() {
+            return toString(new StringBuilder()).toString();
+        }
+    }
+
+    @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
+    @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
+    @JsonIgnoreProperties(ignoreUnknown=true)
+    @XmlRootElement
+    @XmlAccessorType(XmlAccessType.PROPERTY)
+    public static class AtlasEntityWithExtInfo extends AtlasEntityExtInfo {
+        private static final long serialVersionUID = 1L;
+
+        private AtlasEntity entity;
+
+        public AtlasEntityWithExtInfo() {
+            this(null, null);
+        }
+
+        public AtlasEntityWithExtInfo(AtlasEntity entity) {
+            this(entity, null);
+        }
+
+        public AtlasEntityWithExtInfo(AtlasEntity entity, AtlasEntityExtInfo extInfo) {
+            super(extInfo);
+
+            this.entity = entity;
+        }
+
+        public AtlasEntity getEntity() { return entity; }
+
+        public void setEntity(AtlasEntity entity) { this.entity = entity; }
+
+        @JsonIgnore
+        @Override
+        public AtlasEntity getEntity(String guid) {
+            AtlasEntity ret = super.getEntity(guid);
+
+            if (ret == null && entity != null) {
+                if (StringUtils.equals(guid, entity.getGuid())) {
+                    ret = entity;
+                }
+            }
+
+            return ret;
+        }
+
+        @JsonIgnore
+        @Override
+        public void compact() {
+            super.compact();
+
+            // remove 'entity' from referredEntities
+            if (entity != null) {
+                removeEntity(entity.getGuid());
+            }
+        }
+
+        @Override
+        public StringBuilder toString(StringBuilder sb) {
+            if (sb == null) {
+                sb = new StringBuilder();
+            }
+
+            sb.append("AtlasEntityWithExtInfo{");
+            sb.append("entity=").append(entity).append(",");
+            super.toString(sb);
+            sb.append("}");
+
+            return sb;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+
+            AtlasEntityWithExtInfo that = (AtlasEntityWithExtInfo) o;
+            return Objects.equals(entity, that.entity);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(super.hashCode(), entity);
+        }
+    }
+
+    @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
+    @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
+    @JsonIgnoreProperties(ignoreUnknown=true)
+    @XmlRootElement
+    @XmlAccessorType(XmlAccessType.PROPERTY)
+    public static class AtlasEntitiesWithExtInfo extends AtlasEntityExtInfo {
+        private static final long serialVersionUID = 1L;
+
+        private List<AtlasEntity> entities;
+
+        public AtlasEntitiesWithExtInfo() {
+            this(null, null);
+        }
+
+        public AtlasEntitiesWithExtInfo(List<AtlasEntity> entities) {
+            this(entities, null);
+        }
+
+        public AtlasEntitiesWithExtInfo(List<AtlasEntity> entities, AtlasEntityExtInfo extInfo) {
+            super(extInfo);
+
+            this.entities = entities;
+        }
+
+        public List<AtlasEntity> getEntities() { return entities; }
+
+        public void setEntities(List<AtlasEntity> entities) { this.entities = entities; }
+
+        @JsonIgnore
+        @Override
+        public AtlasEntity getEntity(String guid) {
+            AtlasEntity ret = super.getEntity(guid);
+
+            if (ret == null && CollectionUtils.isNotEmpty(entities)) {
+                for (AtlasEntity entity : entities) {
+                    if (StringUtils.equals(guid, entity.getGuid())) {
+                        ret = entity;
+
+                        break;
+                    }
+                }
+            }
+
+            return ret;
+        }
+
+        public void addEntity(AtlasEntity entity) {
+            List<AtlasEntity> entities = this.entities;
+
+            if (entities == null) {
+                entities = new ArrayList<>();
+
+                this.entities = entities;
+            }
+
+            entities.add(entity);
+        }
+
+        public void removeEntity(AtlasEntity entity) {
+            List<AtlasEntity> entities = this.entities;
+
+            if (entity != null && entities != null) {
+                entities.remove(entity);
+            }
+        }
+
+        @Override
+        public void compact() {
+            super.compact();
+
+            // remove 'entities' from referredEntities
+            if (CollectionUtils.isNotEmpty(entities)) {
+                for (AtlasEntity entity : entities) {
+                    removeReferredEntity(entity.getGuid());
+                }
+            }
+        }
+
+        @Override
+        public StringBuilder toString(StringBuilder sb) {
+            if (sb == null) {
+                sb = new StringBuilder();
+            }
+
+            sb.append("AtlasEntitiesWithExtInfo{");
+            sb.append("entities=[");
+            AtlasBaseTypeDef.dumpObjects(entities, sb);
+            sb.append("],");
+            super.toString(sb);
+            sb.append("}");
+
+            return sb;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+
+            AtlasEntitiesWithExtInfo that = (AtlasEntitiesWithExtInfo) o;
+            return Objects.equals(entities, that.entities);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(super.hashCode(), entities);
+        }
+    }
+
+    /**
      * REST serialization friendly list.
      */
     @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)

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

@@ -18,7 +18,6 @@
 package org.apache.atlas.model.instance;
 
 import java.io.Serializable;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -52,6 +51,7 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
     private String             guid                = null;
     private AtlasEntity.Status status              = AtlasEntity.Status.ACTIVE;
     private String             displayText         = null;
+    private List<String>       classificationNames = null;
 
     public AtlasEntityHeader() {
         this(null, null);
@@ -67,12 +67,15 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
 
     public AtlasEntityHeader(String typeName, Map<String, Object> attributes) {
         super(typeName, attributes);
+
+        setClassificationNames(null);
     }
 
 
     public AtlasEntityHeader(String typeName, String guid,  Map<String, Object> attributes) {
         super(typeName, attributes);
         setGuid(guid);
+        setClassificationNames(null);
     }
 
 
@@ -83,6 +86,7 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
             setGuid(other.getGuid());
             setStatus(other.getStatus());
             setDisplayText(other.getDisplayText());
+            setClassificationNames(other.getClassificationNames());
         }
     }
 
@@ -110,6 +114,14 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
         this.displayText = displayText;
     }
 
+    public List<String> getClassificationNames(){
+        return classificationNames;
+    }
+
+    public void setClassificationNames(List<String> classificationNames) {
+        this.classificationNames = classificationNames;
+    }
+
     @Override
     public StringBuilder toString(StringBuilder sb) {
         if (sb == null) {
@@ -120,6 +132,9 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
         sb.append("guid='").append(guid).append('\'');
         sb.append(", status=").append(status);
         sb.append(", displayText=").append(displayText);
+        sb.append(", classificationNames=[");
+        dumpObjects(classificationNames, sb);
+        sb.append("],");
         sb.append(", ");
         super.toString(sb);
         sb.append('}');
@@ -135,12 +150,13 @@ public class AtlasEntityHeader extends AtlasStruct implements Serializable {
         AtlasEntityHeader that = (AtlasEntityHeader) o;
         return Objects.equals(guid, that.guid) &&
                 status == that.status &&
-                Objects.equals(displayText, that.displayText);
+                Objects.equals(displayText, that.displayText) &&
+                Objects.equals(classificationNames, that.classificationNames);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(super.hashCode(), guid, status, displayText);
+        return Objects.hash(super.hashCode(), guid, status, displayText, classificationNames);
     }
 
     @Override

intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityHeaderWithAssociations.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.model.instance;
-
-import org.apache.atlas.model.PList;
-import org.apache.atlas.model.SearchFilter;
-import org.apache.atlas.model.typedef.AtlasEntityDef;
-import org.codehaus.jackson.annotate.JsonAutoDetect;
-import org.codehaus.jackson.annotate.JsonIgnoreProperties;
-import org.codehaus.jackson.map.annotate.JsonSerialize;
-
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlRootElement;
-import javax.xml.bind.annotation.XmlSeeAlso;
-import java.io.Serializable;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-
-import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.NONE;
-import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONLY;
-
-/**
- * An instance of an entity and its associations - like hive_table, hive_database.
- */
-@JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
-@JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
-@JsonIgnoreProperties(ignoreUnknown=true)
-@XmlRootElement
-@XmlAccessorType(XmlAccessType.PROPERTY)
-public class AtlasEntityHeaderWithAssociations extends AtlasEntityHeader implements Serializable{
-    private static final long serialVersionUID = 1L;
-
-    private List<String> classificationNames;
-
-    public AtlasEntityHeaderWithAssociations(){
-        this(null, null);
-    }
-
-    public AtlasEntityHeaderWithAssociations(AtlasEntityDef entityDef) {
-        this(entityDef != null ? entityDef.getName() : null, null);
-    }
-
-    public AtlasEntityHeaderWithAssociations(String typeName, Map<String, Object> attributes) {
-        super(typeName, attributes);
-        setClassificationNames(null);
-    }
-
-    public AtlasEntityHeaderWithAssociations(AtlasEntityHeaderWithAssociations other) {
-        super(other);
-
-        if (other != null) {
-            setClassificationNames(other.getClassificationNames());
-        }
-    }
-
-    public List<String> getClassificationNames(){
-        return classificationNames;
-    }
-
-    public void setClassificationNames(List<String> classificationNames) {
-        this.classificationNames = classificationNames;
-    }
-
-
-    @Override
-    public StringBuilder toString(StringBuilder sb) {
-        if (sb == null) {
-            sb = new StringBuilder();
-        }
-
-        sb.append("AtlasEntityHeaderwithAssociations{");
-        sb.append(", classificationNames=[");
-        dumpObjects(classificationNames, sb);
-        sb.append("],");
-        super.toString(sb);
-        sb.append('}');
-
-        return sb;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-        AtlasEntityHeaderWithAssociations that = (AtlasEntityHeaderWithAssociations) o;
-        return Objects.equals(classificationNames, that.classificationNames);
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(super.hashCode(), classificationNames);
-    }
-
-    @Override
-    public String toString() {
-        return toString(new StringBuilder()).toString();
-    }
-
-    /**
-     * REST serialization friendly list.
-     */
-    @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
-    @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
-    @JsonIgnoreProperties(ignoreUnknown=true)
-    @XmlRootElement
-    @XmlAccessorType(XmlAccessType.PROPERTY)
-    @XmlSeeAlso(AtlasEntity.class)
-    public static class  AtlasEntityHeadersWithAssociations extends PList<AtlasEntityHeaderWithAssociations> {
-        private static final long serialVersionUID = 1L;
-
-        public AtlasEntityHeadersWithAssociations() {
-            super();
-        }
-
-        public AtlasEntityHeadersWithAssociations(List<AtlasEntityHeaderWithAssociations> list) {
-            super(list);
-        }
-
-        public AtlasEntityHeadersWithAssociations(List list, long startIndex, int pageSize, long totalCount,
-                                  SearchFilter.SortType sortType, String sortBy) {
-            super(list, startIndex, pageSize, totalCount, sortType, sortBy);
-        }
-    }
-
-}

intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityWithAssociations.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.model.instance;
-
-import org.apache.atlas.model.PList;
-import org.apache.atlas.model.SearchFilter.SortType;
-import org.apache.atlas.model.typedef.AtlasEntityDef;
-import org.codehaus.jackson.annotate.JsonAutoDetect;
-import org.codehaus.jackson.annotate.JsonIgnoreProperties;
-import org.codehaus.jackson.map.annotate.JsonSerialize;
-
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlRootElement;
-import javax.xml.bind.annotation.XmlSeeAlso;
-import java.io.Serializable;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-
-import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.NONE;
-import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONLY;
-
-
-/**
- * An instance of an entity - like hive_table, hive_database along with its assictaed classifications, terms etc.
- */
-@JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
-@JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
-@JsonIgnoreProperties(ignoreUnknown=true)
-@XmlRootElement
-@XmlAccessorType(XmlAccessType.PROPERTY)
-public class AtlasEntityWithAssociations extends AtlasEntity implements Serializable {
-    private static final long serialVersionUID = 1L;
-
-    List<AtlasClassification> classifications;
-
-    public AtlasEntityWithAssociations() {
-        this(null, null);
-    }
-
-    public AtlasEntityWithAssociations(String typeName) {
-        this(typeName, null);
-    }
-
-    public AtlasEntityWithAssociations(AtlasEntityDef entityDef) {
-        this(entityDef != null ? entityDef.getName() : null, null);
-    }
-
-    public AtlasEntityWithAssociations(String typeName, Map<String, Object> attributes) {
-        super(typeName, attributes);
-        setClassifications(null);
-    }
-
-    public AtlasEntityWithAssociations(AtlasEntityWithAssociations other) {
-        super(other);
-
-        setClassifications(other != null ? other.getClassifications() : null);
-    }
-
-    public AtlasEntityWithAssociations(AtlasEntity other) {
-        super(other);
-    }
-
-    @Override
-    public StringBuilder toString(StringBuilder sb) {
-        if (sb == null) {
-            sb = new StringBuilder();
-        }
-
-        sb.append("AtlasEntityWithAssociations{");
-        sb.append("classifications='").append(classifications).append('\'');
-        sb.append(", ");
-        super.toString(sb);
-        sb.append('}');
-
-        return sb;
-    }
-
-    @Override
-    public boolean equals(Object o) {
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        if (!super.equals(o)) return false;
-        AtlasEntityWithAssociations that = (AtlasEntityWithAssociations) o;
-        return Objects.equals(classifications, that.classifications);
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(super.hashCode(), classifications);
-    }
-
-    public List<AtlasClassification> getClassifications() {
-        return classifications;
-    }
-
-    public void setClassifications(final List<AtlasClassification> classifications) {
-        this.classifications = classifications;
-    }
-
-    @Override
-    public String toString() {
-        return toString(new StringBuilder()).toString();
-    }
-
-    /**
-     * REST serialization friendly list.
-     */
-    @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
-    @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
-    @JsonIgnoreProperties(ignoreUnknown=true)
-    @XmlRootElement
-    @XmlAccessorType(XmlAccessType.PROPERTY)
-    @XmlSeeAlso(AtlasEntityWithAssociations.class)
-    public static class AtlasEntitiesWithAssociations extends PList<AtlasEntityWithAssociations> {
-        private static final long serialVersionUID = 1L;
-
-        public AtlasEntitiesWithAssociations() {
-            super();
-        }
-
-        public AtlasEntitiesWithAssociations(List<AtlasEntityWithAssociations> list) {
-            super(list);
-        }
-
-        public AtlasEntitiesWithAssociations(List list, long startIndex, int pageSize, long totalCount,
-                             SortType sortType, String sortBy) {
-            super(list, startIndex, pageSize, totalCount, sortType, sortBy);
-        }
-    }
-}

repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java

@@ -26,7 +26,7 @@ import org.apache.atlas.discovery.graph.DefaultGraphPersistenceStrategy;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.discovery.AtlasSearchResult;
 import org.apache.atlas.model.instance.AtlasEntity.Status;
-import org.apache.atlas.model.instance.AtlasEntityHeaderWithAssociations;
+import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.query.Expressions.AliasExpression;
 import org.apache.atlas.query.Expressions.Expression;
@@ -97,7 +97,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
                 if (firstElement instanceof AtlasVertex) {
                     for (Object element : queryResult) {
                         if (element instanceof AtlasVertex) {
-                            ret.addEntity(toAtlasEntityHeaderwithAssociations((AtlasVertex)element));
+                            ret.addEntity(toAtlasEntityHeader((AtlasVertex)element));
                         } else {
                             LOG.warn("searchUsingDslQuery({}): expected an AtlasVertex; found unexpected entry in result {}", dslQuery, element);
                         }
@@ -115,7 +115,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
                                     Object entry = ((List)value).get(0);
 
                                     if (entry instanceof AtlasVertex) {
-                                        ret.addEntity(toAtlasEntityHeaderwithAssociations((AtlasVertex)entry));
+                                        ret.addEntity(toAtlasEntityHeader((AtlasVertex)entry));
                                     }
                                 }
                             }
@@ -159,7 +159,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
             String guid = vertex != null ? vertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class) : null;
 
             if (guid != null) {
-                AtlasEntityHeaderWithAssociations entity = toAtlasEntityHeaderwithAssociations(vertex);
+                AtlasEntityHeader entity = toAtlasEntityHeader(vertex);
                 Double score = idxQueryResult.getScore();
                 ret.add(new AtlasFullTextResult(entity, score));
             }
@@ -204,12 +204,12 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
         return new QueryParams(limit, offset);
     }
 
-    private AtlasEntityHeaderWithAssociations toAtlasEntityHeaderwithAssociations(AtlasVertex vertex) {
+    private AtlasEntityHeader toAtlasEntityHeader(AtlasVertex vertex) {
         if (vertex == null) {
             return null;
         }
 
-        AtlasEntityHeaderWithAssociations ret = new AtlasEntityHeaderWithAssociations();
+        AtlasEntityHeader ret = new AtlasEntityHeader();
 
         String typeName = vertex.getProperty(Constants.TYPE_NAME_PROPERTY_KEY, String.class);
 

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

@@ -18,23 +18,13 @@
 
 package org.apache.atlas.repository.graph;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Set;
-import java.util.Stack;
-import java.util.UUID;
-
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.BiMap;
+import com.google.common.collect.HashBiMap;
 import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.RequestContext;
+import org.apache.atlas.model.instance.AtlasEntity.Status;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.RepositoryException;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
@@ -43,6 +33,7 @@ import org.apache.atlas.repository.graphdb.AtlasElement;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.typesystem.IReferenceableInstance;
 import org.apache.atlas.typesystem.ITypedInstance;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
@@ -70,9 +61,19 @@ import org.codehaus.jettison.json.JSONArray;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.BiMap;
-import com.google.common.collect.HashBiMap;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.Stack;
+import java.util.UUID;
 
 /**
  * Utility class for graph operations.
@@ -649,10 +650,15 @@ public final class GraphHelper {
         return element.getProperty(Constants.VERSION_PROPERTY_KEY, Integer.class);
     }
 
+
+
     public static String getStateAsString(AtlasElement element) {
         return element.getProperty(Constants.STATE_PROPERTY_KEY, String.class);
     }
 
+    public static Status getStatus(AtlasElement element) {
+        return (getState(element) == Id.EntityState.DELETED) ? Status.DELETED : Status.ACTIVE;
+    }
 
     //Added conditions in fetching system attributes to handle test failures in GremlinTest where these properties are not set
     public static String getCreatedByAsString(AtlasElement element){
@@ -812,6 +818,7 @@ public final class GraphHelper {
         }
         return Collections.emptyList();
     }
+
     /**
      * Guid and AtlasVertex combo
      */
@@ -974,6 +981,12 @@ public final class GraphHelper {
 
     }
 
+    public static boolean isReference(AtlasType type) {
+        return ((type.getTypeCategory() == org.apache.atlas.model.TypeCategory.STRUCT) ||
+                (type.getTypeCategory() == org.apache.atlas.model.TypeCategory.ENTITY));
+
+    }
+
     public static void setArrayElementsProperty(IDataType elementType, AtlasVertex instanceVertex, String propertyName, List<Object> values) {
         String actualPropertyName = GraphHelper.encodePropertyKey(propertyName);
         if(GraphHelper.isReference(elementType)) {
@@ -1004,6 +1017,27 @@ public final class GraphHelper {
         }
     }
 
+    public static Object getMapValueProperty(AtlasType elementType, AtlasVertex instanceVertex, String propertyName) {
+        String vertexPropertyName = GraphHelper.encodePropertyKey(propertyName);
+
+        if (GraphHelper.isReference(elementType)) {
+            return instanceVertex.getProperty(vertexPropertyName, AtlasEdge.class);
+        } else {
+            return instanceVertex.getProperty(vertexPropertyName, Object.class).toString();
+        }
+    }
+
+    // newly added
+    public static List<Object> getArrayElementsProperty(AtlasType elementType, AtlasVertex instanceVertex, String propertyName) {
+        String encodedPropertyName = GraphHelper.encodePropertyKey(propertyName);
+        if(GraphHelper.isReference(elementType)) {
+            return (List)instanceVertex.getListProperty(encodedPropertyName, AtlasEdge.class);
+        }
+        else {
+            return (List)instanceVertex.getListProperty(encodedPropertyName);
+        }
+    }
+
     public static List<Object> getArrayElementsProperty(IDataType elementType, AtlasVertex instanceVertex, String propertyName) {
         String actualPropertyName = GraphHelper.encodePropertyKey(propertyName);
         if(GraphHelper.isReference(elementType)) {
@@ -1014,7 +1048,6 @@ public final class GraphHelper {
         }
     }
 
-
     public static void dumpToLog(final AtlasGraph<?,?> graph) {
         LOG.debug("*******************Graph Dump****************************");
         LOG.debug("Vertices of {}", graph);

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

@@ -19,12 +19,12 @@ package org.apache.atlas.repository.store.graph;
 
 
 import org.apache.atlas.exception.AtlasBaseException;
-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.AtlasEntityWithAssociations;
-import org.apache.atlas.model.instance.EntityMutations;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
 import org.apache.atlas.model.instance.EntityMutationResponse;
+import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 
 import java.util.List;
@@ -45,9 +45,9 @@ public interface AtlasEntityStore {
      *
      * Get entity definition by its guid
      * @param guid
-     * @return
+     * @return AtlasEntity
      */
-    AtlasEntity   getById(String guid);
+    AtlasEntityWithExtInfo getById(String guid) throws AtlasBaseException;
 
     /**
      * Delete an entity by its guid
@@ -71,15 +71,7 @@ public interface AtlasEntityStore {
      * @return
      * @throws AtlasBaseException
      */
-    AtlasEntity.AtlasEntities getByIds(List<String> guid) throws AtlasBaseException;
-
-    /**
-     * Batch GET to retrieve entities and their associations by their ID
-     * @param guid
-     * @return
-     * @throws AtlasBaseException
-     */
-    AtlasEntityWithAssociations getWithAssociationsByIds(List<String> guid) throws AtlasBaseException;
+    AtlasEntitiesWithExtInfo getByIds(List<String> guid) throws AtlasBaseException;
 
     /*
      * Return list of deleted entity guids
@@ -89,12 +81,12 @@ public interface AtlasEntityStore {
     /**
      *
      * Get an eneity by its unique attribute
-     * @param typeName
-     * @param attrName
-     * @param attrValue
-     * @return
+     * @param entityType
+     * @param uniqAttributes
+     * @return AtlasEntity
      */
-    AtlasEntity  getByUniqueAttribute(String typeName, String attrName, String attrValue);
+    AtlasEntityWithExtInfo getByUniqueAttribute(AtlasEntityType entityType, Map<String, Object> uniqAttributes)
+                                                                                             throws AtlasBaseException;
 
     /**
      * @deprecated
@@ -134,5 +126,4 @@ public interface AtlasEntityStore {
      * Delete classification(s)
      */
     void deleteClassifications(String guid, List<String> classificationNames) throws AtlasBaseException;
-
 }

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

@@ -18,9 +18,7 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 
-import java.util.ArrayList;
-import java.util.List;
-
+import com.google.inject.Inject;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.GraphTransaction;
 import org.apache.atlas.RequestContextV1;
@@ -28,9 +26,16 @@ import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.Status;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.EntityMutationResponse;
+import org.apache.atlas.repository.graph.AtlasGraphProvider;
+import org.apache.atlas.repository.graph.GraphHelper;
+import org.apache.atlas.repository.Constants;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscovery;
@@ -41,22 +46,28 @@ import org.apache.atlas.type.AtlasTypeRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 
-import com.google.inject.Inject;
-
 
 public class AtlasEntityStoreV1 implements AtlasEntityStore {
 
     protected AtlasTypeRegistry typeRegistry;
+    protected final GraphHelper graphHelper = GraphHelper.getInstance();
 
-    private EntityGraphMapper graphMapper;
+    private final EntityGraphMapper graphMapper;
+    private final GraphEntityMapper entityMapper;
+    private final AtlasGraph        graph;
 
     private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityStoreV1.class);
 
     @Inject
-    public AtlasEntityStoreV1(EntityGraphMapper vertexMapper) {
+    public AtlasEntityStoreV1(EntityGraphMapper vertexMapper, GraphEntityMapper entityMapper) {
         this.graphMapper  = vertexMapper;
+        this.entityMapper = entityMapper;
+        this.graph        = AtlasGraphProvider.getGraphInstance();
     }
 
     @Inject
@@ -65,8 +76,43 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Override
-    public AtlasEntity getById(final String guid) {
-        return null;
+    public AtlasEntityWithExtInfo getById(final String guid) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Retrieving entity with guid={}", guid);
+        }
+
+        return entityMapper.toAtlasEntity(guid, true);
+    }
+
+    @Override
+    public AtlasEntityWithExtInfo getByUniqueAttribute(AtlasEntityType entityType, Map<String, Object> uniqAttributes) throws AtlasBaseException {
+        String entityTypeName = entityType.getTypeName();
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Retrieving entity with type={} and attributes={}: values={}", entityTypeName, uniqAttributes);
+        }
+
+        AtlasGraphQuery query = graph.query();
+
+        for (Map.Entry<String, Object> e : uniqAttributes.entrySet()) {
+            String attrName = e.getKey();
+            Object attrValue = e.getValue();
+
+            query = query.has(entityType.getQualifiedAttributeName(attrName), attrValue);
+        }
+
+        Iterator<AtlasVertex> result = query.has(Constants.ENTITY_TYPE_PROPERTY_KEY, entityTypeName)
+                                            .has(Constants.STATE_PROPERTY_KEY, Status.ACTIVE.name())
+                                            .vertices().iterator();
+        AtlasVertex entityVertex = result.hasNext() ? result.next() : null;
+
+        if (entityVertex == null) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, entityTypeName, uniqAttributes.keySet().toString(), uniqAttributes.values().toString());
+        }
+
+        String guid = GraphHelper.getGuid(entityVertex);
+
+        return entityMapper.toAtlasEntity(guid, true);
     }
 
     @Override
@@ -96,12 +142,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Override
-    public AtlasEntity.AtlasEntities getByIds(final List<String> guid) throws AtlasBaseException {
-        return null;
-    }
-
-    @Override
-    public AtlasEntityWithAssociations getWithAssociationsByIds(final List<String> guid) throws AtlasBaseException {
+    public AtlasEntitiesWithExtInfo getByIds(final List<String> guids) throws AtlasBaseException {
         return null;
     }
 
@@ -111,11 +152,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Override
-    public AtlasEntity getByUniqueAttribute(final String typeName, final String attrName, final String attrValue) {
-        return null;
-    }
-
-    @Override
     public EntityMutationResponse updateByUniqueAttribute(final String typeName, final String attributeName, final String attributeValue, final AtlasEntity entity) throws AtlasBaseException {
         return null;
     }

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

@@ -21,11 +21,9 @@ import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 
 import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.AtlasException;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasStruct;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graph.GraphHelper;
@@ -34,14 +32,10 @@ import org.apache.atlas.repository.graphdb.AtlasElement;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasStructType;
 import org.apache.atlas.type.AtlasType;
-import org.apache.atlas.typesystem.persistence.Id;
-import org.apache.atlas.typesystem.types.DataTypes;
-import org.apache.atlas.typesystem.types.IDataType;
 import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Date;
 import java.util.HashMap;

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

+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas.repository.store.graph.v1;
+
+import com.google.inject.Inject;
+import com.google.inject.Singleton;
+import org.apache.atlas.AtlasErrorCode;
+import org.apache.atlas.AtlasException;
+import org.apache.atlas.exception.AtlasBaseException;
+import org.apache.atlas.model.instance.AtlasClassification;
+import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
+import org.apache.atlas.model.instance.AtlasObjectId;
+import org.apache.atlas.model.instance.AtlasStruct;
+import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
+import org.apache.atlas.repository.Constants;
+import org.apache.atlas.repository.graph.GraphHelper;
+import org.apache.atlas.repository.graphdb.AtlasEdge;
+import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
+import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.type.AtlasArrayType;
+import org.apache.atlas.type.AtlasMapType;
+import org.apache.atlas.type.AtlasStructType;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
+import org.apache.atlas.type.AtlasType;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BIGDECIMAL;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BIGINTEGER;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BOOLEAN;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BYTE;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_DATE;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_DOUBLE;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_FLOAT;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_INT;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_LONG;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_SHORT;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_STRING;
+import static org.apache.atlas.repository.graph.GraphHelper.EDGE_LABEL_PREFIX;
+
+@Singleton
+public final class GraphEntityMapper {
+    private static final Logger LOG = LoggerFactory.getLogger(GraphEntityMapper.class);
+    private static final GraphHelper graphHelper = GraphHelper.getInstance();
+
+    @Inject
+    protected AtlasTypeRegistry typeRegistry;
+
+    public AtlasEntityWithExtInfo toAtlasEntity(String guid, boolean includeReferences) throws AtlasBaseException {
+        AtlasEntityExtInfo     entityExtInfo = includeReferences ? new AtlasEntityExtInfo() : null;
+        AtlasEntity            entity        = mapVertexToAtlasEntity(guid, entityExtInfo);
+        AtlasEntityWithExtInfo ret           = new AtlasEntityWithExtInfo(entity, entityExtInfo);
+
+        ret.compact();
+
+        return ret;
+    }
+
+    private AtlasEntity mapVertexToAtlasEntity(String guid, AtlasEntityExtInfo entityExtInfo) throws AtlasBaseException {
+        AtlasEntity entity = entityExtInfo != null ? entityExtInfo.getEntity(guid) : null;
+
+        if (entity != null) {
+            return entity;
+        }
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Mapping graph vertex to atlas entity for guid {}", guid);
+        }
+
+        try {
+            AtlasVertex entityVertex = graphHelper.getVertexForGUID(guid);
+            entity = new AtlasEntity();
+
+            if (entityExtInfo != null) {
+                entityExtInfo.addReferredEntity(guid, entity);
+            }
+
+            mapSystemAttributes(entityVertex, entity);
+
+            mapAttributes(entityVertex, entity, entityExtInfo);
+
+            mapClassifications(entityVertex, entity, entityExtInfo);
+
+        } catch (AtlasException e) {
+            throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
+        }
+
+        return entity;
+    }
+
+    private AtlasEntity mapSystemAttributes(AtlasVertex entityVertex, AtlasEntity entity) {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Mapping system attributes for type {}", entity.getTypeName());
+        }
+
+        entity.setGuid(GraphHelper.getGuid(entityVertex));
+        entity.setTypeName(GraphHelper.getSingleValuedProperty(entityVertex, Constants.ENTITY_TYPE_PROPERTY_KEY, String.class));
+        entity.setStatus(GraphHelper.getStatus(entityVertex));
+        entity.setVersion(GraphHelper.getVersion(entityVertex).longValue());
+
+        entity.setCreatedBy(GraphHelper.getCreatedByAsString(entityVertex));
+        entity.setUpdatedBy(GraphHelper.getModifiedByAsString(entityVertex));
+
+        entity.setCreateTime(new Date(GraphHelper.getCreatedTime(entityVertex)));
+        entity.setUpdateTime(new Date(GraphHelper.getModifiedTime(entityVertex)));
+
+        return entity;
+    }
+
+    private void mapAttributes(AtlasVertex entityVertex, AtlasStruct struct, AtlasEntityExtInfo entityExtInfo) throws AtlasBaseException {
+        AtlasType objType  = typeRegistry.getType(struct.getTypeName());
+
+        if (!(objType instanceof AtlasStructType)) {
+            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, struct.getTypeName());
+        }
+
+        AtlasStructType            structType = (AtlasStructType) objType;
+        Collection<AtlasAttribute> attributes = structType.getAllAttributes().values();
+
+        if (CollectionUtils.isNotEmpty(attributes)) {
+            for (AtlasAttribute attribute : attributes) {
+                Object attrValue = mapVertexToAttribute(entityVertex, attribute, entityExtInfo);
+
+                struct.setAttribute(attribute.getName(), attrValue);
+            }
+        }
+    }
+
+    private void mapClassifications(AtlasVertex entityVertex, AtlasEntity entity, AtlasEntityExtInfo entityExtInfo) throws AtlasBaseException {
+        List<String> classficationNames = GraphHelper.getTraitNames(entityVertex);
+
+        if (CollectionUtils.isNotEmpty(classficationNames)) {
+            List<AtlasClassification> classifications = new ArrayList<>();
+
+            for (String classficationName : classficationNames) {
+                if (LOG.isDebugEnabled()) {
+                    LOG.debug("mapping classification {} to atlas entity", classficationName);
+                }
+
+                Iterable<AtlasEdge> edges = entityVertex.getEdges(AtlasEdgeDirection.OUT, classficationName);
+
+                AtlasEdge edge = (edges != null && edges.iterator().hasNext()) ? edges.iterator().next() : null;
+
+                if (edge != null) {
+                    AtlasClassification classification = new AtlasClassification(classficationName);
+
+                    mapAttributes(edge.getInVertex(), classification, entityExtInfo);
+
+                    classifications.add(classification);
+                }
+            }
+
+            entity.setClassifications(classifications);
+        }
+    }
+
+    private Object mapVertexToAttribute(AtlasVertex entityVertex, AtlasAttribute attribute, AtlasEntityExtInfo entityExtInfo) throws AtlasBaseException {
+        Object    ret                = null;
+        AtlasType attrType           = attribute.getAttributeType();
+        String    vertexPropertyName = attribute.getQualifiedName();
+        String    edgeLabel          = EDGE_LABEL_PREFIX + vertexPropertyName;
+        boolean   isOwnedAttribute   = attribute.isOwnedRef();
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Mapping vertex {} to atlas entity {}.{}", entityVertex, attribute.getDefinedInDef().getName(), attribute.getName());
+        }
+
+        switch (attrType.getTypeCategory()) {
+            case PRIMITIVE:
+                ret = mapVertexToPrimitive(entityVertex, vertexPropertyName, attribute.getAttributeDef());
+                break;
+            case ENUM:
+                ret = GraphHelper.getProperty(entityVertex, vertexPropertyName);
+                break;
+            case STRUCT:
+                ret = mapVertexToStruct(entityVertex, edgeLabel, null, entityExtInfo);
+                break;
+            case ENTITY:
+                ret = mapVertexToObjectId(entityVertex, edgeLabel, null, entityExtInfo, isOwnedAttribute);
+                break;
+            case ARRAY:
+                ret = mapVertexToArray(entityVertex, (AtlasArrayType) attrType, vertexPropertyName, entityExtInfo, isOwnedAttribute);
+                break;
+            case MAP:
+                ret = mapVertexToMap(entityVertex, (AtlasMapType) attrType, vertexPropertyName, entityExtInfo, isOwnedAttribute);
+                break;
+            case CLASSIFICATION:
+                // do nothing
+                break;
+        }
+
+        return ret;
+    }
+
+    private Map<String, Object> mapVertexToMap(AtlasVertex entityVertex, AtlasMapType atlasMapType, final String propertyName,
+                                               AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute) throws AtlasBaseException {
+        Map<String, Object> ret       = new HashMap<>();
+        List<String>        mapKeys   = GraphHelper.getListProperty(entityVertex, propertyName);
+        AtlasType           mapValueType = atlasMapType.getValueType();
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Mapping map attribute {} for vertex {}", atlasMapType.getTypeName(), entityVertex);
+        }
+
+        if (CollectionUtils.isEmpty(mapKeys)) {
+            return null;
+        }
+
+        for (String mapKey : mapKeys) {
+            final String keyPropertyName = propertyName + "." + mapKey;
+            final String edgeLabel       = EDGE_LABEL_PREFIX + keyPropertyName;
+            final Object keyValue        = GraphHelper.getMapValueProperty(mapValueType, entityVertex, keyPropertyName);
+
+            Object mapValue = mapVertexToCollectionEntry(entityVertex, mapValueType, keyValue, edgeLabel, entityExtInfo, isOwnedAttribute);
+            if (mapValue != null) {
+                ret.put(mapKey, mapValue);
+            }
+        }
+
+        return ret;
+    }
+
+    private List<Object> mapVertexToArray(AtlasVertex entityVertex, AtlasArrayType arrayType, String propertyName,
+                                          AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute) throws AtlasBaseException {
+        AtlasType    arrayElementType = arrayType.getElementType();
+        List<Object> arrayElements    = GraphHelper.getArrayElementsProperty(arrayElementType, entityVertex, propertyName);
+        List         arrValues        = new ArrayList();
+        String       edgeLabel        = EDGE_LABEL_PREFIX + propertyName;
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("Mapping array attribute {} for vertex {}", arrayElementType.getTypeName(), entityVertex);
+        }
+
+        if (CollectionUtils.isEmpty(arrayElements)) {
+            return null;
+        }
+
+        for (Object element : arrayElements) {
+            Object arrValue = mapVertexToCollectionEntry(entityVertex, arrayElementType, element,
+                                                         edgeLabel, entityExtInfo, isOwnedAttribute);
+
+            arrValues.add(arrValue);
+        }
+
+        return arrValues;
+    }
+
+    private Object mapVertexToCollectionEntry(AtlasVertex entityVertex, AtlasType arrayElement, Object value, String edgeLabel,
+                                              AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute) throws AtlasBaseException {
+        Object ret = null;
+
+        switch (arrayElement.getTypeCategory()) {
+            case PRIMITIVE:
+            case ENUM:
+                ret = value;
+                break;
+
+            case ARRAY:
+            case MAP:
+            case CLASSIFICATION:
+                break;
+
+            case STRUCT:
+                ret = mapVertexToStruct(entityVertex, edgeLabel, (AtlasEdge) value, entityExtInfo);
+                break;
+
+            case ENTITY:
+                ret = mapVertexToObjectId(entityVertex, edgeLabel, (AtlasEdge) value, entityExtInfo, isOwnedAttribute);
+                break;
+
+            default:
+                break;
+        }
+
+        return ret;
+    }
+
+    private Object mapVertexToPrimitive(AtlasVertex entityVertex, final String vertexPropertyName, AtlasAttributeDef attrDef) {
+        Object ret = null;
+
+        if (GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Object.class) == null) {
+            return null;
+        }
+
+        switch (attrDef.getTypeName().toLowerCase()) {
+            case ATLAS_TYPE_STRING:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, String.class);
+                break;
+            case ATLAS_TYPE_SHORT:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Short.class);
+                break;
+            case ATLAS_TYPE_INT:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Integer.class);
+                break;
+            case ATLAS_TYPE_BIGINTEGER:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, BigInteger.class);
+                break;
+            case ATLAS_TYPE_BOOLEAN:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Boolean.class);
+                break;
+            case ATLAS_TYPE_BYTE:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Byte.class);
+                break;
+            case ATLAS_TYPE_LONG:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Long.class);
+                break;
+            case ATLAS_TYPE_FLOAT:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Float.class);
+                break;
+            case ATLAS_TYPE_DOUBLE:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Double.class);
+                break;
+            case ATLAS_TYPE_BIGDECIMAL:
+                ret = GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, BigDecimal.class);
+                break;
+            case ATLAS_TYPE_DATE:
+                ret = new Date(GraphHelper.getSingleValuedProperty(entityVertex, vertexPropertyName, Long.class));
+                break;
+            default:
+                break;
+        }
+
+        return ret;
+    }
+
+    private AtlasObjectId mapVertexToObjectId(AtlasVertex entityVertex, final String edgeLabel, final AtlasEdge optionalEdge,
+                                              AtlasEntityExtInfo entityExtInfo, boolean isOwnedAttribute) throws AtlasBaseException {
+        AtlasObjectId ret = new AtlasObjectId();
+        AtlasEdge     edge;
+
+        if (optionalEdge == null) {
+            edge = graphHelper.getEdgeForLabel(entityVertex, edgeLabel);
+        } else {
+            edge = optionalEdge;
+        }
+
+        if (GraphHelper.elementExists(edge)) {
+            final AtlasVertex referenceVertex = edge.getInVertex();
+            final String      guid            = GraphHelper.getSingleValuedProperty(referenceVertex, Constants.GUID_PROPERTY_KEY, String.class);
+            final String      typeName        = GraphHelper.getTypeName(referenceVertex);
+
+            if (entityExtInfo != null && isOwnedAttribute) {
+                mapVertexToAtlasEntity(guid, entityExtInfo);
+            }
+
+            ret = new AtlasObjectId(typeName, guid);
+        }
+
+        return ret;
+    }
+
+    private AtlasStruct mapVertexToStruct(AtlasVertex entityVertex, final String edgeLabel, final AtlasEdge optionalEdge,
+                                          AtlasEntityExtInfo entityExtInfo) throws AtlasBaseException {
+        AtlasStruct ret = null;
+        AtlasEdge   edge;
+
+        if (optionalEdge == null) {
+            edge = graphHelper.getEdgeForLabel(entityVertex, edgeLabel);
+        } else {
+            edge = optionalEdge;
+        }
+
+        if (GraphHelper.elementExists(edge)) {
+            final AtlasVertex referenceVertex = edge.getInVertex();
+            ret = new AtlasStruct(GraphHelper.getSingleValuedProperty(referenceVertex, Constants.ENTITY_TYPE_PROPERTY_KEY, String.class));
+
+            mapAttributes(referenceVertex, ret, entityExtInfo);
+        }
+
+        return ret;
+    }
+}
\ No newline at end of file

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

@@ -35,8 +35,6 @@ import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
-import org.apache.atlas.repository.store.graph.EntityGraphDiscovery;
-import org.apache.atlas.repository.store.graph.EntityResolver;
 import org.apache.atlas.services.MetadataService;
 import org.apache.atlas.store.AtlasTypeDefStore;
 import org.apache.atlas.type.AtlasArrayType;
@@ -135,7 +133,7 @@ public class AtlasEntityStoreV1Test {
         MapVertexMapper mapVertexMapper = new MapVertexMapper(deleteHandler);
 
 
-        entityStore = new AtlasEntityStoreV1(new EntityGraphMapper(arrVertexMapper, mapVertexMapper, deleteHandler));
+        entityStore = new AtlasEntityStoreV1(new EntityGraphMapper(arrVertexMapper, mapVertexMapper, deleteHandler), new GraphEntityMapper());
         entityStore.init(typeRegistry);
 
         RequestContextV1.clear();

server-api/src/main/java/org/apache/atlas/services/MetadataService.java

@@ -29,7 +29,6 @@ import org.apache.atlas.typesystem.Referenceable;
 import org.apache.atlas.typesystem.Struct;
 import org.apache.atlas.typesystem.IStruct;
 import org.apache.atlas.typesystem.types.cache.TypeCache;
-import org.apache.atlas.utils.ParamChecker;
 import org.codehaus.jettison.json.JSONObject;
 
 import java.util.List;
@@ -49,7 +48,7 @@ public interface MetadataService {
      */
     JSONObject createType(String typeDefinition) throws AtlasException;
 
-    /**
+    /**z
      * Updates the given types in the type definition
      * @param typeDefinition
      * @return

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

@@ -34,12 +34,8 @@ import org.apache.atlas.model.SearchFilter;
 import org.apache.atlas.model.discovery.AtlasSearchResult;
 import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasFullTextResult;
 import org.apache.atlas.model.discovery.AtlasSearchResult.AttributeSearchResult;
-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.AtlasEntityHeaderWithAssociations;
-import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
-import org.apache.atlas.model.instance.EntityMutationResponse;
+import org.apache.atlas.model.instance.*;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
 import org.apache.atlas.model.lineage.AtlasLineageInfo;
 import org.apache.atlas.model.lineage.AtlasLineageInfo.LineageDirection;
@@ -57,11 +53,7 @@ import org.apache.commons.configuration.Configuration;
 import org.apache.commons.lang.ArrayUtils;
 
 import javax.ws.rs.core.MultivaluedMap;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef.CONSTRAINT_TYPE_INVERSE_REF;
 import static org.apache.atlas.model.typedef.AtlasStructDef.AtlasConstraintDef.CONSTRAINT_TYPE_OWNED_REF;
@@ -227,7 +219,7 @@ public class QuickStartV2 {
                                   AtlasTypeUtil.createOptionalAttrDef("dataType", "string"),
                                   AtlasTypeUtil.createOptionalAttrDef("comment", "string"),
                                   AtlasTypeUtil.createOptionalAttrDefWithConstraint("table", TABLE_TYPE, CONSTRAINT_TYPE_INVERSE_REF,
-                                          new HashMap<String, Object>() {{ put(CONSTRAINT_PARAM_ATTRIBUTE, "table"); }}));
+                                          new HashMap<String, Object>() {{ put(CONSTRAINT_PARAM_ATTRIBUTE, "columns"); }}));
 
         colType.setOptions(new HashMap<String, String>() {{ put("schemaAttributes", "[\"name\", \"description\", \"owner\", \"type\", \"comment\", \"position\"]"); }});
 
@@ -353,8 +345,8 @@ public class QuickStartV2 {
         List<AtlasEntityHeader> entities = response.getEntitiesByOperation(EntityOperation.CREATE);
 
         if (CollectionUtils.isNotEmpty(entities)) {
-            List<AtlasEntityWithAssociations> getByGuidResponse = entitiesClient.getEntityByGuid(entities.get(0).getGuid()); 
-            ret = getByGuidResponse.get(0);
+            AtlasEntityWithExtInfo getByGuidResponse = entitiesClient.getEntityByGuid(entities.get(0).getGuid());
+            ret = getByGuidResponse.getEntity();
             System.out.println("Created entity of type [" + ret.getTypeName() + "], guid: " + ret.getGuid());
         }
 
@@ -367,7 +359,7 @@ public class QuickStartV2 {
 
     AtlasEntity createDatabase(String name, String description, String owner, String locationUri, String... traitNames)
             throws Exception {
-        AtlasEntityWithAssociations entity = new AtlasEntityWithAssociations(DATABASE_TYPE);
+        AtlasEntity entity = new AtlasEntity(DATABASE_TYPE);
 
         entity.setClassifications(toAtlasClassifications(traitNames));
         entity.setAttribute("name", name);
@@ -406,7 +398,7 @@ public class QuickStartV2 {
 
     AtlasEntity createColumn(String name, String dataType, String comment, String... traitNames) throws Exception {
 
-        AtlasEntityWithAssociations entity = new AtlasEntityWithAssociations(COLUMN_TYPE);
+        AtlasEntity entity = new AtlasEntity(COLUMN_TYPE);
         entity.setClassifications(toAtlasClassifications(traitNames));
         entity.setAttribute("name", name);
         entity.setAttribute("dataType", dataType);
@@ -417,7 +409,7 @@ public class QuickStartV2 {
 
     AtlasEntity createTable(String name, String description, AtlasEntity db, AtlasEntity sd, String owner, String tableType,
                             List<AtlasEntity> columns, String... traitNames) throws Exception {
-        AtlasEntityWithAssociations entity = new AtlasEntityWithAssociations(TABLE_TYPE);
+        AtlasEntity entity = new AtlasEntity(TABLE_TYPE);
 
         entity.setClassifications(toAtlasClassifications(traitNames));
         entity.setAttribute("name", name);
@@ -428,16 +420,16 @@ public class QuickStartV2 {
         entity.setAttribute("createTime", System.currentTimeMillis());
         entity.setAttribute("lastAccessTime", System.currentTimeMillis());
         entity.setAttribute("retention", System.currentTimeMillis());
-        entity.setAttribute("db", db);
-        entity.setAttribute("sd", sd);
-        entity.setAttribute("columns", columns);
+        entity.setAttribute("db", db.getAtlasObjectId());
+        entity.setAttribute("sd", sd.getAtlasObjectId());
+        entity.setAttribute("columns", getObjectIds(columns));
 
         return createInstance(entity, traitNames);
     }
 
     AtlasEntity createProcess(String name, String description, String user, List<AtlasEntity> inputs, List<AtlasEntity> outputs,
             String queryText, String queryPlan, String queryId, String queryGraph, String... traitNames) throws Exception {
-        AtlasEntityWithAssociations entity = new AtlasEntityWithAssociations(LOAD_PROCESS_TYPE);
+        AtlasEntity entity = new AtlasEntity(LOAD_PROCESS_TYPE);
 
         entity.setClassifications(toAtlasClassifications(traitNames));
         entity.setAttribute(AtlasClient.NAME, name);
@@ -457,7 +449,7 @@ public class QuickStartV2 {
     }
 
     AtlasEntity createView(String name, AtlasEntity db, List<AtlasEntity> inputTables, String... traitNames) throws Exception {
-        AtlasEntityWithAssociations entity = new AtlasEntityWithAssociations(VIEW_TYPE);
+        AtlasEntity entity = new AtlasEntity(VIEW_TYPE);
 
         entity.setClassifications(toAtlasClassifications(traitNames));
         entity.setAttribute("name", name);
@@ -531,7 +523,7 @@ public class QuickStartV2 {
             AtlasSearchResult results = discoveryClient.dslSearchWithParams(dslQuery, 10, 0);
 
             if (results != null) {
-                List<AtlasEntityHeaderWithAssociations> entitiesResult  = results.getEntities();
+                List<AtlasEntityHeader>   entitiesResult  = results.getEntities();
                 List<AtlasFullTextResult> fullTextResults = results.getFullTextResult();
                 AttributeSearchResult     attribResult    = results.getAttributes();
 
@@ -565,7 +557,24 @@ public class QuickStartV2 {
     }
 
     private String getTableId(String tableName) throws AtlasServiceException {
-        AtlasEntity tableEntity = entitiesClient.getEntityByAttribute(TABLE_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName).get(0);
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName);
+
+        AtlasEntity tableEntity = entitiesClient.getEntityByAttribute(TABLE_TYPE, attributes).getEntity();
         return tableEntity.getGuid();
     }
+
+    private Collection<AtlasObjectId> getObjectIds(Collection<AtlasEntity> entities) {
+        List<AtlasObjectId> ret = new ArrayList<>();
+
+        if (CollectionUtils.isNotEmpty(entities)) {
+            for (AtlasEntity entity : entities) {
+                if (entity != null) {
+                    ret.add(entity.getAtlasObjectId());
+                }
+            }
+        }
+
+        return ret;
+    }
 }

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

@@ -24,7 +24,6 @@ import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntity.Status;
-import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasType;
@@ -32,7 +31,6 @@ import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.typesystem.IReferenceableInstance;
 import org.apache.atlas.typesystem.IStruct;
 import org.apache.atlas.typesystem.Referenceable;
-import org.apache.atlas.typesystem.persistence.AtlasSystemAttributes;
 import org.apache.atlas.typesystem.persistence.Id;
 import org.apache.atlas.typesystem.persistence.Id.EntityState;
 import org.apache.commons.collections.CollectionUtils;
@@ -77,7 +75,7 @@ public class AtlasEntityFormatConverter extends AtlasStructFormatConverter {
                         LOG.error("IReferenceableInstance.getValuesMap() failed", excp);
                     }
 
-                    AtlasEntityWithAssociations entity = new AtlasEntityWithAssociations(entRef.getTypeName(),
+                    AtlasEntity entity = new AtlasEntity(entRef.getTypeName(),
                                                          super.fromV1ToV2(entityType, v1Attribs, context));
                     entity.setGuid(entRef.getId()._getId());
                     entity.setStatus(convertState(entRef.getId().getState()));

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

@@ -21,7 +21,6 @@ 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;
@@ -36,24 +35,16 @@ public interface AtlasFormatConverter {
 
     public static class ConverterContext {
 
-        private Map<String, AtlasEntityWithAssociations> entities = null;
+        private Map<String, AtlasEntity> entities = null;
 
-        public void addEntity(AtlasEntityWithAssociations entity) {
+        public void addEntity(AtlasEntity entity) {
             if (entities == null) {
                 entities = new HashMap<>();
             }
             entities.put(entity.getGuid(), entity);
         }
 
-        public void addEntity(AtlasEntity entity) {
-            if (entity instanceof AtlasEntityWithAssociations) {
-                this.addEntity((AtlasEntityWithAssociations)entity);
-            } else {
-                this.addEntity(new AtlasEntityWithAssociations(entity));
-            }
-        }
-
-        public AtlasEntityWithAssociations getById(String guid) {
+        public AtlasEntity getById(String guid) {
             if( entities != null) {
                 return entities.get(guid);
             }
@@ -63,7 +54,7 @@ public interface AtlasFormatConverter {
 
         public boolean entityExists(String guid) { return entities != null && entities.containsKey(guid); }
 
-        public Map<String, AtlasEntityWithAssociations> getEntities() {
+        public Map<String, AtlasEntity> getEntities() {
             return entities;
         }
     }

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

@@ -17,8 +17,6 @@
  */
 package org.apache.atlas.web.adapters;
 
-import java.util.List;
-
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.AtlasException;
@@ -28,7 +26,6 @@ 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.AtlasEntityHeader;
-import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
 import org.apache.atlas.model.instance.GuidMapping;
@@ -50,7 +47,6 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.Iterator;
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
@@ -129,7 +125,7 @@ public class AtlasInstanceRestAdapters {
         return ret;
     }
 
-    public Map<String, AtlasEntityWithAssociations> getAtlasEntity(IReferenceableInstance referenceable) throws AtlasBaseException {
+    public Map<String, AtlasEntity> getAtlasEntity(IReferenceableInstance referenceable) throws AtlasBaseException {
 
         AtlasFormatConverter converter  = instanceFormatters.getConverter(TypeCategory.ENTITY);
         AtlasEntityType      entityType = typeRegistry.getEntityTypeByName(referenceable.getTypeName());

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

@@ -22,11 +22,8 @@ import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.CreateUpdateEntitiesResult;
 import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.model.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;
@@ -54,9 +51,7 @@ import javax.ws.rs.QueryParam;
 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;
 
@@ -139,7 +134,7 @@ public class EntitiesREST {
         for (String guid : guids) {
             try {
                ITypedReferenceableInstance ref       = metadataService.getEntityDefinition(guid);
-               Map<String, AtlasEntityWithAssociations> entityRet = restAdapters.getAtlasEntity(ref);
+               Map<String, AtlasEntity>    entityRet = restAdapters.getAtlasEntity(ref);
 
                addToEntityList(entityList, entityRet.values());
 
@@ -152,8 +147,8 @@ public class EntitiesREST {
         return entities;
     }
 
-    private void addToEntityList(final List<AtlasEntity> entityList, final Collection<AtlasEntityWithAssociations> values) {
-        for (AtlasEntityWithAssociations val : values) {
+    private void addToEntityList(final List<AtlasEntity> entityList, final Collection<AtlasEntity> values) {
+        for (AtlasEntity val : values) {
             if ( !entityList.contains(val)) {
                 entityList.add(val);
             }

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

@@ -25,10 +25,12 @@ import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
+import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.services.MetadataService;
+import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.type.AtlasTypeRegistry;
@@ -39,15 +41,19 @@ import org.apache.atlas.typesystem.Referenceable;
 import org.apache.atlas.web.adapters.AtlasFormatConverter;
 import org.apache.atlas.web.adapters.AtlasInstanceRestAdapters;
 import org.apache.atlas.web.util.Servlets;
+import org.apache.commons.collections.MapUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import javax.inject.Inject;
 import javax.inject.Singleton;
+import javax.servlet.http.HttpServletRequest;
 import javax.ws.rs.*;
+import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -63,49 +69,54 @@ public class EntityREST {
 
     private static final Logger LOG = LoggerFactory.getLogger(EntityREST.class);
 
-    private final AtlasTypeRegistry typeRegistry;
+    public static final String PREFIX_ATTR = "attr:";
 
+    private final AtlasTypeRegistry         typeRegistry;
     private final AtlasInstanceRestAdapters restAdapters;
-
     private final MetadataService           metadataService;
+    private final AtlasEntityStore          entitiesStore;
 
     @Inject
-    public EntityREST(AtlasTypeRegistry typeRegistry, AtlasInstanceRestAdapters restAdapters, MetadataService metadataService) {
+    public EntityREST(AtlasTypeRegistry typeRegistry, AtlasInstanceRestAdapters restAdapters, MetadataService metadataService, AtlasEntityStore entitiesStore) {
         this.typeRegistry    = typeRegistry;
         this.restAdapters    = restAdapters;
         this.metadataService = metadataService;
+        this.entitiesStore   = entitiesStore;
     }
 
     /**
      * Fetch the complete definition of an entity given its GUID.
      *
      * @param guid GUID for the entity
+     * @return AtlasEntity
+     * @throws AtlasBaseException
      */
     @GET
     @Path("/guid/{guid}")
+    @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    public List<AtlasEntityWithAssociations> getById(@PathParam("guid") String guid) throws AtlasBaseException {
-        try {
-            ITypedReferenceableInstance              ref      = metadataService.getEntityDefinition(guid);
-            Map<String, AtlasEntityWithAssociations> entities = restAdapters.getAtlasEntity(ref);
-
-            return getOrderedEntityList(entities, guid);
-        } catch (AtlasException e) {
-            throw toAtlasBaseException(e);
-        }
+    public AtlasEntityWithExtInfo getById(@PathParam("guid") String guid) throws AtlasBaseException {
+        return entitiesStore.getById(guid);
     }
 
     /**
-     * Fetch the complete definition of an entity given its GUID including its associations
-     * like classifications, terms etc.
-     *
-     * @param guid GUID for the entity
+     * Get entity information using entity type and unique attribute.
+     * @param typeName
+     * @return
+     * @throws AtlasBaseException
      */
     @GET
-    @Path("/guid/{guid}/associations")
+    @Path("/uniqueAttribute/type/{typeName}")
+    @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    public List<AtlasEntityWithAssociations> getWithAssociationsByGuid(@PathParam("guid") String guid) throws AtlasBaseException {
-        return this.getById(guid);
+    public AtlasEntityWithExtInfo getByUniqueAttributes(@PathParam("typeName") String typeName,
+                                                        @Context HttpServletRequest servletRequest) throws AtlasBaseException {
+        AtlasEntityType     entityType = ensureEntityType(typeName);
+        Map<String, Object> attributes = getAttributes(servletRequest);
+
+        validateAttributes(entityType, attributes);
+
+        return entitiesStore.getByUniqueAttribute(entityType, attributes);
     }
 
     /**
@@ -161,14 +172,15 @@ public class EntityREST {
     @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
     @Path("/uniqueAttribute/type/{typeName}/attribute/{attrName}")
-    public EntityMutationResponse deleteByUniqueAttribute(@PathParam("typeName") String entityType,
+    public EntityMutationResponse deleteByUniqueAttribute(@PathParam("typeName") String typeName,
         @PathParam("attrName") String attribute,
         @QueryParam("value") String value) throws Exception {
 
-        AtlasEntityType type = (AtlasEntityType) validateType(entityType, TypeCategory.ENTITY);
-        validateUniqueAttribute(type, attribute);
+        AtlasEntityType entityType = ensureEntityType(typeName);
+
+        validateUniqueAttribute(entityType, attribute);
 
-        final AtlasClient.EntityResult result = metadataService.deleteEntityByUniqueAttribute(entityType, attribute, value);
+        final AtlasClient.EntityResult result = metadataService.deleteEntityByUniqueAttribute(typeName, attribute, value);
         return toEntityMutationResponse(result);
     }
 
@@ -181,17 +193,18 @@ public class EntityREST {
     @Consumes({Servlets.JSON_MEDIA_TYPE, MediaType.APPLICATION_JSON})
     @Produces(Servlets.JSON_MEDIA_TYPE)
     @Path("/uniqueAttribute/type/{typeName}/attribute/{attrName}")
-    public List<AtlasEntityWithAssociations> getByUniqueAttribute(@PathParam("typeName") String entityType,
+    public List<AtlasEntity> getByUniqueAttribute(@PathParam("typeName") String typeName,
         @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);
+        List<AtlasEntity> entityList = new ArrayList<>();
+        AtlasEntityType                   entityType = ensureEntityType(typeName);
+
+        validateUniqueAttribute(entityType, attribute);
 
         try {
-            final ITypedReferenceableInstance entityDefinitionReference = metadataService.getEntityDefinitionReference(entityType, attribute, value);
-            Map<String, AtlasEntityWithAssociations> entityRet = restAdapters.getAtlasEntity(entityDefinitionReference);
+            final ITypedReferenceableInstance entityDefinitionReference = metadataService.getEntityDefinitionReference(typeName, attribute, value);
+            Map<String, AtlasEntity> entityRet = restAdapters.getAtlasEntity(entityDefinitionReference);
             entityList.addAll(entityRet.values());
         } catch (AtlasException e) {
             throw toAtlasBaseException(e);
@@ -209,16 +222,16 @@ public class EntityREST {
     @GET
     @Path("/guid/{guid}/classification/{classificationName}")
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    public AtlasClassification getClassification(@PathParam("guid") String guid, @PathParam("classificationName") String classificationName) throws AtlasBaseException {
+    public AtlasClassification getClassification(@PathParam("guid") String guid, @PathParam("classificationName") String typeName) throws AtlasBaseException {
 
         if (StringUtils.isEmpty(guid)) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
         }
 
-        validateType(classificationName, TypeCategory.CLASSIFICATION);
+        ensureClassificationType(typeName);
 
         try {
-            IStruct trait = metadataService.getTraitDefinition(guid, classificationName);
+            IStruct trait = metadataService.getTraitDefinition(guid, typeName);
             return restAdapters.getClassification(trait);
 
         } catch (AtlasException e) {
@@ -313,23 +326,23 @@ public class EntityREST {
      * Deletes a given classification from an existing entity represented by a guid.
      *
      * @param guid      globally unique identifier for the entity
-     * @param classificationName name of the trait
+     * @param typeName name of the trait
      */
     @DELETE
     @Path("/guid/{guid}/classification/{classificationName}")
     @Consumes({Servlets.JSON_MEDIA_TYPE, MediaType.APPLICATION_JSON})
     @Produces(Servlets.JSON_MEDIA_TYPE)
     public void deleteClassification(@PathParam("guid") String guid,
-        @PathParam("classificationName") String classificationName) throws AtlasBaseException {
+        @PathParam("classificationName") String typeName) throws AtlasBaseException {
 
         if (StringUtils.isEmpty(guid)) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
         }
 
-        validateType(classificationName, TypeCategory.CLASSIFICATION);
+        ensureClassificationType(typeName);
 
         try {
-            metadataService.deleteTrait(guid, classificationName);
+            metadataService.deleteTrait(guid, typeName);
         } catch (AtlasException e) {
             throw toAtlasBaseException(e);
         }
@@ -348,7 +361,28 @@ public class EntityREST {
         return type;
     }
 
+    private AtlasEntityType ensureEntityType(String typeName) throws AtlasBaseException {
+        AtlasEntityType ret = typeRegistry.getEntityTypeByName(typeName);
+
+        if (ret == null) {
+            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), typeName);
+        }
+
+        return ret;
+    }
+
+    private AtlasClassificationType ensureClassificationType(String typeName) throws AtlasBaseException {
+        AtlasClassificationType ret = typeRegistry.getClassificationTypeByName(typeName);
+
+        if (ret == null) {
+            throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.CLASSIFICATION.name(), typeName);
+        }
+
+        return ret;
+    }
+
     /**
+     * Deprecated method - not used
      * Validate that attribute is unique attribute
      * @param entityType     the entity type
      * @param attributeName  the name of the attribute
@@ -361,10 +395,10 @@ public class EntityREST {
         }
     }
 
-    private List<AtlasEntityWithAssociations> getOrderedEntityList(Map<String, AtlasEntityWithAssociations> entities, String firstItemGuid) {
-        List<AtlasEntityWithAssociations> ret = new ArrayList<>(entities.size());
+    private List<AtlasEntity> getOrderedEntityList(Map<String, AtlasEntity> entities, String firstItemGuid) {
+        List<AtlasEntity> ret = new ArrayList<>(entities.size());
 
-        for (AtlasEntityWithAssociations entity : entities.values()) {
+        for (AtlasEntity entity : entities.values()) {
             if (StringUtils.equals(entity.getGuid(), firstItemGuid)) {
                 ret.add(0, entity);
             } else {
@@ -374,4 +408,39 @@ public class EntityREST {
 
         return ret;
     }
+
+    private void validateAttributes(AtlasEntityType entityType, Map<String, Object> attributes)
+            throws AtlasBaseException {
+
+        if (MapUtils.isEmpty(attributes)) {
+            throw new AtlasBaseException(AtlasErrorCode.ATTRIBUTE_UNIQUE_INVALID, entityType.getTypeName(), "");
+        }
+
+        for (String attrName : attributes.keySet()) {
+            AtlasAttributeDef attrDef = entityType.getAttributeDef(attrName);
+
+            if (attrDef == null || !attrDef.getIsUnique()) {
+                throw new AtlasBaseException(AtlasErrorCode.ATTRIBUTE_UNIQUE_INVALID, entityType.getTypeName(), attrName);
+            }
+        }
+    }
+
+    private Map<String, Object> getAttributes(HttpServletRequest request) {
+        Map<String, Object> attributes = new HashMap<>();
+
+        if (MapUtils.isNotEmpty(request.getParameterMap())) {
+            for (Map.Entry<String, String[]> e : request.getParameterMap().entrySet()) {
+                String key = e.getKey();
+
+                if (key != null && key.startsWith(PREFIX_ATTR)) {
+                    String[] values = e.getValue();
+                    String   value  = values != null && values.length > 0 ? values[0] : null;
+
+                    attributes.put(key.substring(PREFIX_ATTR.length()), value);
+                }
+            }
+        }
+
+        return attributes;
+    }
 }

webapp/src/test/java/org/apache/atlas/examples/QuickStartV2IT.java

@@ -32,6 +32,7 @@ import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -56,7 +57,9 @@ public class QuickStartV2IT extends BaseResourceIT {
     }
 
     private AtlasEntity getDB(String dbName) throws AtlasServiceException, JSONException {
-        AtlasEntity dbEntity = entitiesClientV2.getEntityByAttribute(QuickStartV2.DATABASE_TYPE, "name", dbName).get(0);
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put("name", dbName);
+        AtlasEntity dbEntity = entitiesClientV2.getEntityByAttribute(QuickStartV2.DATABASE_TYPE, attributes).getEntity();
         return dbEntity;
     }
 
@@ -73,12 +76,16 @@ public class QuickStartV2IT extends BaseResourceIT {
     }
 
     private AtlasEntity getTable(String tableName) throws AtlasServiceException {
-        AtlasEntity tableEntity = entitiesClientV2.getEntityByAttribute(QuickStartV2.TABLE_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName).get(0);
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName);
+        AtlasEntity tableEntity = entitiesClientV2.getEntityByAttribute(QuickStartV2.TABLE_TYPE, attributes).getEntity();
         return tableEntity;
     }
 
     private AtlasEntity getProcess(String processName) throws AtlasServiceException {
-        AtlasEntity processEntity = entitiesClientV2.getEntityByAttribute(QuickStartV2.LOAD_PROCESS_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, processName).get(0);
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, processName);
+        AtlasEntity processEntity = entitiesClientV2.getEntityByAttribute(QuickStartV2.LOAD_PROCESS_TYPE, attributes).getEntity();
         return processEntity;
     }
      
@@ -115,8 +122,9 @@ public class QuickStartV2IT extends BaseResourceIT {
 
     @Test
     public void testProcessIsAdded() throws AtlasServiceException, JSONException {
-        AtlasEntity loadProcess = entitiesClientV2.getEntityByAttribute(QuickStartV2.LOAD_PROCESS_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
-                QuickStartV2.LOAD_SALES_DAILY_PROCESS).get(0);
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, QuickStartV2.LOAD_SALES_DAILY_PROCESS);
+        AtlasEntity loadProcess = entitiesClientV2.getEntityByAttribute(QuickStartV2.LOAD_PROCESS_TYPE, attributes).getEntity();
 
         Map loadProcessAttribs = loadProcess.getAttributes();
         assertEquals(QuickStartV2.LOAD_SALES_DAILY_PROCESS, loadProcessAttribs.get(AtlasClient.NAME));
@@ -169,7 +177,9 @@ public class QuickStartV2IT extends BaseResourceIT {
 
     @Test
     public void testViewIsAdded() throws AtlasServiceException, JSONException {
-        AtlasEntity view = entitiesClientV2.getEntityByAttribute(QuickStartV2.VIEW_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, QuickStartV2.PRODUCT_DIM_VIEW).get(0);
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, QuickStartV2.PRODUCT_DIM_VIEW);
+        AtlasEntity view = entitiesClientV2.getEntityByAttribute(QuickStartV2.VIEW_TYPE, attributes).getEntity();
         Map<String, Object> viewAttributes = view.getAttributes();
         assertEquals(QuickStartV2.PRODUCT_DIM_VIEW, viewAttributes.get(AtlasClient.NAME));
 

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

@@ -22,8 +22,8 @@ import org.apache.atlas.RequestContext;
 import org.apache.atlas.TestUtilsV2;
 import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
-import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
@@ -98,10 +98,11 @@ public class TestEntityREST {
     @Test
     public void testGetEntityById() throws Exception {
         createOrUpdateEntity();
-        final List<AtlasEntityWithAssociations> response = entityREST.getById(dbGuid);
+        AtlasEntityWithExtInfo response = entityREST.getById(dbGuid);
 
         Assert.assertNotNull(response);
-        TestEntitiesREST.verifyAttributes(response.get(0).getAttributes(), dbEntity.getAttributes());
+        Assert.assertNotNull(response.getEntity());
+        TestEntitiesREST.verifyAttributes(response.getEntity().getAttributes(), dbEntity.getAttributes());
     }
 
     @Test
@@ -130,8 +131,8 @@ public class TestEntityREST {
     @Test(dependsOnMethods = "testAddAndGetClassification")
     public void  testGetEntityWithAssociations() throws Exception {
 
-        List<AtlasEntityWithAssociations> entity = entityREST.getWithAssociationsByGuid(dbGuid);
-        final List<AtlasClassification> retrievedClassifications = entity.get(0).getClassifications();
+        AtlasEntityWithExtInfo entity = entityREST.getById(dbGuid);
+        final List<AtlasClassification> retrievedClassifications = entity.getEntity().getClassifications();
 
         Assert.assertNotNull(retrievedClassifications);
         Assert.assertEquals(new ArrayList<AtlasClassification>() {{ add(testClassification); }}, retrievedClassifications);
@@ -172,7 +173,7 @@ public class TestEntityREST {
         Assert.assertTrue(AtlasEntity.isAssigned(dbGuid));
 
         //Get By unique attribute
-        List<AtlasEntityWithAssociations> entities = entityREST.getByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, TestUtilsV2.NAME, updatedDBName);
+        List<AtlasEntity> entities = entityREST.getByUniqueAttribute(TestUtilsV2.DATABASE_TYPE, TestUtilsV2.NAME, updatedDBName);
         Assert.assertNotNull(entities);
         Assert.assertNotNull(entities.get(0).getGuid());
         Assert.assertEquals(entities.get(0).getGuid(), dbGuid);

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

@@ -40,7 +40,6 @@ import org.apache.atlas.AtlasTypedefClientV2;
 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.AtlasStruct;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
@@ -561,9 +560,8 @@ public abstract class BaseResourceIT {
         return tableInstance;
     }
 
-    protected AtlasEntityWithAssociations createHiveTableInstanceV2(AtlasEntity databaseInstance, String tableName) throws Exception {
-        AtlasEntityWithAssociations tableInstance =
-                new AtlasEntityWithAssociations(HIVE_TABLE_TYPE_V2);
+    protected AtlasEntity createHiveTableInstanceV2(AtlasEntity databaseInstance, String tableName) throws Exception {
+        AtlasEntity tableInstance = new AtlasEntity(HIVE_TABLE_TYPE_V2);
         tableInstance.setClassifications(
                 Arrays.asList(new AtlasClassification("classification"),
                         new AtlasClassification("pii"),

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

@@ -26,7 +26,7 @@ import org.apache.atlas.model.discovery.AtlasSearchResult;
 import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasFullTextResult;
 import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasQueryType;
 import org.apache.atlas.model.instance.AtlasEntity.Status;
-import org.apache.atlas.model.instance.AtlasEntityHeaderWithAssociations;
+import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.typesystem.TypesDef;
 import org.apache.atlas.typesystem.types.ClassType;
 import org.apache.atlas.typesystem.types.DataTypes;
@@ -68,11 +68,11 @@ public class EntityDiscoveryJerseyResourceIT extends BaseResourceIT {
         assertEquals(searchResult.getQueryText(), dslQuery);
         assertEquals(searchResult.getQueryType(), AtlasQueryType.DSL);
 
-        List<AtlasEntityHeaderWithAssociations> entities = searchResult.getEntities();
+        List<AtlasEntityHeader> entities = searchResult.getEntities();
         assertNotNull(entities);
         assertEquals(entities.size(), 1);
 
-        AtlasEntityHeaderWithAssociations dbEntity = entities.get(0);
+        AtlasEntityHeader dbEntity = entities.get(0);
         assertEquals(dbEntity.getTypeName(), DATABASE_TYPE_BUILTIN);
         assertEquals(dbEntity.getDisplayText(), dbName);
         assertEquals(dbEntity.getStatus(), Status.ACTIVE);
@@ -130,11 +130,11 @@ public class EntityDiscoveryJerseyResourceIT extends BaseResourceIT {
 
         assertEquals(searchResult.getQueryText(), query);
         assertEquals(searchResult.getQueryType(), AtlasQueryType.DSL);
-        List<AtlasEntityHeaderWithAssociations> entities = searchResult.getEntities();
+        List<AtlasEntityHeader> entities = searchResult.getEntities();
         assertNotNull(entities);
         assertEquals(entities.size(), 1);
 
-        AtlasEntityHeaderWithAssociations dbEntity = entities.get(0);
+        AtlasEntityHeader dbEntity = entities.get(0);
         assertEquals(dbEntity.getTypeName(), DATABASE_TYPE_BUILTIN);
         assertEquals(dbEntity.getDisplayText(), dbName);
         assertEquals(dbEntity.getStatus(), Status.ACTIVE);

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

@@ -37,7 +37,6 @@ import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasClassification.AtlasClassifications;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
-import org.apache.atlas.model.instance.AtlasEntityWithAssociations;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations;
 import org.apache.atlas.model.typedef.AtlasClassificationDef;
@@ -80,7 +79,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
     private String traitName;
 
     private AtlasEntity dbEntity;
-    private AtlasEntityWithAssociations tableEntity;
+    private AtlasEntity tableEntity;
     @Inject
     private NotificationInterface notificationInterface;
     private NotificationConsumer<EntityNotification> notificationConsumer;
@@ -187,7 +186,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
 
         //Test the same across references
         final String tableName = randomString();
-        AtlasEntityWithAssociations hiveTableInstanceV2 = createHiveTableInstanceV2(hiveDBInstanceV2, tableName);
+        AtlasEntity hiveTableInstanceV2 = createHiveTableInstanceV2(hiveDBInstanceV2, tableName);
         hiveTableInstanceV2.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName);
 
         EntityMutationResponse entity = entitiesClientV2.createEntity(hiveTableInstanceV2);
@@ -273,7 +272,9 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         AtlasEntity hiveDB = createHiveDB();
         String qualifiedName = (String) hiveDB.getAttribute(NAME);
         //get entity by attribute
-        AtlasEntity byAttribute = entitiesClientV2.getEntityByAttribute(DATABASE_TYPE_V2, NAME, qualifiedName).get(0);
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put(NAME, qualifiedName);
+        AtlasEntity byAttribute = entitiesClientV2.getEntityByAttribute(DATABASE_TYPE_V2, attributes).getEntity();
         assertEquals(byAttribute.getTypeName(), DATABASE_TYPE_V2);
         assertEquals(byAttribute.getAttribute(NAME), qualifiedName);
     }
@@ -427,13 +428,13 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         return hiveDBInstanceV2;
     }
 
-    private TypeUtils.Pair<AtlasEntity, AtlasEntityWithAssociations> createDBAndTable() throws Exception {
+    private TypeUtils.Pair<AtlasEntity, AtlasEntity> createDBAndTable() throws Exception {
         AtlasEntity dbInstanceV2 = createHiveDB();
-        AtlasEntityWithAssociations hiveTableInstanceV2 = createHiveTable();
+        AtlasEntity hiveTableInstanceV2 = createHiveTable();
         return TypeUtils.Pair.of(dbInstanceV2, hiveTableInstanceV2);
     }
 
-    private AtlasEntityWithAssociations createHiveTable() throws Exception {
+    private AtlasEntity createHiveTable() throws Exception {
         if (tableEntity == null) {
             tableEntity = createHiveTable(createHiveDB(), TABLE_NAME);
         }
@@ -441,8 +442,8 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
 
     }
 
-    private AtlasEntityWithAssociations createHiveTable(AtlasEntity dbInstanceV2, String tableName) throws Exception {
-        AtlasEntityWithAssociations hiveTableInstanceV2 = createHiveTableInstanceV2(dbInstanceV2, tableName);
+    private AtlasEntity createHiveTable(AtlasEntity dbInstanceV2, String tableName) throws Exception {
+        AtlasEntity hiveTableInstanceV2 = createHiveTableInstanceV2(dbInstanceV2, tableName);
         AtlasEntityHeader createdHeader = createEntity(hiveTableInstanceV2);
         assertNotNull(createdHeader);
         assertNotNull(createdHeader.getGuid());
@@ -478,7 +479,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         AtlasClassificationDef classificationByName = typedefClientV2.getClassificationByName(traitName);
         assertNotNull(classificationByName);
 
-        AtlasEntityWithAssociations hiveTable = createHiveTable();
+        AtlasEntity hiveTable = createHiveTable();
         assertEquals(hiveTable.getClassifications().size(), 7);
 
         AtlasClassification piiClassification = new AtlasClassification(piiTrait.getName());
@@ -509,7 +510,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         entitiesClientV2.addClassifications(guid, ImmutableList.of(traitInstance));
 
         // verify the response
-        AtlasEntityWithAssociations withAssociationByGuid = entitiesClientV2.getEntityWithAssociationByGuid(guid).get(0);
+        AtlasEntity withAssociationByGuid = entitiesClientV2.getEntityByGuid(guid).getEntity();
         assertNotNull(withAssociationByGuid);
         assertFalse(withAssociationByGuid.getClassifications().isEmpty());
 
@@ -620,8 +621,8 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
 
         AtlasEntity ref = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values);
         columns.add(ref);
-        AtlasEntityWithAssociations hiveTable = createHiveTable();
-        AtlasEntityWithAssociations tableUpdated = hiveTable;
+        AtlasEntity hiveTable = createHiveTable();
+        AtlasEntity tableUpdated = hiveTable;
 
         hiveTable.setAttribute("columns", columns);
 
@@ -656,7 +657,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
     }
 
     private AtlasEntity getEntityByGuid(String guid) throws AtlasServiceException {
-        return entitiesClientV2.getEntityByGuid(guid).get(0);
+        return entitiesClientV2.getEntityByGuid(guid).getEntity();
     }
 
     @Test(dependsOnMethods = "testSubmitEntity")
@@ -678,7 +679,7 @@ public class EntityV2JerseyResourceIT extends BaseResourceIT {
         AtlasEntity ref2 = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values2);
         columns.add(ref1);
         columns.add(ref2);
-        AtlasEntityWithAssociations hiveTable = createHiveTable();
+        AtlasEntity hiveTable = createHiveTable();
         hiveTable.setAttribute("columns", columns);
         EntityMutationResponse updateEntityResult = entitiesClientV2.updateEntity(hiveTable);
         assertNotNull(updateEntityResult);