feat: update suggestion response

atlanhq · Aug 19, 2024 · 6e8f8a6 · 6e8f8a6
1 parent a0dbc6c
commit 6e8f8a6
Showing 1 changed file with 11 additions and 5 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
@@ -51,7 +51,7 @@ public SuggestionResponse searchSuggestions(String queryStr) throws IOException
         LinkedHashMap suggestions = suggestionMap.get(0);
         List<LinkedHashMap> options = (List<LinkedHashMap>) suggestions.get("options");
         for (LinkedHashMap option : options) {
-            suggestionResponse.addSuggestion((String) option.get("text"));
+            suggestionResponse.addSuggestion((LinkedHashMap) option.get("_source"));
         }
 
         return suggestionResponse;
@@ -63,19 +63,25 @@ public SuggestionResponse searchSuggestions(String queryStr) throws IOException
     public class SuggestionResponse {
 
         public SuggestionResponse() { }
-        private List<String> suggestions = new ArrayList<>();
+        private List<LinkedHashMap> suggestions = new ArrayList<>();
+        private final List<LinkedHashMap> sources = new ArrayList<>();
 
-        public List<String> getSuggestions() {
+        public List<LinkedHashMap> getSuggestions() {
             return suggestions;
         }
 
-        public void addSuggestion(String suggestion) {
+        public void addSuggestion(LinkedHashMap suggestion) {
             this.suggestions.add(suggestion);
         }
 
-        public void setSuggestions(List<String> suggestions) {
+        public void addSource(LinkedHashMap source) {
+            this.sources.add(source);
+        }
+
+        public void setSuggestions(List<LinkedHashMap> suggestions) {
             this.suggestions = suggestions;
         }
+
     }