diff --git a/repository/src/main/java/org/apache/atlas/discovery/ESBasedSuggestionService.java b/repository/src/main/java/org/apache/atlas/discovery/ESBasedSuggestionService.java new file mode 100644 index 0000000000..e5b67ab3f1 --- /dev/null +++ b/repository/src/main/java/org/apache/atlas/discovery/ESBasedSuggestionService.java @@ -0,0 +1,82 @@ +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 SuggestionResponse searchSuggestions(String queryStr) throws IOException { + // Build the suggestor query for ES + String suggestorName = "suggest_keyword"; + Request queryRequest = new Request("POST", "/autocomplete/_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 responseMap = AtlasType.fromJson(esResponseString, Map.class); + Map suggestMap = (Map) responseMap.get("suggest"); + ArrayList suggestionMap = (ArrayList) suggestMap.get(suggestorName); + LinkedHashMap suggestions = suggestionMap.get(0); + List options = (List) 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 suggestions = new ArrayList<>(); + + public List getSuggestions() { + return suggestions; + } + + public void addSuggestion(String suggestion) { + this.suggestions.add(suggestion); + } + + public void setSuggestions(List suggestions) { + this.suggestions = suggestions; + } + } + + +} diff --git a/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java b/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java index 590e3cb0bf..7223ab0c06 100644 --- a/webapp/src/main/java/org/apache/atlas/web/rest/DiscoveryREST.java +++ b/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,9 +35,11 @@ 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; +import org.apache.atlas.type.AtlasType; import org.apache.atlas.type.AtlasTypeRegistry; import org.apache.atlas.utils.AtlasPerfMetrics; import org.apache.atlas.utils.AtlasPerfTracer; @@ -92,6 +95,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 TRACKING_UTM_TAGS = new HashSet<>(Arrays.asList("ui_main_list", "ui_popup_searchbar")); private static final String UTM_TAG_FROM_PRODUCT = "project_webapp"; @@ -827,17 +832,19 @@ public AtlasQuickSearchResult quickSearch(QuickSearchParameters quickSearchParam } @Path("suggestions") - @GET + @POST @Timed - public AtlasSuggestionsResult getSuggestions(@QueryParam("prefixString") String prefixString, @QueryParam("fieldName") String fieldName) { + public ESBasedSuggestionService.SuggestionResponse getSuggestions(Object queryStr) { AtlasPerfTracer perf = null; try { if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) { - perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.getSuggestions(" + prefixString + "," + fieldName + ")"); + perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "DiscoveryREST.getSuggestions(" + queryStr + ")"); } - return discoveryService.getSuggestions(prefixString, fieldName); + return esBasedSuggestionService.searchSuggestions(AtlasType.toJson(queryStr)); + } catch (IOException e) { + throw new RuntimeException(e); } finally { AtlasPerfTracer.log(perf); }