ATLAS-3876 : Relationship Search API not showing correct approximateCount

intg/src/main/java/org/apache/atlas/type/AtlasStructType.java

@@ -670,7 +670,7 @@ public class AtlasStructType extends AtlasType {
         throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_ATTRIBUTE, attrName, structDef.getName());
     }
 
-    AtlasEntityType getReferencedEntityType(AtlasType type) {
+    static AtlasEntityType getReferencedEntityType(AtlasType type) {
         if (type instanceof AtlasArrayType) {
             type = ((AtlasArrayType)type).getElementType();
         }
@@ -980,6 +980,11 @@ public class AtlasStructType extends AtlasType {
             return (relationshipLabel == null) ? getEdgeLabel(qualifiedName) : relationshipLabel;
         }
 
+        public AtlasEntityType getReferencedEntityType(AtlasTypeRegistry typeRegistry) throws AtlasBaseException {
+            AtlasType type = typeRegistry.getType(attributeDef.getTypeName());
+            return AtlasStructType.getReferencedEntityType(type);
+        }
+
         public static String getQualifiedAttributeName(AtlasStructDef structDef, String attrName) {
             if (isRootType(structDef)) {
                 return attrName;

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

@@ -87,11 +87,12 @@ public interface AtlasDiscoveryService {
      * @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.
      * @return AtlasSearchResult
      */
-    AtlasSearchResult searchRelatedEntities(String guid, String relation, Set<String> attributes, 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, boolean getApproximateCount, int limit, int offset) throws AtlasBaseException;
 
     /**
      *

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

@@ -75,7 +75,7 @@ public abstract class SearchProcessor {
     public static final String  CUSTOM_ATTR_SEARCH_FORMAT  = "\"\\\"%s\\\":\\\"%s\\\"\"";
     public static final String  CUSTOM_ATTR_SEARCH_FORMAT_GRAPH = "\"%s\":\"%s\"";
     private static final Map<SearchParameters.Operator, String>                            OPERATOR_MAP           = new HashMap<>();
-    private static final Map<SearchParameters.Operator, VertexAttributePredicateGenerator> OPERATOR_PREDICATE_MAP = new HashMap<>();
+    private static final Map<SearchParameters.Operator, ElementAttributePredicateGenerator> OPERATOR_PREDICATE_MAP = new HashMap<>();
 
     static
     {
@@ -697,7 +697,7 @@ public abstract class SearchProcessor {
         Predicate ret = null;
 
         AtlasAttribute                    attribute = type.getAttribute(attrName);
-        VertexAttributePredicateGenerator predicate = OPERATOR_PREDICATE_MAP.get(op);
+        ElementAttributePredicateGenerator predicate = OPERATOR_PREDICATE_MAP.get(op);
 
         if (attribute != null && predicate != null) {
             final AtlasType attrType = attribute.getAttributeType();

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java

@@ -201,6 +201,10 @@ public class EntityGraphRetriever {
         return toAtlasEntityHeader(getEntityVertex(guid));
     }
 
+    public AtlasEntityHeader toAtlasEntityHeader(String guid, Set<String> attributes) throws AtlasBaseException {
+        return toAtlasEntityHeader(getEntityVertex(guid), attributes);
+    }
+
     public AtlasEntityHeader toAtlasEntityHeader(AtlasVertex entityVertex) throws AtlasBaseException {
         return toAtlasEntityHeader(entityVertex, Collections.<String>emptySet());
     }

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("getApproximateCount")    boolean     getApproximateCount,
                                                    @QueryParam("limit")                  int         limit,
                                                    @QueryParam("offset")                 int         offset) throws AtlasBaseException {
         Servlets.validateQueryParamLength("guid", guid);
@@ -380,11 +381,11 @@ public class DiscoveryREST {
 
         try {
             if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
-                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.relatedEntitiesSearchUsingGremlin(" + guid +
-                        ", " + relation + ", " + sortByAttribute + ", " + sortOrder + ", " + excludeDeletedEntities + ", " + ", " + limit + ", " + offset + ")");
+                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.relatedEntitiesSearch(" + guid +
+                        ", " + relation + ", " + sortByAttribute + ", " + sortOrder + ", " + excludeDeletedEntities + ", " + getApproximateCount + ", " + limit + ", " + offset + ")");
             }
 
-            return discoveryService.searchRelatedEntities(guid, relation, attributes, sortByAttribute, sortOrder, excludeDeletedEntities, limit, offset);
+            return discoveryService.searchRelatedEntities(guid, relation, attributes, sortByAttribute, sortOrder, excludeDeletedEntities, getApproximateCount, limit, offset);
         } finally {
             AtlasPerfTracer.log(perf);
         }