ATLAS-3517: Enhance suggestions REST API to have optional 'fieldName' parameter…

graphdb/api/src/main/java/org/apache/atlas/repository/graphdb/AtlasGraphIndexClient.java

@@ -18,12 +18,9 @@
 package org.apache.atlas.repository.graphdb;
 
 import org.apache.atlas.model.discovery.AtlasAggregationEntry;
-import org.apache.atlas.model.discovery.SearchParameters;
-import org.apache.atlas.type.AtlasEntityType;
 
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 /**
  * Represents a graph client work with indices used by Jansgraph.
@@ -39,9 +36,10 @@ public interface AtlasGraphIndexClient {
     /**
      * Returns top 5 suggestions for the given prefix string.
      * @param prefixString the prefix string whose value needs to be retrieved.
+     * @param indexFieldName the indexed field name from which to retrieve suggestions
      * @return top 5 suggestion strings with prefix String
      */
-    List<String> getSuggestions(String prefixString);
+    List<String> getSuggestions(String prefixString, String indexFieldName);
 
     /**
      *  The implementers should apply the search weights for the passed in properties.

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusGraphIndexClient.java

@@ -28,6 +28,7 @@ import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.collections.MapUtils;
 import org.apache.commons.configuration.Configuration;
+import org.apache.commons.lang.StringUtils;
 import org.apache.solr.client.solrj.SolrClient;
 import org.apache.solr.client.solrj.SolrQuery;
 import org.apache.solr.client.solrj.SolrRequest;
@@ -55,6 +56,8 @@ public class AtlasJanusGraphIndexClient implements AtlasGraphIndexClient {
     private static final FreqComparator FREQ_COMPARATOR          = new FreqComparator();
     private static final int            DEFAULT_SUGGESTION_COUNT = 5;
     private static final int            MIN_FACET_COUNT_REQUIRED = 1;
+    private static final String         TERMS_PREFIX             = "terms.prefix";
+    private static final String         TERMS_FIELD              = "terms.fl";
 
     private final Configuration configuration;
 
@@ -269,7 +272,7 @@ public class AtlasJanusGraphIndexClient implements AtlasGraphIndexClient {
     }
 
     @Override
-    public List<String> getSuggestions(String prefixString) {
+    public List<String> getSuggestions(String prefixString, String indexFieldName) {
         SolrClient solrClient = null;
 
         try {
@@ -284,9 +287,13 @@ public class AtlasJanusGraphIndexClient implements AtlasGraphIndexClient {
             SolrQuery solrQuery = new SolrQuery();
 
             solrQuery.setRequestHandler(Constants.TERMS_REQUEST_HANDLER)
-                     .setParam("terms.prefix", prefixString)
+                     .setParam(TERMS_PREFIX, prefixString)
                      .setParam(CommonParams.OMIT_HEADER, true);
 
+            if (StringUtils.isNotEmpty(indexFieldName)) {
+                solrQuery.setParam(TERMS_FIELD, indexFieldName);
+            }
+
             QueryResponse queryResponse = solrClient.query(VERTEX_INDEX, solrQuery);
             TermsResponse termsResponse = queryResponse == null? null: queryResponse.getTermsResponse();
 

intg/src/main/java/org/apache/atlas/model/discovery/AtlasSuggestionsResult.java

@@ -33,9 +33,11 @@ import static com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility.PUBLIC_
 public class AtlasSuggestionsResult {
     private List<String> suggestions;
     private String       prefixString;
+    private String       fieldName;
 
-    public AtlasSuggestionsResult(String prefixString) {
+    public AtlasSuggestionsResult(String prefixString, String fieldName) {
         this.prefixString = prefixString;
+        this.fieldName    = fieldName;
     }
 
     public List<String> getSuggestions() {
@@ -53,4 +55,12 @@ public class AtlasSuggestionsResult {
     public void setPrefixString(String prefixString) {
         this.prefixString = prefixString;
     }
+
+    public String getFieldName() {
+        return fieldName;
+    }
+
+    public void setFieldName(String fieldName) {
+        this.fieldName = fieldName;
+    }
 }

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

@@ -151,7 +151,8 @@ public interface AtlasDiscoveryService {
     /**
      * Should return top 5 suggestion strings for the given prefix.
      * @param prefixString the prefix string
+     * @param fieldName field from which to retrieve suggestions
      * @return top 5 suggestion strings for the given prefix.
      */
-    AtlasSuggestionsResult getSuggestions(String prefixString);
+    AtlasSuggestionsResult getSuggestions(String prefixString, String fieldName);
 }

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

@@ -109,7 +109,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
         this.maxTagsLengthInIdxQuery  = ApplicationProperties.get().getInt(Constants.INDEX_SEARCH_TAGS_MAX_QUERY_STR_LENGTH, 512);
         this.indexSearchPrefix        = AtlasGraphUtilsV2.getIndexSearchPrefix();
         this.userProfileService       = userProfileService;
-        this.suggestionsProvider      = new SuggestionsProviderImpl(graph);
+        this.suggestionsProvider      = new SuggestionsProviderImpl(graph, typeRegistry);
     }
 
     @Override
@@ -448,8 +448,8 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
 
     @Override
     @GraphTransaction
-    public AtlasSuggestionsResult getSuggestions(String prefixString) {
-        return suggestionsProvider.getSuggestions(prefixString);
+    public AtlasSuggestionsResult getSuggestions(String prefixString, String fieldName) {
+        return suggestionsProvider.getSuggestions(prefixString, fieldName);
     }
 
     @Override

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

@@ -20,5 +20,5 @@ package org.apache.atlas.discovery;
 import org.apache.atlas.model.discovery.AtlasSuggestionsResult;
 
 public interface SuggestionsProvider {
-    AtlasSuggestionsResult getSuggestions(String prefixString);
+    AtlasSuggestionsResult getSuggestions(String prefixString, String indexFieldName);
 }

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

@@ -21,6 +21,7 @@ import org.apache.atlas.AtlasException;
 import org.apache.atlas.model.discovery.AtlasSuggestionsResult;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasGraphIndexClient;
+import org.apache.atlas.type.AtlasTypeRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -31,19 +32,22 @@ public class SuggestionsProviderImpl implements SuggestionsProvider {
     private static final Logger LOG = LoggerFactory.getLogger(SuggestionsProviderImpl.class);
 
     private final AtlasGraph        graph;
+    private final AtlasTypeRegistry typeRegistry;
 
-    public SuggestionsProviderImpl(AtlasGraph graph) {
+    public SuggestionsProviderImpl(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
         this.graph        = graph;
+        this.typeRegistry = typeRegistry;
     }
 
     @Override
-    public AtlasSuggestionsResult getSuggestions(String prefixString) {
-        AtlasSuggestionsResult result = new AtlasSuggestionsResult(prefixString);
+    public AtlasSuggestionsResult getSuggestions(String prefixString, String fieldName) {
+        AtlasSuggestionsResult result = new AtlasSuggestionsResult(prefixString, fieldName);
 
         try {
             AtlasGraphIndexClient graphIndexClient = graph.getGraphIndexClient();
+            String                indexFieldName   = (fieldName == null) ? null : typeRegistry.getIndexFieldName(fieldName);
 
-            result.setSuggestions(graphIndexClient.getSuggestions(prefixString));
+            result.setSuggestions(graphIndexClient.getSuggestions(prefixString, indexFieldName));
         } catch (AtlasException e) {
             LOG.error("Error encountered in performing quick suggestions. Will return no suggestions.", e);
 

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

@@ -650,15 +650,15 @@ public class DiscoveryREST {
 
     @Path("suggestions")
     @GET
-    public AtlasSuggestionsResult getSuggestions(@QueryParam("prefixString") String prefixString) throws AtlasBaseException {
+    public AtlasSuggestionsResult getSuggestions(@QueryParam("prefixString") String prefixString, @QueryParam("fieldName") String fieldName) {
         AtlasPerfTracer perf = null;
 
         try {
             if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
-                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.getSuggestions(" + prefixString + ")");
+                perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.getSuggestions(" + prefixString + "," + fieldName + ")");
             }
 
-            return discoveryService.getSuggestions(prefixString);
+            return discoveryService.getSuggestions(prefixString, fieldName);
         } finally {
             AtlasPerfTracer.log(perf);
         }