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/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/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/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);