feat: add direct query option

atlanhq · Aug 12, 2024 · 3d6913a · 3d6913a
1 parent 408e5ca
commit 3d6913a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 25 deletions.
diff --git a/repository/src/main/java/org/apache/atlas/discovery/ESBasedSuggestionService.java b/repository/src/main/java/org/apache/atlas/discovery/ESBasedSuggestionService.java
@@ -33,29 +33,11 @@ public class ESBasedSuggestionService {
     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 {
+    public SuggestionResponse searchSuggestions(String queryStr) 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");
+        Request queryRequest = new Request("POST", "/autocomplete/_search");
         queryRequest.setJsonEntity(queryStr);
         Response response = esRestClient.performRequest(queryRequest);
 

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
@@ -39,6 +39,7 @@
 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;
@@ -831,18 +832,17 @@ public AtlasQuickSearchResult quickSearch(QuickSearchParameters quickSearchParam
     }
 
     @Path("suggestions")
-    @GET
+    @POST
     @Timed
-    public ESBasedSuggestionService.SuggestionResponse getSuggestions(@QueryParam("prefixString") String prefixString, @QueryParam("fieldName") String fieldName,
-                                                                      @QueryParam("fuzziness") int fuzziness) {
+    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 esBasedSuggestionService.searchSuggestions(prefixString, fieldName, fuzziness);
+            return esBasedSuggestionService.searchSuggestions(AtlasType.toJson(queryStr));
         } catch (IOException e) {
             throw new RuntimeException(e);
         } finally {