You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@atlas.apache.org by ma...@apache.org on 2017/01/17 03:51:48 UTC

incubator-atlas git commit: ATLAS-1460: v2 search API updated to return name/description/owner and classification names in result

Repository: incubator-atlas
Updated Branches:
  refs/heads/master 53f88222a -> e0c6b98ef


ATLAS-1460: v2 search API updated to return name/description/owner and classification names in result

Signed-off-by: Madhan Neethiraj <ma...@apache.org>


Project: http://git-wip-us.apache.org/repos/asf/incubator-atlas/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-atlas/commit/e0c6b98e
Tree: http://git-wip-us.apache.org/repos/asf/incubator-atlas/tree/e0c6b98e
Diff: http://git-wip-us.apache.org/repos/asf/incubator-atlas/diff/e0c6b98e

Branch: refs/heads/master
Commit: e0c6b98ef274e984ae3b1a61bddf79c5e21aa17c
Parents: 53f8822
Author: Vimal Sharma <sv...@apache.org>
Authored: Mon Jan 16 18:07:31 2017 -0800
Committer: Madhan Neethiraj <ma...@apache.org>
Committed: Mon Jan 16 19:51:30 2017 -0800

----------------------------------------------------------------------
 .../model/discovery/AtlasSearchResult.java      |  36 ++---
 .../AtlasEntityHeaderWithAssociations.java      | 143 +++++++++++++++++++
 release-log.txt                                 |   1 +
 .../atlas/discovery/EntityDiscoveryService.java | 133 +++++++++--------
 .../EntityDiscoveryJerseyResourceIT.java        |  10 +-
 5 files changed, 231 insertions(+), 92 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/e0c6b98e/intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java
----------------------------------------------------------------------
diff --git a/intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java b/intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java
index c4c0063..e0f24b1 100644
--- a/intg/src/main/java/org/apache/atlas/model/discovery/AtlasSearchResult.java
+++ b/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.AtlasEntityHeader;
+import org.apache.atlas.model.instance.AtlasEntityHeaderWithAssociations;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang.StringUtils;
 import org.codehaus.jackson.annotate.JsonAutoDetect;
@@ -42,11 +42,11 @@ import static org.codehaus.jackson.annotate.JsonAutoDetect.Visibility.PUBLIC_ONL
 @XmlRootElement
 @XmlAccessorType(XmlAccessType.PROPERTY)
 public class AtlasSearchResult implements Serializable {
-    private String                    queryText;
-    private AtlasQueryType            queryType;
-    private List<AtlasEntityHeader>   entities;
-    private AttributeSearchResult     attributes;
-    private List<AtlasFullTextResult> fullTextResult;
+    private String                                  queryText;
+    private AtlasQueryType                          queryType;
+    private List<AtlasEntityHeaderWithAssociations> entities;
+    private AttributeSearchResult                   attributes;
+    private List<AtlasFullTextResult>               fullTextResult;
 
     public AtlasSearchResult() {}
 
@@ -66,9 +66,9 @@ public class AtlasSearchResult implements Serializable {
 
     public void setQueryType(AtlasQueryType queryType) { this.queryType = queryType; }
 
-    public List<AtlasEntityHeader> getEntities() { return entities; }
+    public List<AtlasEntityHeaderWithAssociations> getEntities() { return entities; }
 
-    public void setEntities(List<AtlasEntityHeader> entities) { this.entities = entities; }
+    public void setEntities(List<AtlasEntityHeaderWithAssociations> entities) { this.entities = entities; }
 
     public AttributeSearchResult getAttributes() { return attributes; }
 
@@ -104,7 +104,7 @@ public class AtlasSearchResult implements Serializable {
                 '}';
     }
 
-    public void addEntity(AtlasEntityHeader newEntity) {
+    public void addEntity(AtlasEntityHeaderWithAssociations newEntity) {
         if (entities == null) {
             entities = new ArrayList<>();
         }
@@ -117,13 +117,13 @@ public class AtlasSearchResult implements Serializable {
         }
     }
 
-    public void removeEntity(AtlasEntityHeader entity) {
-        List<AtlasEntityHeader> entities = this.entities;
+    public void removeEntity(AtlasEntityHeaderWithAssociations entity) {
+        List<AtlasEntityHeaderWithAssociations> entities = this.entities;
 
         if (CollectionUtils.isNotEmpty(entities)) {
-            Iterator<AtlasEntityHeader> iter = entities.iterator();
+            Iterator<AtlasEntityHeaderWithAssociations> iter = entities.iterator();
             while (iter.hasNext()) {
-                AtlasEntityHeader currEntity = iter.next();
+                AtlasEntityHeaderWithAssociations 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 {
-        AtlasEntityHeader entity;
-        Double            score;
+        AtlasEntityHeaderWithAssociations entity;
+        Double                            score;
 
         public AtlasFullTextResult() {}
 
-        public AtlasFullTextResult(AtlasEntityHeader entity, Double score) {
+        public AtlasFullTextResult(AtlasEntityHeaderWithAssociations entity, Double score) {
             this.entity = entity;
             this.score  = score;
         }
 
-        public AtlasEntityHeader getEntity() { return entity; }
+        public AtlasEntityHeaderWithAssociations getEntity() { return entity; }
 
-        public void setEntity(AtlasEntityHeader entity) { this.entity = entity; }
+        public void setEntity(AtlasEntityHeaderWithAssociations entity) { this.entity = entity; }
 
         public Double getScore() { return score; }
 

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/e0c6b98e/intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityHeaderWithAssociations.java
----------------------------------------------------------------------
diff --git a/intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityHeaderWithAssociations.java b/intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityHeaderWithAssociations.java
new file mode 100644
index 0000000..ed1adb4
--- /dev/null
+++ b/intg/src/main/java/org/apache/atlas/model/instance/AtlasEntityHeaderWithAssociations.java
@@ -0,0 +1,143 @@
+/**
+ * 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);
+        }
+    }
+
+}

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/e0c6b98e/release-log.txt
----------------------------------------------------------------------
diff --git a/release-log.txt b/release-log.txt
index 2a52fcc..369a5c1 100644
--- a/release-log.txt
+++ b/release-log.txt
@@ -9,6 +9,7 @@ ATLAS-1060 Add composite indexes for exact match performance improvements for al
 ATLAS-1127 Modify creation and modification timestamps to Date instead of Long(sumasai)
 
 ALL CHANGES:
+ATLAS-1460 v2 search API updated to return name/description/owner and classification names in result (vimalsharma via mneethiraj)
 ATLAS-1434 fixed unit test to use correct type names; updated error message per review comments (ashutoshm via mneethiraj)
 ATLAS-1391 Add exclusion mechanism for Atlas audit mechanism (guptaneeru via svimal2106)
 ATLAS-1280 Atlas changes to support Hive hook for Hive2 (mneethiraj via svimal2106)

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/e0c6b98e/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
index be9fed2..2be9a2d 100644
--- a/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
+++ b/repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java
@@ -17,6 +17,7 @@
  */
 package org.apache.atlas.discovery;
 
+import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasConfiguration;
 import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasFullTextResult;
 import org.apache.atlas.model.discovery.AtlasSearchResult.AtlasQueryType;
@@ -25,7 +26,8 @@ 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.AtlasEntityHeader;
+import org.apache.atlas.model.instance.AtlasEntityHeaderWithAssociations;
+import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.query.Expressions.AliasExpression;
 import org.apache.atlas.query.Expressions.Expression;
 import org.apache.atlas.query.Expressions.SelectExpression;
@@ -44,7 +46,6 @@ import org.apache.atlas.repository.graphdb.AtlasIndexQuery.Result;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.collections.MapUtils;
-import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import scala.Option;
@@ -54,10 +55,10 @@ import scala.util.parsing.combinator.Parsers.NoSuccess;
 import javax.inject.Inject;
 import javax.script.ScriptException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 import static org.apache.atlas.AtlasErrorCode.DISCOVERY_QUERY_FAILED;
 
@@ -67,6 +68,10 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
     private final DefaultGraphPersistenceStrategy graphPersistenceStrategy;
     private static final Logger LOG = LoggerFactory.getLogger(EntityDiscoveryService.class);
 
+    private final static String PROPERTY_KEY_NAME        = AtlasBaseTypeDef.ATLAS_TYPE_ASSET + "." + AtlasClient.NAME;
+    private final static String PROPERTY_KEY_DESCRIPTION = AtlasBaseTypeDef.ATLAS_TYPE_ASSET + "." + AtlasClient.DESCRIPTION;
+    private final static String PROPERTY_KEY_OWNER       = AtlasBaseTypeDef.ATLAS_TYPE_ASSET + "." + AtlasClient.OWNER;
+
     @Inject
     EntityDiscoveryService(MetadataRepository metadataRepository) {
         this.graph = AtlasGraphProvider.getGraphInstance();
@@ -85,16 +90,39 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
 
             Object result = graph.executeGremlinScript(gremlinQuery.queryStr(), false);
 
-            if (result instanceof List) {
-                List queryResult = (List) result;
-
-                if (isAtlasVerticesList(queryResult)) {
-                    for (Object entity : queryResult) {
-                        ret.addEntity(toAtlasEntityHeader(entity));
+            if (result instanceof List && CollectionUtils.isNotEmpty((List)result)) {
+                List   queryResult  = (List) result;
+                Object firstElement = queryResult.get(0);
+
+                if (firstElement instanceof AtlasVertex) {
+                    for (Object element : queryResult) {
+                        if (element instanceof AtlasVertex) {
+                            ret.addEntity(toAtlasEntityHeaderwithAssociations((AtlasVertex)element));
+                        } else {
+                            LOG.warn("searchUsingDslQuery({}): expected an AtlasVertex; found unexpected entry in result {}", dslQuery, element);
+                        }
+                    }
+                } else if (firstElement instanceof Map &&
+                           (((Map)firstElement).containsKey("theInstance") || ((Map)firstElement).containsKey("theTrait"))) {
+                    for (Object element : queryResult) {
+                        if (element instanceof Map) {
+                            Map map = (Map)element;
+
+                            if (map.containsKey("theInstance")) {
+                                Object value = map.get("theInstance");
+
+                                if (value instanceof List && CollectionUtils.isNotEmpty((List)value)) {
+                                    Object entry = ((List)value).get(0);
+
+                                    if (entry instanceof AtlasVertex) {
+                                        ret.addEntity(toAtlasEntityHeaderwithAssociations((AtlasVertex)entry));
+                                    }
+                                }
+                            }
+                        } else {
+                            LOG.warn("searchUsingDslQuery({}): expected a trait result; found unexpected entry in result {}", dslQuery, element);
+                        }
                     }
-                } else if (isTraitList(queryResult)) {
-                    ret.setEntities(toTraitResult(queryResult));
-
                 } else if (gremlinQuery.hasSelectList()) {
                     ret.setAttributes(toAttributesResult(queryResult, gremlinQuery));
                 }
@@ -128,10 +156,10 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
         while (iter.hasNext() && ret.size() < params.limit()) {
             Result idxQueryResult = iter.next();
             AtlasVertex vertex = idxQueryResult.getVertex();
-            String guid = vertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class);
+            String guid = vertex != null ? vertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class) : null;
 
             if (guid != null) {
-                AtlasEntityHeader entity = toAtlasEntityHeader(idxQueryResult.getVertex());
+                AtlasEntityHeaderWithAssociations entity = toAtlasEntityHeaderwithAssociations(vertex);
                 Double score = idxQueryResult.getScore();
                 ret.add(new AtlasFullTextResult(entity, score));
             }
@@ -176,73 +204,40 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
         return new QueryParams(limit, offset);
     }
 
-    private AtlasEntityHeader toAtlasEntityHeader(Object vertexObj) {
-        AtlasEntityHeader ret = new AtlasEntityHeader();
-
-        if (vertexObj instanceof AtlasVertex) {
-            AtlasVertex vertex = (AtlasVertex) vertexObj;
-            ret.setTypeName(vertex.getProperty(Constants.TYPE_NAME_PROPERTY_KEY, String.class));
-            ret.setGuid(vertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class));
-            ret.setDisplayText(vertex.getProperty(Constants.QUALIFIED_NAME, String.class));
-
-            String state = vertex.getProperty(Constants.STATE_PROPERTY_KEY, String.class);
-            if (state != null) {
-                Status status = (state.equalsIgnoreCase("ACTIVE") ? Status.STATUS_ACTIVE : Status.STATUS_DELETED);
-                ret.setStatus(status);
-            }
-
+    private AtlasEntityHeaderWithAssociations toAtlasEntityHeaderwithAssociations(AtlasVertex vertex) {
+        if (vertex == null) {
+            return null;
         }
 
-        return ret;
-    }
-
-    private AtlasIndexQuery toAtlasIndexQuery(String fullTextQuery) {
-        String graphQuery = String.format("v.\"%s\":(%s)", Constants.ENTITY_TEXT_PROPERTY_KEY, fullTextQuery);
-        return graph.indexQuery(Constants.FULLTEXT_INDEX, graphQuery);
-    }
+        AtlasEntityHeaderWithAssociations ret = new AtlasEntityHeaderWithAssociations();
 
-    private boolean isAtlasVerticesList(List list) {
-        boolean ret = false;
-
-        if (CollectionUtils.isNotEmpty(list)) {
-            ret = list.get(0) instanceof AtlasVertex;
-        }
+        String typeName = vertex.getProperty(Constants.TYPE_NAME_PROPERTY_KEY, String.class);
 
-        return ret;
-    }
+        ret.setTypeName(typeName);
+        ret.setGuid(vertex.getProperty(Constants.GUID_PROPERTY_KEY, String.class));
+        ret.setDisplayText(vertex.getProperty(Constants.QUALIFIED_NAME, String.class));
+        ret.setAttribute(AtlasClient.NAME, vertex.getProperty(PROPERTY_KEY_NAME, String.class));
+        ret.setAttribute(AtlasClient.DESCRIPTION, vertex.getProperty(PROPERTY_KEY_DESCRIPTION, String.class));
+        ret.setAttribute(AtlasClient.OWNER, vertex.getProperty(PROPERTY_KEY_OWNER, String.class));
 
-    private boolean isTraitList(List list) {
-        boolean ret = false;
+        Collection<String> classificationNames = vertex.getPropertyValues(Constants.TRAIT_NAMES_PROPERTY_KEY, String.class);
 
-        if (CollectionUtils.isNotEmpty(list)) {
-            Object firstObj = list.get(0);
+        if (CollectionUtils.isNotEmpty(classificationNames)) {
+            ret.setClassificationNames(new ArrayList<>(classificationNames));
+        }
 
-            if (firstObj instanceof Map) {
-                Map map  = (Map) firstObj;
-                Set keys = map.keySet();
-                ret = (keys.contains("theInstance") || keys.contains("theTrait"));
-            }
+        String state = vertex.getProperty(Constants.STATE_PROPERTY_KEY, String.class);
+        if (state != null) {
+            Status status = (state.equalsIgnoreCase("ACTIVE") ? Status.STATUS_ACTIVE : Status.STATUS_DELETED);
+            ret.setStatus(status);
         }
 
         return ret;
     }
 
-    private List<AtlasEntityHeader> toTraitResult(List list) {
-        List<AtlasEntityHeader> ret = new ArrayList();
-
-        for (Object mapObj : list) {
-            Map map = (Map) mapObj;
-            if (MapUtils.isNotEmpty(map)) {
-                for (Object key : map.keySet()) {
-                    List values = (List) map.get(key);
-                    if (StringUtils.equals(key.toString(), "theInstance") && isAtlasVerticesList(values)) {
-                        ret.add(toAtlasEntityHeader(values.get(0)));
-                    }
-                }
-            }
-        }
-
-        return ret;
+    private AtlasIndexQuery toAtlasIndexQuery(String fullTextQuery) {
+        String graphQuery = String.format("v.\"%s\":(%s)", Constants.ENTITY_TEXT_PROPERTY_KEY, fullTextQuery);
+        return graph.indexQuery(Constants.FULLTEXT_INDEX, graphQuery);
     }
 
     private AttributeSearchResult toAttributesResult(List list, GremlinQuery query) {

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/e0c6b98e/webapp/src/test/java/org/apache/atlas/web/resources/EntityDiscoveryJerseyResourceIT.java
----------------------------------------------------------------------
diff --git a/webapp/src/test/java/org/apache/atlas/web/resources/EntityDiscoveryJerseyResourceIT.java b/webapp/src/test/java/org/apache/atlas/web/resources/EntityDiscoveryJerseyResourceIT.java
index a576dd0..2ade5b0 100755
--- a/webapp/src/test/java/org/apache/atlas/web/resources/EntityDiscoveryJerseyResourceIT.java
+++ b/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.AtlasEntityHeader;
+import org.apache.atlas.model.instance.AtlasEntityHeaderWithAssociations;
 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<AtlasEntityHeader> entities = searchResult.getEntities();
+        List<AtlasEntityHeaderWithAssociations> entities = searchResult.getEntities();
         assertNotNull(entities);
         assertEquals(entities.size(), 1);
 
-        AtlasEntityHeader dbEntity = entities.get(0);
+        AtlasEntityHeaderWithAssociations dbEntity = entities.get(0);
         assertEquals(dbEntity.getTypeName(), DATABASE_TYPE);
         assertEquals(dbEntity.getDisplayText(), dbName);
         assertEquals(dbEntity.getStatus(), Status.STATUS_ACTIVE);
@@ -130,11 +130,11 @@ public class EntityDiscoveryJerseyResourceIT extends BaseResourceIT {
 
         assertEquals(searchResult.getQueryText(), query);
         assertEquals(searchResult.getQueryType(), AtlasQueryType.DSL);
-        List<AtlasEntityHeader> entities = searchResult.getEntities();
+        List<AtlasEntityHeaderWithAssociations> entities = searchResult.getEntities();
         assertNotNull(entities);
         assertEquals(entities.size(), 1);
 
-        AtlasEntityHeader dbEntity = entities.get(0);
+        AtlasEntityHeaderWithAssociations dbEntity = entities.get(0);
         assertEquals(dbEntity.getTypeName(), DATABASE_TYPE);
         assertEquals(dbEntity.getDisplayText(), dbName);
         assertEquals(dbEntity.getStatus(), Status.STATUS_ACTIVE);