feat: create the suggestions API for suggestions

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

@@ -0,0 +1,100 @@
+package org.apache.atlas.discovery;
+
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import org.apache.atlas.type.AtlasType;
+import org.apache.http.util.EntityUtils;
+import org.codehaus.jackson.annotate.JsonAutoDetect;
+import org.codehaus.jackson.annotate.JsonIgnoreProperties;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.RestClient;
+import org.elasticsearch.common.unit.Fuzziness;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.suggest.SuggestBuilder;
+import org.elasticsearch.search.suggest.SuggestBuilders;
+import org.elasticsearch.search.suggest.completion.CompletionSuggestionBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+public class ESBasedSuggestionService {
+    private static final Logger LOG = LoggerFactory.getLogger(ESBasedSuggestionService.class);
+
+    private RestClient esRestClient;
+
+    public static final String INDEX_NAME = "suggest";
+
+    public ESBasedSuggestionService(RestClient lowLevelClient) {
+        this.esRestClient = lowLevelClient;
+    }
+    public static SearchRequest buildSuggestQuery(String userQuery, String field, String suggestion_name, int fuzziness) {
+        SearchRequest searchRequest = new SearchRequest(INDEX_NAME);
+        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
+
+        Fuzziness fuzzinessObj = Fuzziness.build(fuzziness);
+        CompletionSuggestionBuilder completionSuggestionBuilder = SuggestBuilders.completionSuggestion(field)
+                .prefix(userQuery, fuzzinessObj);
+
+        SuggestBuilder suggestBuilder = new SuggestBuilder()
+                .addSuggestion(suggestion_name, completionSuggestionBuilder);
+
+        searchSourceBuilder.suggest(suggestBuilder);
+        searchRequest.source(searchSourceBuilder);
+
+        return searchRequest;
+    }
+
+    public SuggestionResponse searchSuggestions(String query, String fieldName, int fuzziness) throws IOException {
+        // Build the suggestor query for ES
+        String suggestorName = "suggest_keyword";
+        SearchRequest searchRequest = buildSuggestQuery(query, fieldName, suggestorName, fuzziness);
+        String queryStr = searchRequest.source().toString();
+        Request queryRequest = new Request("POST", "/suggest/_search");
+        queryRequest.setJsonEntity(queryStr);
+        Response response = esRestClient.performRequest(queryRequest);
+
+        SuggestionResponse suggestionResponse = new SuggestionResponse();
+        String esResponseString =  EntityUtils.toString(response.getEntity());
+
+        // Parse the response and return the suggestions
+        Map<String, Object> responseMap = AtlasType.fromJson(esResponseString, Map.class);
+        Map<String, Object> suggestMap = (Map<String, Object>) responseMap.get("suggest");
+        ArrayList<LinkedHashMap> suggestionMap = (ArrayList<LinkedHashMap>) suggestMap.get(suggestorName);
+        LinkedHashMap suggestions = suggestionMap.get(0);
+        List<LinkedHashMap> options = (List<LinkedHashMap>) suggestions.get("options");
+        for (LinkedHashMap option : options) {
+            suggestionResponse.addSuggestion((String) option.get("text"));
+        }
+
+        return suggestionResponse;
+    }
+
+    @JsonAutoDetect(getterVisibility=JsonAutoDetect.Visibility.PUBLIC_ONLY, setterVisibility=JsonAutoDetect.Visibility.PUBLIC_ONLY, fieldVisibility=JsonAutoDetect.Visibility.PUBLIC_ONLY)
+    @JsonSerialize(include=JsonSerialize.Inclusion.ALWAYS)
+    @JsonIgnoreProperties(ignoreUnknown=true)
+    public class SuggestionResponse {
+
+        public SuggestionResponse() { }
+        private List<String> suggestions = new ArrayList<>();
+
+        public List<String> getSuggestions() {
+            return suggestions;
+        }
+
+        public void addSuggestion(String suggestion) {
+            this.suggestions.add(suggestion);
+        }
+
+        public void setSuggestions(List<String> suggestions) {
+            this.suggestions = suggestions;
+        }
+    }
+
+
+}

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

@@ -25,6 +25,7 @@
 import org.apache.atlas.annotation.Timed;
 import org.apache.atlas.authorize.AtlasAuthorizationUtils;
 import org.apache.atlas.discovery.AtlasDiscoveryService;
+import org.apache.atlas.discovery.ESBasedSuggestionService;
 import org.apache.atlas.discovery.EntityDiscoveryService;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.discovery.*;
@@ -34,6 +35,7 @@
 import org.apache.atlas.model.searchlog.SearchLogSearchResult;
 import org.apache.atlas.model.searchlog.SearchRequestLogData.SearchRequestLogDataBuilder;
 import org.apache.atlas.repository.Constants;
+import org.apache.atlas.repository.graphdb.janus.AtlasElasticsearchDatabase;
 import org.apache.atlas.searchlog.SearchLoggingManagement;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasStructType;
@@ -92,6 +94,8 @@ public class DiscoveryREST {
     private final AtlasDiscoveryService discoveryService;
     private final SearchLoggingManagement loggerManagement;
 
+    private final ESBasedSuggestionService esBasedSuggestionService = new ESBasedSuggestionService(AtlasElasticsearchDatabase.getLowLevelClient());
+
     private static final String INDEXSEARCH_TAG_NAME = "indexsearch";
     private static final Set<String> TRACKING_UTM_TAGS = new HashSet<>(Arrays.asList("ui_main_list", "ui_popup_searchbar"));
     private static final String UTM_TAG_FROM_PRODUCT = "project_webapp";
@@ -829,15 +833,18 @@ public AtlasQuickSearchResult quickSearch(QuickSearchParameters quickSearchParam
     @Path("suggestions")
     @GET
     @Timed
-    public AtlasSuggestionsResult getSuggestions(@QueryParam("prefixString") String prefixString, @QueryParam("fieldName") String fieldName) {
+    public ESBasedSuggestionService.SuggestionResponse getSuggestions(@QueryParam("prefixString") String prefixString, @QueryParam("fieldName") String fieldName,
+                                                                      @QueryParam("fuzziness") int fuzziness) {
         AtlasPerfTracer perf = null;
 
         try {
             if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
                 perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.getSuggestions(" + prefixString + "," + fieldName + ")");
             }
 
-            return discoveryService.getSuggestions(prefixString, fieldName);
+            return esBasedSuggestionService.searchSuggestions(prefixString, fieldName, fuzziness);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
         } finally {
             AtlasPerfTracer.log(perf);
         }