ATLAS-3949 : Relationship search API, add parameter to get classification…

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

@@ -83,16 +83,11 @@ public interface AtlasDiscoveryService {
      *
      * @param guid unique ID of the entity.
      * @param relation relation name.
-     * @param attributes set of attributes in search result.
-     * @param sortByAttribute sort the result using this attribute name, default value is 'name'
-     * @param sortOrder sorting order
-     * @param excludeDeletedEntities exclude deleted entities in search result.
      * @param getApproximateCount
-     * @param limit number of resultant rows (for pagination). [ limit > 0 ] and [ limit < maxlimit ]. -1 maps to atlas.search.defaultlimit property.
-     * @param offset offset to the results returned (for pagination). [ offset >= 0 ]. -1 maps to offset 0.
+     * @param searchParameters
      * @return AtlasSearchResult
      */
-    AtlasSearchResult searchRelatedEntities(String guid, String relation, Set<String> attributes, String sortByAttribute, SortOrder sortOrder, boolean excludeDeletedEntities, boolean getApproximateCount, int limit, int offset) throws AtlasBaseException;
+     AtlasSearchResult searchRelatedEntities(String guid, String relation, boolean getApproximateCount, SearchParameters searchParameters) throws AtlasBaseException;
 
     /**
      *

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

@@ -569,8 +569,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
 
     @Override
     @GraphTransaction
-    public AtlasSearchResult searchRelatedEntities(String guid, String relation, Set<String> attributes, String sortBy, SortOrder sortOrder,
-                                                   boolean excludeDeletedEntities, boolean getApproximateCount, int limit, int offset) throws AtlasBaseException {
+    public AtlasSearchResult searchRelatedEntities(String guid, String relation, boolean getApproximateCount, SearchParameters searchParameters) throws AtlasBaseException {
         AtlasSearchResult ret = new AtlasSearchResult(AtlasQueryType.RELATIONSHIP);
 
         if (StringUtils.isEmpty(guid) || StringUtils.isEmpty(relation)) {
@@ -605,6 +604,10 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
         }
 
         //validate sortBy attribute
+        String    sortBy           = searchParameters.getSortBy();
+        SortOrder sortOrder        = searchParameters.getSortOrder();
+        int       offset           = searchParameters.getOffset();
+        int       limit            = searchParameters.getLimit();
         String sortByAttributeName = DEFAULT_SORT_ATTRIBUTE_NAME;
 
         if (StringUtils.isNotEmpty(sortBy)) {
@@ -642,7 +645,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
         //get relationship(end vertices) vertices
         GraphTraversal gt = graph.V(entityVertex.getId()).bothE(relation).otherV();
 
-        if (excludeDeletedEntities) {
+        if (searchParameters.getExcludeDeletedEntities()) {
             gt.has(Constants.STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
         }
 
@@ -662,7 +665,13 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
 
             if (v != null && v.property(Constants.GUID_PROPERTY_KEY).isPresent()) {
                 String endVertexGuid     = v.property(Constants.GUID_PROPERTY_KEY).value().toString();
-                resultList.add(entityRetriever.toAtlasEntityHeader(endVertexGuid, attributes));
+                AtlasVertex vertex       = entityRetriever.getEntityVertex(endVertexGuid);
+                AtlasEntityHeader entity = entityRetriever.toAtlasEntityHeader(vertex, searchParameters.getAttributes());
+
+                if (searchParameters.getIncludeClassificationAttributes()) {
+                    entity.setClassifications(entityRetriever.getAllClassifications(vertex));
+                }
+                resultList.add(entity);
             }
         }
 
@@ -677,7 +686,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
         if (getApproximateCount) {
             Iterator<AtlasEdge> edges = GraphHelper.getAdjacentEdgesByLabel(entityVertex, AtlasEdgeDirection.BOTH, relation);
 
-            if (excludeDeletedEntities) {
+            if (searchParameters.getExcludeDeletedEntities()) {
                 List<AtlasEdge> edgeList = new ArrayList<>();
                 edges.forEachRemaining(edgeList::add);
 

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

@@ -370,6 +370,7 @@ public class DiscoveryREST {
                                                    @QueryParam("sortBy")                          String      sortByAttribute,
                                                    @QueryParam("sortOrder")                       SortOrder   sortOrder,
                                                    @QueryParam("excludeDeletedEntities")          boolean     excludeDeletedEntities,
+                                                   @QueryParam("includeClassificationAttributes") boolean     includeClassificationAttributes,
                                                    @QueryParam("getApproximateCount")             boolean     getApproximateCount,
                                                    @QueryParam("limit")                           int         limit,
                                                    @QueryParam("offset")                          int         offset) throws AtlasBaseException {
@@ -385,7 +386,16 @@ public class DiscoveryREST {
                         ", " + relation + ", " + sortByAttribute + ", " + sortOrder + ", " + excludeDeletedEntities + ", " + getApproximateCount + ", " + limit + ", " + offset + ")");
             }
 
-            return discoveryService.searchRelatedEntities(guid, relation, attributes, sortByAttribute, sortOrder, excludeDeletedEntities, getApproximateCount, limit, offset);
+            SearchParameters parameters = new SearchParameters();
+            parameters.setAttributes(attributes);
+            parameters.setSortBy(sortByAttribute);
+            parameters.setSortOrder(sortOrder);
+            parameters.setExcludeDeletedEntities(excludeDeletedEntities);
+            parameters.setLimit(limit);
+            parameters.setOffset(offset);
+            parameters.setIncludeClassificationAttributes(includeClassificationAttributes);
+            return discoveryService.searchRelatedEntities(guid, relation, getApproximateCount, parameters);
+
         } finally {
             AtlasPerfTracer.log(perf);
         }