ATLAS-3867 : Relationship search API should have a provision to fetch custom…

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

@@ -19,13 +19,13 @@
 package org.apache.atlas.discovery;
 
 
-import com.sun.xml.bind.v2.model.annotation.Quick;
 import org.apache.atlas.SortOrder;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.discovery.*;
 import org.apache.atlas.model.profile.AtlasUserSavedSearch;
 
 import java.util.List;
+import java.util.Set;
 
 public interface AtlasDiscoveryService {
     /**
@@ -83,6 +83,7 @@ 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.
@@ -90,7 +91,7 @@ public interface AtlasDiscoveryService {
      * @param offset offset to the results returned (for pagination). [ offset >= 0 ]. -1 maps to offset 0.
      * @return AtlasSearchResult
      */
-    AtlasSearchResult searchRelatedEntities(String guid, String relation, String sortByAttribute, SortOrder sortOrder, boolean excludeDeletedEntities, int limit, int offset) throws AtlasBaseException;
+    AtlasSearchResult searchRelatedEntities(String guid, String relation, Set<String> attributes, String sortByAttribute, SortOrder sortOrder, boolean excludeDeletedEntities, int limit, int offset) throws AtlasBaseException;
 
     /**
      *

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

@@ -563,7 +563,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
 
     @Override
     @GraphTransaction
-    public AtlasSearchResult searchRelatedEntities(String guid, String relation, String sortBy, SortOrder sortOrder,
+    public AtlasSearchResult searchRelatedEntities(String guid, String relation, Set<String> attributes, String sortBy, SortOrder sortOrder,
                                                    boolean excludeDeletedEntities, int limit, int offset) throws AtlasBaseException {
         AtlasSearchResult ret = new AtlasSearchResult(AtlasQueryType.RELATIONSHIP);
 
@@ -685,7 +685,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
                     List<AtlasEntityHeader> resultList = new ArrayList<>(vertices.size());
 
                     for (AtlasVertex vertex : vertices) {
-                        resultList.add(entityRetriever.toAtlasEntityHeader(vertex));
+                        resultList.add(entityRetriever.toAtlasEntityHeader(vertex, attributes));
                     }
 
                     ret.setEntities(resultList);

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

@@ -55,6 +55,7 @@ import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import java.io.IOException;
 import java.util.List;
+import java.util.Set;
 
 /**
  * REST interface for data discovery using dsl or full text search
@@ -344,6 +345,7 @@ public class DiscoveryREST {
      *
      * @param guid            Attribute name
      * @param relation        relationName
+     * @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 limit           limit the result set to only include the specified number of entries
@@ -357,6 +359,7 @@ public class DiscoveryREST {
     @Path("relationship")
     public AtlasSearchResult searchRelatedEntities(@QueryParam("guid")                   String      guid,
                                                    @QueryParam("relation")               String      relation,
+                                                   @QueryParam("attributes")             Set<String> attributes,
                                                    @QueryParam("sortBy")                 String      sortByAttribute,
                                                    @QueryParam("sortOrder")              SortOrder   sortOrder,
                                                    @QueryParam("excludeDeletedEntities") boolean     excludeDeletedEntities,
@@ -374,7 +377,7 @@ public class DiscoveryREST {
                         ", " + relation + ", " + sortByAttribute + ", " + sortOrder + ", " + excludeDeletedEntities + ", " + ", " + limit + ", " + offset + ")");
             }
 
-            return discoveryService.searchRelatedEntities(guid, relation, sortByAttribute, sortOrder, excludeDeletedEntities, limit, offset);
+            return discoveryService.searchRelatedEntities(guid, relation, attributes, sortByAttribute, sortOrder, excludeDeletedEntities, limit, offset);
         } finally {
             AtlasPerfTracer.log(perf);
         }